finish indoor_pipeline

tests/test_indoor_pipeline.py

@@ -21,14 +21,13 @@ def test_scannet_pipeline():
             type='IndoorGlobalRotScale',
             use_height=True,
             rot_range=[-np.pi * 1 / 36, np.pi * 1 / 36],
-            scale_range=None),
+            scale_range=None)
     ]
     pipeline = Compose(pipelines)
-    info = mmcv.load('./tests/data/scannet/scannet_infos.pkl')
+    info = mmcv.load('./tests/data/scannet/scannet_infos.pkl')[0]
     results = dict()
     data_path = './tests/data/scannet/scannet_train_instance_data'
     results['data_path'] = data_path
-    info = info[0]
     scan_name = info['point_cloud']['lidar_idx']
     results['pts_filename'] = osp.join(data_path, f'{scan_name}_vert.npy')
     if info['annos']['gt_num'] != 0:
@@ -77,3 +76,68 @@ def test_scannet_pipeline():
     assert np.all(gt_labels.flatten() == expected_gt_labels)
     assert np.all(pts_semantic_mask == expected_pts_semantic_mask)
     assert np.all(pts_instance_mask == expected_pts_instance_mask)
+
+
+def test_sunrgbd_pipeline():
+    np.random.seed(0)
+    pipelines = [
+        dict(
+            type='IndoorLoadPointsFromFile',
+            use_height=True,
+            load_dim=6,
+            use_dim=[0, 1, 2]),
+        dict(type='IndoorFlipData', flip_ratio_yz=1.0),
+        dict(
+            type='IndoorGlobalRotScale',
+            use_height=True,
+            rot_range=[-np.pi / 6, np.pi / 6],
+            scale_range=[0.85, 1.15]),
+        dict(type='IndoorPointSample', num_points=5),
+    ]
+    pipeline = Compose(pipelines)
+    results = dict()
+    info = mmcv.load('./tests/data/sunrgbd/sunrgbd_infos.pkl')[0]
+    data_path = './tests/data/sunrgbd/sunrgbd_trainval'
+    scan_name = info['point_cloud']['lidar_idx']
+    results['pts_filename'] = osp.join(data_path, 'lidar',
+                                       f'{scan_name:06d}.npy')
+
+    if info['annos']['gt_num'] != 0:
+        gt_bboxes_3d = info['annos']['gt_boxes_upright_depth']
+        gt_labels = info['annos']['class'].reshape(-1, 1)
+        gt_bboxes_3d_mask = np.ones_like(gt_labels)
+    else:
+        gt_bboxes_3d = np.zeros((1, 6), dtype=np.float32)
+        gt_labels = np.zeros((1, 1))
+        gt_bboxes_3d_mask = np.zeros((1, 1))
+    results['gt_bboxes_3d'] = gt_bboxes_3d
+    results['gt_labels'] = gt_labels
+    results['gt_bboxes_3d_mask'] = gt_bboxes_3d_mask
+    results = pipeline(results)
+    points = results['points']
+    gt_bboxes_3d = results['gt_bboxes_3d']
+    gt_labels = results['gt_labels']
+    expected_points = np.array(
+        [[0.6570105, 1.5538014, 0.24514851, 1.0165423],
+         [0.656101, 1.558591, 0.21755838, 0.98895216],
+         [0.6293659, 1.5679953, -0.10004003, 0.67135376],
+         [0.6068739, 1.5974995, -0.41063973, 0.36075398],
+         [0.6464709, 1.5573514, 0.15114647, 0.9225402]])
+    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
+                                      ]])
+    expected_gt_labels = np.array([0, 7, 6])
+    assert np.allclose(gt_bboxes_3d, expected_gt_bboxes_3d)
+    assert np.allclose(gt_labels.flatten(), expected_gt_labels)
+    assert np.allclose(points, expected_points)
\ No newline at end of file