test_indoor_pipeline.py 6.97 KB
Newer Older
liyinhao's avatar
liyinhao committed
1
2
3
4
5
6
7
8
9
import os.path as osp

import mmcv
import numpy as np

from mmdet3d.datasets.pipelines import Compose


def test_scannet_pipeline():
10
11
12
13
14
    class_names = ('cabinet', 'bed', 'chair', 'sofa', 'table', 'door',
                   'window', 'bookshelf', 'picture', 'counter', 'desk',
                   'curtain', 'refrigerator', 'showercurtrain', 'toilet',
                   'sink', 'bathtub', 'garbagebin')

liyinhao's avatar
liyinhao committed
15
16
17
    np.random.seed(0)
    pipelines = [
        dict(
zhangwenwei's avatar
zhangwenwei committed
18
19
            type='LoadPointsFromFile',
            shift_height=True,
liyinhao's avatar
liyinhao committed
20
21
            load_dim=6,
            use_dim=[0, 1, 2]),
zhangwenwei's avatar
zhangwenwei committed
22
23
24
25
26
27
        dict(
            type='LoadAnnotations3D',
            with_bbox_3d=True,
            with_label_3d=True,
            with_mask_3d=True,
            with_seg_3d=True),
liyinhao's avatar
liyinhao committed
28
29
30
31
        dict(type='IndoorPointSample', num_points=5),
        dict(type='IndoorFlipData', flip_ratio_yz=1.0, flip_ratio_xz=1.0),
        dict(
            type='IndoorGlobalRotScale',
zhangwenwei's avatar
zhangwenwei committed
32
            shift_height=True,
liyinhao's avatar
liyinhao committed
33
            rot_range=[-1 / 36, 1 / 36],
34
35
36
37
38
            scale_range=None),
        dict(type='DefaultFormatBundle3D', class_names=class_names),
        dict(
            type='Collect3D',
            keys=[
liyinhao's avatar
liyinhao committed
39
                'points', 'gt_bboxes_3d', 'gt_labels_3d', 'pts_semantic_mask',
40
41
                'pts_instance_mask'
            ]),
liyinhao's avatar
liyinhao committed
42
43
    ]
    pipeline = Compose(pipelines)
yinchimaoliang's avatar
yinchimaoliang committed
44
    info = mmcv.load('./tests/data/scannet/scannet_infos.pkl')[0]
liyinhao's avatar
liyinhao committed
45
    results = dict()
liyinhao's avatar
liyinhao committed
46
47
    data_path = './tests/data/scannet'
    results['pts_filename'] = osp.join(data_path, info['pts_path'])
liyinhao's avatar
liyinhao committed
48
    if info['annos']['gt_num'] != 0:
liyinhao's avatar
liyinhao committed
49
50
51
        scannet_gt_bboxes_3d = info['annos']['gt_boxes_upright_depth'].astype(
            np.float32)
        scannet_gt_labels_3d = info['annos']['class'].astype(np.long)
liyinhao's avatar
liyinhao committed
52
53
    else:
        scannet_gt_bboxes_3d = np.zeros((1, 6), dtype=np.float32)
liyinhao's avatar
liyinhao committed
54
        scannet_gt_labels_3d = np.zeros((1, ), dtype=np.long)
zhangwenwei's avatar
zhangwenwei committed
55
56
    results['ann_info'] = dict()
    results['ann_info']['pts_instance_mask_path'] = osp.join(
liyinhao's avatar
liyinhao committed
57
        data_path, info['pts_instance_mask_path'])
zhangwenwei's avatar
zhangwenwei committed
58
    results['ann_info']['pts_semantic_mask_path'] = osp.join(
liyinhao's avatar
liyinhao committed
59
        data_path, info['pts_semantic_mask_path'])
zhangwenwei's avatar
zhangwenwei committed
60
61
    results['ann_info']['gt_bboxes_3d'] = scannet_gt_bboxes_3d
    results['ann_info']['gt_labels_3d'] = scannet_gt_labels_3d
62

zhangwenwei's avatar
zhangwenwei committed
63
64
65
    results['bbox3d_fields'] = []
    results['pts_mask_fields'] = []
    results['pts_seg_fields'] = []
66

liyinhao's avatar
liyinhao committed
67
    results = pipeline(results)
68
69
70

    points = results['points']._data
    gt_bboxes_3d = results['gt_bboxes_3d']._data
liyinhao's avatar
liyinhao committed
71
    gt_labels_3d = results['gt_labels_3d']._data
72
73
    pts_semantic_mask = results['pts_semantic_mask']._data
    pts_instance_mask = results['pts_instance_mask']._data
liyinhao's avatar
liyinhao committed
74
75
76
77
78
79
80
81
82
83
84
85
86
    expected_points = np.array(
        [[-2.9078157, -1.9569951, 2.3543026, 2.389488],
         [-0.71360034, -3.4359822, 2.1330001, 2.1681855],
         [-1.332374, 1.474838, -0.04405887, -0.00887359],
         [2.1336637, -1.3265059, -0.02880373, 0.00638155],
         [0.43895668, -3.0259454, 1.5560012, 1.5911865]])
    expected_gt_bboxes_3d = np.array([
        [-1.5005362, -3.512584, 1.8565295, 1.7457027, 0.24149807, 0.57235193],
        [-2.8848705, 3.4961755, 1.5268247, 0.66170084, 0.17433672, 0.67153597],
        [-1.1585636, -2.192365, 0.61649567, 0.5557011, 2.5375574, 1.2144762],
        [-2.930457, -2.4856408, 0.9722377, 0.6270478, 1.8461524, 0.28697443],
        [3.3114715, -0.00476722, 1.0712197, 0.46191898, 3.8605113, 2.1603441]
    ])
liyinhao's avatar
liyinhao committed
87
    expected_gt_labels_3d = np.array([
liyinhao's avatar
liyinhao committed
88
89
90
91
92
93
94
        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
    ])
    expected_pts_semantic_mask = np.array([3, 1, 2, 2, 15])
    expected_pts_instance_mask = np.array([44, 22, 10, 10, 57])
    assert np.allclose(points, expected_points)
    assert np.allclose(gt_bboxes_3d[:5, :], expected_gt_bboxes_3d)
liyinhao's avatar
liyinhao committed
95
    assert np.all(gt_labels_3d.numpy() == expected_gt_labels_3d)
96
97
    assert np.all(pts_semantic_mask.numpy() == expected_pts_semantic_mask)
    assert np.all(pts_instance_mask.numpy() == expected_pts_instance_mask)
yinchimaoliang's avatar
yinchimaoliang committed
98
99
100


def test_sunrgbd_pipeline():
101
102
    class_names = ('bed', 'table', 'sofa', 'chair', 'toilet', 'desk',
                   'dresser', 'night_stand', 'bookshelf', 'bathtub')
yinchimaoliang's avatar
yinchimaoliang committed
103
104
105
    np.random.seed(0)
    pipelines = [
        dict(
zhangwenwei's avatar
zhangwenwei committed
106
107
            type='LoadPointsFromFile',
            shift_height=True,
yinchimaoliang's avatar
yinchimaoliang committed
108
109
            load_dim=6,
            use_dim=[0, 1, 2]),
zhangwenwei's avatar
zhangwenwei committed
110
        dict(type='LoadAnnotations3D'),
yinchimaoliang's avatar
yinchimaoliang committed
111
112
113
        dict(type='IndoorFlipData', flip_ratio_yz=1.0),
        dict(
            type='IndoorGlobalRotScale',
zhangwenwei's avatar
zhangwenwei committed
114
            shift_height=True,
liyinhao's avatar
liyinhao committed
115
            rot_range=[-1 / 6, 1 / 6],
yinchimaoliang's avatar
yinchimaoliang committed
116
117
            scale_range=[0.85, 1.15]),
        dict(type='IndoorPointSample', num_points=5),
118
        dict(type='DefaultFormatBundle3D', class_names=class_names),
liyinhao's avatar
liyinhao committed
119
120
        dict(
            type='Collect3D', keys=['points', 'gt_bboxes_3d', 'gt_labels_3d']),
yinchimaoliang's avatar
yinchimaoliang committed
121
122
123
124
    ]
    pipeline = Compose(pipelines)
    results = dict()
    info = mmcv.load('./tests/data/sunrgbd/sunrgbd_infos.pkl')[0]
liyinhao's avatar
liyinhao committed
125
126
    data_path = './tests/data/sunrgbd'
    results['pts_filename'] = osp.join(data_path, info['pts_path'])
yinchimaoliang's avatar
yinchimaoliang committed
127
128

    if info['annos']['gt_num'] != 0:
liyinhao's avatar
liyinhao committed
129
130
131
        gt_bboxes_3d = info['annos']['gt_boxes_upright_depth'].astype(
            np.float32)
        gt_labels_3d = info['annos']['class'].astype(np.long)
yinchimaoliang's avatar
yinchimaoliang committed
132
133
    else:
        gt_bboxes_3d = np.zeros((1, 6), dtype=np.float32)
liyinhao's avatar
liyinhao committed
134
        gt_labels_3d = np.zeros((1, ), dtype=np.long)
135

zhangwenwei's avatar
zhangwenwei committed
136
137
138
139
140
141
142
143
    # prepare input of pipeline
    results['ann_info'] = dict()
    results['ann_info']['gt_bboxes_3d'] = gt_bboxes_3d
    results['ann_info']['gt_labels_3d'] = gt_labels_3d
    results['bbox3d_fields'] = []
    results['pts_mask_fields'] = []
    results['pts_seg_fields'] = []

yinchimaoliang's avatar
yinchimaoliang committed
144
    results = pipeline(results)
145
146
    points = results['points']._data
    gt_bboxes_3d = results['gt_bboxes_3d']._data
liyinhao's avatar
liyinhao committed
147
    gt_labels_3d = results['gt_labels_3d']._data
liyinhao's avatar
liyinhao committed
148
149
150
151
152
    expected_points = np.array([[0.6512, 1.5781, 0.0710, 0.0499],
                                [0.6473, 1.5701, 0.0657, 0.0447],
                                [0.6464, 1.5635, 0.0826, 0.0616],
                                [0.6453, 1.5603, 0.0849, 0.0638],
                                [0.6488, 1.5786, 0.0461, 0.0251]])
yinchimaoliang's avatar
yinchimaoliang committed
153
154
155
156
157
158
159
160
161
162
163
164
165
166
    expected_gt_bboxes_3d = np.array([[
        -2.012483, 3.9473376, -0.25446942, 2.3730404, 1.9457763, 2.0303352,
        1.2205974
    ],
                                      [
                                          -3.7036808, 4.2396426, -0.81091917,
                                          0.6032123, 0.91040343, 1.003341,
                                          1.2662518
                                      ],
                                      [
                                          0.6528646, 2.1638472, -0.15228128,
                                          0.7347852, 1.6113238, 2.1694272,
                                          2.81404
                                      ]])
liyinhao's avatar
liyinhao committed
167
    expected_gt_labels_3d = np.array([0, 7, 6])
yinchimaoliang's avatar
yinchimaoliang committed
168
    assert np.allclose(gt_bboxes_3d, expected_gt_bboxes_3d)
liyinhao's avatar
liyinhao committed
169
    assert np.allclose(gt_labels_3d.flatten(), expected_gt_labels_3d)
liyinhao's avatar
liyinhao committed
170
    assert np.allclose(points, expected_points, 1e-2)