def forward_refinement_head(self, datum, rpn_rois, rpn_rois_rotation, fpn_outputs, is_train):
if is_train:
ref_rois, ref_rots, target_class_ids, target_bbox, target_rots, roi_gt_box_assignment \
= self._get_detection_target(rpn_rois, rpn_rois_rotation, datum['bboxes_cls'],
datum['bboxes_normalized'], datum['bboxes_rotations'])
unnorm_ref_rois = [
detection_utils.unnormalize_boxes(roi.cpu().numpy(), self.config.max_ptc_size)
for roi in ref_rois
]
unnorm_ref_posrois = [r[:len(b)] for r, b in zip(unnorm_ref_rois, roi_gt_box_assignment)]
mask_coords, mask_insts, _ = self._get_mask_target(unnorm_ref_posrois, datum['coords'],
roi_gt_box_assignment,
datum['instance_target'], is_train=True)
ref_class_logits, ref_class_probs, ref_bbox, ref_rotdelta = self.fpn_classifier_network(
unnorm_ref_rois, ref_rots, fpn_outputs)
output = {
'ref_rois': ref_rois,
'ref_rots': ref_rots,
'unnorm_ref_rois': unnorm_ref_rois,
'unnorm_ref_posrois': unnorm_ref_posrois,
'ref_target_class_ids': target_class_ids,
'ref_target_bbox': target_bbox,
'ref_target_rots': target_rots,
'ref_mask_coords': mask_coords,
'ref_mask_insts': mask_insts,
'ref_class_logits': ref_class_logits,
'ref_class_probs': ref_class_probs,
'ref_bbox': ref_bbox,
'ref_rotdelta': ref_rotdelta,
}
else:
rpn_rois = [roi[~torch.all(roi == 0, 1)] for roi in rpn_rois]
unnorm_ref_rois = [
detection_utils.unnormalize_boxes(roi.cpu().numpy(), self.config.max_ptc_size)
for roi in rpn_rois
]
ref_rots = rpn_rois_rotation
if ref_rots is not None and self.config.normalize_rotation2:
ref_rots = [rot / 2 + np.pi / 2 for rot in ref_rots]
_, ref_class_probs, ref_bbox, ref_rotdelta = self.fpn_classifier_network(
unnorm_ref_rois, ref_rots, fpn_outputs)
detection = self.detection_refinement(
ref_class_probs, rpn_rois, ref_bbox, rpn_rois_rotation, ref_rotdelta)
mask_coords, _, mask_coordmasks = self._get_mask_target(detection, datum['coords'],
is_train=False)
output = {
'detection': detection,
'mask_coords': mask_coords,
'mask_coordmasks': mask_coordmasks,
}
return output
def visualize_predictions(self, datum, output, iteration):
coords = datum['coords'].numpy()
batch_size = coords[:, 0].max() + 1
output_path = pathlib.Path(self.config.visualize_path)
output_path.mkdir(exist_ok=True)
for i in range(batch_size):
# Visualize RGB input.
coords_mask = coords[:, 0] == i
coords_b = coords[coords_mask, 1:]
if datum['input'].shape[1] > 3:
rgb_b = ((datum['input'][:, :3] + 0.5) * 255).numpy()
else:
rgb_b = ((datum['input'].repeat(1, 3) - datum['input'].min())
/ (datum['input'].max() - datum['input'].min()) * 255).numpy()
rgb_ply_dest = output_path / ('visualize_%04d_rgb.ply' % iteration)
pc_utils.save_point_cloud(np.hstack((coords_b, rgb_b)), rgb_ply_dest)
# Visualize positive anchor centers.
positive_anchors_mask = F.softmax(output['rpn_class_logits'].cpu(), dim=2)[0, :, 1] > 0.5
positive_anchors = datum['anchors'][torch.where(positive_anchors_mask)]
if len(positive_anchors.shape) == 1:
positive_anchors = np.array([positive_anchors])
if positive_anchors.shape[0] > 0:
positive_anchors = np.unique((positive_anchors[:, 3:] + positive_anchors[:, :3]) / 2,
axis=0)
anchors_all_ply_dest = output_path / ('visualize_%04d_anchors_pred.ply' % iteration)
pc_utils.save_point_cloud(positive_anchors, anchors_all_ply_dest)
# Visualize region proposals.
rpn_rois = output['rpn_rois'][i].cpu().numpy()
rpn_mask = ~np.all(rpn_rois == 0, 1)
rpn_rois = rpn_rois[rpn_mask][:100]
rpn_rois = detection_utils.unnormalize_boxes(rpn_rois, self.config.max_ptc_size)
if output.get('rpn_rois_rotation') is None:
rpn_rois_ptc = pc_utils.visualize_bboxes(rpn_rois)
else:
rpn_rois_rotation = output['rpn_rois_rotation'][i].cpu().numpy()[rpn_mask][:100]
rpn_rois = np.hstack((rpn_rois, rpn_rois_rotation[:, None]))
rpn_rois_ptc = pc_utils.visualize_bboxes(rpn_rois, bbox_param='xyzxyzr')
rpn_rois_ply_dest = output_path / ('visualize_%04d_rpn_rois_pred.ply' % iteration)
pc_utils.save_point_cloud(rpn_rois_ptc, rpn_rois_ply_dest)
# Visualize final detection.
detection_pred = output['detection'][i]
if detection_pred.shape[1] == 8:
detection_ptc = pc_utils.visualize_bboxes(detection_pred[:, :6], detection_pred[:, 6])
elif detection_pred.shape[1] == 9:
detection_ptc = pc_utils.visualize_bboxes(detection_pred[:, :7], detection_pred[:, 7],
bbox_param='xyzxyzr')
else:
raise ValueError('Unknown bounding box output.')
detection_dest = output_path / ('visualize_%04d_detection_pred.ply' % iteration)
pc_utils.save_point_cloud(detection_ptc, detection_dest)
def visualize_groundtruth(self, datum, iteration):
coords = datum['coords'].numpy()
batch_size = coords[:, 0].max() + 1
output_path = pathlib.Path(self.config.visualize_path)
output_path.mkdir(exist_ok=True)
for i in range(batch_size):
# Visualize ground-truth positive anchors.
anchors_gt = datum['anchors'][torch.where(datum['rpn_match'].cpu() == 1)[1]]
anchors_gt_ptc = pc_utils.visualize_bboxes(anchors_gt)
anchors_gt_ply_dest = output_path / ('visualize_%04d_anchors_gt.ply' % iteration)
pc_utils.save_point_cloud(anchors_gt_ptc, anchors_gt_ply_dest)
# Visualize center location of all ground-truth anchors.
anchors_all = np.unique((datum['anchors'][:, 3:] + datum['anchors'][:, :3]) / 2, axis=0)
anchors_all_ply_dest = output_path / ('visualize_%04d_all_anchors_centers.ply' % iteration)
pc_utils.save_point_cloud(anchors_all, anchors_all_ply_dest)
# Visualize ground-truth positive anchors.
if datum.get('bboxes_rotations') is None:
bboxes_gt = pc_utils.visualize_bboxes(datum['bboxes_coords'][i], datum['bboxes_cls'][i])
else:
bboxes_gt = np.hstack((datum['bboxes_coords'][i], datum['bboxes_rotations'][i][:, None]))
bboxes_gt = pc_utils.visualize_bboxes(bboxes_gt, datum['bboxes_cls'][i],
bbox_param='xyzxyzr')
bboxes_gt_ply_dest = output_path / ('visualize_%04d_bboxes_gt.ply' % iteration)
pc_utils.save_point_cloud(bboxes_gt, bboxes_gt_ply_dest)
# Visualize reconstructed ground-truth rpn targets.
rpn_bbox_anchors = datum['anchors'][(datum['rpn_match'].flatten() == 1).cpu().numpy()]
rpn_bbox_anchors = detection_utils.normalize_boxes(rpn_bbox_anchors, self.config.max_ptc_size)
rpn_bbox_target = datum['rpn_bbox'].reshape(-1, 6)
rpn_bbox_mask = ~torch.all(rpn_bbox_target == 0, 1)
rpn_bbox_target = rpn_bbox_target[rpn_bbox_mask].cpu().numpy()
rpn_bbox_target *= np.reshape(self.config.rpn_bbox_std, (1, len(self.config.rpn_bbox_std)))
rpn_bbox_target = detection_utils.apply_box_deltas(torch.from_numpy(rpn_bbox_anchors),
torch.from_numpy(rpn_bbox_target),
self.config.normalize_bbox)
rpn_bbox_target = detection_utils.unnormalize_boxes(rpn_bbox_target.numpy(),
self.config.max_ptc_size)
if datum.get('rpn_rotation') is None:
bboxes_gt_recon = pc_utils.visualize_bboxes(rpn_bbox_target)
else:
rpn_rot_target = datum['rpn_rotation'][i][rpn_bbox_mask].cpu().numpy()
bboxes_gt_recon = np.hstack((rpn_bbox_target, rpn_rot_target[:, None]))
bboxes_gt_recon = pc_utils.visualize_bboxes(bboxes_gt_recon, bbox_param='xyzxyzr')
bboxes_gt_recon_ply_dest = output_path / ('visualize_%04d_bboxes_gt_recon.ply' % iteration)
pc_utils.save_point_cloud(bboxes_gt_recon, bboxes_gt_recon_ply_dest)
def _get_rpnonly_detection(self, rpn_rois, rpn_rois_score, rpn_rois_rotation):
detections = []
if rpn_rois is not None:
for i, (rpn_roi, rpn_score) in enumerate(zip(rpn_rois, rpn_rois_score)):
rpn_roi_unnorm = detection_utils.unnormalize_boxes(
rpn_roi[:rpn_score.shape[0], :6].cpu().numpy(), self.config.max_ptc_size)
if self.EVAL_PERCLASS_MAP:
rpn_classes = rpn_roi[:, 6, None].cpu().numpy()
else:
rpn_classes = np.zeros((rpn_roi.shape[0], 1))
detection = np.hstack((rpn_roi_unnorm, rpn_classes, rpn_score.cpu().numpy()[:, None]))
if rpn_rois_rotation is not None:
detection_rot = rpn_rois_rotation[i].cpu().numpy()[:, None]
if self.config.normalize_rotation2:
detection_rot = detection_rot / 2 + np.pi / 2
detection = np.hstack((detection[:, :6], detection_rot, detection[:, 6:]))
detection_mask = detection[:, -1] > self.config.post_nms_min_confidence
detection = detection[detection_mask]
detections.append(detection)
else:
detections = [np.zeros((0, 9 if self.dataset.IS_ROTATION_BBOX else 8))]
return detections
def detection_refinement(self, b_probs, b_rois, b_deltas, b_rots, b_rotdeltas):
if b_probs is None:
num_channel = 9 if b_rots is None else 8
return [np.zeros((0, num_channel))]
num_batch = [rois.shape[0] for rois in b_rois]
num_samples = sum(num_batch)
assert num_samples == b_probs.shape[0] == b_deltas.shape[0]
if b_rots is not None:
assert num_samples == sum(rots.shape[0] for rots in b_rots) == b_rotdeltas.shape[0]
batch_split = [(sum(num_batch[:i]), sum(num_batch[:(i + 1)])) for i in range(len(num_batch))]
b_probs = [b_probs[i:j] for (i, j) in batch_split]
b_deltas = [b_deltas[i:j] for (i, j) in batch_split]
if b_rots is not None:
b_rotdeltas = [b_rotdeltas[i:j] for (i, j) in batch_split]
b_nms = []
b_nms_rot = None if b_rots is None else []
for i, (probs, rois, deltas) in enumerate(zip(b_probs, b_rois, b_deltas)):
rois = rois.reshape(-1, rois.shape[-1])
class_ids = torch.argmax(probs, dim=1)
batch_slice = range(probs.shape[0])
class_scores = probs[batch_slice, class_ids]
class_deltas = deltas[batch_slice, class_ids - 1]
class_deltas *= torch.tensor(self.config.rpn_bbox_std).to(deltas)
refined_rois = detection_utils.apply_box_deltas(rois, class_deltas,
self.config.normalize_bbox)
if b_rots is not None:
class_rot_deltas = b_rotdeltas[i][batch_slice, class_ids - 1]
class_rot_deltas = self.ref_rotation_criterion.pred(class_rot_deltas)
refined_rots = detection_utils.normalize_rotation(b_rots[i] + class_rot_deltas)
keep = torch.where(class_ids > 0)[0].cpu().numpy()
if self.config.detection_min_confidence:
conf_keep = torch.where(class_scores > self.config.detection_min_confidence)[0]
keep = np.array(list(set(conf_keep.cpu().numpy()).intersection(keep)))
if keep.size == 0:
b_nms.append(np.zeros((0, 8)))
if b_rots is not None:
b_nms_rot.append(np.zeros(0))
else:
pre_nms_class_ids = class_ids[keep] - 1
pre_nms_scores = class_scores[keep]
pre_nms_rois = refined_rois[keep]
if b_rots is not None:
pre_nms_rots = refined_rots[keep]
nms_scores = []
nms_rois = []
nms_classes = []
nms_rots = []
for class_id in torch.unique(pre_nms_class_ids):
class_nms_mask = pre_nms_class_ids == class_id
class_nms_scores = pre_nms_scores[class_nms_mask]
class_nms_rois = pre_nms_rois[class_nms_mask]
pre_nms_class_rots = None
if b_rots is not None:
pre_nms_class_rots = pre_nms_rots[class_nms_mask]
nms_roi, nms_rot, nms_score = detection_utils.non_maximum_suppression(
class_nms_rois, pre_nms_class_rots, class_nms_scores,
self.config.detection_nms_threshold, self.config.detection_max_instances,
self.config.detection_rot_nms, self.config.detection_aggregate_overlap)
nms_rois.append(nms_roi)
nms_scores.append(nms_score)
nms_classes.append(torch.ones(len(nms_score)).to(class_nms_rois) * class_id)
if b_rots is not None:
if self.config.normalize_rotation2:
nms_rot = nms_rot / 2 + np.pi / 2
nms_rots.append(nms_rot)
nms_scores = torch.cat(nms_scores)
nms_rois = torch.cat(nms_rois)
nms_classes = torch.cat(nms_classes)
detection_max_instances = min(self.config.detection_max_instances, nms_scores.shape[0])
ix = torch.topk(nms_scores, detection_max_instances)[1]
nms_rois_unnorm = detection_utils.unnormalize_boxes(
nms_rois[ix].cpu().numpy(), self.config.max_ptc_size)
nms_bboxes = np.hstack((nms_rois_unnorm, nms_classes[ix, None].cpu().numpy(),
nms_scores[ix, None].cpu().numpy()))
if b_rots is not None:
nms_rots = torch.cat(nms_rots)[ix, None].cpu().numpy()
nms_bboxes = np.hstack((nms_bboxes[:, :6], nms_rots, nms_bboxes[:, 6:]))
b_nms.append(nms_bboxes)
return b_nms