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

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

网络结构

上图是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_conf，dets_coor，dets_orien。

其中dets_conf的尺寸是：18207*80。这里$18207=3\ast (17^2+34^2+68^2)$，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) / 2final_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 suppressionRefer to `torchvision.ops.boxes.batched_nms` for details.Args:dets (n, 6): bounding boxes with x, y, width, height and scorecats (n,): cls indices of bounding boxesthreshold (float): IoU threshold for NMSnormalized (bool): use normalized coordinates or notReturns: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() / 2offsets = cats.float().view(-1, 1) * (max_coordinate + 0.5)dets_for_nms = dets.clone()dets_for_nms[:, :2] += offsetsprint("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], keepdef 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_xymasks = ((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的过程。