def _get_bbox_output(self, bbox, trans_output, height, width):
bbox = self._coco_box_to_bbox(bbox).copy()
rect = np.array([[bbox[0], bbox[1]], [bbox[0], bbox[3]],
[bbox[2], bbox[3]], [bbox[2], bbox[1]]],
dtype=np.float32)
for t in range(4):
rect[t] = dataset_utils.affine_transform(rect[t], trans_output)
bbox[:2] = rect[:, 0].min(), rect[:, 1].min()
bbox[2:] = rect[:, 0].max(), rect[:, 1].max()
bbox_amodal = copy.deepcopy(bbox)
bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, self.opt.output_res.width - 1)
bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, self.opt.output_res.height - 1)
h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]
return bbox, bbox_amodal
def _add_hps(self, ret, k, ann, gt_det, trans_output, ct_int, bbox, h, w):
num_joints = self.num_joints
pts = np.array(ann['keypoints'], np.float32).reshape(num_joints, 3) \
if 'keypoints' in ann else np.zeros((self.num_joints, 3), np.float32)
hp_radius = dataset_utils.gaussian_radius((math.ceil(h), math.ceil(w)))
hp_radius = max(0, int(hp_radius))
for j in range(num_joints):
pts[j, :2] = dataset_utils.affine_transform(
pts[j, :2], trans_output)
if pts[j, 2] > 0:
if pts[j, 0] >= 0 and pts[j, 0] < self.opt.output_res.width and pts[j, 1] >= 0 and \
pts[j, 1] < self.opt.output_res.height:
ret['hps'][k, j * 2:j * 2 + 2] = pts[j, :2] - ct_int
ret['hps_mask'][k, j * 2:j * 2 + 2] = 1
pt_int = pts[j, :2].astype(np.int32)
ret['hp_offset'][k * num_joints + j] = pts[j, :2] - pt_int
ret['hp_ind'][
k * num_joints +
j] = pt_int[1] * self.opt.output_res.width + pt_int[0]
ret['hp_offset_mask'][k * num_joints + j] = 1
ret['hm_hp_mask'][k * num_joints + j] = 1
ret['joint'][k * num_joints + j] = j
dataset_utils.draw_umich_gaussian(ret['hm_hp'][j], pt_int,
hp_radius)
if pts[j, 2] == 1:
ret['hm_hp'][j, pt_int[1], pt_int[0]] = self.ignore_val
ret['hp_offset_mask'][k * num_joints + j] = 0
ret['hm_hp_mask'][k * num_joints + j] = 0
else:
pts[j, :2] *= 0
else:
pts[j, :2] *= 0
self._ignore_region(ret['hm_hp'][j,
int(bbox[1]):int(bbox[3]) + 1,
int(bbox[0]):int(bbox[2]) +
1])
gt_det['hps'].append(pts[:, :2].reshape(num_joints * 2))
def test_affine_transform(self):
point = np.array([10., 10.])
transformed_point = dataset_utils.affine_transform(point, np.eye(2, 3))
np.testing.assert_allclose(point, transformed_point)
def _get_pre_dets(self, anns, trans_input, trans_output):
hm_h, hm_w = self.opt.input_res.height, self.opt.input_res.width
down_ratio = self.opt.down_ratio
trans = trans_input
reutrn_hm = self.opt.pre_hm
pre_hm = np.zeros(
(1, hm_h, hm_w), dtype=np.float32) if reutrn_hm else None
pre_cts, track_ids = [], []
for ann in anns:
cls_id = ann['category_id']
if cls_id > self.num_classes or cls_id <= IGNORED_CLASS_ID or (
'iscrowd' in ann and ann['iscrowd'] > 0):
continue
bbox = self._coco_box_to_bbox(ann['bbox'])
bbox[:2] = dataset_utils.affine_transform(bbox[:2], trans)
bbox[2:] = dataset_utils.affine_transform(bbox[2:], trans)
bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, hm_w - 1)
bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, hm_h - 1)
h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]
max_rad = 1
if (h > 0 and w > 0):
radius = dataset_utils.gaussian_radius(
(math.ceil(h), math.ceil(w)))
radius = max(0, int(radius))
max_rad = max(max_rad, radius)
ct = np.array([(bbox[0] + bbox[2]) / 2,
(bbox[1] + bbox[3]) / 2],
dtype=np.float32)
ct0 = ct.copy()
conf = 1
ct[0] = ct[0] + np.random.randn() * self.opt.hm_disturb * w
ct[1] = ct[1] + np.random.randn() * self.opt.hm_disturb * h
conf = 1 if np.random.random() > self.opt.lost_disturb else 0
ct_int = ct.astype(np.int32)
if conf == 0:
pre_cts.append(ct / down_ratio)
else:
pre_cts.append(ct0 / down_ratio)
track_ids.append(ann['track_id'] if 'track_id' in ann else -1)
if reutrn_hm:
dataset_utils.draw_umich_gaussian(pre_hm[0],
ct_int,
radius,
k=conf)
if np.random.random() < self.opt.fp_disturb and reutrn_hm:
ct2 = ct0.copy()
# Hard code heatmap disturb ratio, haven't tried other
# numbers.
ct2[0] = ct2[0] + np.random.randn() * 0.05 * w
ct2[1] = ct2[1] + np.random.randn() * 0.05 * h
ct2_int = ct2.astype(np.int32)
dataset_utils.draw_umich_gaussian(pre_hm[0],
ct2_int,
radius,
k=conf)
return pre_hm, pre_cts, track_ids
def _add_instance(self,
ret,
gt_det,
k,
cls_id,
bbox,
bbox_amodal,
ann,
trans_output,
aug_s,
calib,
pre_cts=None,
track_ids=None):
h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]
if h <= 0 or w <= 0:
return
radius = dataset_utils.gaussian_radius((math.ceil(h), math.ceil(w)))
radius = max(0, int(radius))
ct = np.array([(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2],
dtype=np.float32)
ct_int = ct.astype(np.int32)
ret['cat'][k] = cls_id - 1
ret['mask'][k] = 1
if 'wh' in ret:
ret['wh'][k] = 1. * w, 1. * h
ret['wh_mask'][k] = 1
ret['ind'][k] = ct_int[1] * self.opt.output_res.width + ct_int[0]
ret['reg'][k] = ct - ct_int
ret['reg_mask'][k] = 1
dataset_utils.draw_umich_gaussian(ret['hm'][cls_id - 1], ct_int,
radius)
gt_det['bboxes'].append(
np.array(
[ct[0] - w / 2, ct[1] - h / 2, ct[0] + w / 2, ct[1] + h / 2],
dtype=np.float32))
gt_det['scores'].append(1)
gt_det['clses'].append(cls_id - 1)
gt_det['cts'].append(ct)
if 'tracking' in self.opt.heads:
if ann['track_id'] in track_ids:
pre_ct = pre_cts[track_ids.index(ann['track_id'])]
ret['tracking_mask'][k] = 1
ret['tracking'][k] = pre_ct - ct_int
gt_det['tracking'].append(ret['tracking'][k])
else:
gt_det['tracking'].append(np.zeros(2, np.float32))
if 'ltrb' in self.opt.heads:
ret['ltrb'][k] = bbox[0] - ct_int[0], bbox[1] - ct_int[1], \
bbox[2] - ct_int[0], bbox[3] - ct_int[1]
ret['ltrb_mask'][k] = 1
if 'ltrb_amodal' in self.opt.heads:
ret['ltrb_amodal'][k] = \
bbox_amodal[0] - ct_int[0], bbox_amodal[1] - ct_int[1], \
bbox_amodal[2] - ct_int[0], bbox_amodal[3] - ct_int[1]
ret['ltrb_amodal_mask'][k] = 1
gt_det['ltrb_amodal'].append(bbox_amodal)
if 'velocity' in self.opt.heads:
if ('velocity' in ann) and min(ann['velocity']) > -1000:
ret['velocity'][k] = np.array(ann['velocity'], np.float32)[:3]
ret['velocity_mask'][k] = 1
gt_det['velocity'].append(ret['velocity'][k])
if 'hps' in self.opt.heads:
self._add_hps(ret, k, ann, gt_det, trans_output, ct_int, bbox, h,
w)
if 'rot' in self.opt.heads:
self._add_rot(ret, ann, k, gt_det)
if 'dep' in self.opt.heads:
if 'depth' in ann:
ret['dep_mask'][k] = 1
ret['dep'][k] = ann['depth'] * aug_s
gt_det['dep'].append(ret['dep'][k])
else:
gt_det['dep'].append(2)
if 'dim' in self.opt.heads:
if 'dim' in ann:
ret['dim_mask'][k] = 1
ret['dim'][k] = ann['dim']
gt_det['dim'].append(ret['dim'][k])
else:
gt_det['dim'].append([1, 1, 1])
if 'amodel_offset' in self.opt.heads:
if 'amodel_center' in ann:
amodel_center = dataset_utils.affine_transform(
ann['amodel_center'], trans_output)
ret['amodel_offset_mask'][k] = 1
ret['amodel_offset'][k] = amodel_center - ct_int
gt_det['amodel_offset'].append(ret['amodel_offset'][k])
else:
gt_det['amodel_offset'].append([0, 0])