test_s3dis_dataset.py

# Copyright (c) OpenMMLab. All rights reserved.
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]
        # assert the path should contains data_prefix and data_root
        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)

        # assert the keys in ann_info and the type
        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)

        # assert the keys in ann_info and the type
        self.assertIn('gt_labels_3d', ann_info)
        # assert mapping to -1 or 1
        assert (ann_info['gt_labels_3d'] <= 0).all()
        self.assertEqual(ann_info['gt_labels_3d'].dtype, np.int64)
        # all instance have been filtered by classes
        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())