|
|
|
|
|
import unittest |
|
|
|
|
|
import numpy as np |
|
|
import torch |
|
|
from mmengine.testing import assert_allclose |
|
|
|
|
|
from mmdet3d.datasets import ScanNetDataset, ScanNetSegDataset |
|
|
from mmdet3d.structures import DepthInstance3DBoxes |
|
|
from mmdet3d.utils import register_all_modules |
|
|
|
|
|
|
|
|
def _generate_scannet_seg_dataset_config(): |
|
|
data_root = './tests/data/scannet/' |
|
|
ann_file = 'scannet_infos.pkl' |
|
|
classes = ('wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', |
|
|
'door', 'window', 'bookshelf', 'picture', 'counter', 'desk', |
|
|
'curtain', 'refrigerator', 'showercurtrain', 'toilet', 'sink', |
|
|
'bathtub', 'otherfurniture') |
|
|
palette = [ |
|
|
[174, 199, 232], |
|
|
[152, 223, 138], |
|
|
[31, 119, 180], |
|
|
[255, 187, 120], |
|
|
[188, 189, 34], |
|
|
[140, 86, 75], |
|
|
[255, 152, 150], |
|
|
[214, 39, 40], |
|
|
[197, 176, 213], |
|
|
[148, 103, 189], |
|
|
[196, 156, 148], |
|
|
[23, 190, 207], |
|
|
[247, 182, 210], |
|
|
[219, 219, 141], |
|
|
[255, 127, 14], |
|
|
[158, 218, 229], |
|
|
[44, 160, 44], |
|
|
[112, 128, 144], |
|
|
[227, 119, 194], |
|
|
[82, 84, 163], |
|
|
] |
|
|
scene_idxs = [0] |
|
|
modality = dict(use_lidar=True, use_camera=False) |
|
|
pipeline = [ |
|
|
dict( |
|
|
type='LoadPointsFromFile', |
|
|
coord_type='DEPTH', |
|
|
shift_height=False, |
|
|
use_color=True, |
|
|
load_dim=6, |
|
|
use_dim=[0, 1, 2, 3, 4, 5]), |
|
|
dict( |
|
|
type='LoadAnnotations3D', |
|
|
with_bbox_3d=False, |
|
|
with_label_3d=False, |
|
|
with_mask_3d=False, |
|
|
with_seg_3d=True), |
|
|
dict(type='PointSegClassMapping'), |
|
|
dict( |
|
|
type='IndoorPatchPointSample', |
|
|
num_points=5, |
|
|
block_size=1.5, |
|
|
ignore_index=len(classes), |
|
|
use_normalized_coord=True, |
|
|
enlarge_size=0.2, |
|
|
min_unique_num=None), |
|
|
dict(type='NormalizePointsColor', color_mean=None), |
|
|
dict(type='Pack3DDetInputs', keys=['points', 'pts_semantic_mask']) |
|
|
] |
|
|
|
|
|
data_prefix = dict( |
|
|
pts='points', |
|
|
pts_instance_mask='instance_mask', |
|
|
pts_semantic_mask='semantic_mask') |
|
|
return (data_root, ann_file, classes, palette, scene_idxs, data_prefix, |
|
|
pipeline, modality) |
|
|
|
|
|
|
|
|
def _generate_scannet_dataset_config(): |
|
|
data_root = 'tests/data/scannet' |
|
|
ann_file = 'scannet_infos.pkl' |
|
|
classes = ('cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window', |
|
|
'bookshelf', 'picture', 'counter', 'desk', 'curtain', |
|
|
'refrigerator', 'showercurtrain', 'toilet', 'sink', 'bathtub', |
|
|
'garbagebin') |
|
|
modality = dict(use_lidar=True, use_camera=False) |
|
|
pipeline = [ |
|
|
dict( |
|
|
type='LoadPointsFromFile', |
|
|
coord_type='DEPTH', |
|
|
shift_height=True, |
|
|
load_dim=6, |
|
|
use_dim=[0, 1, 2]), |
|
|
dict( |
|
|
type='LoadAnnotations3D', |
|
|
with_bbox_3d=True, |
|
|
with_label_3d=True, |
|
|
with_mask_3d=True, |
|
|
with_seg_3d=True), |
|
|
dict(type='GlobalAlignment', rotation_axis=2), |
|
|
dict(type='PointSegClassMapping'), |
|
|
dict(type='PointSample', num_points=5), |
|
|
dict( |
|
|
type='RandomFlip3D', |
|
|
sync_2d=False, |
|
|
flip_ratio_bev_horizontal=1.0, |
|
|
flip_ratio_bev_vertical=1.0), |
|
|
dict( |
|
|
type='GlobalRotScaleTrans', |
|
|
rot_range=[-0.087266, 0.087266], |
|
|
scale_ratio_range=[1.0, 1.0], |
|
|
shift_height=True), |
|
|
dict( |
|
|
type='Pack3DDetInputs', |
|
|
keys=[ |
|
|
'points', 'pts_semantic_mask', 'gt_bboxes_3d', 'gt_labels_3d', |
|
|
'pts_instance_mask' |
|
|
]) |
|
|
] |
|
|
data_prefix = dict( |
|
|
pts='points', |
|
|
pts_instance_mask='instance_mask', |
|
|
pts_semantic_mask='semantic_mask') |
|
|
return data_root, ann_file, classes, data_prefix, pipeline, modality |
|
|
|
|
|
|
|
|
class TestScanNetDataset(unittest.TestCase): |
|
|
|
|
|
def test_scannet(self): |
|
|
np.random.seed(0) |
|
|
data_root, ann_file, classes, data_prefix, \ |
|
|
pipeline, modality, = _generate_scannet_dataset_config() |
|
|
register_all_modules() |
|
|
scannet_dataset = ScanNetDataset( |
|
|
data_root, |
|
|
ann_file, |
|
|
data_prefix=data_prefix, |
|
|
pipeline=pipeline, |
|
|
metainfo=dict(classes=classes), |
|
|
modality=modality) |
|
|
|
|
|
scannet_dataset.prepare_data(0) |
|
|
input_dict = scannet_dataset.get_data_info(0) |
|
|
scannet_dataset[0] |
|
|
|
|
|
self.assertIn(data_prefix['pts'], |
|
|
input_dict['lidar_points']['lidar_path']) |
|
|
self.assertIn(data_root, input_dict['lidar_points']['lidar_path']) |
|
|
|
|
|
ann_info = scannet_dataset.parse_ann_info(input_dict) |
|
|
|
|
|
|
|
|
except_label = np.array([ |
|
|
6, 6, 4, 9, 11, 11, 10, 0, 15, 17, 17, 17, 3, 12, 4, 4, 14, 1, 0, |
|
|
0, 0, 0, 0, 0, 5, 5, 5 |
|
|
]) |
|
|
|
|
|
self.assertEqual(ann_info['gt_labels_3d'].dtype, np.int64) |
|
|
assert_allclose(ann_info['gt_labels_3d'], except_label) |
|
|
self.assertIsInstance(ann_info['gt_bboxes_3d'], DepthInstance3DBoxes) |
|
|
assert len(ann_info['gt_bboxes_3d']) == 27 |
|
|
assert torch.allclose(ann_info['gt_bboxes_3d'].tensor.sum(), |
|
|
torch.tensor([107.7353])) |
|
|
|
|
|
no_class_scannet_dataset = ScanNetDataset( |
|
|
data_root, ann_file, metainfo=dict(classes=['cabinet'])) |
|
|
|
|
|
input_dict = no_class_scannet_dataset.get_data_info(0) |
|
|
ann_info = no_class_scannet_dataset.parse_ann_info(input_dict) |
|
|
|
|
|
|
|
|
self.assertIn('gt_labels_3d', ann_info) |
|
|
|
|
|
assert (ann_info['gt_labels_3d'] <= 0).all() |
|
|
self.assertEqual(ann_info['gt_labels_3d'].dtype, np.int64) |
|
|
|
|
|
self.assertEqual(len(ann_info['gt_labels_3d']), 27) |
|
|
self.assertEqual(len(no_class_scannet_dataset.metainfo['classes']), 1) |
|
|
|
|
|
def test_scannet_seg(self): |
|
|
data_root, ann_file, classes, palette, scene_idxs, data_prefix, \ |
|
|
pipeline, modality, = _generate_scannet_seg_dataset_config() |
|
|
|
|
|
register_all_modules() |
|
|
np.random.seed(0) |
|
|
scannet_seg_dataset = ScanNetSegDataset( |
|
|
data_root, |
|
|
ann_file, |
|
|
metainfo=dict(classes=classes, palette=palette), |
|
|
data_prefix=data_prefix, |
|
|
pipeline=pipeline, |
|
|
modality=modality, |
|
|
scene_idxs=scene_idxs) |
|
|
|
|
|
input_dict = scannet_seg_dataset.prepare_data(0) |
|
|
|
|
|
points = input_dict['inputs']['points'] |
|
|
data_sample = input_dict['data_samples'] |
|
|
pts_semantic_mask = data_sample.gt_pts_seg.pts_semantic_mask |
|
|
|
|
|
expected_points = torch.tensor([[ |
|
|
0.0000, 0.0000, 1.2427, 0.6118, 0.5529, 0.4471, -0.6462, -1.0046, |
|
|
0.4280 |
|
|
], |
|
|
[ |
|
|
0.1553, -0.0074, 1.6077, 0.5882, |
|
|
0.6157, 0.5569, -0.6001, -1.0068, |
|
|
0.5537 |
|
|
], |
|
|
[ |
|
|
0.1518, 0.6016, 0.6548, 0.1490, |
|
|
0.1059, 0.0431, -0.6012, -0.8309, |
|
|
0.2255 |
|
|
], |
|
|
[ |
|
|
-0.7494, 0.1033, 0.6756, 0.5216, |
|
|
0.4353, 0.3333, -0.8687, -0.9748, |
|
|
0.2327 |
|
|
], |
|
|
[ |
|
|
-0.6836, -0.0203, 0.5884, 0.5765, |
|
|
0.5020, 0.4510, -0.8491, -1.0105, |
|
|
0.2027 |
|
|
]]) |
|
|
expected_pts_semantic_mask = np.array([13, 13, 12, 2, 0]) |
|
|
|
|
|
assert torch.allclose(points, expected_points, 1e-2) |
|
|
self.assertTrue( |
|
|
(pts_semantic_mask.numpy() == expected_pts_semantic_mask).all()) |
|
|
|
