密集型目标检测网络_目标检测网络CenterNet详解(四)

"""Transforms described in https://arxiv.org/abs/1904.07850."""# pylint: disable=too-many-function-argsfrom__future__ import absolute_import

import numpyasnp

import mxnetasmxfrom gluoncv.data.transforms import bbox astbboxfrom gluoncv.data.transforms import image astimagefromgluoncv.data.transforms import experimentalfromgluoncv.utils.filesystem import try_import_cv2frommxnet import ndfrommxnet import gluonclassCenterNetTargetGenerator(gluon.Block):"""Target generator for CenterNet.

Parameters----------num_class :intNumber of categories.

output_width :intWidth of the network output.

output_height :intHeight of the network output."""def __init__(self, num_class, output_width, output_height):

super(CenterNetTargetGenerator, self).__init__()

self._num_class=num_class

self._output_width= int(output_width)

self._output_height= int(output_height)

def forward(self, gt_boxes, gt_ids):"""Target generation"""# pylint: disable=arguments-differ

h_scale= 1.0 # already scaled inaffinetransform

w_scale= 1.0 # already scaled inaffinetransform

heatmap=np.zeros((self._num_class, self._output_height, self._output_width),

dtype=np.float32)

wh_target= np.zeros((2, self._output_height, self._output_width), dtype=np.float32)

wh_mask= np.zeros((2, self._output_height, self._output_width), dtype=np.float32)

center_reg= np.zeros((2, self._output_height, self._output_width), dtype=np.float32)

center_reg_mask= np.zeros((2, self._output_height, self._output_width), dtype=np.float32)for bbox, cid inzip(gt_boxes, gt_ids):

cid= int(cid)

box_h, box_w= bbox[3] - bbox[1], bbox[2] - bbox[0]if box_h > 0 and box_w > 0:

radius=_gaussian_radius((np.ceil(box_h), np.ceil(box_w)))

radius= max(0, int(radius))

center=np.array(

[(bbox[0] + bbox[2]) / 2 * w_scale, (bbox[1] + bbox[3]) / 2 *h_scale],

dtype=np.float32)

center_int=center.astype(np.int32) #浮点数变为整数

center_x, center_y=center_int

assert center_x output_width: {}'.format(center_x, self._output_width)

assert center_y output_height: {}'.format(center_y, self._output_height)

_draw_umich_gaussian(heatmap[cid], center_int, radius)

print(radius)

wh_target[0, center_y, center_x] = box_w *w_scale

wh_target[1, center_y, center_x] = box_h *h_scale

wh_mask[:, center_y, center_x]= 1.0center_reg[:, center_y, center_x]= center -center