【视频分割】【深度学习】MiVOS官方Pytorch代码–Propagation模块解析

MiVOS模型将交互到掩码和掩码传播分离，从而实现更高的泛化性和更好的性能。单独训练的交互模块将用户交互转换为对象掩码，传播模块使用一种新的top-k过滤策略在读取时空存储器时进行临时传播，本博客将讲解Propagation(用户交互产生分割图)。

【解析代码地址】

前言

在详细解析MiVOS代码之前,首要任务是成功运行MiVOS代码【win10下参考教程】,后续学习才有意义。
本博客讲解Propagation(掩码传播)的功能模块，之前的博文已经讲解了S2M(用户交互产生分割图)的功能模块，因此interactive_gui.py文件只详细讲解与Propagation功能相关的代码。
【代码:用interactive_gui_4.1.py代替interactive_gui.py】

用户界面新增Propagation模块

主函数新增代码

在主函数中，实例化Propagation对象和Fusion对象并加载权重。

prop_saved = torch.load(args.prop_model)
prop_model = PropagationNetwork().cuda().eval()
prop_model.load_state_dict(prop_saved)

fusion_saved = torch.load(args.fusion_model)
fusion_model = FusionNet().cuda().eval()
fusion_model.load_state_dict(fusion_saved)

InferenceCore封装了prop_net模型和fuse_net模型为一个推理器，对应掩码传播和融合过程。

self.processor = InferenceCore(prop_net, fuse_net, images_to_torch(images, device='cpu'),
                                       num_objects, mem_freq=mem_freq, mem_profile=mem_profile)

__init__函数新增代码

新增Propagate按钮和Propagate进度条

self.run_button = QPushButton('Propagate')
self.run_button.clicked.connect(self.on_run)

# 功能：显示progress进度
self.progress = QProgressBar(self)
self.progress.setGeometry(0, 0, 300, 25)
self.progress.setMinimumWidth(300)
self.progress.setMinimum(0)
self.progress.setMaximum(100)
self.progress.setFormat('Idle')
self.progress.setStyleSheet("QProgressBar{color: black;}")
self.progress.setAlignment(Qt.AlignCenter)

on_run函数

获取关键图片mask后，通过interact完成掩膜传播得到所有图片的mask。

def on_run(self):
    if self.interacted_mask is None:
        return
    # 掩码传播
    self.current_mask = self.processor.interact(self.interacted_mask, self.cursur,
                        self.progress_total_cb, self.progress_step_cb)
    self.interacted_mask = None
    # clear scribble and reset
    self.show_current_frame()
    # 清除当前图像所有交互对象的交互记录
    self.reset_this_interaction()
    self.progress.setFormat('Idle')
    self.progress.setValue(0)

progress_total_cb/progress_step_cb函数

直观的显示当前传播的进度。

def progress_step_cb(self):
    self.progress_num += 1
    ratio = self.progress_num/self.progress_max
    self.progress.setValue(int(ratio*100))
    self.progress.setFormat('%2.1f%%' % (ratio*100))
    QApplication.processEvents()

def progress_total_cb(self, total):
    self.progress_max = total
    self.progress_num = -1
    self.progress_step_cb()

InferenceCore类关键代码讲解

在inference_core.py文件中。

__init__函数

InferenceCore类封装了prop_net模型和fuse_net模型，对输入模型的图片做了简单的预处理，设置了程序运行时GPU资源的使用权限等等。

def __init__(self, prop_net : PropagationNetwork, fuse_net : FusionNet, images, num_objects,
                mem_profile=0, mem_freq=5, device='cuda:0'):
    self.prop_net = prop_net.to(device, non_blocking=True)          # 掩码传播
    if fuse_net is not None:
        self.fuse_net = fuse_net.to(device, non_blocking=True)      # 掩码融合
    self.mem_profile = mem_profile                                  # 使用GPU资源的权限,0~3级 数字越高权限越低
    self.mem_freq = mem_freq                                        # 使用内存资源的时间,数字越高内存使用越少
    self.device = device

    # 使用GPU资源的不同权限
    if mem_profile == 0:
        self.data_dev = device
        self.result_dev = device
        self.q_buf_size = 105
        self.i_buf_size = -1    # 无需缓冲的图像(-1表示都加载)
    elif mem_profile == 1:
        self.data_dev = 'cpu'
        self.result_dev = device
        self.q_buf_size = 105
        self.i_buf_size = 105   # 无需缓冲的图像 105张
    elif mem_profile == 2:
        self.data_dev = 'cpu'
        self.result_dev = 'cpu'
        self.q_buf_size = 3
        self.i_buf_size = 3
    else:
        self.data_dev = 'cpu'
        self.result_dev = 'cpu'
        self.q_buf_size = 1
        self.i_buf_size = 1

    # 图片数量
    t = images.shape[1]
    h, w = images.shape[-2:]
    # 目标数
    self.k = num_objects

    # 用补padding的方式将输入的图片高宽变为16的整数倍
    self.images, self.pad = pad_divide_by(images, 16, images.shape[-2:])
    # 补过padding后新图片的高宽
    nh, nw = self.images.shape[-2:]
    # 图片是否加载到GPU
    self.images = self.images.to(self.data_dev, non_blocking=False)

    # 这两者以不同的格式存储相同的信息
    # masks是为了输入到模型中,np_masks是为了在主界面展示
    self.masks = torch.zeros((t, 1, nh, nw), dtype=torch.uint8, device=self.result_dev)
    self.np_masks = np.zeros((t, h, w), dtype=np.uint8)

    # prob表示包括背景在内目标的概率
    self.prob = torch.zeros((self.k+1, t, 1, nh, nw), dtype=torch.float32, device=self.result_dev)   # k+1:object_num t:batch
    # 防止背景概率为0
    self.prob[0] = 1e-7

    # 记录图片batchsize 原始高宽
    self.t, self.h, self.w = t, h, w
    # 记录padding处理后的高宽
    self.nh, self.nw = nh, nw
    # padding处理后图像的比
    self.kh = self.nh//16
    self.kw = self.nw//16

    # 查询缓冲区
    self.query_buf = {}
    # 图片缓冲区
    self.image_buf = {}
    # 存储有过交互的图片序号
    self.interacted = set()
    # 存储key/value
    self.certain_mem_k = None
    self.certain_mem_v = None

关于Propagation模块中深度学习网络会单出一期详细讲解，现在暂不展开讲解，避免内容过于杂乱，只需要读者知道流程。

interact函数

以用户当前帧作为起点，对mask进行双向传播，通过用户标注，获取某些帧的mask，再通过双向传播获得所有帧的mask。
过去带有对象掩码的帧视为内存帧self.image_buf(get_image_buffered)，计算memory的key和value存入self.certain_mem_k和self.certain_mem_v中，用于预测查询帧的对象掩码。通过do_pass完成包括传播和融合的完整过程。

def interact(self, mask, idx, total_cb=None, step_cb=None):
    """
    Interact -> Propagate -> Fuse

    mask - One-hot mask of the interacted frame, background included
    idx - Frame index of the interacted frame
    total_cb, step_cb - Callback functions for the GUI

    Return: all mask results in np format for DAVIS evaluation
    """
    # 记录有过交互的图片序号
    self.interacted.add(idx)
    mask = mask.to(self.device)
    # 用补padding的方式将输入的mask高宽变为16的整数倍
    mask, _ = pad_divide_by(mask, 16, mask.shape[-2:])

    self.mask_diff = mask - self.prob[:, idx].to(self.device)
    # idx正差异
    self.pos_mask_diff = self.mask_diff.clamp(0, 1)
    # idx负差异
    self.neg_mask_diff = (-self.mask_diff).clamp(0, 1)
    # 对于图片idx,直接用mask更新self.prob
    self.prob[:, idx] = mask
    # mask是包括背景的各目标可能的概率-->mask[1:]除去背景
    # 获得图像idx的存储key/value
    key_k, key_v = self.prop_net.memorize(self.get_image_buffered(idx), mask[1:])
    # key_k 用于评估当前帧和之前帧的相似性 [num_objects, 512, 1, H/16, W/16]
    # key_v 用来生成最后mask精细结果信息 [num_objects, 512, 1, H/16, W/16]

    # 记录所有有过交互的图片idex的存储key/value
    if self.certain_mem_k is None:
        self.certain_mem_k = key_k
        self.certain_mem_v = key_v
    else:
        # shape的第2维就代表certain_mem_k存储的图片key/value的数量
        self.certain_mem_k = torch.cat([self.certain_mem_k, key_k], 2)
        self.certain_mem_v = torch.cat([self.certain_mem_v, key_v], 2)

    # 统计当前图片前后传播给其他图片的范围

    # 这部分与GUI界面中进度条有关
    if total_cb is not None:
        # 前向传播找最近的(min)  idx---->
        front_limit = min([ti for ti in self.interacted if ti > idx] + [self.t])
        # 反向传播找最近的(max)  <----idx
        back_limit = max([ti for ti in self.interacted if ti < idx] + [-1])
        # max<----idx---->min
        # 传播不包括max和min以及自身idx
        total_num = front_limit - back_limit - 2    # -1 for shift, -1 for center frame
        if total_num > 0:
            total_cb(total_num)
     # 这部分与GUI界面中进度条有关

    # 这部分与模型真在的传播过程有关
    # True 正向传播 False 反向传播
    self.do_pass(key_k, key_v, idx, True, step_cb=step_cb)
    self.do_pass(key_k, key_v, idx, False, step_cb=step_cb)
    # 这部分与模型真在的传播过程有关

    # This is a more memory-efficient argmax
    # mask对应的原始片像素位属于概率最大的obj
    for ti in range(self.t):
        self.masks[ti] = torch.argmax(self.prob[:, ti], dim=0)
    out_masks = self.masks

    # 去除padding部分
    if self.pad[2]+self.pad[3] > 0:
        out_masks = out_masks[:, :, self.pad[2]:-self.pad[3], :]
    if self.pad[0]+self.pad[1] > 0:
        out_masks = out_masks[:, :, :, self.pad[0]:-self.pad[1]]

    # np_masks是为了GUI界面显示
    self.np_masks = (out_masks.detach().cpu().numpy()[:, 0]).astype(np.uint8)
    return self.np_masks

产生memory的key/value示意图：

get_image_buffered函数

将当前带有对象掩码的帧放入到内存self.image_buf中。

def get_image_buffered(self, idx):
    # 模型的加载模式(cpu/GPU)与数据的加载模式一致,数据不再处理
    if self.data_dev == self.device:
        return self.images[:, idx]
    # 图片是否在图片缓冲区
    if idx not in self.image_buf:
        # 是否超过了图像缓冲区支持的最大缓冲值
        if len(self.image_buf) > self.i_buf_size:
            self.image_buf = {}     # 重置图片缓冲区

    # 设置当前图片加载模式与模型的一致,并放入图像缓冲区,
    self.image_buf[idx] = self.images[:, idx].to(self.device)
    # 返回图片缓冲区
    result = self.image_buf[idx]
    return result

do_pass函数

传播和融合的完整过程。
双向传播操作，forward标志用判断是前向传播还是反向传播，一旦遇到结尾帧或者交互过的帧就停止传播。max<----idx---->min

差异感知与线性结果作为共同输入来预测最终的输出(fuse_one_frame)。

这里是传播和深度学习训练过程的传播不是一个概念

def do_pass(self, key_k, key_v, idx, forward=True, step_cb=None):
    """
    Do a complete pass that includes propagation and fusion
    key_k/key_v -  memory feature of the starting frame
    idx - Frame index of the starting frame
    forward - forward/backward propagation
    step_cb - Callback function used for GUI (progress bar) only
    """
    # Pointer in the memory bank
    # 已经存在的关键帧数量
    num_certain_keys = self.certain_mem_k.shape[2]
    # 关键帧指针
    m_front = num_certain_keys

    # Determine the required size of the memory bank
    # 双向传播操作,一旦遇到结尾帧或者交互过的帧就停止传播
    if forward:
        #  前向传播找最近的(min)  idx---->min
        closest_ti = min([ti for ti in self.interacted if ti > idx] + [self.t])
        # closest_ti - idx - 1:不包括idx和min; mem_freq是内存容量
        total_m = (closest_ti - idx - 1)//self.mem_freq + 1 + num_certain_keys
    else:
        # 反向传播找最近的(max)  max<----idx
        closest_ti = max([ti for ti in self.interacted if ti < idx] + [-1])
        # 不包括idx和max
        total_m = (idx - closest_ti - 1)//self.mem_freq + 1 + num_certain_keys
    K, CK, _, H, W = key_k.shape
    _, CV, _, _, _ = key_v.shape

    # Pre-allocate keys/values memory
    keys = torch.empty((K, CK, total_m, H, W), dtype=torch.float32, device=self.device)
    values = torch.empty((K, CV, total_m, H, W), dtype=torch.float32, device=self.device)

    # Initial key/value passed in
    # 初始化关键帧的key和value
    keys[:, :, 0:num_certain_keys] = self.certain_mem_k
    values[:, :, 0:num_certain_keys] = self.certain_mem_v
    prev_in_mem = True
    last_ti = idx    # 在传播过程中保证在mem_freq范围内

    # Note that we never reach closest_ti, just the frame before it
    if forward:
        # 前向传播的范围
        this_range = range(idx+1, closest_ti)
        step = +1
        end = closest_ti - 1
    else:
        # 反向传播的范围
        this_range = range(idx-1, closest_ti, -1)
        step = -1
        end = closest_ti + 1

    for ti in this_range:
        if prev_in_mem:     # 包不包括idx的key/value
            this_k = keys[:, :, :m_front]
            this_v = values[:, :, :m_front]
        else:
            this_k = keys[:, :, :m_front+1]
            this_v = values[:, :, :m_front+1]
        query = self.get_query_kv_buffered(ti)
        # 获得mask:正确区分背景和多个目标前景
        out_mask = self.prop_net.segment_with_query(this_k, this_v, *query)
        out_mask = aggregate_wbg(out_mask, keep_bg=True)
        if ti != end:
            keys[:, :, m_front:m_front+1], values[:, :, m_front:m_front+1] = self.prop_net.memorize(
                    self.get_image_buffered(ti), out_mask[1:])
            if abs(ti-last_ti) >= self.mem_freq:
                # Memorize the frame
                m_front += 1
                last_ti = ti
                prev_in_mem = True
            else:
                prev_in_mem = False

        # In-place fusion, maximizes the use of queried buffer
        # esp. for long sequence where the buffer will be flushed
        if (closest_ti != self.t) and (closest_ti != -1):
            # 差异感知融合
            self.prob[:, ti] = self.fuse_one_frame(closest_ti, idx, ti, self.prob[:, ti], out_mask,
                                    key_k, query[3]).to(self.result_dev)
        else:
            self.prob[:, ti] = out_mask.to(self.result_dev)

        # Callback function for the GUI
        if step_cb is not None:
            step_cb()
    return closest_ti

产生query的key/value示意图：

计算query和memory的关系:

get_query_kv_buffered函数

def get_query_kv_buffered(self, idx):
    # 图片是否在查询缓冲区
    if idx not in self.query_buf:
        # 是否超过了查询缓冲区支持的最大缓冲值
        if len(self.query_buf) > self.q_buf_size:
            self.query_buf = {}     # 重置查询缓存区
        #  返回f16, f8, f4, k16, v16
        #  k16 [num_objects, 128, 1, H/16, W/16] v16 [num_objects, 512, 1, H/16, W/16]
        self.query_buf[idx] = self.prop_net.get_query_values(self.get_image_buffered(idx))
    # 返回查询缓冲区
    result = self.query_buf[idx]
    return result

fuse_one_frame函数

结合了差异感知融合方法与线性融合方法，将差异感知与线性结果作为共同输入，再通过一个简单的网络来预测最终的输出。

def fuse_one_frame(self, tc, tr, ti, prev_mask, curr_mask, mk16, qk16):
    assert(tc<ti<tr or tr<ti<tc)    # 必须在符合的传播范围内
    prob = torch.zeros((self.k, 1, self.nh, self.nw), dtype=torch.float32, device=self.device)
    nc = abs(tc-ti) / abs(tc-tr)
    nr = abs(tr-ti) / abs(tc-tr)
    # 时间相关
    dist = torch.FloatTensor([nc, nr]).to(self.device).unsqueeze(0)
    for k in range(1, self.k+1):
        # 注意力位置
        attn_map = self.prop_net.get_attention(mk16[k-1:k], self.pos_mask_diff[k:k+1], self.neg_mask_diff[k:k+1], qk16)
        # 融合过程
        w = torch.sigmoid(self.fuse_net(self.get_image_buffered(ti), 
                prev_mask[k:k+1].to(self.device), curr_mask[k:k+1].to(self.device), attn_map, dist))
        prob[k-1] = w 
    return aggregate_wbg(prob, keep_bg=True)

融合的完整过程示意图。

总结

尽可能简单、详细的介绍MiVOS中Propagation模块的代码。后续会讲解Propagation的网络原理和代码(PropagationNetwork和FusionNet)以及MiVOS的训练。