CenterNet代码解析

这篇文章主要就是介绍一些用到的重要的函数,只介绍detection部分。

0.网站

https://github.com/xingyizhou/CenterNet

install:

https://github.com/xingyizhou/CenterNet/blob/master/readme/INSTALL.md

dataset:

https://github.com/xingyizhou/CenterNet/blob/master/readme/DATA.md

1.ctdet_decode:作用是将heat_map解码成b-box

将输出转化成det的函数是lib\models\decode.py中的ctdet_decode。

1.1 首先经过_nms:

def _nms(heat, kernel=3):
    pad = (kernel - 1) // 2

    hmax = nn.functional.max_pool2d(
        heat, (kernel, kernel), stride=1, padding=pad)
    keep = (hmax == heat).float()
    return heat * keep

hmax用来寻找8-近邻极大值点,keep为h极大值点的位置,返回heat*keep,筛选出极大值点,为原值,其余为0。

 

2.1 之后经过_topk

def _topk(scores, K=40):
    batch, cat, height, width = scores.size()
      
    topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), K)

    topk_inds = topk_inds % (height * width)
    topk_ys   = (topk_inds / width).int().float()
    topk_xs   = (topk_inds % width).int().float()
      
    topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), K)
    topk_clses = (topk_ind / K).int()
    topk_inds = _gather_feat(
        topk_inds.view(batch, -1, 1), topk_ind).view(batch, K)
    topk_ys = _gather_feat(topk_ys.view(batch, -1, 1), topk_ind).view(batch, K)
    topk_xs = _gather_feat(topk_xs.view(batch, -1, 1), topk_ind).view(batch, K)

    return topk_score, topk_inds, topk_clses, topk_ys, topk_xs

topk_scores:  batch * cat * K,    batch代表batchsize,cat代表类别数,K代表K个最大值。

topk_indsbatch * cat * K,        index取值:[0, W x H - 1]

topk_scorestopk_inds分别为每个batch每张heatmap(每个类别)中前K个最大的score和id。

之后对topk_inds使用取余和除法得到横纵坐标top_ys、top_xs。

然后在每个batch中取所有heatmap的前K个最大score以及id,不考虑类别的影响。

topk_scorebatch * K

topk_indbatch * K                        index取值:[0, cat x K - 1]

之后对topk_inds(view后)topk_ind调用了_gather_feat函数,在utils文件中:

2.2 _gather_feat

def _gather_feat(feat, ind, mask=None):
    dim  = feat.size(2)
    ind  = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)
    feat = feat.gather(1, ind)
    if mask is not None:
        mask = mask.unsqueeze(2).expand_as(feat)
        feat = feat[mask]
        feat = feat.view(-1, dim)
    return feat

 输入:

feat(topk_inds):  batch * (cat x K) * 1          (假设输入的是topk_indstopk_ind)

ind(topk_ind)batch * K

首先将ind扩展一个指标,变为 batch * K * 1

之后使用gather,将ind对应的值取出来。

返回的是index:

feat:  batch * K * 1              取值:[0, cat x K - 1]

更一般的情况如下:

feat :  A * B * C      

indA * D

首先将ind扩展一个指标,并且expand为dim的大小,变为 A * D * C,其中对于任意的i, j, 数组ind[i, j, :]中所有的元素均相同,等于原来A * D shape的ind[i, j]。

之后使用gather,将ind对应的值取出来。

得到的feat: A * D * C

 

2.3 返回值

最后返回有四个:topk_score, topk_inds, topk_clses, topk_ys, topk_xs

topk_score:batch * K。每张图片中最大的K个值

topk_inds:batch * K 。没张图片中最大的K个值对应的index,这个index在[0, W x H - 1]之间。

后两个类似。

 

3.3 _tranpose_and_gather_feat,将_topk得到的index用于取值。

_tranpose_and_gather_feat的输入有reg,也有wh,前者应该是回归offset的,后者应该是得到bbox的W和H的。

    scores, inds, clses, ys, xs = _topk(heat, K=K)
    if reg is not None:
      reg = _tranpose_and_gather_feat(reg, inds)
wh = _tranpose_and_gather_feat(wh, inds)

以下是_tranpose_and_gather_feat的定义:

def _tranpose_and_gather_feat(feat, ind):
    feat = feat.permute(0, 2, 3, 1).contiguous()
    feat = feat.view(feat.size(0), -1, feat.size(3))
    feat = _gather_feat(feat, ind)
    return feat

输入:

featbatch * C(channel) * W * H

indbatch * K

首先将feat中各channel的元素放到最后一个index中,并且使用contiguous将内存变为连续的,用于后面的view。

之后将feat变为batch * (W x H) * C的形状,使用_gather_feat根据ind取出feat中对应的元素

返回:

featbatch * K * C

feat[i, j, k]为第i个batch,第k个channel的第j个最大值。

 

总体来说有点复杂,直接把它的逻辑用图来描述出来

假设输入是:\begin{bmatrix} [1 & 2 & 3\\ 1 & 2 & 3\\ 1 & 2 & 6]\\ [3 & 4 & 5\\ 3 & 4 & 7\\ 3 & 4 & 5]\\ \end{bmatrix},shape为batch * C * W * H(batch size直接设为1,忽略),对应于图中就是1 * 2 * 3 * 3,假设K=2。则经过以下两步之后

    scores, inds, clses, ys, xs = _topk(heat, K=K)
    if reg is not None:
      reg = _tranpose_and_gather_feat(reg, inds)

最终得到的是:\begin{bmatrix} [3 & 7] [6 & 3] \end{bmatrix},shape为batch * K * C。[3, 7]中的7是所有channel中最大的元素,6则是第二大的元素,将所有channel对应对应位置的元素取出来就得到了最终的结果。

其中__gather_feat起到的作用是消除各个channel区别的作用,最终得到的inds是对于所有channel而言的。

_tranpose_and_gather_feat的作用则是解码获得的inds,取得最终的结果。

_topk输入的feat就是定位的heat_map,在这上面获得inds后,这个inds就可以应用到offset_heat_map、size_heat_map上面。

 

下面用图示详细解释这两行代码的过程:

CenterNet代码解析_第1张图片

CenterNet代码解析_第2张图片

ctdet_decode的代码解释如下:

def ctdet_decode(heat, wh, reg=None, cat_spec_wh=False, K=100):
    batch, cat, height, width = heat.size()

    # heat = torch.sigmoid(heat)
    # perform nms on heatmaps
    heat = _nms(heat)
    
    scores, inds, clses, ys, xs = _topk(heat, K=K)
    # xs、ys是inds转化成在heat_map上面的行、列

    if reg is not None:
      reg = _tranpose_and_gather_feat(reg, inds)
      reg = reg.view(batch, K, 2)
      xs = xs.view(batch, K, 1) + reg[:, :, 0:1]
      ys = ys.view(batch, K, 1) + reg[:, :, 1:2]
    else:
      xs = xs.view(batch, K, 1) + 0.5
      ys = ys.view(batch, K, 1) + 0.5

    # xs、ys都加上一个偏移

    wh = _tranpose_and_gather_feat(wh, inds)
    # 取wh中对应与inds的元素,就像上面的例子中一样。


    if cat_spec_wh:
      wh = wh.view(batch, K, cat, 2)
      clses_ind = clses.view(batch, K, 1, 1).expand(batch, K, 1, 2).long()
      wh = wh.gather(2, clses_ind).view(batch, K, 2)
    else:
      wh = wh.view(batch, K, 2)
    clses  = clses.view(batch, K, 1).float()
    scores = scores.view(batch, K, 1)
    bboxes = torch.cat([xs - wh[..., 0:1] / 2, 
                        ys - wh[..., 1:2] / 2,
                        xs + wh[..., 0:1] / 2, 
                        ys + wh[..., 1:2] / 2], dim=2)
    # bbox就这样获得了。
    detections = torch.cat([bboxes, scores, clses], dim=2)
      
    return detections

2.后处理

上面根据heatmap得到了dets,但是还需要进一步处理:

1. demo中的line 30:ret = detector.run(img),detector为ctdet

2. base_detector中的line 82:run函数:

images -> output、dets

dets->      dets = self.post_process(dets, meta, scale) ->    detections.append(dets)

detections ->    results = self.merge_outputs(detections) ->results

3.上面的两个过程:post_process和merge_outputs在ctdet中进行了定义

4.post_process:

    dets = dets.detach().cpu().numpy()
    dets = dets.reshape(1, -1, dets.shape[2])
    dets = ctdet_post_process(
        dets.copy(), [meta['c']], [meta['s']],
        meta['out_height'], meta['out_width'], self.opt.num_classes)
    for j in range(1, self.num_classes + 1):
      dets[0][j] = np.array(dets[0][j], dtype=np.float32).reshape(-1, 5)
      dets[0][j][:, :4] /= scale
    return dets[0]

做的应该就是尺度变换之类的吧。

5.merge_outputs:

  def merge_outputs(self, detections):
    results = {}
    for j in range(1, self.num_classes + 1):
      results[j] = np.concatenate(
        [detection[j] for detection in detections], axis=0).astype(np.float32)
      if len(self.scales) > 1 or self.opt.nms:
         soft_nms(results[j], Nt=0.5, method=2)
    scores = np.hstack(
      [results[j][:, 4] for j in range(1, self.num_classes + 1)])
    if len(scores) > self.max_per_image:
      kth = len(scores) - self.max_per_image
      thresh = np.partition(scores, kth)[kth]
      for j in range(1, self.num_classes + 1):
        keep_inds = (results[j][:, 4] >= thresh)
        results[j] = results[j][keep_inds]
    return results

大致上是先做soft_nms,之后scores进行一个筛选,将那些大于最多检测框(100)剔除掉。

你可能感兴趣的:(目标检测,Object,detection,代码解析)