add seed and more data for test unit, change indoor sample

mmdet3d/datasets/pipelines/indoor_sample.py

@@ -3,10 +3,26 @@ import numpy as np
 from mmdet.datasets.registry import PIPELINES
 
 
-def points_random_sampling(points,
+@PIPELINES.register_module
+class PointSample(object):
+    """Point Sample.
+
+    Sampling data to a certain number.
+
+    Args:
+        name (str): Name of the dataset.
+        num_points (int): Number of points to be sampled.
+    """
+
+    def __init__(self, num_points):
+        self.num_points = num_points
+
+    def points_random_sampling(self,
+                               points,
                                num_samples,
                                replace=None,
-                           return_choices=False):
+                               return_choices=False,
+                               seed=None):
         """Points Random Sampling.
 
         Sample points to a certain number.
@@ -21,39 +37,25 @@ def points_random_sampling(points,
             points (ndarray): 3D Points.
             choices (ndarray): The generated random samples
         """
+        if seed is not None:
+            np.random.seed(seed)
         if replace is None:
             replace = (points.shape[0] < num_samples)
-    choices = np.random.choice(points.shape[0], num_samples, replace=replace)
+        choices = np.random.choice(
+            points.shape[0], num_samples, replace=replace)
         if return_choices:
             return points[choices], choices
         else:
             return points[choices]
 
-
-@PIPELINES.register_module
-class PointSample(object):
-    """Point Sample.
-
-    Sampling data to a certain number.
-
-    Args:
-        name (str): Name of the dataset.
-        num_points (int): Number of points to be sampled.
-    """
-
-    def __init__(self, name, num_points):
-        assert name in ['scannet', 'sunrgbd']
-        self.name = name
-        self.num_points = num_points
-
-    def __call__(self, results):
-        points = results.get('points', None)
+    def __call__(self, results, seed=None):
+        point_cloud = results.get('point_cloud', None)
         pcl_color = results.get('pcl_color', None)
-        points, choices = points_random_sampling(
-            points, self.num_points, return_choices=True)
-        results['points'] = points
+        point_cloud, choices = self.points_random_sampling(
+            point_cloud, self.num_points, return_choices=True, seed=seed)
+        results['point_cloud'] = point_cloud
 
-        if self.name == 'scannet':
+        if pcl_color is not None:
             pcl_color = pcl_color[choices]
             instance_labels = results.get('instance_labels', None)
             semantic_labels = results.get('semantic_labels', None)
@@ -67,6 +69,5 @@ class PointSample(object):
 
     def __repr__(self):
         repr_str = self.__class__.__name__
-        repr_str += '(dataset_name={})'.format(self.name)
         repr_str += '(num_points={})'.format(self.num_points)
         return repr_str

tests/test_indoor_sample.py

@@ -4,35 +4,65 @@ from mmdet3d.datasets.pipelines.indoor_sample import PointSample
 
 
 def test_indoor_sample():
-    scannet_sample_points = PointSample('scannet', 1)
+    scannet_sample_points = PointSample(5)
     scannet_results = dict()
-    scannet_results['points'] = np.array(
-        [[1.6110241e+00, -1.6903955e-01, 5.8115810e-01, 5.9897250e-01],
-         [1.3978075e+00, 4.2035791e-01, 3.8729519e-01, 4.0510958e-01]])
-    scannet_results['instance_labels'] = np.array([9, 3])
-    scannet_results['pcl_color'] = np.array([[29.0, 142.0, 122.0],
-                                             [21., 17., 16.]])
-    scannet_results['semantic_labels'] = np.array([1, 5])
+    scannet_point_cloud = np.array(
+        [[1.0719866, -0.7870435, 0.8408122, 0.9196809],
+         [1.103661, 0.81065744, 2.6616862, 2.7405548],
+         [1.0276475, 1.5061463, 2.6174362, 2.6963048],
+         [-0.9709588, 0.6750515, 0.93901765, 1.0178864],
+         [1.0578915, 1.1693821, 0.87503505, 0.95390373],
+         [0.05560996, -1.5688863, 1.2440368, 1.3229055],
+         [-0.15731563, -1.7735453, 2.7535574, 2.832426],
+         [1.1188195, -0.99211365, 2.5551798, 2.6340485],
+         [-0.9186557, -1.7041215, 2.0562649, 2.1351335],
+         [-1.0128691, -1.3394243, 0.040936, 0.1198047]])
+    scannet_results['point_cloud'] = scannet_point_cloud
+    scannet_instance_labels = np.array([15, 12, 11, 38, 0, 18, 17, 12, 17, 0])
+    scannet_results['instance_labels'] = scannet_instance_labels
+    scannet_pcl_color = np.array([[77., 74., 63.], [176., 166., 156.],
+                                  [178., 164., 153.], [240., 235., 237.],
+                                  [58., 58., 59.], [245., 236., 229.],
+                                  [158., 148., 141.], [195., 184., 178.],
+                                  [193., 181., 174.], [105., 102., 97.]])
+    scannet_results['pcl_color'] = scannet_pcl_color
+    scannet_semantic_labels = np.array([38, 1, 1, 40, 0, 40, 1, 1, 1, 0])
+    scannet_results['semantic_labels'] = scannet_semantic_labels
 
-    scannet_results = scannet_sample_points(scannet_results)
-    points = scannet_results.get('points', None)
-    pcl_color = scannet_results.get('pcl_color', None)
-    instance_labels = scannet_results.get('instance_labels', None)
-    semantic_labels = scannet_results.get('semantic_labels', None)
-    assert points.shape == (1, 4)
-    assert pcl_color.shape == (1, 3)
-    assert instance_labels.shape == (1, )
-    assert semantic_labels.shape == (1, )
+    scannet_results = scannet_sample_points(scannet_results, 0)
+    scannet_point_cloud_result = scannet_results.get('point_cloud', None)
+    scannet_pcl_color_result = scannet_results.get('pcl_color', None)
+    scannet_instance_labels_result = scannet_results.get(
+        'instance_labels', None)
+    scannet_semantic_labels_result = scannet_results.get(
+        'semantic_labels', None)
+    scannet_choices = np.array([2, 8, 4, 9, 1])
+    assert np.allclose(scannet_point_cloud[scannet_choices],
+                       scannet_point_cloud_result)
+    assert np.allclose(scannet_pcl_color[scannet_choices],
+                       scannet_pcl_color_result)
+    assert np.all(scannet_instance_labels[scannet_choices] ==
+                  scannet_instance_labels_result)
+    assert np.all(scannet_semantic_labels[scannet_choices] ==
+                  scannet_semantic_labels_result)
 
-    sunrgbd_sample_points = PointSample('sunrgbd', 1)
+    sunrgbd_sample_points = PointSample(5)
     sunrgbd_results = dict()
-    sunrgbd_results['points'] = np.array(
-        [[1.2113925, 2.8755326, -1.1801991, 0.01056887],
-         [3.6554186, 4.5093756, 0.33279705, 1.523565]])
+    sunrgbd_point_cloud = np.array(
+        [[-1.8135729e-01, 1.4695230e+00, -1.2780589e+00, 7.8938007e-03],
+         [1.2581362e-03, 2.0561588e+00, -1.0341064e+00, 2.5184631e-01],
+         [6.8236995e-01, 3.3611867e+00, -9.2599887e-01, 3.5995382e-01],
+         [-2.9432583e-01, 1.8714852e+00, -9.0929651e-01, 3.7665617e-01],
+         [-0.5024875, 1.8032674, -1.1403012, 0.14565146],
+         [-0.520559, 1.6324949, -0.9896099, 0.2963428],
+         [0.95929825, 2.9402404, -0.8746674, 0.41128528],
+         [-0.74624217, 1.5244724, -0.8678476, 0.41810507],
+         [0.56485355, 1.5747732, -0.804522, 0.4814307],
+         [-0.0913099, 1.3673826, -1.2800645, 0.00588822]])
+    sunrgbd_results['point_cloud'] = sunrgbd_point_cloud
 
-    sunrgbd_results = sunrgbd_sample_points(sunrgbd_results)
-    point_cloud = sunrgbd_results.get('points', None)
-    assert point_cloud.shape == (1, 4)
-
-
-test_indoor_sample()
+    sunrgbd_results = sunrgbd_sample_points(sunrgbd_results, 0)
+    sunrgbd_choices = np.array([2, 8, 4, 9, 1])
+    sunrgbd_point_cloud_result = sunrgbd_results.get('point_cloud', None)
+    assert np.allclose(sunrgbd_point_cloud[sunrgbd_choices],
+                       sunrgbd_point_cloud_result)