def test_load_segmentation_mask():
# Test loading semantic segmentation mask on ScanNet dataset
scannet_info = mmcv.load('./tests/data/scannet/scannet_infos.pkl')[0]
scannet_load_annotations3D = LoadAnnotations3D(
with_bbox_3d=False,
with_label_3d=False,
with_mask_3d=False,
with_seg_3d=True)
scannet_results = dict()
data_path = './tests/data/scannet'
# prepare input of loading pipeline
scannet_results['ann_info'] = dict()
scannet_results['ann_info']['pts_semantic_mask_path'] = osp.join(
data_path, scannet_info['pts_semantic_mask_path'])
scannet_results['pts_seg_fields'] = []
scannet_results = scannet_load_annotations3D(scannet_results)
scannet_pts_semantic_mask = scannet_results['pts_semantic_mask']
assert scannet_pts_semantic_mask.shape == (100, )
# Convert class_id to label and assign ignore_index
scannet_seg_class_mapping = \
PointSegClassMapping((1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16,
24, 28, 33, 34, 36, 39))
scannet_results = scannet_seg_class_mapping(scannet_results)
scannet_pts_semantic_mask = scannet_results['pts_semantic_mask']
assert np.all(scannet_pts_semantic_mask == np.array([
13, 20, 1, 2, 6, 2, 13, 1, 13, 2, 0, 20, 5, 20, 2, 0, 1, 13, 0, 0, 0,
20, 6, 20, 13, 20, 2, 20, 20, 2, 16, 5, 13, 5, 13, 0, 20, 0, 0, 1, 7,
20, 20, 20, 20, 20, 20, 20, 0, 1, 2, 13, 16, 1, 1, 1, 6, 2, 12, 20, 3,
20, 20, 14, 1, 20, 2, 1, 7, 2, 0, 5, 20, 5, 20, 20, 3, 6, 5, 20, 0, 13,
12, 2, 20, 0, 0, 13, 20, 1, 20, 5, 3, 0, 13, 1, 2, 2, 2, 1
]))
def test_load_annotations3D():
# Test scannet LoadAnnotations3D
scannet_info = mmcv.load('./tests/data/scannet/scannet_infos.pkl')[0]
scannet_load_annotations3D = LoadAnnotations3D(
with_bbox_3d=True,
with_label_3d=True,
with_mask_3d=True,
with_seg_3d=True)
scannet_results = dict()
data_path = './tests/data/scannet'
if scannet_info['annos']['gt_num'] != 0:
scannet_gt_bboxes_3d = scannet_info['annos']['gt_boxes_upright_depth']
scannet_gt_labels_3d = scannet_info['annos']['class']
else:
scannet_gt_bboxes_3d = np.zeros((1, 6), dtype=np.float32)
scannet_gt_labels_3d = np.zeros((1, ))
# prepare input of loading pipeline
scannet_results['ann_info'] = dict()
scannet_results['ann_info']['pts_instance_mask_path'] = osp.join(
data_path, scannet_info['pts_instance_mask_path'])
scannet_results['ann_info']['pts_semantic_mask_path'] = osp.join(
data_path, scannet_info['pts_semantic_mask_path'])
scannet_results['ann_info']['gt_bboxes_3d'] = DepthInstance3DBoxes(
scannet_gt_bboxes_3d, box_dim=6, with_yaw=False)
scannet_results['ann_info']['gt_labels_3d'] = scannet_gt_labels_3d
scannet_results['bbox3d_fields'] = []
scannet_results['pts_mask_fields'] = []
scannet_results['pts_seg_fields'] = []
scannet_results = scannet_load_annotations3D(scannet_results)
scannet_gt_boxes = scannet_results['gt_bboxes_3d']
scannet_gt_labels = scannet_results['gt_labels_3d']
scannet_pts_instance_mask = scannet_results['pts_instance_mask']
scannet_pts_semantic_mask = scannet_results['pts_semantic_mask']
repr_str = repr(scannet_load_annotations3D)
expected_repr_str = 'LoadAnnotations3D(\n with_bbox_3d=True, ' \
'with_label_3d=True, with_attr_label=False, ' \
'with_mask_3d=True, with_seg_3d=True, ' \
'with_bbox=False, with_label=False, ' \
'with_mask=False, with_seg=False, ' \
'with_bbox_depth=False, poly2mask=True)'
assert repr_str == expected_repr_str
assert scannet_gt_boxes.tensor.shape == (27, 7)
assert scannet_gt_labels.shape == (27, )
assert scannet_pts_instance_mask.shape == (100, )
assert scannet_pts_semantic_mask.shape == (100, )