标签：Real orien Segmentation Orientation pred grid dets anchor nms
目录
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 ∗ ( 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的过程。
标签：Real,orien,Segmentation,Orientation,pred,grid,dets,anchor,nms 来源： https://blog.csdn.net/weixin_47062350/article/details/121125988
本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享； 2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关； 3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关； 4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除； 5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。