【從零開始學CenterNet】7. CenterNet測試推理過程

【GiantPandaCV導語】這是CenterNet系列的最后一篇。本文主要講CenterNet在推理過程中的數據加載和后處理部分代碼。最后提供了一個已經配置好的數據集供大家使用。

代碼注釋在：https://github.com/pprp/SimpleCVReproduction/tree/master/CenterNet

1. eval部分數據加載

由於CenterNet是生成了一個heatmap進行的目標檢測，而不是傳統的基於anchor的方法，所以訓練時候的數據加載和測試時的數據加載結果是不同的。並且在測試的過程中使用到了Test Time Augmentation（TTA），使用到了多尺度測試，翻轉等。

在CenterNet中由於不需要非極大抑制，速度比較快。但是CenterNet如果在測試的過程中加入了多尺度測試，那就回調用soft nms將不同尺度的返回的框進行抑制。

class PascalVOC_eval(PascalVOC):
    def __init__(self, data_dir, split, test_scales=(1,), test_flip=False, fix_size=True, **kwargs):
        super(PascalVOC_eval, self).__init__(data_dir, split, **kwargs)
        # test_scale = [0.5,0.75,1,1.25,1.5]
        self.test_flip = test_flip
        self.test_scales = test_scales
        self.fix_size = fix_size

    def __getitem__(self, index):
        img_id = self.images[index]
        img_path = os.path.join(
            self.img_dir, self.coco.loadImgs(ids=[img_id])[0]['file_name'])
        image = cv2.imread(img_path)
        height, width = image.shape[0:2]

        out = {}
        for scale in self.test_scales:
            # 得到多個尺度的圖片大小
            new_height = int(height * scale)
            new_width = int(width * scale)

            if self.fix_size:
                # fix size代表根據參數固定圖片大小
                img_height, img_width = self.img_size['h'], self.img_size['w']
                center = np.array(
                    [new_width / 2., new_height / 2.], dtype=np.float32)
                scaled_size = max(height, width) * 1.0
                scaled_size = np.array(
                    [scaled_size, scaled_size], dtype=np.float32)
            else:
                # self.padding = 31  # 127 for hourglass
                img_height = (new_height | self.padding) + 1
                img_width = (new_width | self.padding) + 1
                # 按位或運算，找到最接近的[32,64,128,256,512]
                center = np.array(
                    [new_width // 2, new_height // 2], dtype=np.float32)
                scaled_size = np.array(
                    [img_width, img_height], dtype=np.float32)

            img = cv2.resize(image, (new_width, new_height))
            trans_img = get_affine_transform(
                center, scaled_size, 0, [img_width, img_height])
            img = cv2.warpAffine(img, trans_img, (img_width, img_height))

            img = img.astype(np.float32) / 255.
            img -= self.mean
            img /= self.std
            # from [H, W, C] to [1, C, H, W]
            img = img.transpose(2, 0, 1)[None, :, :, :]

            if self.test_flip: # 橫向翻轉
                img = np.concatenate((img, img[:, :, :, ::-1].copy()), axis=0)

            out[scale] = {'image': img,
                          'center': center,
                          'scale': scaled_size,
                          'fmap_h': img_height // self.down_ratio, # feature map的大小
                          'fmap_w': img_width // self.down_ratio}

        return img_id, out

以上是eval過程的數據加載部分的代碼，主要有兩個需要關注的點：

1. 如果是多尺度會根據test_scale的值返回不同尺度的結果，每個尺度都有img,center等信息。這部分代碼可以和test.py代碼的多尺度處理一塊理解。
2. 尺度處理部分，有一個padding參數

img_height = (new_height | self.padding) + 1
img_width = (new_width | self.padding) + 1

這部分代碼作用就是通過按位或運算，找到最接近的2的倍數-1作為最終的尺度。

'''
>>> 10 | 31
31
>>> 20 | 31
31
>>> 510 | 31
511
>>> 256 | 31
287
>>> 510 | 127
511
>>> 1000 | 127
1023
'''

例如：輸入512，多尺度開啟：0.5,0.7,1.5，那最終的結果是

512 x 0.5 | 31 = 287

512 x 0.7 | 31 = 383

512 x 1.5 | 31 = 799

2. 推理過程

上圖是CenterNet的結構圖，使用的是PlotNeuralNet工具繪制。在推理階段，輸入圖片通過骨干網絡進行特征提取，然后對下采樣得到的特征圖進行預測，得到三個頭，分別是offset head、wh head、heatmap head。

推理過程核心工作就是從heatmap提取得到需要的bounding box，具體的提取方法是使用了一個3x3的最大化池化，檢查當前熱點的值是否比周圍8個臨近點的值都大。然后取100個這樣的點，再做篩選。

以上過程的核心函數是：

output = model(inputs[scale]['image'])[-1]
dets = ctdet_decode(*output, K=cfg.test_topk)

ctdet_decode這個函數功能就是將heatmap轉化成bbox:

def ctdet_decode(hmap, regs, w_h_, K=100):
    '''
    hmap提取中心點位置為xs,ys
    regs保存的是偏置，需要加在xs,ys上，代表精確的中心位置
    w_h_保存的是對應目標的寬和高
    '''
    # dets = ctdet_decode(*output, K=cfg.test_topk)
    batch, cat, height, width = hmap.shape
    hmap = torch.sigmoid(hmap) # 歸一化到0-1

    # if flip test
    if batch > 1: # batch > 1代表使用了翻轉
        # img = np.concatenate((img, img[:, :, :, ::-1].copy()), axis=0)
        hmap = (hmap[0:1] + flip_tensor(hmap[1:2])) / 2
        w_h_ = (w_h_[0:1] + flip_tensor(w_h_[1:2])) / 2
        regs = regs[0:1]
    batch = 1
    # 這里的nms和帶anchor的目標檢測方法中的不一樣，這里使用的是3x3的maxpool篩選
    hmap = _nms(hmap)  # perform nms on heatmaps
    # 找到前K個極大值點代表存在目標
    scores, inds, clses, ys, xs = _topk(hmap, K=K)

    regs = _tranpose_and_gather_feature(regs, inds)

    regs = regs.view(batch, K, 2)

    xs = xs.view(batch, K, 1) + regs[:, :, 0:1]
    ys = ys.view(batch, K, 1) + regs[:, :, 1:2]

    w_h_ = _tranpose_and_gather_feature(w_h_, inds)
    w_h_ = w_h_.view(batch, K, 2)

    clses = clses.view(batch, K, 1).float()
    scores = scores.view(batch, K, 1)

    # xs,ys是中心坐標，w_h_[...,0:1]是w,1:2是h
    bboxes = torch.cat([xs - w_h_[..., 0:1] / 2,
                        ys - w_h_[..., 1:2] / 2,
                        xs + w_h_[..., 0:1] / 2,
                        ys + w_h_[..., 1:2] / 2], dim=2)
    detections = torch.cat([bboxes, scores, clses], dim=2)
    return detections

第一步

將hmap歸一化，使用了sigmoid函數

hmap = torch.sigmoid(hmap) # 歸一化到0-1

第二步

進入_nms函數：

def _nms(heat, kernel=3):
    hmax = F.max_pool2d(heat, kernel, stride=1, padding=(kernel - 1) // 2)
    keep = (hmax == heat).float() # 找到極大值點
    return heat * keep

hmax代表特征圖經過3x3卷積以后的結果，keep為極大點的位置，返回的結果是篩選后的極大值點，其余不符合8-近鄰極大值點的都歸為0。

這時候通過heatmap得到了滿足8近鄰極大值點的所有值。

這里的nms曾經在群里討論過，有群友認為僅通過3x3的並不合理，可以嘗試使用3x3,5x5,7x7這樣的maxpooling，相當於也進行了多尺度測試，據說能提高一點點mAP。

第三步

進入_topk函數，這里K是一個超參數，CenterNet中設置K=100

def _topk(scores, K=40):
    # score shape : [batch, class , h, w]
    batch, cat, height, width = scores.size()

    # to shape: [batch , class, h * w] 分類別，每個class channel統計最大值
    # topk_scores和topk_inds分別是前K個score和對應的id
    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()

    # to shape: [batch , class * h * w] 這樣的結果是不分類別的，全體class中最大的100個
    topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), K)

    # 所有類別中找到最大值
    topk_clses = (topk_ind / K).int()

    topk_inds = _gather_feature(topk_inds.view(
        batch, -1, 1), topk_ind).view(batch, K)

    topk_ys = _gather_feature(topk_ys.view(
        batch, -1, 1), topk_ind).view(batch, K)

    topk_xs = _gather_feature(topk_xs.view(
        batch, -1, 1), topk_ind).view(batch, K)

    return topk_score, topk_inds, topk_clses, topk_ys, topk_xs

torch.topk的一個demo如下：

>>> x
array([[0.11530714, 0.014376  , 0.23392263, 0.48629663],
       [0.59611302, 0.83697236, 0.27330404, 0.17728915],
       [0.36443852, 0.46562404, 0.73033529, 0.44751189]])
>>> torch.topk(torch.from_numpy(x), 3)
torch.return_types.topk(
    values=tensor([[0.4863, 0.2339, 0.1153],
                   [0.8370, 0.5961, 0.2733],
                   [0.7303, 0.4656, 0.4475]], dtype=torch.float64),
    indices=tensor([[3, 2, 0],
                    [1, 0, 2],
                    [2, 1, 3]]))

topk_scores和topk_inds分別是前K個score和對應的id。

• topk_scores 形狀【batch, class, K】K代表得分最高的前100個點， 其保存的內容是每個類別前100個最大的score。

• topk_inds 形狀 【batch, class, K】class代表80個類別channel，其保存的是每個類別對應100個score的下角標。

• topk_score 形狀 【batch, K】,通過gather feature 方法獲取，其保存的是全部類別前100個最大的score。

• topk_ind 形狀 【batch , K】，代表通過topk調用結果的下角標, 其保存的是全部類別對應的100個score的下角標。

• topk_inds、topk_ys、topk_xs三個變量都經過gather feature函數，其主要功能是從對應張量中根據下角標提取結果，具體函數如下：

def _gather_feature(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) # 按照dim=1獲取ind
  if mask is not None:
    mask = mask.unsqueeze(2).expand_as(feat)
    feat = feat[mask]
    feat = feat.view(-1, dim)
  return feat

以topk_inds為例(K=100,class=80)

feat (topk_inds) 形狀為：【batch, 80x100, 1】

ind (topk_ind) 形狀為：【batch，100】

ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)擴展一個位置，ind形狀變為：【batch, 100, 1】

feat = feat.gather(1, ind)按照dim=1獲取ind,為了方便理解和回憶，這里舉一個例子：

>>> import torch
>>> a = torch.randn(1, 10)
>>> b = torch.tensor([[3,4,5]])
>>> a.gather(1, b)
tensor([[ 0.7257, -0.4977,  1.2522]])
>>> a
tensor([[ 1.0684, -0.9655,  0.7381,  0.7257, -0.4977,  1.2522,  1.5084,  0.2669,
         -0.5471,  0.5998]])

相當於是feat根據ind的角標的值獲取到了對應feat位置上的結果。最終feat形狀為【batch，100,1】

第四步

經過topk函數，得到了四個返回值，topk_score、topk_inds、topk_ys、topk_xs四個參數的形狀都是【batch, 100】,其中topk_inds是每張圖片的前100個最大的值對應的index。

regs = _tranpose_and_gather_feature(regs, inds)

w_h_ = _tranpose_and_gather_feature(w_h_, inds)

transpose_and_gather_feat函數功能是將topk得到的index取值，得到對應前100的regs和wh的值。

def _tranpose_and_gather_feature(feat, ind):
  # ind代表的是ground truth中設置的存在目標點的下角標
  feat = feat.permute(0, 2, 3, 1).contiguous()# from [bs c h w] to [bs, h, w, c] 
  feat = feat.view(feat.size(0), -1, feat.size(3)) # to [bs, wxh, c]
  feat = _gather_feature(feat, ind) # 從中取得ind對應值
  return feat

到這一步為止，可以將top100的score、wh、regs等值提取，並且得到對應的bbox，最終ctdet_decode返回了detections變量。

3. 數據集

之前在CenterNet系列第一篇PyTorch版CenterNet訓練自己的數據集中講解了如何配置數據集，為了更方便學習和調試這部分代碼，筆者從github上找到了一個浣熊數據集，這個數據集僅有200張圖片，方便大家快速訓練和debug。

鏈接：https://pan.baidu.com/s/1unK-QZKDDaGwCrHrOFCXEA
提取碼：pdcv

以上數據集已經制作好了，只要按照第一篇文章中將DCN、NMS等編譯好，就可以直接使用。

5. 參考

https://blog.csdn.net/fsalicealex/article/details/91955759

https://zhuanlan.zhihu.com/p/66048276

https://zhuanlan.zhihu.com/p/85194783