Skip to content

About

No description, website, or topics provided.

Resources

Stars

1 star

Watchers

0 watching

Forks

Repository files navigation

impainting

面向机器人采集视频的绿色手套移除流水线。项目使用 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,确认显存余量后再增加并发。

1. 克隆项目

git clone git@github.com:MrSecant/impainting.git
cd impainting

如果服务器使用本文配置过的 SSH 别名:

git clone git@github.com-zyh:MrSecant/impainting.git
cd impainting

2. 创建 Conda 环境

先确保当前 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。

3. 验证环境

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)'

4. 准备模型

代码使用完全本地的 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

如果计算节点不能访问外网,请提前从其他机器复制这三个文件。

5. 单视频运行

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

6. 批量运行

批处理器默认搜索以下目录结构:

<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

常见问题

conda: 未找到命令

source /mnt/data/lcy/miniconda3/etc/profile.d/conda.sh
conda activate impainting

如果 Miniconda 安装在其他位置,请替换为对应的 etc/profile.d/conda.sh。

traceback 指向了错误的 Python 环境

不要只看 shell 提示符,直接检查:

which python
python -c 'import sys; print(sys.executable)'

批处理时可显式指定:

--python /absolute/path/to/env/bin/python

找不到 DINO 或 SAM

确认:

test -f model_ckpt/grounding-dino-base/model.safetensors
test -f model_ckpt/sam-vit-huge/model.safetensors

模型不会在推理时自动从 Hugging Face 下载。

ProPainter 首次启动很慢

首次运行可能正在下载权重。检查:

ls -lh third_party/ProPainter/weights/

CUDA 显存不足

  • 将 --workers-per-gpu 降为 1;
  • 减小 --subvideo-length;
  • 减小 --neighbor-length;
  • 确认没有其他进程占用 GPU。

FFmpeg 编码失败

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

About

No description, website, or topics provided.

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages