|
|
|
import unittest |
|
|
|
import numpy as np |
|
import torch |
|
from mmengine.testing import assert_allclose |
|
|
|
from mmdet3d.datasets import S3DISDataset, S3DISSegDataset |
|
from mmdet3d.structures import DepthInstance3DBoxes |
|
from mmdet3d.utils import register_all_modules |
|
|
|
|
|
def _generate_s3dis_seg_dataset_config(): |
|
data_root = './tests/data/s3dis/' |
|
ann_file = 's3dis_infos.pkl' |
|
classes = ('ceiling', 'floor', 'wall', 'beam', 'column', 'window', 'door', |
|
'table', 'chair', 'sofa', 'bookcase', 'board', 'clutter') |
|
palette = [[0, 255, 0], [0, 0, 255], [0, 255, 255], [255, 255, 0], |
|
[255, 0, 255], [100, 100, 255], [200, 200, 100], |
|
[170, 120, 200], [255, 0, 0], [200, 100, 100], [10, 200, 100], |
|
[200, 200, 200], [50, 50, 50]] |
|
scene_idxs = [0 for _ in range(20)] |
|
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.0, |
|
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_s3dis_dataset_config(): |
|
data_root = 'tests/data/s3dis' |
|
ann_file = 's3dis_infos.pkl' |
|
classes = ('table', 'chair', 'sofa', 'bookcase', 'board') |
|
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=True, |
|
with_label_3d=True, |
|
with_mask_3d=True, |
|
with_seg_3d=True), |
|
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]), |
|
dict(type='NormalizePointsColor', color_mean=None), |
|
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 TestS3DISDataset(unittest.TestCase): |
|
|
|
def test_s3dis(self): |
|
np.random.seed(0) |
|
data_root, ann_file, classes, data_prefix, \ |
|
pipeline, modality = _generate_s3dis_dataset_config() |
|
register_all_modules() |
|
s3dis_dataset = S3DISDataset( |
|
data_root, |
|
ann_file, |
|
data_prefix=data_prefix, |
|
pipeline=pipeline, |
|
metainfo=dict(classes=classes), |
|
modality=modality) |
|
|
|
s3dis_dataset.prepare_data(0) |
|
input_dict = s3dis_dataset.get_data_info(0) |
|
s3dis_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 = s3dis_dataset.parse_ann_info(input_dict) |
|
|
|
|
|
except_label = np.array([1, 1, 3, 1, 2, 0, 0, 0, 3]) |
|
|
|
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']) == 9 |
|
assert torch.allclose(ann_info['gt_bboxes_3d'].tensor.sum(), |
|
torch.tensor([63.0455])) |
|
|
|
no_class_s3dis_dataset = S3DISDataset( |
|
data_root, ann_file, metainfo=dict(classes=['table'])) |
|
|
|
input_dict = no_class_s3dis_dataset.get_data_info(0) |
|
ann_info = no_class_s3dis_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']), 9) |
|
self.assertEqual(len(no_class_s3dis_dataset.metainfo['classes']), 1) |
|
|
|
def test_s3dis_seg(self): |
|
data_root, ann_file, classes, palette, scene_idxs, data_prefix, \ |
|
pipeline, modality, = _generate_s3dis_seg_dataset_config() |
|
|
|
register_all_modules() |
|
np.random.seed(0) |
|
|
|
s3dis_seg_dataset = S3DISSegDataset( |
|
data_root, |
|
ann_file, |
|
metainfo=dict(classes=classes, palette=palette), |
|
data_prefix=data_prefix, |
|
pipeline=pipeline, |
|
modality=modality, |
|
scene_idxs=scene_idxs) |
|
|
|
input_dict = s3dis_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, 3.1720, 0.4706, 0.4431, 0.3725, 0.4624, 0.7502, |
|
0.9543 |
|
], |
|
[ |
|
0.2880, -0.5900, 0.0650, 0.3451, |
|
0.3373, 0.3490, 0.5119, 0.5518, |
|
0.0196 |
|
], |
|
[ |
|
0.1570, 0.6000, 3.1700, 0.4941, |
|
0.4667, 0.3569, 0.4893, 0.9519, |
|
0.9537 |
|
], |
|
[ |
|
-0.1320, 0.3950, 0.2720, 0.3216, |
|
0.2863, 0.2275, 0.4397, 0.8830, |
|
0.0818 |
|
], |
|
[ |
|
-0.4860, -0.0640, 3.1710, 0.3843, |
|
0.3725, 0.3059, 0.3789, 0.7286, |
|
0.9540 |
|
]]) |
|
|
|
expected_pts_semantic_mask = np.array([0, 1, 0, 8, 0]) |
|
|
|
assert torch.allclose(points, expected_points, 1e-2) |
|
self.assertTrue( |
|
(pts_semantic_mask.numpy() == expected_pts_semantic_mask).all()) |
|
|
