ProPainter——实现视频消除特定对象、去水印、视频修复

github：https://github.com/sczhou/ProPainter

paper：ICCV2023 | ProPainter: Improving Propagation and Transformer for Video Inpainting

Demo Video：https://www.youtube.com/watch?v=92EHfgCO5-Q

首先确定下自己的电脑GPU显存是否满足要求，下方展示了不同尺寸、帧数的视频在fp32/fp16两种精确度下的显存需求

分辨率	50帧	80帧
1280 x 720	28G / 19G	OOM / 25G
720 x 480	11G / 7G	13G / 8G
640 x 480	10G / 6G	12G / 7G
320 x 240	3G / 2G	4G / 3G

1. 安装

---

1.1 克隆项目

git clone https://github.com/sczhou/ProPainter.git

1.2 创建虚拟环境和安装依赖库

# create new anaconda env
conda create -n propainter python=3.8 -y
conda activate propainter

# install python dependencies
pip3 install -r requirements.txt

• CUDA >= 9.2
• PyTorch >= 1.7.1
• Torchvision >= 0.8.2
• 其他库的要求查看requrements.txt

1.3 下载权重文件

从Releases V0.1.0下载权重文件保存到下方路径

weights
   |- ProPainter.pth
   |- recurrent_flow_completion.pth
   |- raft-things.pth
   |- i3d_rgb_imagenet.pt (for evaluating VFID metric)
   |- README.md

2. 自带示例

---

2.1 消除物体

python inference_propainter.py --video inputs/object_removal/bmx-trees --mask inputs/object_removal/bmx-trees_mask

inputs/object_removal/bmx-trees目录下是视频的所有帧保存的图片，inputs/object_removal/bmx-trees_mask是每帧图片中需要消除的对象的掩码。运行代码后结果保存在results/bmx-trees下。

2.2 视频修复

python inference_propainter.py --video inputs/video_completion/running_car.mp4 --mask inputs/video_completion/mask_square.png --height 240 --width 432

inputs/video_completion/running_car.mp4是输入的视频，inputs/video_completion/mask_square.png是输入的掩码。

这个跟上面的区别就是，一个对像素点位置在变化的区域进行修复，一个是对像素点固定区域进行修复。对于待修复区域位置一直在变化的，需要输入每一帧的掩码，待修复区域位置不变的，则只需要一帧的掩码即可。

3. 实际应用案例

---

实际操作，用它试试给视频去水印。

• 首先抓取视频中的一帧保存为图片

import cv2

# 打开视频文件
cap = cv2.VideoCapture('F:/ProPainter/inputs/self/11.mp4')

# 获取视频总帧数和帧率和高宽
count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
fps = int(cap.get(cv2.CAP_PROP_FPS))
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
print('Video info w: {}, h: {}, count: {}, fps: {}'.format(w, h, count, fps))

while cap.isOpened():
    ret, frame = cap.read()
    # 选第一帧保存为图片
    cv2.imwrite("F:/ProPainter/inputs/self/frame.jpg", frame)
    break

# 释放资源
cap.release()

• 使用win10自带的画图软件获得水印位置，使用脚本生成掩码图

from PIL import Image, ImageDraw

# 创建一张与视频同样尺寸全黑的图片
img = Image.new('RGB', (1280, 578), color='black')
draw = ImageDraw.Draw(img)

# 三个矩形区域设为白色，分别为三个举行的左上点和右下点的坐标
rectangles = [
    [(301, 33), (413, 86)],
    [(945, 29), (1046, 68)],
    [(959, 194), (1007, 392)]
]

for rect in rectangles:
    draw.rectangle(rect, fill='white')

# 保存图片到本地
img.save('F:/ProPainter/inputs/self/mask.png')

• 运行程序

python inference_propainter.py --video F:/ProPainter/inputs/self/11.mp4 --mask F:/ProPainter/inputs/self/mask.png --resize_ratio 0.3 --fp16

视频修复需要大量的GPU显存，为了避免内存不足（OOM）错误。可以选用下方参数进一步减少内存使用量：

• --neighbor_length默认10，通过减少相邻帧数的数量；
• --ref_stride默认10，通过增加stride；
• --resize_ratio默认1.0，通过调小处理视频的大小；
• 通过指定--width和--height来设置较小的视频尺寸；
• --fp16，在推理期间使用半精度；
• --subvideo_length默认80，减少子视频的帧数；

4.训练自己的数据集

请看官方自己的文档，需要将视频每帧转为图片，然后每张图片获取一张掩码图，所有图片需要resize到432x240。