Real-time Instance Segmentation with Discriminative Orientation Maps

目录

OrienMask是2021年新出的一篇实时的实例分割模型。下图是OrienMask在coco数据集上的测试结果。在单张2080ti上可以达到42.7fps的速度。

Real-time Instance Segmentation with Discriminative Orientation Maps

这篇文章的主体思路是先使用yolo检测出目标的bbox,然后再检测出bbox中的前景来完成实例分割任务。

网络结构

Real-time Instance Segmentation with Discriminative Orientation Maps
上图是OrienMask的网络结构,主要分为3块:Backbone、FPN、Heads。

使用darknet53作为Backbone,进行了32倍下采样,分别在4x,8x,16x,32x下采样的进行FPN来进行特征融合;最后是检测heads,32x、16x、8倍下采样分支上输出bbox结果,4倍下采样的分支上输出Orientation Maps。可以看出来除了在4倍下采样的分支上输出Orientation Maps之外只是一个yolo。

网络输出

这里详细介绍网络的输出和Orientation Maps的作用。

在后处理之前,模型的输出是:dets_confdets_coordets_orien

其中dets_conf的尺寸是:18207*80。这里 18207 = 3 ∗ ( 1 7 2 + 3 4 2 + 6 8 2 ) 18207=3*(17^2+34^2+68^2) 18207=3∗(172+342+682),3是表示每个尺度下有3个不同比例的anchor,17、34、68分别是3个尺度的grid size;80代表类的数量。dets_conf里存放的是所有尺度的grid的置信度。

dets_coor的尺寸是:18207*4。18207的定义同上,4代表bbox的长度即xywh。这里存放的所有尺度的grid的xywh的偏移量。

dets_orien的尺寸是9*2*544*544。这里9代表9个anchor(3个尺度*3个不同长宽比),2代表x,y(向量),544是图像输入尺寸。即dets_orien里存了输入图像在9个anchor上的方向向量。

后处理

怕表达不清楚,后处理这块照着代码聊吧。这里提供的代码是我转tensorrt时没能转进onnx的部分,与源码中的后处理代码可能会有一点出入,但是最终结果是对的。后面会介绍转tensorrt的全部流程和代码。

def post_processing(dets_coord, dets_conf, dets_orien, dets_anchor_idx, grid_sizes, base_xy, grid_anchors, pre_post_num=400, conf_thresh=0.005):
    selected_inds, dets_cls = torch.nonzero(dets_conf > conf_thresh, as_tuple=True)
    selected_inds = selected_inds.view(-1)
    dets_cls = dets_cls.view(-1)
    dets_conf = dets_conf[selected_inds, dets_cls]
    # 选择固定数量的候选框留给 NMS 操作
    if selected_inds.numel() > pre_post_num:
        dets_conf, topk_inds = dets_conf.topk(pre_post_num)
        selected_inds = selected_inds[topk_inds]
        dets_cls = dets_cls[topk_inds]

    dets_coord = dets_coord[selected_inds]
    dets_anchor_idx = dets_anchor_idx[selected_inds]
    dets_orien = dets_orien * grid_anchors.reshape(3, 6, 1, 1) / 2
    final_pred = multi_class_nms(base_xy, grid_sizes, dets_coord, dets_conf, dets_cls, dets_anchor_idx, dets_orien)
    return final_pred

第2到5行是选择出置信度大于阈值的,selected_inds是anchor的索引,dets_cls是类别索引,这里的阈值设的非常小,会影响速率,这个问题在讲nms部分时会分析。第7行的if是防止太多的数据进入nms,即满足置信度条件的数量如果大于pre_post_num就取个topk。直到第15行,进入multi_class_nms,这里不仅做了bb的nms还根据dets_orien求出mask图。

下面展开multi_class_nms代码:

def batched_nms(dets, cats, threshold=0.5, normalized=True):
    """Batched non maximum suppression

    Refer to `torchvision.ops.boxes.batched_nms` for details.

    Args:
        dets (n, 6): bounding boxes with x, y, width, height and score
        cats (n,): cls indices of bounding boxes
        threshold (float): IoU threshold for NMS
        normalized (bool): use normalized coordinates or not

    Returns:
        remains dets, cats and their indices
    """
    if dets.size(0) == 0:
        keep = dets.new_zeros(0, dtype=torch.long)
    else:
        max_coordinate = 1.5 if normalized else dets[:, :2].max() + dets[:, 2:4].max() / 2
        offsets = cats.float().view(-1, 1) * (max_coordinate + 0.5)
        dets_for_nms = dets.clone()
        dets_for_nms[:, :2] += offsets
        print("dets_for_nms: ", dets_for_nms.shape)
        print(dets_for_nms)
        print("*"*88)
        if dets.is_cuda:
            keep = nms_cuda.nms(dets_for_nms, threshold)
        else:
            keep = nms_cpu.nms(dets_for_nms, threshold)

    return dets[keep], cats[keep], keep


def multi_class_nms(base_xy, grid_sizes, pred_coord, pred_conf, pred_cls, pred_anchor_idx, pred_orien, post_nms_num=100, orien_thresh=0.3):
    dets = torch.cat([pred_coord, pred_conf.unsqueeze(-1)], dim=1)
    # 使用Batch_NMS求取结果
    dets, cats, keep = batched_nms(dets, pred_cls)
    # 求取前100个结果
    if keep.numel() > post_nms_num:
        _, topk_inds = dets[:, -1].topk(post_nms_num)
        dets = dets[topk_inds]
        cats = cats[topk_inds]
        keep = keep[topk_inds]

    anchor_idx = pred_anchor_idx[keep]
    x_centers = (grid_sizes[anchor_idx, 0] * dets[:, 0]).view(-1, 1, 1)
    y_centers = (grid_sizes[anchor_idx, 1] * dets[:, 1]).view(-1, 1, 1)
    det_width = dets[:, 2].view(-1, 1, 1)
    det_height = dets[:, 3].view(-1, 1, 1)

    pred_orien = pred_orien.reshape(9, 2, 544, 544)
    pred_orien = pred_orien + base_xy

    masks = ((torch.abs(pred_orien[anchor_idx, 0] - x_centers) < orien_thresh * det_width * grid_sizes[anchor_idx, 0].view(-1, 1, 1)) &
             (torch.abs(pred_orien[anchor_idx, 1] - y_centers) < orien_thresh * det_height * grid_sizes[anchor_idx, 1].view(-1, 1, 1)))

    return {'bbox': dets, 'mask': masks, 'cls': cats}

multi_class_nms最开始将xywh偏移量(pred_coord)和置信度(pred_conf)cat到一起送到batched_nms里。batched_nms里主要做了bbox的nms,在这里又一次卡了置信的的阈值,这次阈值比较大也是最终的bbox的阈值,个人感觉这样做是不太合理的,卡一次最终的阈值就可以了,因为在C++复现后处理的时候会有速率的影响,并且我在C++那里把卡阈值的操作做成了函数,连续调用两次相同的函数看起来就很呆。batched_nms是非常常用的东西就不展开介绍了,重点是根据pred_orien生成mask图:由于预测出的orien是单位向量,所以在post_processing中先乘上了grid_anchors,再在multi_class_nms里加上了base_xy。

这里grid_anchors是anchor的尺寸,在这个模型里设置的是:

grid_anchors = torch.tensor([[1.5000, 2.0000],
                             [2.3750, 4.5000],
                             [5.0000, 3.5000],
                             [2.2500, 4.6875],
                             [4.7500, 3.4375],
                             [4.5000, 9.1250],
                             [4.4375, 3.4375],
                             [6.0000, 7.5938],
                             [14.3438, 12.5312]])/2

base_xy是在不同尺度grid上的伪像素坐标,生成方式我也摘出来了:

grid_size = [[17, 17],
             [34, 34],
             [68, 68]]
nHs = [size[0] for size in grid_size]
nWs = [size[1] for size in grid_size]
num_anchors = [len(m) for m in anchor_mask]
base_xy = torch.zeros(9, 2, 544, 544)
grid_x, grid_y, anchor_idx = [], [], []
for m, nA, nH, nW in zip(anchor_mask, num_anchors, nHs, nWs):
    # 构建一个基础的方向坐标 (unit size=grid)
    base_y, base_x = torch.meshgrid([torch.arange(544, dtype=torch.float) / 544 * nH,
                                     torch.arange(544, dtype=torch.float) / 544 * nW])
    # 第一个维度上进行拼接操作
    base_xy[m] = torch.stack([base_x, base_y], dim=0)

最后orien向量的x<bbox.width*threshold 且 y<bbox.height*threshold时认为时bbox的前景。这样就完成了实例分割。

总结

这个模型的训练工作是其他同事来做的,我的主要工作是把这个模型移植到tensorrt上所以更多关注了模型的后处理部分,模型架构部分的代码没去细看。整体看下来这个模型架构比较清晰,后处理过程也比较简单。我的初版代码在tensorrt上infer时间是18ms后处理大部分放到了cuda上是3ms,预处理使用opencv写的会比较耗时大概20ms,但是预处理可以和infer+后处理放在不同线程去跑。综合帧率还是能达到40fps左右的(速度上没体现出来tensorrt的优势,惭愧)。代码优化后并且把模型转fp16或者量化到int8应该会更快的。后续有时间的话再记录转tensorrt的过程。

上一篇:图像分割基础


下一篇:Semi-supervised semantic segmentation needs strong, varied perturbations