使用 YOLOv8x-seg 为 ProPainter 生成遮罩

一、概述

YOLOv8x-seg 是 Ultralytics 提供的实例分割模型，可以自动识别并分割视频中的物体。结合 ProPainter，可以实现自动化的物体移除工作流。

工作流程：

使用 YOLOv8x-seg 检测并分割视频中的目标物体
将分割结果转换为 ProPainter 所需的黑白遮罩
使用 ProPainter 进行物体移除

二、环境准备

1. 安装依赖

pip install ultralytics opencv-python numpy pillow

2. YOLOv8x-seg 模型

首次运行时会自动下载 yolov8x-seg.pt 模型文件（约 130MB）

三、ProPainter 遮罩要求

根据 ProPainter 的规范：

格式：PNG 图片
颜色模式：灰度图或 RGB（只使用黑白）
白色（255）：需要移除的区域
黑色（0）：保留的区域
命名：与视频帧对应（如 0001.png, 0002.png...）
尺寸：与视频帧尺寸一致

四、使用方法

方法一：自动生成所有帧的遮罩

使用提供的脚本 generate_masks_yolov8.py，可以：

自动提取视频帧
使用 YOLOv8x-seg 检测指定类别的物体
生成对应的遮罩文件
输出可直接用于 ProPainter 的文件结构

命令示例：

# 移除视频中的所有人（class_id=0）- 输出到当前目录
python generate_masks_yolov8.py --video input.mp4 --classes 0

# 自定义输出名称
python generate_masks_yolov8.py --video input.mp4 --name bmx-trees --classes 0

# 移除多个类别（人和自行车）
python generate_masks_yolov8.py --video input.mp4 --classes 0 1

# 指定输出到其他目录
python generate_masks_yolov8.py --video input.mp4 --output_dir outputs --classes 0

输出结构（默认在当前目录，符合 ProPainter 规范）：

./
├── input/           # 视频帧（用于 ProPainter --video 参数）
│   ├── 0001.jpg
│   ├── 0002.jpg
│   └── ...
└── input_mask/      # 遮罩文件（用于 ProPainter --mask 参数）
    ├── 0001.png
    ├── 0002.png
    └── ...

然后运行 ProPainter：

python inference_propainter.py --video input --mask input_mask

五、YOLO 常用类别 ID

YOLOv8 使用 COCO 数据集的 80 个类别：

ID	类别名称	ID	类别名称	ID	类别名称
0	person	1	bicycle	2	car
3	motorcycle	5	bus	7	truck
16	dog	17	cat	24	backpack
26	handbag	28	suitcase	32	sports ball

完整列表见：COCO 类别

六、高级用法

1. 只处理特定帧范围

修改脚本中的帧提取逻辑，或手动删除不需要的帧和遮罩

2. 手动修正遮罩

生成的遮罩可能不完美，可以使用图像编辑软件（如 Photoshop、GIMP）手动修正：

用白色画笔补充漏检区域
用黑色画笔去除误检区域

3. 遮罩膨胀处理

为了更好的移除效果，可以对遮罩进行膨胀操作，扩大白色区域：

import cv2
mask = cv2.imread('mask.png', 0)
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
dilated = cv2.dilate(mask, kernel, iterations=1)
cv2.imwrite('mask_dilated.png', dilated)

七、注意事项

显存占用：YOLOv8x-seg 是大模型，需要较多显存（建议 8GB+）
检测精度：置信度阈值（--conf）影响检测结果，默认 0.25
遮罩质量：YOLO 的分割边缘可能不够精细，必要时需手动修正
处理速度：视频越长，处理时间越久，建议先用短视频测试

八、故障排查

问题 1：YOLO 没有检测到目标

降低 --conf 阈值
确认目标类别 ID 正确
检查视频质量和目标清晰度

问题 2：遮罩边缘不准确

使用遮罩膨胀处理
手动修正关键帧
考虑使用更精细的分割模型（如 SAM）

问题 3：ProPainter 效果不佳

检查遮罩是否完全覆盖目标
确保遮罩与视频帧尺寸一致
尝试调整 ProPainter 的参数

目录CONTENT

使用YOLOv8x-seg生成遮罩