Faceswap 图像 IO 库深度解析：lib.image 的图像读取、PNG 元数据与后台加载器

本文基于 Faceswap 仓库中的 docs/full/lib/image.rst 文档页展开。该页面通过 Sphinx 的 automodapi 指令自动为 lib/image.py 模块生成完整 API 参考，涵盖图像读写、多格式元数据（PNG iTXt / TIFF ImageDescription）、多线程批量加载以及后台图像 IO 线程体系。读完后你将掌握 Faceswap 如何处理训练数据、提取人脸与视频帧的图像输入输出，并能基于该库实现自己的批量图像处理流程。

1. 文档页与模块定位

docs/full/lib/image.rst 的全部内容是一个 automodapi 指令：

.. automodapi:: lib.image
   :include-all-objects:

它指示 Sphinx 自动抓取 lib.image 模块中所有公开对象（函数与类）的 docstring 并生成 API 参考。因此该文档页的实际内容就是 lib/image.py 的全部公开 API。模块开头自述为 """ Utilities for working with images """，其服务对象包括：

• 提取阶段（scripts/extract.py）、转换阶段（scripts/convert.py）、训练阶段（scripts/train.py）等脚本；
• 数据管线（lib/training/data/data_set.py 通过 from lib.image import read_image 读取带元数据的人脸）；
• 各类 GUI 工具（manual、alignments、preview、sort 等）。

整个模块按 __all__ = get_module_objects(__name__) 导出全部公开对象，下面按「单张读取 → 批量读取 → 元数据 → 编码 → 后台 IO」的顺序逐块剖析。

2. read_image：确保“真的读到了”的图像加载

read_image(filename, raise_error=False, with_metadata=False)（见 lib/image.py#L59-L152）是 Faceswap 最核心的图像入口。其 docstring 明确说明：它扩展了 cv2.imread() 的功能，确保图像确实被加载——因为 cv2.imread 在失败时会静默返回 None，这会在批量处理中埋下隐患。

参数与返回值（继承自 docstring）：

参数	说明	默认
filename	图像完整路径	必填
raise_error	True 时任何失败（含返回 None）都会抛出异常；False 时仅记录错误日志、不抛错，流程可继续	False
with_metadata	仅对提取出来的 Faceswap 人脸有效；True 时额外返回存储在 PNG EXIF/iTXt 头中的 Faceswap 元数据	False

实现上有几个值得注意的细节（对应源码逻辑）：

1. 绕过 cv2.imread 直接按字节解码：先 open(filename, "rb") 读入原始字节，再用 cv2.imdecode(np.frombuffer(raw_file, dtype=np.uint8), cv2.IMREAD_UNCHANGED) 解码。这样做既保留原始位深，又为后续按同一份字节流解析 PNG 元数据（png_read_meta(raw_file)）做了铺垫。
2. 通道归一化：灰度图（ndim == 2）被 cv2.cvtColor 转成 BGR；带 4 通道（RGBA）的图像直接 image[:, :, :3] 剥离 alpha 掩码。
3. 位深归一化到 UINT8：整数型图像若 iinfo.max != 255（例如 16 位图），按 image / max * 255.0 缩放；浮点图像则 np.clip(image, 0.0, 1.0) * 255 粗暴裁剪后转 np.float32；最终统一 np.clip(..., 0, 255).astype(np.uint8)。
4. 三类异常处理：TypeError、ValueError（docstring 特别提示可能是文件名含特殊字符或文件损坏）以及兜底 Exception，均按 raise_error 决定抛错或仅 logger.error 后返回 None。

返回值在开启 with_metadata 时为二元组 (image, PNGHeader)；类型重载（@T.overload，见 lib/image.py#L34-L56）把这四种组合精确标注了出来，便于静态类型检查。

3. read_image_batch：多线程批量读图

read_image_batch(filenames, with_metadata=False)（lib/image.py#L165-L238）在 docstring 中说明：利用多线程同时从磁盘读图，大幅降低批量读取耗时。实现要点：

• 内部用 concurrent.futures.ThreadPoolExecutor 提交所有 read_image(filename, raise_error=True, with_metadata=...) 任务，并以 {future: idx} 字典记录序号；
• 通过 futures.as_completed 收集结果并按 idx 回填到预分配的 batch 列表，保证返回顺序与输入 filenames 一致；
• 返回 np.array(batch)，with_metadata=True 时返回 (batch, metadata_list)。

docstring 的 Notes 明确警告：批量图像应尺寸一致，否则会退化为非齐次的 object 数组；其示例展示了典型的批量形状 (3, 64, 64, 3)，这正是 Faceswap 训练/转换时 64×64 人脸小图的典型尺寸。

4. PNG 元数据：人脸信息随图像一起旅行

Faceswap 提取的人脸 PNG 不仅是一张小图，还携带完整对齐信息。这一机制由 lib/align/objects.py#L250-L274 中的两个 dataclass 定义：

• PNGSource：alignments_version（对齐文件版本）、original_filename、face_index（该脸在帧内的索引）、source_filename、source_is_video、source_frame_dims（原帧的高宽）；
• PNGHeader：聚合 PNGAlignments（人脸对齐信息）与 PNGSource 两块，可经 to_dict() / from_dict() 与字典互转。

写入与读取由 lib/image.py 中的一组函数完成：

4.1 pack_to_itxt / png_write_meta

pack_to_itxt(metadata)（L348-L370）把元数据打包成合法的 PNG iTXt 块：以 faceswap 为关键词（key），内容为 str(metadata) 的 UTF-8 编码字节；按 关键词 + b"\0\0\0\0\0" + 数据 组装块体（5 个零字节是 iTXt 的压缩标志、语言标签与翻译标记的约定写法），再计算 crc32（以 b"iTXt" 为初值）拼上 4 字节大端长度与 CRC，得到完整块。

png_write_meta(image, data)（L466-L488）在编码好的 PNG 字节流中定位 b"IDAT"，用 split = image.find(b"IDAT") - 4 找到插入点，把 iTXt 块直接拼接到 IDAT 之前。docstring 坦承这是一个“精简且不追求健壮”的写头器：因为 OpenCV 本身不写 iTXt 头，所以可以假设文件中唯一的 iTXt 就是 Faceswap 自己写的那个。

4.2 png_read_meta

png_read_meta(image)（L599-L635）从字节流中循环 find(b"iTXt") 定位每个 iTXt 块，解析长度、关键词与值；跳过非 faceswap 关键词的块，命中后以 literal_eval 还原为字典并经 PNGHeader.from_dict 返回强类型对象。注释同样强调它是“精简、非健壮、非安全”的定向读取器。

4.3 read_image_meta / read_image_meta_batch

read_image_meta(filename)（L241-L304）是按文件流解析的轻量版本，适合只关心尺寸与元数据、不想解码整张图时使用：

• 非 .png 文件：直接 cv2.imread 取 height/width 返回；
• PNG 文件：先校验 8 字节签名 b"\x89PNG\r\n\x1a\n"（否则抛 ValueError），然后逐块解析——IHDR 块取前 8 字节大端无符号整数得到宽高；iTXt 块拆分 keyword\0value，关键词为 faceswap 时 literal_eval(value[4:].decode("utf-8")) 存入 itxt 并结束；遇到 IDAT 或空块即停止，其余块用 in_file.seek(length + 4, 1) 跳过。

read_image_meta_batch(filenames)（L307-L345）把上述逻辑放进 ThreadPoolExecutor 并行执行，以生成器形式产出 (filename, metadata) 二元组；docstring 特意注明产出顺序是非确定的（as_completed 顺序），不保证与输入一致。

4.4 update_existing_metadata

update_existing_metadata(filename, metadata)（L373-L422）用于原地更新已提取人脸 PNG 头的元数据（例如调整框选后写回对齐数据）。它把文件复制到 filename + "~" 临时文件，逐块搬运：非 iTXt 块原样拷贝（含 4 字节 CRC）；非 faceswap 的 iTXt 块也原样保留；命中 faceswap 块时改为写入 pack_to_itxt(metadata) 的新块并跳过旧 CRC；到达 IDAT 后把剩余图像数据整体写出，最后 os.replace 原子替换原文件。整个流程完全基于 struct 手工解析，不依赖任何第三方 PNG 库。

5. TIFF 元数据：手工改写 IFD

转换阶段（convert writer）会把矩阵等信息写入 TIFF 的 ImageDescription 字段，由 tiff_write_meta / tiff_read_meta 实现（lib/image.py#L491-L596）：

• 两个函数都断言文件为小端（image[:2] == b"II"）且版本号为 42，明确仅支持单页 TIFF（下一个 IFD 指针必须为 0）；
• 写入时读取 IFD 起始处的 tag 数量，num_tags + 1 重新打包，遍历每个 12 字节 tag：为偏移量类 tag 的 offset 统一 + 12（为新增条目腾出空间），按 tag id 排序规则找到 tag 270（ImageDescription）的插入位，追加 struct.pack("HH", 270, 2) + (长度, 数据偏移) 的新 tag，并把 JSON 序列化后的数据（json.dumps(..., ensure_ascii=True).encode("ascii")）追加到文件尾部；
• 读取时（tiff_read_meta）按同样方式遍历 IFD，找到 tag 270 后按类型（{2: "1s", 3: "1H", 4: "1I", 7: "1B"}）计算大小、按 offset 取数据并 json.loads 返回。

docstring 同样声明这是“为非常具体任务服务”的实现，不适用于其他用途。

6. encode_image 与 generate_thumbnail

encode_image(image, extension, encoding_args=None, metadata=None)（L425-L463）把 np.ndarray（BGR）编码为指定格式的字节流：

• 元数据仅支持 .png 与 .tif，传入其他扩展名会抛 ValueError("Metadata is only supported for .png and .tif images")；
• 核心是 cv2.imencode(extension, image, args)[1].tobytes()，encoding_args 直接透传给 cv2.imencode（例如 JPEG 质量参数）；
• 带 metadata 时按扩展名分派到 png_write_meta 或 tiff_write_meta。

generate_thumbnail(image, size=96, quality=60)（L638-L663）生成 96×96 JPEG 缩略图，默认质量 60。当原图大于目标尺寸时用 cv2.INTER_AREA 下采样（抗混叠），小于时改用 cv2.INTER_CUBIC 放大，最后以 cv2.IMWRITE_JPEG_QUALITY 编码。这个 96px 缩略图正是 lib/align/objects.py 中 FileAlignments.thumb 字段所描述的“96px JPEG thumbnail”。

7. 颜色空间与杂项工具

• batch_convert_color(batch, color_space)（L666-L708）：把整批 (n, h, w, c) 图像一次性 reshape 成 (n*h, w, c) 再调用一次 cv2.cvtColor(batch, getattr(cv2, f"COLOR_{color_space}"))，避免了逐张循环。docstring 说明这只适用于源/目标颜色空间形状相同的转换，并建议需要完整色彩范围时使用 32 位图像（8 位转换有信息损失）；
• hex_to_rgb("#0d25ac") / rgb_to_hex((0, 255, 255))：十六进制色值与 RGB 元组互转，供 GUI 配色使用。

8. ImageIO 体系：后台线程 + 队列的图像 IO

模块后半部分（L749 起）是一套生产级 IO 抽象，父类 ImageIO 与两个子类构成「后台读」「后台写」骨架。

8.1 ImageIO（基类）

ImageIO(path, queue_size, args=None)：

• path 可以是目录、视频文件或图像文件列表；保存场景必须是已存在的目录；
• _check_location_exists 对不存在的路径（或列表中任一不存在的路径）抛 FaceswapError；
• 内部维护一个 queue.Queue(maxsize=queue_size) 缓冲与一个 FSThread（来自 lib/multithreading），线程名即类名，并共享 FSThread 的 error_state 用于跨线程错误传播；
• _process(queue) 是抽象方法，子类实现具体的读/写循环；close() 负责 join 线程。

8.2 ImagesLoader：后台预读

ImagesLoader(path, queue_size=8, fast_count=True, skip_list=None, count=None, pts=None, keyframes=None)（lib/image.py#L834-L1109）：

参数	说明	默认
path	图像目录 / 视频文件 / 图像文件列表	必填
queue_size	内部缓冲区持有的图像数量	8
fast_count	视频计数：True 快速但不保证精确，False 慢但精确	True
skip_list	需要跳过的帧/图索引集合	None
count	已知数量时传入以跳过计数步骤	None
pts / keyframes	视频的显示时间戳与关键帧（用于加速 seek）	None

关键机制：

1. 确定性顺序：目录加载通过 get_image_paths 生成有序文件列表；视频加载构造虚拟帧名 _dummy_video_frame_name(index)，格式为 <视频名>_<帧号:06d>.<扩展名>（帧号从 1 开始），以便与 alignments 文件按文件名对齐；
2. 后台线程 + 有界队列：_process 在后台迭代帧，每读出一帧先检查全黑/失败帧（image is None or (not image.any() and image.ndim not in (2, 3))，注释解释全黑帧 any() 为 False 但仍是合法维度所以双重判断），再 queue.put(retval, timeout=0.2) 循环重试直到 QueueFull 消失或检测到线程错误；
3. EOF 哨兵：迭代结束后向队列放入字符串 "EOF"，主线程 load() 生成器收到后 break 并 close()；
4. 错误传播：load() 每次取队列前检查 error_state.has_error——若当前是主线程则 re_raise() 抛出，否则静默中断迭代。

load() 产出 (filename, image)（FacesLoader 为三元组），典型用法即 docstring 示例：

loader = ImagesLoader('/path/to/video.mp4')
for filename, image in loader.load():
    # do processing

8.3 FacesLoader 与 SingleFrameLoader

• FacesLoader(path, skip_list=None, count=None)（L1112-L1173）：固定 queue_size=8，重写 _get_count_and_filelist 只保留 .png 文件（提取人脸都存为 PNG），重写 _from_folder 用 read_image(filename, raise_error=False, with_metadata=True) 读图，产出 (filename, face, metadata) 三元组——训练管线正是借此把对齐信息随每张脸一起喂给数据加载器；
• SingleFrameLoader(path, video_meta_data=None)（L1176-L1258）：面向 GUI 的随机访问加载器。因为“反正要等这一帧读完”，故 queue_size=1、fast_count=False 且不走后台线程；video_meta_data 可传入含 pts_time 与 keyframes 的缓存避免重复扫描视频；image_from_index(index) 对视频走 VideoReader.get(index)（配合 pts/keyframes 实现快速 seek），对图片列表则直接 read_image。

8.4 ImagesSaver：后台并发保存

ImagesSaver(path, queue_size=8, as_bytes=False)（L1261-L1399）：

• 保存目录必须是已存在的文件夹（_check_location_exists 额外校验 os.path.isdir 与字符串类型）；
• 后台线程 _process 从队列取 (filename, image, sub_folder)，提交给内部 ThreadPoolExecutor 并发落盘——即“保存线程调度 + 线程池执行”两层结构；收到 "EOF" 后 executor.shutdown() 退出；
• _save 中文件名会被 os.path.basename 去路径化后拼到 location（或其 sub_folder 子目录，目录不存在时 os.makedirs），as_bytes=True 时直接 open(..., "wb").write(image)，否则 cv2.imwrite；保存失败仅 logger.error 不中断流程；
• 公共 API 为 save(filename, image, sub_folder=None) 与 close()（close 先放 "EOF" 再 join 线程）。docstring 强调：所有保存操作完成后必须调用 close() 才能确保队列排空。

典型用法：

saver = ImagesSaver('/path/to/save/folder')
for filename, image in <image_iterator>:
    saver.save(filename, image)
saver.close()

9. 在训练管线中的实际用法

以 lib/training/data/data_set.py 为例，数据加载器直接依赖本模块：from lib.image import read_image，并在构造训练样本时调用 image, meta = read_image(filename, ...)（该文件 L415、L514 附近），把人脸像素与其 PNGHeader 元数据（框选、掩码来源、原始帧信息等）一并取出，供数据增强与身份标签使用。这印证了第 4 节所述设计：PNG iTXt 元数据是 Faceswap 把对齐结果“贴”在图像文件上的唯一载体，read_image(with_metadata=True) + FacesLoader 则是消费端的标准入口。

10. 小结

lib/image.py 是 Faceswap 图像数据流的基石：read_image 解决“可靠读取与位深/通道归一化”，read_image_batch 与 ImageIO 家族解决“后台并发、顺序确定、错误可控”的吞吐问题，pack_to_itxt/png_write_meta/tiff_write_meta 等函数则用纯标准库（struct + zlib.crc32）实现了对 PNG iTXt 与 TIFF IFD 的字节级读写，使 64×64 的人脸小图能携带完整对齐上下文在提取、转换、训练各环节流转。对希望扩展 Faceswap 数据管线的开发者而言，理解这套“有界队列 + EOF 哨兵 + ErrorState”的线程模式与手写图像头解析代码，是复用或改造其 IO 层的关键。