centernet笔记 - inference阶段后处理

预备知识，mxnet.ndarray的一些操作
shape_array

shape_array([[1,2,3,4], [5,6,7,8]]) = [2,4]  # 获取ndarry的shape
# Returns a 1D int64 array containing the shape of data.

split

Splits an array along a particular axis into multiple sub-arrays.
cc = [  1  80 128 128]
N, C, H, W = cc.split(num_outputs=4, axis=0)
# N = [1], C = [80], H = [128], W = [128]

topk

topk(data=None, axis=_Null, k=_Null, ret_typ=_Null, is_ascend=_Null, dtype=_Null, out=None, name=None, **kwargs)
# data : The input array
# axis : 从哪个维度选取topk，默认-1，即最后一个维度
# ret_typ : return type，可以选择返回值，或者index，或者both
# is_ascend : 是否逆序，默认为0，返回 top k largest

cast

Casts all elements of the input to a new type.
# 将array中元素转换为指定的类型
cast([0.9, 1.3], dtype='int32') = [0, 1]

slice_like

# Slices a region of the array like the shape of another array. 
x = [[  1.,   2.,   3.,   4.],
     [  5.,   6.,   7.,   8.],
     [  9.,  10.,  11.,  12.]]
y = [[  0.,   0.,   0.],
     [  0.,   0.,   0.]]
slice_like(x, y) = [[ 1.,  2.,  3.]
                    [ 5.,  6.,  7.]]
slice_like(x, y, axes=(0, 1)) = [[ 1.,  2.,  3.]
                                 [ 5.,  6.,  7.]]
slice_like(x, y, axes=(0)) = [[ 1.,  2.,  3.,  4.]
                              [ 5.,  6.,  7.,  8.]]
slice_like(x, y, axes=(-1)) = [[  1.,   2.,   3.]
                               [  5.,   6.,   7.]
                               [  9.,  10.,  11.]]

expand_dims

# Inserts a new axis of size 1 into the array shape For example, given x with shape (2,3,4),
# then expand_dims(x, axis=1) will return a new array with shape (2,1,3,4).
mxnet.ndarray.expand_dims(data=None, axis=_Null, out=None, name=None, **kwargs)
# axis=-1, 则在末尾增加一个维度

tile

# 将array复制n次
x = [[1, 2],
     [3, 4]]
tile(x, reps=(2,3)) = [[ 1.,  2.,  1.,  2.,  1.,  2.],
                       [ 3.,  4.,  3.,  4.,  3.,  4.],
                       [ 1.,  2.,  1.,  2.,  1.,  2.],
                       [ 3.,  4.,  3.,  4.,  3.,  4.]]

gather_nd
参考：https://discuss.gluon.ai/t/topic/2413/3

# 从array中，按照输入的索引取出其中的元素
data = [[0, 1], [2, 3]]
indices = [[1, 1, 0], [0, 1, 0]]
gather_nd(data, indices) = [2, 3, 0] # 取出(1,0),(1,1),(0,0)的元素

centernet训练与预测过程中，图片的几何变化
训练阶段：原图 -> (512,512) -> (128,128)
预测阶段：(128,128) -> (512,512) -> 原图

在inference阶段，模型的输出是在128x128的feature shape上，在feature map上得到的结果，需要乘以scale=4.0，才能映射到512x512上，接着，还必须经过512x512到原图的一个仿射变换，才能得到模型最终输出值

核心操作及解码
核心操作

# self.heatmap_nms = nn.MaxPool2D(pool_size=3, strides=1, padding=1)
heatmap = outs[0]
keep = F.broadcast_equal(self.heatmap_nms(heatmap), heatmap)
results = self.decoder(keep * heatmap, outs[1], outs[2])

outs为模型输出，在heatmap所在维度上做3x3的max pooling，把峰值点找出来
解码

    def hybrid_forward(self, F, x, wh, reg):
        """Forward of decoder"""
        import pdb
        pdb.set_trace()
        # 这里假设batch_size = 4, resize w和h为512x512
        _, _, out_h, out_w = x.shape_array().split(num_outputs=4, axis=0) # 获取feature map的H，W
        scores, indices = x.reshape((0, -1)).topk(k=self._topk, ret_typ='both') # 获取top_100的scores和indices，x nx80x128x128 -> nx(80x128x128), 某类的128x128元素连在一起
        indices = F.cast(indices, 'int64')
        topk_classes = F.cast(
            F.broadcast_div(
                indices,
                (out_h * out_w)),
            'float32')  #  根据indices，获取top_100对应的类别，0-79
        topk_indices = F.broadcast_mod(indices, (out_h * out_w))  # 获取在128x128平面拉平后的索引
        topk_ys = F.broadcast_div(topk_indices, out_w) # 1维索引恢复出2维平面上的y
        topk_xs = F.broadcast_mod(topk_indices, out_w) # 1维索引恢复出2维平面上的x
        center = reg.transpose((0, 2, 3, 1)).reshape((0, -1, 2)) # shape: (4, 16384, 2)
        wh = wh.transpose((0, 2, 3, 1)).reshape((0, -1, 2)) # shape: (4, 16384, 2)
        batch_indices = F.cast(F.arange(256).slice_like(center, axes=(
            0)).expand_dims(-1).tile(reps=(1, self._topk)), 'int64') # shape: (4, 100), 依次为全0，全1，全2，全3
        reg_xs_indices = F.zeros_like(batch_indices, dtype='int64') # shape: (4, 100)，全0
        reg_ys_indices = F.ones_like(batch_indices, dtype='int64') # shape: (4, 100)，全1
        reg_xs = F.concat(
            batch_indices, topk_indices, reg_xs_indices, dim=0).reshape(
            (3, -1)) # shape： (3, 400)
        ‘’‘reg_xs = 
        [[   0    0    0 ...    3    3    3]
 		 [9664 8207 9425 ... 9639 5593 9044]
 		 [   0    0    0 ...    0    0    0]]
        ’‘’
        reg_ys = F.concat(
            batch_indices, topk_indices, reg_ys_indices, dim=0).reshape(
            (3, -1))
        ‘’‘reg_ys = 
        [[   0    0    0 ...    3    3    3]
 		 [9664 8207 9425 ... 9639 5593 9044]
 		 [   1    1    1 ...    1    1    1]]
        ’‘’
        xs = F.cast(
            F.gather_nd(
                center, reg_xs).reshape(
                (-1, self._topk)), 'float32') # shape : (4, 100)
        ys = F.cast(
            F.gather_nd(
                center, reg_ys).reshape(
                (-1, self._topk)), 'float32') # shape : (4, 100)
        topk_xs = F.cast(topk_xs, 'float32') + xs
        topk_ys = F.cast(topk_ys, 'float32') + ys # feature map上坐标+偏移 = 真实坐标
        w = F.cast(F.gather_nd(wh, reg_xs).reshape((-1, self._topk)), 'float32')  # noqa
        h = F.cast(F.gather_nd(wh, reg_ys).reshape((-1, self._topk)), 'float32')  # noqa
        half_w = w / 2
        half_h = h / 2
        results = [
            topk_xs - half_w,
            topk_ys - half_h,
            topk_xs + half_w,
            topk_ys + half_h]
        results = F.concat(*[tmp.expand_dims(-1) for tmp in results], dim=-1) # shape: (4, 100, 4)
        return topk_classes, scores, results * self._scale # (4, 100),(4, 100),(4, 100, 4)

由512x512映射到原图

affine_mat = get_post_transform(orig_width, orig_height, 512, 512) # 获取仿射变换矩阵
bbox[0:2] = affine_transform(bbox[0:2], affine_mat) # 对x1,y1做仿射变换
bbox[2:4] = affine_transform(bbox[2:4], affine_mat) # 对x2,y2做仿射变换