面向机器人采集视频的绿色手套移除流水线。项目使用 Grounding DINO 检测手部候选区域、SAM 生成像素级掩码,再由 ProPainter 完成时序视频修复。
设计目标是保守地移除手套:宁可保留少量边缘,也尽量不修改操作物体、夹爪和掩码外区域。
两个示例视频来自同一段 640 × 480、30 FPS、约 37 秒的输入,便于直接比较处理前后的效果。
输入视频
→ Grounding DINO 检测手套/手/前臂
→ SAM 生成候选掩码
→ 颜色、连通域和时序轨迹约束
→ ProPainter 视频修复
→ 恢复掩码外原始像素
→ H.264 / yuv420p MP4
impainting/
├── example/ # 输入与处理结果示例
├── remove_hand.py # 单视频入口
├── batch_remove_hand.py # 多任务、多 GPU 批处理入口
├── remove.sh # 当前服务器的批处理示例
├── _grounded_sam.py # Grounding DINO、SAM 和视频编码
├── _mask_utils.py # 掩码后处理
├── _propainter_utils.py # ProPainter 调用与状态文件
├── requirements.txt # Python 依赖(不包含 PyTorch)
├── model_ckpt/ # 本地 DINO/SAM 模型,不进入 Git
└── third_party/ProPainter/
└── weights/ # ProPainter 权重,不进入 Git
推荐并已验证的环境:
- Linux
- Python 3.10
- NVIDIA GPU 和可用的
nvidia-smi - PyTorch 2.6.0 + CUDA 12.4
- torchvision 0.21.0
- FFmpeg(需要 H.264 编码支持)
模型推理需要较多显存。批处理时建议先使用每张 GPU 一个 worker,确认显存余量后再增加并发。
git clone git@github.com:MrSecant/impainting.git
cd impainting如果服务器使用本文配置过的 SSH 别名:
git clone git@github.com-zyh:MrSecant/impainting.git
cd impainting先确保当前 shell 中存在 conda 命令:
source /mnt/data/lcy/miniconda3/etc/profile.d/conda.sh创建独立环境:
conda create -n impainting python=3.10 pip -y
conda activate impainting
conda install -c conda-forge ffmpeg -y安装 CUDA 12.4 对应的 PyTorch:
python -m pip install \
torch==2.6.0 \
torchvision==0.21.0 \
--index-url https://download.pytorch.org/whl/cu124安装其余依赖:
python -m pip install --upgrade pip
python -m pip install -r requirements.txt如果机器使用其他 CUDA/PyTorch 组合,只替换上面的 PyTorch 安装命令;不要同时安装多个不同 CUDA 版本的 torch。
which python
python -c 'import sys; print(sys.executable)'
python -c 'import cv2, torch, torchvision, transformers; print("torch:", torch.__version__); print("CUDA:", torch.cuda.is_available())'
nvidia-smi
ffmpeg -version预期:
sys.executable位于 Conda 的impainting环境中;- PyTorch 版本包含
2.6.0; torch.cuda.is_available()输出True。
不要使用 Conda base 的 Python 或其他项目的 Python 启动脚本。需要确认时始终执行:
python -c 'import sys; print(sys.executable)'代码使用完全本地的 Hugging Face 模型,运行时设置了 local_files_only=True。目录必须是:
model_ckpt/
├── grounding-dino-base/
│ ├── config.json
│ ├── model.safetensors
│ ├── preprocessor_config.json
│ └── tokenizer...
└── sam-vit-huge/
├── config.json
├── model.safetensors
└── preprocessor_config.json
可以使用安装依赖后提供的 hf 命令下载:
mkdir -p model_ckpt
hf download IDEA-Research/grounding-dino-base \
--local-dir model_ckpt/grounding-dino-base
hf download facebook/sam-vit-huge \
--local-dir model_ckpt/sam-vit-huge也可以直接从已有机器复制完整目录:
cp -a /path/to/model_ckpt/grounding-dino-base model_ckpt/
cp -a /path/to/model_ckpt/sam-vit-huge model_ckpt/ProPainter 首次运行会自动把以下权重下载到 third_party/ProPainter/weights/:
raft-things.pth
recurrent_flow_completion.pth
ProPainter.pth
如果计算节点不能访问外网,请提前从其他机器复制这三个文件。
conda activate impainting
cd /path/to/impainting
python remove_hand.py \
--input /path/to/base_0_rgb.mp4 \
--models-root "$PWD/model_ckpt" \
--propainter-root "$PWD/third_party/ProPainter" \
--output-dir /path/to/output默认输出:
/path/to/output/base_0_rgb_remove_hand.mp4
/path/to/output/metadata.json
/path/to/output/status.json
/path/to/output/timing.json
一组完整的保守参数示例:
python remove_hand.py \
--input /path/to/base_0_rgb.mp4 \
--models-root "$PWD/model_ckpt" \
--propainter-root "$PWD/third_party/ProPainter" \
--output-dir /path/to/output \
--output-name base_0_rgb_remove_hand.mp4 \
--box-threshold 0.22 \
--text-threshold 0.22 \
--roi-bottom-ratio 0.38 \
--dino-interval 1 \
--box-expand-ratio 0.20 \
--min-box-margin 12 \
--min-area 100 \
--mask-dilation 5 \
--temporal-flow-width 320 \
--mask-mode strict_color \
--neighbor-length 8 \
--ref-stride 5 \
--subvideo-length 20 \
--raft-iter 10 \
--propainter-mask-dilation 1 \
--progress-every 25批处理器默认搜索以下目录结构:
<task-root>/<episode>/camera/base_0_rgb.mp4
建议先 dry-run,确认任务发现是否正确:
python batch_remove_hand.py \
--task-list /path/to/task-root \
--gpus 0,1,2,3 \
--workers-per-gpu 1 \
--dry-run正式运行:
python batch_remove_hand.py \
--task-list /path/to/task-root-a \
--task-list /path/to/task-root-b \
--python "$(command -v python)" \
--models-root "$PWD/model_ckpt" \
--propainter-root "$PWD/third_party/ProPainter" \
--gpus 0,1,2,3 \
--workers-per-gpu 1也可以使用任务列表文件,每行一个 task root,空行和 # 注释会被忽略:
python batch_remove_hand.py \
--task-list-file /path/to/tasks.txt \
--python "$(command -v python)" \
--gpus auto \
--workers-per-gpu 1每条输入的结果写回对应的 camera 目录:
<episode>/camera/base_0_rgb.mp4
<episode>/camera/base_0_rgb_remove_hand.mp4
<episode>/camera/remove_hand/metadata.json
<episode>/camera/remove_hand/status.json
<episode>/camera/remove_hand/timing.json
<episode>/camera/remove_hand/*.log
已存在结果默认跳过;确实需要重跑时添加 --overwrite。
remove.sh 是当前服务器任务路径的示例脚本。执行前应检查其中的 --task-list 和 GPU 设置:
conda activate impainting
bash remove.sh| 参数 | 默认值 | 作用 |
|---|---|---|
--mask-mode |
strict_color |
strict_color 更保守;color_shadow 覆盖更多绿色阴影 |
--mask-dilation |
5 |
最终掩码膨胀;越大越少残边,但越容易影响物体 |
--propainter-mask-dilation |
1 |
ProPainter 内部额外膨胀 |
--box-threshold |
0.22 |
Grounding DINO 检测阈值 |
--text-threshold |
0.22 |
文本匹配阈值 |
--roi-bottom-ratio |
0.38 |
只保留画面较下方的候选框 |
--dino-interval |
1 |
DINO 调用间隔;1 的质量最稳定 |
--neighbor-length |
8 |
ProPainter 邻近参考帧数 |
--subvideo-length |
20 |
ProPainter 分段长度;降低可减少显存占用 |
--raft-iter |
10 |
RAFT 光流迭代次数 |
如果操作物体被误伤,优先尝试:
--mask-mode strict_color
--mask-dilation 3
--propainter-mask-dilation 0
--box-threshold 0.26
--text-threshold 0.26
如果绿色手套边缘残留明显,优先尝试:
--mask-mode color_shadow
--mask-dilation 6
--propainter-mask-dilation 2
--box-threshold 0.20
--text-threshold 0.20
source /mnt/data/lcy/miniconda3/etc/profile.d/conda.sh
conda activate impainting如果 Miniconda 安装在其他位置,请替换为对应的 etc/profile.d/conda.sh。
不要只看 shell 提示符,直接检查:
which python
python -c 'import sys; print(sys.executable)'批处理时可显式指定:
--python /absolute/path/to/env/bin/python确认:
test -f model_ckpt/grounding-dino-base/model.safetensors
test -f model_ckpt/sam-vit-huge/model.safetensors模型不会在推理时自动从 Hugging Face 下载。
首次运行可能正在下载权重。检查:
ls -lh third_party/ProPainter/weights/- 将
--workers-per-gpu降为1; - 减小
--subvideo-length; - 减小
--neighbor-length; - 确认没有其他进程占用 GPU。
ffmpeg -version项目也安装了 imageio-ffmpeg 作为后备,但推荐使用系统或 Conda 安装的 FFmpeg。
metadata.json:输入输出、帧数、FPS、掩码统计和运行参数;status.json:当前阶段、成功或错误信息;timing.json:端到端及 DINO、SAM、ProPainter 等阶段耗时。
批处理出现失败时,优先查看对应 camera/remove_hand/*.log。
项目包含修改后的 ProPainter 推理代码。其许可证见:
third_party/ProPainter/LICENSE