FaceSwap lib.infer 包详解：从检测到对齐的多线程并行推理管线实现

lib.infer 是 FaceSwap 中负责"从图片和视频中提取人脸"的核心推理库，其 Sphinx API 参考页见 docs/full/lib/infer.rst。读完本文，你能掌握该包 10 个子模块的职责划分（detect / align / mask / identity / runner / iterator / objects / handler / plugin_utils / profile）、Handler-Runner-队列的三级多线程管线架构，以及 re_feeds、re_align、rotation、身份过滤阈值等关键参数的源码级行为，从而理解 FaceSwap 提取管线（extract pipeline）是如何把一帧原始图像变成带关键点、掩码和身份向量的人脸批次的。

1. 包定位与模块全景

包入口 lib/infer/init.py 的模块 docstring 一句话概括了包的定位：Parallel batched inference library for faceswap.py（faceswap.py 的并行批处理推理库）。它对外导出 6 个主要类：Align、Detect、FileHandler（别名 File）、Identity、Mask 和 Profiler。

API 文档页 docs/full/lib/infer.rst 用 automodapi 指令依次列出了包的 10 个子模块，每个模块在 lib/infer/ 目录下都有对应源文件。结合源码，各模块的职责如下：

模块	源文件	核心职责
align	align.py	人脸关键点（landmark）对齐插件的调度、re-feed 抖动、re-align 二次对齐、输入归一化、对齐结果过滤
detect	detect.py	人脸检测插件调度：整图缩放/旋转增强、检测框回映射与尺寸过滤
handler	handler.py	所有插件处理器的抽象基类：模型加载/编译/预热、OOM 处理、图像格式化、文件直通处理器
identity	identity.py	身份识别插件调度与基于余弦相似度的 filter/nFilter 人脸过滤
iterator	iterator.py	管线各队列上的迭代器：分批、合并批次、处理 FLUSH/SHUTDOWN 信号与零检测帧
mask	mask.py	人脸掩码插件调度：带 ROI alpha 通道的对齐人脸输入、越界裁切
objects	objects.py	管线数据对象：ExtractBatch、ExtractBatchAligned、ExtractBatchMask、ExtractSignal
plugin_utils	plugin_utils.py	通用工具：随机输入构造、Torch 模块递归发现、插件预热、模型编译
profile	profile.py	GPU 性能剖析（单插件基准测试与整管线吞吐优化）
runner	runner.py	ExtractRunner 执行器：每插件 3 个线程 + 4 个队列的线程/队列编排

2. 管线架构：Handler、Runner 与队列

2.1 ExtractHandler：插件的三段式抽象

所有插件处理器都继承 ExtractHandler（抽象基类），其核心约定是三段式处理器元组：

processors = ("pre_process", "process", "post_process")

基类完成与具体插件类型无关的工作：

• 插件加载：构造函数中通过 PluginLoader.get_extractor(self.plugin_type, plugin) 按类型（detect/align/mask/identity）加载插件实例，并用 _is_overridden() 检查插件是否重写了 pre_process/process/post_process，从而决定哪些环节需要走插件自定义逻辑；
• 模型初始化 handler.py：init_model() 加载模型；若请求编译则调用 compile_models()，否则发送一批 warmup 输入；
• 推理入口 _predict() handler.py：两个值得注意的工程细节——
  1. 若启用模型编译且本批数据量小于插件 batch size，会用空数组把批次补齐到 batch size 再送入模型，避免动态形状触发重编译，推理后再裁回原尺寸；
  2. 捕获 torch.cuda.OutOfMemoryError 并抛出带操作建议的 FaceswapError，提示用户关闭占用 GPU 的程序、降低 faceswap/config/extract.ini 中的 batch size、或改用轻量插件（见 handler.py 的 OOM_MESSAGE）；
• 图像格式化 _format_images()：把 0–255 的 UINT8 图像按插件声明的 dtype 和 scale（如 (0,1) 或 (-1,1)）做数值缩放；
• 运行器构建 __call__()：创建 ExtractRunner 并挂接上游 runner 的输出队列，返回 runner。

此外还有两个派生类：

• ExtractHandlerFace：面向"以对齐人脸为输入"的插件（mask/identity 使用）。它读取插件的 centering（legacy/face/head 三种居中方式）和 EXTRACT_RATIOS 计算 padding，并提供 _get_faces()（从整帧做仿射对齐裁剪）与 _get_faces_aligned()（从已对齐人脸做子裁剪/缩放）两条取图路径；
• FileHandler：伪处理器。当整条管线完全由 alignments 文件驱动（不加载任何插件）时，数据从输入直通输出，processors 为空元组，即不启动任何线程。

2.2 ExtractRunner：每插件 3 线程、4 队列的编排

ExtractRunner 是真正驱动数据流动的对象，从源码结构看其拓扑为：

[上游 out 队列 / 本插件入队] → pre_process 线程 → 队列 → process 线程 → 队列 → post_process 线程 → 队列(out) → 下游

• 线程：_get_threads() 按 handler.processors 为每个处理器注册一个 FSThread，由 PluginThreads 统一 start/join 并共享 ErrorState。任一线程出错时，其他线程在 queue.put/迭代检查时发现 error_state.has_error 即停止推进；FileHandler 无线程，可通过 register_external_error_state() 借用调用方线程的错误状态，避免管线挂死；
• 队列：_get_queues() 为每个处理阶段以及输出各建一个 Queue(maxsize=1)，入队队列直接复用上游 runner 的 out_queue。容量为 1 的队列天然形成背压：下游处理不过来时上游 queue.put(timeout=0.2) 超时重试，实现生产者/消费者节流；
• 主循环 _process_batches() runner.py：从本阶段队列迭代取批次 → 调用对应 handler 方法 → 放入下一队列。非编译模型选择在 process 线程内部才 init_model()，以加快线程启动；结束时向下游发送 ExtractSignal.SHUTDOWN；
• 零检测短路 _handle_zero_detections()：对非 detect 插件，若批次里没有任何检测框（batch.frame_ids 为空），则跳过 handler 直接透传到下一队列，避免对"没有人脸"的帧做无谓推理；
• 直通批次 _process_passthrough()：passthrough 批次的框数可能超过插件 batch size，会被切分成多份分别推理后合并；
• 输出清理 _clean_output()：管线出口处删除无检测框的图像内存（替换为空数组），并删除 matrices、data 等仅内部使用的属性，降低内存占用；
• 迭代接口 __next__()：只有链上最后一个 runner 可以 next()，返回 FrameFaces（帧 + 已检测到的人脸）对象；非主线程中检测到线程错误时返回 None，由主线程 re_raise() 抛出。

模块级还有 _PLUGIN_REGISTER 字典记录 runner 链的 uuid 顺序，用于断言管线构建顺序与追踪多插件管线；runner.py 末尾还提供 get_pipeline() 供 profile 模块构建整条管线。

2.3 Iterator：批次装配与信号处理

iterator.py 定义了抽象基类 ExtractIterator 及 4 个子类（InputIterator、InboundIterator、InterimIterator、OutputIterator）。基类要点：

• 每次 queue.get(timeout=0.2) 非阻塞取数，配合 FIFO 列表在本地攒批；
• _zero_detect_threshold = batch_size * 2：当输入类迭代器发现 FIFO 中积压的"零检测帧"达到该阈值时提前放行批次，防止长段无人脸的片源把帧堆进内存并拖长管线空转阶段；
• 管线信号：识别 ExtractSignal（FLUSH=1 冲刷排队项，SHUTDOWN=2 冲刷并停止）。收到 SHUTDOWN 时若 FIFO 中还留有最后一个批次，会先把它吐出再抛 StopIteration，保证收尾数据不丢失；
• 每次取数前检查 _error_state.has_error，出错立即停止迭代。

InputIterator 负责把入口的 FrameFaces 对象组装成 ExtractBatch 并按插件 batch size 切分；InboundIterator 处理"上游已是批次、当前插件 batch size 不同"时的重新分箱；OutputIterator 则把批次还原为按帧组织的 FrameFaces 输出。

3. 数据对象：ExtractBatch 家族

管线中流动的数据封装在 objects.py：

• ExtractSignal（IntEnum）：FLUSH 与 SHUTDOWN 两种控制信号；
• ExtractBatch：核心批次数据类，字段包括 filenames（帧文件名）、images（原始帧图像列表）、sources、is_aligned（images 是整帧还是已对齐人脸）、frame_sizes/frame_metadata（对齐人脸模式下的原始帧信息）、passthrough 等。注释中特别说明了批次语义："检测器之后，批大小不再是图像数，而是一张图可能含 0~N 张人脸，需要拆分/重组以适配各插件的内部 batch size"；
• ExtractBatchAligned：关键点批次的"缓存型"数据类。除 landmarks 与 landmark_type 外，还以惰性缓存属性提供：
  • landmarks_68：把非 68 点关键点（如 51 点）转换为 68 点格式；
  • matrices：由 batch_umeyama(points, MEAN_FACE[...]) 计算的"帧空间 → 归一化(0,1)空间"Umeyama 相似变换矩阵（51 点人脸时取 landmarks[:, 17:] 与 LM_2D_51 均值脸配对）；
  • offsets_face/offsets_head：由 Batch3D.get_offsets() 基于 PnP 估计的 3D 姿态算出的 face/head 居中偏移；
  • rotation/translation：Batch3D.solve_pnp() 求解的 (N,3,1) 旋转/平移向量，供 roll 过滤等使用；
  • 支持切片 __getitem__、append 与 apply_mask（布尔掩码过滤时同步过滤所有缓存数组）。
• ExtractBatchMask：掩码批次的 centering、归一化矩阵、storage_size 与 masks 数组，同样支持切片/追加/掩码过滤。

这套对象设计让每个 handler 只做"就地更新批次字段"（如 batch.data = ...、batch.bboxes = ...），无需自行管理数据传递。

4. Detect 阶段：整图检测与旋转增强

Detect 处理器构造参数及默认值：

参数	默认	含义
plugin	必填	检测插件名
rotation	None	旋转角度列表（0–360 逗号分隔）或单一步进值；None 表示不旋转
min_size / max_size	0	接受检测的最小/最大比例：以"检测框最长边 ÷ 帧最短边"计，0 表示不过滤
compile_model	False	是否编译 PyTorch 模型
config_file	None	自定义 extract 配置路径

预处理（pre_process）：_get_matrices() 为每帧计算"缩放 + 居中 padding"的 3×3 仿射矩阵，把任意宽高比的帧映射到插件的 input_size 正方形输入；_scale_images() 用 cv2.warpAffine 执行变换（缩放因子 < 1 用 INTER_AREA 防止振铃，否则 INTER_CUBIC），再经 _format_images() 做数值格式化后调用 plugin.pre_process()。

推理与旋转增强（process）：未启用旋转时只做一次推理。启用旋转时进入循环——Rotator 预计算每个角度的旋转矩阵及其逆矩阵（旋转矩阵转置 + 反向平移，_pre_compute_inverse_matrices()），然后逐角度用 BORDER_REPLICATE 填充旋转图像，记录"每帧在哪个角度索引被检测到人脸"（indices_angle），已检出的帧从待处理掩码中剔除，直到全部检出或角度用尽。

后处理（post_process）：

1. 非旋转模式在此调用插件 post_process；旋转模式则对每个检测框的四个角点用预计算的逆矩阵 un_rotate() 回原帧坐标，再取四角最小/最大值重建轴对齐框；
2. _stack_boxes() 把各帧的可变长框列表堆叠为单一数组，并用 np.repeat 建立"帧 → 多张脸"的映射 batch.frame_ids；
3. _scale_boxes() 用 batch.matrices 把模型输入尺寸的框反算回原始帧尺寸并取整为 int32；
4. _filter_boxes() 按 min_size/max_size 过滤（比较框最长边与"帧最短边 × 百分比"），并在结束时通过 output_info() 以 info 日志汇总被过滤数量。

5. Align 阶段：关键点精化、re-feed 与 re-align

Align 是参数最多的处理器，构造参数及默认值：

参数	默认	含义
plugin	必填	对齐（landmark）插件名
re_feeds	0	额外推理次数：对检测框做随机微扰后多次推理并取平均
re_align	False	是否基于首轮结果做二次对齐（两遍推理）
normalization	None	对齐输入归一化：none/clahe/hist/mean
filters	False	是否启用对齐结果过滤器
compile_model / config_file	False / None	同 Detect

5.1 常规路径：ROI 裁剪 → 推理 → 回映射

• 首轮 _prepare_data(iteration=1)：plugin.pre_process(boxes) 由检测框生成各插件偏好的正方形 ROI（不同插件的取景范围不同）；_matrices_from_roi() 把 ROI 转为 3×3 仿射矩阵（断言 ROI 必须为正方形），用于把预测的关键点映射回帧空间；_crop_and_resize() 把 ROI 裁出并缩放到 input_size（放大用 INTER_CUBIC、缩小用 INTER_AREA），越界部分由 _clamp_roi()/_get_destinations() 夹取到帧边界并留黑边；
• 推理 process()：无 re_align 时单遍推理，post_process() 先调用 plugin.post_process()（若插件重写），再用 batch_transform(batch.matrices, ...) 把归一化空间的关键点变换回帧空间，得到 batch.landmarks，并按形状推断 LandmarkType（68/98/51 点等）；
• 要求结果为 float32 数组（断言保证）。

5.2 ReFeed：随机抖动平均

ReFeed 通过多次推理抑制单帧预测噪声：total_feeds = re_feeds + 1，除首次外每次对 ROI 做随机扰动——缩放因子在 [1-β, 1+β] 均匀采样、四角平移量也在该范围内随机（β=0.05，即相对框最短边 ±5%）。所有 feed 的预测在 post_process() 中由 merge() 沿 feed 维取平均。re-feed 只在最终一遍（is_final）执行，且当 re_align 启用时会在"对齐后的框"上执行。

5.3 ReAlign：两遍对齐

ReAlign 启用时 iterations=2，原理是从源码注释可归纳为：

1. 第一遍：按检测框 ROI 推理得到首轮关键点；_first_pass() 用 batch_umeyama(landmarks[:, 17:], MEAN_FACE[LM_2D_51]) 算出人脸的标准对齐矩阵，再叠加插件 ROI 相对检测框的 padding/offset 修正与一个固定 30% padding 的 _adjust_matrix（弥补检测框"松紧"），用 cv2.warpAffine 把整帧 warp 成 expanded_size（= input_size × (1 + 2β)）的扩展对齐人脸图存于 ReAlign._images；
2. 第二遍：get_images() 从扩展图中按 re-feed 矩阵裁出 input_size 的子图再次送模型，得到的关键点更贴合真实人脸姿态；_second_pass() 把累积的矩阵（roi_matrices @ inv(align) @ inv(patch)）写回 batch.matrices，供后处理把关键点变换回帧空间。

5.4 输入归一化 Normalize

Normalize 支持 4 种方法：none（直通）、clahe（每通道 CLAHE，clipLimit=2.0、tileGridSize=(4,4)）、hist（每通道 cv2.equalizeHist）、mean（每通道 min-max 归一到 0–255）。可在运行时用 set_normalize_method() 切换（提取界面中可动态调整）。

5.5 AlignedFilter：对齐结果质量过滤

AlignedFilter 的阈值全部取自 extract 配置（cfg.aligner_features()、aligner_min_scale()、aligner_max_scale()、aligner_distance()（百分制，内部除以 100）、aligner_roll()，见 plugins/extract/extract_config.py），只要任一项非零即自动启用。四类过滤器：

• features：检查眼/嘴相对位置——眼睛区域（landmark 索引 17–26 与 36–47）的最低 y 必须高于嘴部（48–67）的最高 y，否则判定为无效脸；
• scale：按对齐矩阵线性分量估算脸相对帧的像素尺度，过滤超出 [min_scale, max_scale] 的脸；
• distance：归一化关键点与 51 点均值脸的逐点平均距离不得超过阈值；
• roll：Batch3D.roll() 估计的 3D 翻滚角绝对值不得超过阈值。

注意该过滤器仅支持 LM_2D_68 与 LM_2D_98，遇到其他关键点类型会打 warning 并自动禁用；每类过滤计数在管线结束时经 output_info() 汇总为一条 [Align filter] info 日志。

6. Mask 与 Identity 阶段

6.1 Mask：带 ROI 通道的掩码推理

Mask 处理器的特殊之处在于输入是 4 通道图像：

• storage_size 取自 cfg.mask_storage_size()，若配置值落在 (0, 64) 会打 warning 并强制抬到 64（掩码存储下限）；
• 预处理：无论输入是整帧还是已对齐人脸，都先经 ExtractHandlerFace._get_faces()/_get_faces_aligned() 得到对齐人脸，再把"原始帧内为白、越界为黑"的 ROI 掩码 warp 到第 4 通道，供模型区分"人脸区"与"出界填充区"；
• 后处理 mask.py：_crop_out_of_bounds() 用 ROI 通道把预测掩码中落在原帧外的部分清零；storage_size 为 0 时按实际输出尺寸初始化；最终 masks × 255 转 uint8 存入 batch.masks[storage_name]。

6.2 Identity：身份嵌入与 filter/nFilter

Identity 处理器接受 threshold 参数（默认 0.4）。预处理复用 ExtractHandlerFace 的双路径取图（整帧对齐 or 已对齐人脸子裁剪），输出嵌入写入 batch.identities[storage_name]。

IdentityFilter 实现基于余弦相似度的白/黑名单过滤：

• add_identities(identities, is_filter)：is_filter=True 为白名单（保留与之相似的脸），False 为 nFilter 黑名单（剔除相似的脸）；
• 相似度计算 _find_cosine_similarity()：两侧嵌入分别 L2 归一后做矩阵乘 t_norms @ s_norms.T，任一参考脸相似度 >= threshold 即判为匹配；
• 白/黑名单同时启用时对两个布尔掩码取 .all() 交集；
• FilterLoader 负责从命令行给定的图片文件或目录构建参考嵌入：目录会扫描其中的图像扩展名文件，不存在的路径打 warning 跳过，而同一文件同时出现在 filter 和 nFilter 中会直接报错退出。

7. 模型工具与性能剖析

7.1 plugin_utils：预热、发现与编译

plugin_utils.py 提供 4 个工具：

• random_input_from_plugin()：按插件的 input_size、scale、dtype 生成随机张量，支持 channels-first/last 两种布局；
• get_torch_modules()：递归遍历插件 model 属性树（含 Mapping/Iterable/__dict__），找出所有 torch.nn.Module，并用 seen 集合防环、用模块名约束不越出插件所在模块；
• warmup_plugin()：用随机批做一次真实推理完成 CUDA 初始化；channels_last=None 时先试 channels-first 再试 channels-last 自动探测布局，期间临时调低 cv2 日志级别以免 dnn 插件误报；
• compile_models()：在进程级锁内对每个模块调用 mod.compile(fullgraph=True, dynamic=False, ...)，选项针对 FaceSwap 的"代码侧处理动态 batch + 静态形状"场景调优：triton.cudagraphs=True、shape_padding（张量对齐 Tensor core）、epilogue_fusion、coordinate_descent_tuning、max_autotune 等；编译前先用单样本 warmup 确认通道顺序，编译后再发一次真实 batch size 的 warmup，最后 torch.cuda.empty_cache() 释放峰值显存。

7.2 profile：GPU 吞吐剖析

profile.py 的模块 docstring 为 GPU profiling for throughput optimization。其中 ModelProfile 对单个 PyTorch GPU 插件做基准测试：max_batch_size 默认 128、每个 batch size 运行 run_time（默认 10 秒）；batch size 从 1 起按 2 的幂递增，先做 2 秒 warmup 并 reset_peak_memory_stats()，然后记录每秒编码数（EG/S）与 VRAM 已分配/已保留峰值（MB），输出固定宽度表格便于抄录；任一 batch size 触发 OutOfMemoryError 即提前终止。包入口导出的 Profiler 则面向整条管线（get_pipeline() 构建后剖析），用于在多插件组合中找吞吐瓶颈。

8. 一次完整提取的数据流

综合以上源码，一次 faceswap extract 调用中的人脸数据流可以概括为：

1. InputIterator 把 FrameFaces（图像或视频帧）装配为 ExtractBatch 放入首个入队；
2. Detect runner：整图缩放 → 检测（可选多角度旋转增强）→ 框回映射、尺寸过滤，输出 bboxes + frame_ids，无脸帧被短路透传；
3. Align runner：ROI 裁剪 → landmark 推理（可选 re-feed 抖动平均、re-align 两遍对齐）→ 关键点回帧空间，可选 AlignedFilter 质量过滤；
4. Mask / Identity runner（可多个插件并行挂接）：分别产出对齐掩码（ExtractBatchMask）与身份嵌入（并按 threshold 做 filter/nFilter 过滤）；
5. 末端 OutputIterator 还原 FrameFaces，主线程 next() 最终 runner 逐帧取结果，管线以 FLUSH/SHUTDOWN 信号有序冲刷、收尾并释放中间属性。

各级队列 maxsize=1 的背压设计加上零检测短路，使 CPU 预处理、GPU 推理、后处理三个阶段能够持续流水化执行，这正是包 docstring 中"Parallel batched inference"的含义。

9. 相关配置与参考路径

• 默认配置文件路径为 faceswap/config/extract.ini（config_file 参数可指定自定义文件，load_config() 在 handler 初始化时加载）；align 过滤阈值、mask 存储尺寸等均从 plugins/extract/extract_config.py 读取；
• 插件侧基类与配置：plugins/extract/base.py（ExtractPlugin）、plugins/extract/extract_config.py；
• 对齐数学工具：lib/align/（batch_umeyama、batch_transform、Batch3D 等，被 infer 各模块大量复用）；
• API 参考文档页：docs/full/lib/infer.rst。

以上所有类、参数与默认值均出自当前仓库 lib/infer/ 源码与包 docstring；管线中线程/队列的具体交互行为（如外部错误状态回注、passthrough 拆分）是从源码结构阅读得出的实现细节，可直接对照上文列出的文件与行号验证。