学习 RobustVideoMatting

源码

https://github.com/PeterL1n/RobustVideoMatting.git

有些笔记看起来很傻，但还是很有必要的，比如下面的推理脚本

python inference.py \
    --variant mobilenetv3 \
    --checkpoint rvm_mobilenetv3.pth \
    --device cuda \
    --input-source "input.mp4" \
    --downsample-ratio 0.25 \
    --output-type video \
    --output-composition "composition.mp4" \
    --output-alpha "alpha.mp4" \
    --output-foreground "foreground.mp4" \
    --output-video-mbps 4 \
    --seq-chunk 12

可以让你拿到项目可以直接跑起来，不至于再去熟悉半天

需要先下载 https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3.pth

然后开始debug其代码，遇到第一个麻烦，就是在forward中打了断点，但是代码跳转不进去，反正还是摸索了很久，发现是如下两行代码导致的：

self.model = torch.jit.script(self.model)
self.model = torch.jit.freeze(self.model)

把这两行代码注释掉，才可以进forward查看

下面开始记流水帐

我是从抖音上下载的视频，进行推理的

源视频分辨率为1280×720

src.shape=[1, 12, 3, 1280, 720]

这里的12是上面的seq-chunk，一次会取12帧画面

首先会进行图像缩放，缩放比例为上面参数downsample-ratio，也就是0.25

缩放后图像尺寸为：

src_sm.shape=[1, 12, 3, 320, 180]

然后是送入backbone(mobilenetv3)中提取特征

得到f1/f2/f3/f4:

f1.shape=[1, 12, 16, 160, 90]

f2.shape=[1, 12, 24, 80, 45]

f3.shape=[1, 12, 40, 40, 23]

f4.shape=[1, 12, 960, 20, 12]

其中f4会送到LRASPP中计算，得到

f4.shape=[1, 12, 128, 20, 12]

然后src_sm和f1/f2/f3/f4/r1/r2/r3/r4全部送入解码

其中r1/r2/r3/r4在最初的时候都为None

解码函数为RecurrentDecoder;

首先会把src_sm进行三次平均池化，得到如下feature

s1.shape=[1, 12, 3, 160, 90]
s2.shape=[1, 12, 3, 80, 45]
s3.shape=[1, 12, 3, 40, 23]

然后把f4/r4送到BottleneckBlock进行计算，然后在这里就会遇到本代码中，最核心的ConvGRU，

 首先会把f4，也就是这里的x切为a/b两半，尺寸均为[1, 12, 64, 20, 12]

然后会把b和r送到ConvGRU去计算，这里的r就是r4,

在ConvGRU中，参数就变为了x和h，这里的x是上面的b,h是r，也就是r4,

但因为最初的时候r4为None，所以在ConvGRU里，会初始化一个和x尺寸一样的全0的tensor，

h.shape=[1, 64, 20, 12]

然后送入到时间序列里去推理，这里的参数x就是上面的b，h就是，啊，h

然后x会按帧进行遍历，每一帧就是xt

xt.shape=[1, 64, 20, 12]

然后送到forward_single_frame里进行计算，这里参数就变为了x/h，x就是xt，h，还是h

然后x和h拼接到一起，[1, 128, 20, 12]

通过一个卷积和sigmoid，变为了[1, 128, 20, 12]

然后split分为了两部分，r和z

r.shape=[1, 64, 20, 12]
z.shape=[1, 64, 20, 12]

然后把r和h相乘，再和x拼接，再送到卷积和Tanh激活函数，得出了c，因为这里输出通道数是输入通道数的一半

c.shape=[1, 64, 20, 12]

h = (1 - z) * h + z * c

然后把h返回

h.shape=[1, 64, 20, 12]

并把h添加到o这个list里去

然后重新进入到下个循环里去，我这里是12帧，所以会循环12次，

o.shape=[1, 12, 64, 20, 12]

当然也会把h返回，这里的h就是最后一次计算出来的h

那么，理论上h就是o的最后一个元素

o[:,-1]==h

然后就返回到b/h，b就是o，r就是h

然后会把a和b再拼接一下，返回回来，得到x4/r4

x4.shape=[1, 12, 128, 20, 12]
r4.shape=[1, 64, 20, 12]

那我们知道，上面的a就是就是输入参数的前一半，所以返回回来之后，就是x4，那么x4的前面一半跟f4的前面一半是一模一样的，

x4/f3/s3/r3会送到 UpsamplingBlock 函数，当然第一次调用的时候r3为None

对应 x/f/s/r

会把x/f/s都展开，

x.shape=[12, 128, 20, 12]
f.shape=[12, 40, 40, 23]
s.shape=[12, 3, 40, 23]

x上采样

x.shape=[12, 128, 40, 24]

但是x要按照s的形状进行一下裁剪，变为了

x.shape=[12, 128, 40, 23]

x/f/s拼接，变为

x.shape=[12, 171, 40, 23]

经过一个卷积变为：

x.shape=[12, 80, 40, 23]

然后又分为a/b两半

还是把b和r送到ConvGRU里去

x3.shape=[1, 12, 80, 40, 23]
r3.shape=[1, 40, 40, 23]

然后decoder1/decoder2/decoder3都是同一个函数

x2.shape=[1, 12, 40, 80, 45]
r2.shape=[1, 20, 80, 45]

x1.shape=[1, 12, 32, 160, 90]
r1.shape=[1, 16, 160, 90]

然后是把x1和s0送到OutputBlock

这一次倒是没有送到ConvGRU，直接拼接一下，卷积一下就得出最终结果了

x0.shape=[1, 12, 16, 320, 180]

然后把x0和r1/r2/r3/r4返回，x0成为了hid

然后这里的hid还要再送到函数Projection里去计算，就是一个卷积运算，

出来的结果又分成两半

fgr_residual.shape=[1, 12, 3, 320, 180]
pha.shape=[1, 12, 1, 320, 180]

然后，因为有下采样，所以src/src_sm/fgr_residual/pha/hid都会送到DeepGuidedFilterRefiner，再经过一通卷积(没仔细看)，得到

fgr_residual.shape=[1, 12, 3, 1280, 720]
pha.shape=[1, 12, 1, 1280, 720]

fgr_residual和src相加，得到fgr,裁剪到0和1之间

然后返回，得到最终结果

rgr/pha和rec

fgr.shape=[1, 12, 3, 1280, 720]
pha.shape=[1, 12, 1, 1280, 720]

fgr就是前景output-foreground

pha就是output-alpha

最终结果就是pha乘以fgr，其余部分，用背景(绿幕)来填充

不过我感觉这个fgr没啥用，直接用pha乘以src 不好么