Skip to content

GitLab

Commit 60371607 authored by liyinhao's avatar liyinhao
Browse files

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

parent 60e3958e
Hide whitespace changes
Inline Side-by-side
Showing with 95 additions and 64 deletions
mmdet3d/datasets/pipelines/indoor_sample.py
View file @ 60371607
......@@ -3,33 +3,6 @@ import numpy as np
from mmdet.datasets.registry import PIPELINES
def points_random_sampling(points,
num_samples,
replace=None,
return_choices=False):
"""Points Random Sampling.
Sample points to a certain number.
Args:
points (ndarray): 3D Points.
num_samples (int): Number of samples to be sampled.
replace (bool): Whether the sample is with or without replacement.
return_choices (bool): Whether return choice.
Returns:
points (ndarray): 3D Points.
choices (ndarray): The generated random samples
"""
if replace is None:
replace = (points.shape[0] < num_samples)
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.
......@@ -41,19 +14,48 @@ class PointSample(object):
num_points (int): Number of points to be sampled.
"""
def __init__(self, name, num_points):
assert name in ['scannet', 'sunrgbd']
self.name = name
def __init__(self, num_points):
self.num_points = num_points
def __call__(self, results):
points = results.get('points', None)
def points_random_sampling(self,
points,
num_samples,
replace=None,
return_choices=False,
seed=None):
"""Points Random Sampling.
Sample points to a certain number.
Args:
points (ndarray): 3D Points.
num_samples (int): Number of samples to be sampled.
replace (bool): Whether the sample is with or without replacement.
return_choices (bool): Whether return choice.
Returns:
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)
if return_choices:
return points[choices], choices
else:
return points[choices]
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
View file @ 60371607
......@@ -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)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment