Commit 603281ae authored by liyinhao's avatar liyinhao
Browse files

add indoor sample points

parent 2f5b24ce
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.
"""
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 IndoorSamplePoints(object):
"""Indoor Sample Points
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)
pcl_color = results.get('pcl_color', None)
points, choices = points_random_sampling(
points, self.num_points, return_choices=True)
results['points'] = points
if self.name == 'scannet':
pcl_color = pcl_color[choices]
instance_labels = results.get('instance_labels', None)
semantic_labels = results.get('semantic_labels', None)
instance_labels = instance_labels[choices]
semantic_labels = semantic_labels[choices]
results['instance_labels'] = instance_labels
results['semantic_labels'] = semantic_labels
results['pcl_color'] = pcl_color
return results
def __repr__(self):
repr_str = self.__class__.__name__
repr_str += '(num_points={})'.format(self.num_points)
return repr_str
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