目标检测 | Anchor free之CenterNet深度解析

def HourglassNetwork(heads, num_stacks, cnv_dim=256, inres=(512, 512), weights='ctdet_coco',
                     dims=[256, 384, 384, 384, 512]):
  """Instantiates the Hourglass architecture.
  Optionally loads weights pre-trained on COCO.
  Note that the data format convention used by the model is
  the one specified in your Keras config at `~/.keras/keras.json`.
  # Arguments
      num_stacks: number of hourglass modules.
      cnv_dim: number of filters after the resolution is decreased.
      inres: network input shape, should be a multiple of 128.
      weights: one of `None` (random initialization),
            'ctdet_coco' (pre-training on COCO for 2D object detection),
            'hpdet_coco' (pre-training on COCO for human pose detection),
            or the path to the weights file to be loaded.
      dims: numbers of channels in the hourglass blocks.
  # Returns
      A Keras model instance.
  # Raises
      ValueError: in case of invalid argument for `weights`,
          or invalid input shape.
  """
  if not (weights in {'ctdet_coco', 'hpdet_coco', None} or os.path.exists(weights)):
    raise ValueError('The `weights` argument should be either '
                     '`None` (random initialization), `ctdet_coco` '
                     '(pre-trained on COCO), `hpdet_coco` (pre-trained on COCO) '
                     'or the path to the weights file to be loaded.')
  input_layer = Input(shape=(inres[0], inres[1], 3), name='HGInput')
  inter = pre(input_layer, cnv_dim)
  prev_inter = None
  outputs = []
  for i in range(num_stacks):
    prev_inter = inter
    _heads, inter = hourglass_module(heads, inter, cnv_dim, i, dims) # return the heads that include  three branchs
    outputs.extend(_heads)
    if i < num_stacks - 1:
      # the joint between the first hourglass module and the second ones
      inter_ = Conv2D(cnv_dim, 1, use_bias=False, name='inter_.%d.0' % i)(prev_inter)
      inter_ = BatchNormalization(epsilon=1e-5, name='inter_.%d.1' % i)(inter_)

      cnv_ = Conv2D(cnv_dim, 1, use_bias=False, name='cnv_.%d.0' % i)(inter)
      cnv_ = BatchNormalization(epsilon=1e-5, name='cnv_.%d.1' % i)(cnv_)

      inter = Add(name='inters.%d.inters.add' % i)([inter_, cnv_])
      inter = Activation('relu', name='inters.%d.inters.relu' % i)(inter)
      inter = residual(inter, cnv_dim, 'inters.%d' % i)

  model = Model(inputs=input_layer, outputs=outputs)
  if weights == 'ctdet_coco':
    weights_path = get_file(
      '%s_hg.hdf5' % weights,
      CTDET_COCO_WEIGHTS_PATH,
      cache_subdir='models',
      file_hash='ce01e92f75b533e3ff8e396c76d55d97ff3ec27e99b1bdac1d7b0d6dcf5d90eb')
    model.load_weights(weights_path)
  elif weights == 'hpdet_coco':
    weights_path = get_file(
      '%s_hg.hdf5' % weights,
      HPDET_COCO_WEIGHTS_PATH,
      cache_subdir='models',
      file_hash='5c562ee22dc383080629dae975f269d62de3a41da6fd0c821085fbee183d555d')
    model.load_weights(weights_path)
  elif weights is not None:
    model.load_weights(weights)

  return model

def _ctdet_decode(hm, reg, wh, k=100, output_stride=4):
    """将网络的输出转换为标准的检测框信息"""
    hm = K.sigmoid(hm)
    hm = _nms(hm)
    hm_shape = K.shape(hm)
    reg_shape = K.shape(reg)
    wh_shape = K.shape(wh)
    # cat为通道数
    batch, width, cat = hm_shape[0], hm_shape[2], hm_shape[3]

    # 对输出的特征图进行铺平
    hm_flat = K.reshape(hm, (batch, -1))
    reg_flat = K.reshape(reg, (reg_shape[0], -1, reg_shape[-1]))
    wh_flat = K.reshape(wh, (wh_shape[0], -1, wh_shape[-1]))

    def _process_sample(args):
        _hm, _reg, _wh = args
        _scores, _inds = tf.math.top_k(_hm, k=k, sorted=True)  # 寻找前k个heatmap的值
        _classes = K.cast(_inds % cat, 'float32')   #获取索引对应的类别
        _inds = K.cast(_inds / cat, 'int32') #在某一类别中的位置（最大长度为 width*width），一维的

        # 一维位置转二维坐标
        _xs = K.cast(_inds % width, 'float32')  #二维坐标中的横坐标
        _ys = K.cast(K.cast(_inds / width, 'int32'), 'float32') #二维坐标的纵坐标
        _wh = K.gather(_wh, _inds) #根据索引获得宽高数据
        _reg = K.gather(_reg, _inds) #根据坐标获得offset

        _xs = _xs + _reg[..., 0]
        _ys = _ys + _reg[..., 1]

        _x1 = _xs - _wh[..., 0] / 2
        _y1 = _ys - _wh[..., 1] / 2
        _x2 = _xs + _wh[..., 0] / 2
        _y2 = _ys + _wh[..., 1] / 2

        # rescale to image coordinates
        _x1 = output_stride * _x1
        _y1 = output_stride * _y1
        _x2 = output_stride * _x2
        _y2 = output_stride * _y2

        _detection = K.stack([_x1, _y1, _x2, _y2, _scores, _classes], -1)
        return _detection

    detections = K.map_fn(_process_sample, [hm_flat, reg_flat, wh_flat], dtype=K.floatx())
    return detections

def _nms(heat, kernel=3):
    hmax = K.pool2d(heat, (kernel, kernel), padding='same', pool_mode='max')
    keep = K.cast(K.equal(hmax, heat), K.floatx())
    return heat * keep