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, )