test_indoor_augment.py 5.96 KB
Newer Older
liyinhao's avatar
liyinhao committed
1
2
import numpy as np

3
4
from mmdet3d.datasets.pipelines.indoor_augment import (IndoorFlipData,
                                                       IndoorGlobalRotScale)
liyinhao's avatar
liyinhao committed
5
6


7
def test_indoor_flip_data():
liyinhao's avatar
liyinhao committed
8
9
    np.random.seed(0)
    sunrgbd_indoor_flip_data = IndoorFlipData()
10
    sunrgbd_results = dict()
liyinhao's avatar
liyinhao committed
11
    sunrgbd_results['points'] = np.array(
12
13
        [[1.02828765e+00, 3.65790772e+00, 1.97294697e-01, 1.61959505e+00],
         [-3.95979017e-01, 1.05465031e+00, -7.49204338e-01, 6.73096001e-01]])
liyinhao's avatar
liyinhao committed
14
    sunrgbd_results['gt_bboxes_3d'] = np.array([[
15
16
        0.213684, 1.036364, -0.982323, 0.61541, 0.572574, 0.872728, 3.07028526
    ],
liyinhao's avatar
liyinhao committed
17
18
19
20
21
22
                                                [
                                                    -0.449953, 1.395455,
                                                    -1.027778, 1.500956,
                                                    1.637298, 0.636364,
                                                    -1.58242359
                                                ]])
23
    sunrgbd_results = sunrgbd_indoor_flip_data(sunrgbd_results)
24
25
    sunrgbd_points = sunrgbd_results['points']
    sunrgbd_gt_bboxes_3d = sunrgbd_results['gt_bboxes_3d']
26

liyinhao's avatar
liyinhao committed
27
    expected_sunrgbd_points = np.array(
28
29
        [[-1.02828765, 3.65790772, 0.1972947, 1.61959505],
         [0.39597902, 1.05465031, -0.74920434, 0.673096]])
liyinhao's avatar
liyinhao committed
30
    expected_sunrgbd_gt_bboxes_3d = np.array([[
31
32
33
34
        -0.213684, 1.036364, -0.982323, 0.61541, 0.572574, 0.872728, 0.07130739
    ], [
        0.449953, 1.395455, -1.027778, 1.500956, 1.637298, 0.636364, 4.72401624
    ]])
liyinhao's avatar
liyinhao committed
35
36
    assert np.allclose(sunrgbd_points, expected_sunrgbd_points)
    assert np.allclose(sunrgbd_gt_bboxes_3d, expected_sunrgbd_gt_bboxes_3d)
37

liyinhao's avatar
liyinhao committed
38
39
    np.random.seed(0)
    scannet_indoor_flip_data = IndoorFlipData()
40
    scannet_results = dict()
liyinhao's avatar
liyinhao committed
41
    scannet_results['points'] = np.array(
42
43
        [[1.6110241e+00, -1.6903955e-01, 5.8115810e-01, 5.9897250e-01],
         [1.3978075e+00, 4.2035791e-01, 3.8729519e-01, 4.0510958e-01]])
liyinhao's avatar
liyinhao committed
44
    scannet_results['gt_bboxes_3d'] = np.array([[
45
46
47
48
49
        0.55903838, 0.48201692, 0.65688646, 0.65370704, 0.60029864, 0.5163464
    ], [
        -0.03226406, 1.70392646, 0.60348618, 0.65165804, 0.72084366, 0.64667457
    ]])
    scannet_results = scannet_indoor_flip_data(scannet_results)
50
51
    scannet_points = scannet_results['points']
    scannet_gt_bboxes_3d = scannet_results['gt_bboxes_3d']
52

liyinhao's avatar
liyinhao committed
53
    expected_scannet_points = np.array(
54
55
        [[-1.6110241, 0.16903955, 0.5811581, 0.5989725],
         [-1.3978075, -0.42035791, 0.38729519, 0.40510958]])
liyinhao's avatar
liyinhao committed
56
    expected_scannet_gt_bboxes_3d = np.array([[
57
58
59
60
        -0.55903838, -0.48201692, 0.65688646, 0.65370704, 0.60029864, 0.5163464
    ], [
        0.03226406, -1.70392646, 0.60348618, 0.65165804, 0.72084366, 0.64667457
    ]])
liyinhao's avatar
liyinhao committed
61
62
    assert np.allclose(scannet_points, expected_scannet_points)
    assert np.allclose(scannet_gt_bboxes_3d, expected_scannet_gt_bboxes_3d)
63
64
65


def test_global_rot_scale():
liyinhao's avatar
liyinhao committed
66
    np.random.seed(0)
liyinhao's avatar
liyinhao committed
67
68
    sunrgbd_augment = IndoorGlobalRotScale(
        True, rot_range=[-np.pi / 6, np.pi / 6], scale_range=[0.85, 1.15])
liyinhao's avatar
liyinhao committed
69
    sunrgbd_results = dict()
liyinhao's avatar
liyinhao committed
70
    sunrgbd_results['points'] = np.array(
liyinhao's avatar
liyinhao committed
71
72
        [[1.02828765e+00, 3.65790772e+00, 1.97294697e-01, 1.61959505e+00],
         [-3.95979017e-01, 1.05465031e+00, -7.49204338e-01, 6.73096001e-01]])
liyinhao's avatar
liyinhao committed
73
    sunrgbd_results['gt_bboxes_3d'] = np.array([[
liyinhao's avatar
liyinhao committed
74
75
        0.213684, 1.036364, -0.982323, 0.61541, 0.572574, 0.872728, 3.07028526
    ],
liyinhao's avatar
liyinhao committed
76
77
78
79
80
81
                                                [
                                                    -0.449953, 1.395455,
                                                    -1.027778, 1.500956,
                                                    1.637298, 0.636364,
                                                    -1.58242359
                                                ]])
liyinhao's avatar
liyinhao committed
82
83

    sunrgbd_results = sunrgbd_augment(sunrgbd_results)
84
85
    sunrgbd_points = sunrgbd_results['points']
    sunrgbd_gt_bboxes_3d = sunrgbd_results['gt_bboxes_3d']
86

liyinhao's avatar
liyinhao committed
87
    expected_sunrgbd_points = np.array(
88
89
        [[0.89427376, 3.94489646, 0.21003141, 1.72415094],
         [-0.47835783, 1.09972989, -0.79757058, 0.71654893]])
liyinhao's avatar
liyinhao committed
90
    expected_sunrgbd_gt_bboxes_3d = np.array([[
91
92
        0.17080999, 1.11345031, -1.04573864, 0.65513891, 0.60953755,
        0.92906854, 3.01916788
liyinhao's avatar
liyinhao committed
93
    ],
liyinhao's avatar
liyinhao committed
94
95
96
97
98
99
100
101
                                              [
                                                  -0.55427876, 1.45912611,
                                                  -1.09412807, 1.59785293,
                                                  1.74299674, 0.67744563,
                                                  -1.63354097
                                              ]])
    assert np.allclose(sunrgbd_points, expected_sunrgbd_points)
    assert np.allclose(sunrgbd_gt_bboxes_3d, expected_sunrgbd_gt_bboxes_3d)
liyinhao's avatar
liyinhao committed
102

liyinhao's avatar
liyinhao committed
103
    np.random.seed(0)
104
    scannet_augment = IndoorGlobalRotScale(
liyinhao's avatar
liyinhao committed
105
        True, rot_range=[-np.pi * 1 / 36, np.pi * 1 / 36], scale_range=None)
liyinhao's avatar
liyinhao committed
106
    scannet_results = dict()
liyinhao's avatar
liyinhao committed
107
    scannet_results['points'] = np.array(
liyinhao's avatar
liyinhao committed
108
109
        [[1.6110241e+00, -1.6903955e-01, 5.8115810e-01, 5.9897250e-01],
         [1.3978075e+00, 4.2035791e-01, 3.8729519e-01, 4.0510958e-01]])
liyinhao's avatar
liyinhao committed
110
    scannet_results['gt_bboxes_3d'] = np.array([[
liyinhao's avatar
liyinhao committed
111
112
113
114
        0.55903838, 0.48201692, 0.65688646, 0.65370704, 0.60029864, 0.5163464
    ], [
        -0.03226406, 1.70392646, 0.60348618, 0.65165804, 0.72084366, 0.64667457
    ]])
liyinhao's avatar
liyinhao committed
115
    scannet_results = scannet_augment(scannet_results)
116
117
    scannet_points = scannet_results['points']
    scannet_gt_bboxes_3d = scannet_results['gt_bboxes_3d']
118

liyinhao's avatar
liyinhao committed
119
    expected_scannet_points = np.array(
120
121
        [[1.61240576, -0.15530836, 0.5811581, 0.5989725],
         [1.39417555, 0.43225122, 0.38729519, 0.40510958]])
liyinhao's avatar
liyinhao committed
122
    expected_scannet_gt_bboxes_3d = np.array([[
123
        0.55491157, 0.48676213, 0.65688646, 0.65879754, 0.60584609, 0.5163464
124
    ], [
125
        -0.04677942, 1.70358975, 0.60348618, 0.65777559, 0.72636927, 0.64667457
126
    ]])
liyinhao's avatar
liyinhao committed
127
128
    assert np.allclose(scannet_points, expected_scannet_points)
    assert np.allclose(scannet_gt_bboxes_3d, expected_scannet_gt_bboxes_3d)