Commit ea7fdbed authored by liyinhao's avatar liyinhao
Browse files

finish indoor load test unit

parent e9ef0711
...@@ -47,7 +47,7 @@ def _rotate_aligned_boxes(input_boxes, rot_mat): ...@@ -47,7 +47,7 @@ def _rotate_aligned_boxes(input_boxes, rot_mat):
class IndoorFlipData(object): class IndoorFlipData(object):
"""Indoor Flip Data """Indoor Flip Data
Flip points and groundtruth boxes. Flip point_cloud and groundtruth boxes.
Args: Args:
name (str): name of the dataset. name (str): name of the dataset.
...@@ -58,12 +58,12 @@ class IndoorFlipData(object): ...@@ -58,12 +58,12 @@ class IndoorFlipData(object):
self.name = name self.name = name
def __call__(self, results): def __call__(self, results):
points = results.get('points', None) point_cloud = results.get('point_cloud', None)
gt_boxes = results.get('gt_boxes', None) gt_boxes = results.get('gt_boxes', None)
if np.random.random() > 0.5: if np.random.random() > 0.5:
# Flipping along the YZ plane # Flipping along the YZ plane
points[:, 0] = -1 * points[:, 0] point_cloud[:, 0] = -1 * point_cloud[:, 0]
gt_boxes[:, 0] = -1 * gt_boxes[:, 0] gt_boxes[:, 0] = -1 * gt_boxes[:, 0]
if self.name == 'sunrgbd': if self.name == 'sunrgbd':
gt_boxes[:, 6] = np.pi - gt_boxes[:, 6] gt_boxes[:, 6] = np.pi - gt_boxes[:, 6]
...@@ -71,10 +71,10 @@ class IndoorFlipData(object): ...@@ -71,10 +71,10 @@ class IndoorFlipData(object):
if self.name == 'scannet' and np.random.random() > 0.5: if self.name == 'scannet' and np.random.random() > 0.5:
# Flipping along the XZ plane # Flipping along the XZ plane
points[:, 1] = -1 * points[:, 1] point_cloud[:, 1] = -1 * point_cloud[:, 1]
gt_boxes[:, 1] = -1 * gt_boxes[:, 1] gt_boxes[:, 1] = -1 * gt_boxes[:, 1]
results['gt_boxes'] = gt_boxes results['gt_boxes'] = gt_boxes
results['points'] = points results['point_cloud'] = point_cloud
return results return results
...@@ -88,7 +88,7 @@ class IndoorFlipData(object): ...@@ -88,7 +88,7 @@ class IndoorFlipData(object):
class IndoorRotateData(object): class IndoorRotateData(object):
"""Indoor Rotate Data """Indoor Rotate Data
Rotate points and groundtruth boxes. Rotate point_cloud and groundtruth boxes.
Args: Args:
name (str): name of the dataset. name (str): name of the dataset.
...@@ -99,12 +99,13 @@ class IndoorRotateData(object): ...@@ -99,12 +99,13 @@ class IndoorRotateData(object):
self.name = name self.name = name
def __call__(self, results): def __call__(self, results):
points = results.get('points', None) point_cloud = results.get('point_cloud', None)
gt_boxes = results.get('gt_boxes', None) gt_boxes = results.get('gt_boxes', None)
rot_angle = (np.random.random() * np.pi / rot_angle = (np.random.random() * np.pi /
3) - np.pi / 6 # -30 ~ +30 degree 3) - np.pi / 6 # -30 ~ +30 degree
rot_mat = _rotz(rot_angle) rot_mat = _rotz(rot_angle)
points[:, 0:3] = np.dot(points[:, 0:3], np.transpose(rot_mat)) point_cloud[:, 0:3] = np.dot(point_cloud[:, 0:3],
np.transpose(rot_mat))
if self.name == 'scannet': if self.name == 'scannet':
gt_boxes = _rotate_aligned_boxes(gt_boxes, rot_mat) gt_boxes = _rotate_aligned_boxes(gt_boxes, rot_mat)
...@@ -112,7 +113,7 @@ class IndoorRotateData(object): ...@@ -112,7 +113,7 @@ class IndoorRotateData(object):
gt_boxes[:, 0:3] = np.dot(gt_boxes[:, 0:3], np.transpose(rot_mat)) gt_boxes[:, 0:3] = np.dot(gt_boxes[:, 0:3], np.transpose(rot_mat))
gt_boxes[:, 6] -= rot_angle gt_boxes[:, 6] -= rot_angle
results['points'] = points results['point_cloud'] = point_cloud
results['gt_boxes'] = gt_boxes results['gt_boxes'] = gt_boxes
return results return results
...@@ -120,24 +121,3 @@ class IndoorRotateData(object): ...@@ -120,24 +121,3 @@ class IndoorRotateData(object):
repr_str = self.__class__.__name__ repr_str = self.__class__.__name__
repr_str += '(dataset_name={})'.format(self.name) repr_str += '(dataset_name={})'.format(self.name)
return repr_str return repr_str
@PIPELINES.register_module()
class PointShuffle(object):
"""Point Shuffle.
Shuffle points.
"""
def __init__(self):
pass
def __call__(self, results):
points = results.get('points')
np.random.shuffle(points)
results['points'] = points
return results
def __repr__(self):
repr_str = self.__class__.__name__
return repr_str
...@@ -49,11 +49,13 @@ class IndoorLoadData(object): ...@@ -49,11 +49,13 @@ class IndoorLoadData(object):
'%06d.npz' % info['point_cloud']['lidar_idx']))['pc'] '%06d.npz' % info['point_cloud']['lidar_idx']))['pc']
if not self.use_color: if not self.use_color:
point_cloud = point_cloud[:, 0:3] # do not use color for now if self.name == 'scannet':
pcl_color = point_cloud[:, 3:6] pcl_color = point_cloud[:, 3:6]
point_cloud = point_cloud[:, 0:3] # do not use color for now
else: else:
point_cloud = point_cloud[:, 0:6] if self.name == 'scannet':
pcl_color = point_cloud[:, 3:6] pcl_color = point_cloud[:, 3:6]
point_cloud = point_cloud[:, 0:6]
point_cloud[:, 3:] = (point_cloud[:, 3:] - point_cloud[:, 3:] = (point_cloud[:, 3:] -
np.array(self.mean_color)) / 256.0 np.array(self.mean_color)) / 256.0
......
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