Skip to content

GitLab

Commit 223c7ae6 authored by liyinhao's avatar liyinhao
Browse files

finish indoor augment

parent 24deecb3
Hide whitespace changes
Inline Side-by-side
Showing with 165 additions and 132 deletions
mmdet3d/datasets/pipelines/indoor_augment.py
View file @ 223c7ae6
...@@ -20,99 +20,143 @@ def _rotz(t): ...@@ -20,99 +20,143 @@ def _rotz(t):
return rot_mat return rot_mat
def _rotate_aligned_boxes(input_boxes, rot_mat): # TODO: merge outdoor indoor transform
centers, lengths = input_boxes[:, 0:3], input_boxes[:, 3:6]
new_centers = np.dot(centers, np.transpose(rot_mat))
dx, dy = lengths[:, 0] / 2.0, lengths[:, 1] / 2.0
new_x = np.zeros((dx.shape[0], 4))
new_y = np.zeros((dx.shape[0], 4))
for i, crnr in enumerate([(-1, -1), (1, -1), (1, 1), (-1, 1)]):
crnrs = np.zeros((dx.shape[0], 3))
crnrs[:, 0] = crnr[0] * dx
crnrs[:, 1] = crnr[1] * dy
crnrs = np.dot(crnrs, np.transpose(rot_mat))
new_x[:, i] = crnrs[:, 0]
new_y[:, i] = crnrs[:, 1]
new_dx = 2.0 * np.max(new_x, 1)
new_dy = 2.0 * np.max(new_y, 1)
new_lengths = np.stack((new_dx, new_dy, lengths[:, 2]), axis=1)
return np.concatenate([new_centers, new_lengths], axis=1)
@PIPELINES.register_module()
class IndoorFlipData(object):
"""Indoor Flip Data
Flip point_cloud and groundtruth boxes.
Args:
name (str): name of the dataset.
"""
def __init__(self, name):
assert name in ['scannet', 'sunrgbd']
self.name = name
def __call__(self, results):
point_cloud = results.get('point_cloud', None)
gt_boxes = results.get('gt_boxes', None)
if np.random.random() > 0.5:
# Flipping along the YZ plane
point_cloud[:, 0] = -1 * point_cloud[:, 0]
gt_boxes[:, 0] = -1 * gt_boxes[:, 0]
if self.name == 'sunrgbd':
gt_boxes[:, 6] = np.pi - gt_boxes[:, 6]
results['gt_boxes'] = gt_boxes
if self.name == 'scannet' and np.random.random() > 0.5:
# Flipping along the XZ plane
point_cloud[:, 1] = -1 * point_cloud[:, 1]
gt_boxes[:, 1] = -1 * gt_boxes[:, 1]
results['gt_boxes'] = gt_boxes
results['point_cloud'] = point_cloud
return results
def __repr__(self):
repr_str = self.__class__.__name__
repr_str += '(dataset_name={})'.format(self.name)
return repr_str
@PIPELINES.register_module() @PIPELINES.register_module()
class IndoorRotateData(object): class IndoorAugment(object):
"""Indoor Rotate Data """Indoor Augment.
Rotate point_cloud and groundtruth boxes. Augment sunrgbd and scannet data.
Args: Args:
name (str): name of the dataset. seed (int): Numpy random seed.
use_flip (bool): Whether to use flip.
use_rotate (bool): Whether to use rotate.
use_shuffle (bool): Whether to use shuffle.
use_color (bool): Whether to use color.
use_height (bool): Whether to use height.
rot_range (float): Range of rotation.
scale_range (float): Range of scale.
(List[float]): Mean color of the point cloud.
""" """
def __init__(self, name): def __init__(self,
assert name in ['scannet', 'sunrgbd'] seed=None,
self.name = name use_flip=True,
self.rot_range = np.pi / 3 use_rotate=True,
use_shuffle=True,
use_color=False,
use_scale=True,
use_height=True,
rot_range=1 / 3,
scale_range=0.3,
mean_color=[0.5, 0.5, 0.5]):
if seed is not None:
np.random.seed(seed)
self.use_flip = use_flip
self.use_rotate = use_rotate
self.use_shuffle = use_shuffle
self.use_color = use_color
self.use_scale = use_scale
self.use_height = use_height
self.rot_range = rot_range
self.scale_range = scale_range
self.mean_color = mean_color
def _rotate_aligned_boxes(self, input_boxes, rot_mat):
"""Rotate Aligned Boxes.
Rotate function for the aligned boxes.
Args:
input_boxes (ndarray): 3D boxes.
rot_mat (ndarray): Rotation matrix.
Returns:
rotated_boxes (ndarry): 3D boxes after rotation.
"""
centers, lengths = input_boxes[:, 0:3], input_boxes[:, 3:6]
new_centers = np.dot(centers, np.transpose(rot_mat))
dx, dy = lengths[:, 0] / 2.0, lengths[:, 1] / 2.0
new_x = np.zeros((dx.shape[0], 4))
new_y = np.zeros((dx.shape[0], 4))
for i, crnr in enumerate([(-1, -1), (1, -1), (1, 1), (-1, 1)]):
crnrs = np.zeros((dx.shape[0], 3))
crnrs[:, 0] = crnr[0] * dx
crnrs[:, 1] = crnr[1] * dy
crnrs = np.dot(crnrs, np.transpose(rot_mat))
new_x[:, i] = crnrs[:, 0]
new_y[:, i] = crnrs[:, 1]
new_dx = 2.0 * np.max(new_x, 1)
new_dy = 2.0 * np.max(new_y, 1)
new_lengths = np.stack((new_dx, new_dy, lengths[:, 2]), axis=1)
return np.concatenate([new_centers, new_lengths], axis=1)
def __call__(self, results): def __call__(self, results):
point_cloud = results.get('point_cloud', 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() * name = 'scannet' if gt_boxes.shape[1] == 6 else 'sunrgbd'
self.rot_range) - self.rot_range / 2 # -30 ~ +30 degree
rot_mat = _rotz(rot_angle) if self.use_flip:
point_cloud[:, 0:3] = np.dot(point_cloud[:, 0:3], if np.random.random() > 0.5:
np.transpose(rot_mat)) # Flipping along the YZ plane
point_cloud[:, 0] = -1 * point_cloud[:, 0]
if self.name == 'scannet': gt_boxes[:, 0] = -1 * gt_boxes[:, 0]
gt_boxes = _rotate_aligned_boxes(gt_boxes, rot_mat) if name == 'sunrgbd':
else: gt_boxes[:, 6] = np.pi - gt_boxes[:, 6]
gt_boxes[:, 0:3] = np.dot(gt_boxes[:, 0:3], np.transpose(rot_mat))
gt_boxes[:, 6] -= rot_angle if name == 'scannet' and np.random.random() > 0.5:
# Flipping along the XZ plane
point_cloud[:, 1] = -1 * point_cloud[:, 1]
gt_boxes[:, 1] = -1 * gt_boxes[:, 1]
if self.use_rotate:
rot_angle = (np.random.random() * self.rot_range * np.pi
) - np.pi * self.rot_range / 2 # -30 ~ +30 degree
rot_mat = _rotz(rot_angle)
point_cloud[:, 0:3] = np.dot(point_cloud[:, 0:3],
np.transpose(rot_mat))
if name == 'scannet':
gt_boxes = self._rotate_aligned_boxes(gt_boxes, rot_mat)
else:
gt_boxes[:, 0:3] = np.dot(gt_boxes[:, 0:3],
np.transpose(rot_mat))
gt_boxes[:, 6] -= rot_angle
if self.use_shuffle:
np.random.shuffle(point_cloud)
# Augment RGB color
if self.use_color:
rgb_color = point_cloud[:, 3:6] + self.mean_color
rgb_color *= (1 + 0.4 * np.random.random(3) - 0.2
) # brightness change for each channel
rgb_color += (0.1 * np.random.random(3) - 0.05
) # color shift for each channel
rgb_color += np.expand_dims(
(0.05 * np.random.random(point_cloud.shape[0]) - 0.025),
-1) # jittering on each pixel
rgb_color = np.clip(rgb_color, 0, 1)
# randomly drop out 30% of the points' colors
rgb_color *= np.expand_dims(
np.random.random(point_cloud.shape[0]) > 0.3, -1)
point_cloud[:, 3:6] = rgb_color - self.mean_color
if self.use_scale:
# Augment point cloud scale: 0.85x-1.15x
scale_ratio = np.random.random(
) * self.scale_range + 1 - self.scale_range / 2
scale_ratio = np.expand_dims(np.tile(scale_ratio, 3), 0)
point_cloud[:, 0:3] *= scale_ratio
gt_boxes[:, 0:3] *= scale_ratio
gt_boxes[:, 3:6] *= scale_ratio
if self.use_height:
point_cloud[:, -1] *= scale_ratio[0, 0]
results['point_cloud'] = point_cloud results['point_cloud'] = point_cloud
results['gt_boxes'] = gt_boxes results['gt_boxes'] = gt_boxes
...@@ -120,5 +164,10 @@ class IndoorRotateData(object): ...@@ -120,5 +164,10 @@ class IndoorRotateData(object):
def __repr__(self): def __repr__(self):
repr_str = self.__class__.__name__ repr_str = self.__class__.__name__
repr_str += '(dataset_name={})'.format(self.name) repr_str += '(use_rotate={})'.format(self.use_rotate)
repr_str += '(use_flip={})'.format(self.use_flip)
repr_str += '(use_rotate={})'.format(self.use_shuffle)
repr_str += '(use_color={})'.format(self.use_color)
repr_str += '(use_scale={})'.format(self.use_scale)
repr_str += '(use_height={})'.format(self.use_height)
return repr_str return repr_str
tests/test_indoor_augment.py
View file @ 223c7ae6
import numpy as np import numpy as np
from mmdet3d.datasets.pipelines.indoor_augment import (IndoorFlipData, from mmdet3d.datasets.pipelines.indoor_augment import IndoorAugment
IndoorRotateData)
def test_indoor_flip_data(): def test_indoor_augment():
sunrgbd_flip_data = IndoorFlipData('sunrgbd') sunrgbd_augment = IndoorAugment(0, True, True, False, False, True, True)
sunrgbd_results = dict() sunrgbd_results = dict()
sunrgbd_results['point_cloud'] = np.array( sunrgbd_results['point_cloud'] = np.array(
[[1.02828765e+00, 3.65790772e+00, 1.97294697e-01, 1.61959505e+00], [[1.02828765e+00, 3.65790772e+00, 1.97294697e-01, 1.61959505e+00],
...@@ -18,13 +17,28 @@ def test_indoor_flip_data(): ...@@ -18,13 +17,28 @@ def test_indoor_flip_data():
1.500956, 1.637298, 0.636364, 1.500956, 1.637298, 0.636364,
-1.58242359 -1.58242359
]]) ]])
sunrgbd_results = sunrgbd_flip_data(sunrgbd_results)
sunrgbd_results = sunrgbd_augment(sunrgbd_results)
sunrgbd_point_cloud = sunrgbd_results.get('point_cloud', None) sunrgbd_point_cloud = sunrgbd_results.get('point_cloud', None)
sunrgbd_gt_boxes = sunrgbd_results.get('gt_boxes', None) sunrgbd_gt_boxes = sunrgbd_results.get('gt_boxes', None)
assert sunrgbd_point_cloud.shape == (2, 4) expected_sunrgbd_point_cloud = np.array(
assert sunrgbd_gt_boxes.shape == (2, 7) [[-1.87572197, 3.4384955, 0.2033771, 1.66952557],
[0.15494677, 1.15088388, -0.77230157, 0.69384689]])
scannet_flip_data = IndoorFlipData('scannet') expected_sunrgbd_gt_boxes = np.array([[
-0.45341026, 0.99208554, -1.01260705, 0.63438248, 0.59022589,
0.89963334, -0.15403838
],
[
0.13067981, 1.50574494,
-1.05946338, 1.54722899,
1.68777428, 0.65598247,
4.49867047
]])
assert np.allclose(sunrgbd_point_cloud, expected_sunrgbd_point_cloud)
assert np.allclose(sunrgbd_gt_boxes, expected_sunrgbd_gt_boxes)
scannet_augment = IndoorAugment(
0, True, True, True, False, False, False, rot_range=1 / 18)
scannet_results = dict() scannet_results = dict()
scannet_results['point_cloud'] = np.array( scannet_results['point_cloud'] = np.array(
[[1.6110241e+00, -1.6903955e-01, 5.8115810e-01, 5.9897250e-01], [[1.6110241e+00, -1.6903955e-01, 5.8115810e-01, 5.9897250e-01],
...@@ -34,46 +48,16 @@ def test_indoor_flip_data(): ...@@ -34,46 +48,16 @@ def test_indoor_flip_data():
], [ ], [
-0.03226406, 1.70392646, 0.60348618, 0.65165804, 0.72084366, 0.64667457 -0.03226406, 1.70392646, 0.60348618, 0.65165804, 0.72084366, 0.64667457
]]) ]])
scannet_results = scannet_flip_data(scannet_results) scannet_results = scannet_augment(scannet_results)
scannet_point_cloud = scannet_results.get('point_cloud', None) scannet_point_cloud = scannet_results.get('point_cloud', None)
scannet_gt_boxes = scannet_results.get('gt_boxes', None) scannet_gt_boxes = scannet_results.get('gt_boxes', None)
assert scannet_point_cloud.shape == (2, 4) expected_scannet_point_cloud = np.array(
assert scannet_gt_boxes.shape == (2, 6) [[-1.61379665, 0.14011924, 0.5811581, 0.5989725],
[-1.39004371, -0.44535946, 0.38729519, 0.40510958]])
expected_scannet_gt_boxes = np.array([[
def test_indoor_rotate_data(): -0.55030367, -0.49196554, 0.65688646, 0.66436803, 0.61192608, 0.5163464
sunrgbd_indoor_rotate_data = IndoorRotateData('sunrgbd')
sunrgbd_results = dict()
sunrgbd_results['point_cloud'] = np.array(
[[1.02828765e+00, 3.65790772e+00, 1.97294697e-01, 1.61959505e+00],
[-3.95979017e-01, 1.05465031e+00, -7.49204338e-01, 6.73096001e-01]])
sunrgbd_results['gt_boxes'] = np.array([[
0.213684, 1.036364, -0.982323, 0.61541, 0.572574, 0.872728, 3.07028526
],
[
-0.449953, 1.395455, -1.027778,
1.500956, 1.637298, 0.636364,
-1.58242359
]])
sunrgbd_results = sunrgbd_indoor_rotate_data(sunrgbd_results)
sunrgbd_point_cloud = sunrgbd_results.get('point_cloud', None)
sunrgbd_gt_boxes = sunrgbd_results.get('gt_boxes', None)
assert sunrgbd_point_cloud.shape == (2, 4)
assert sunrgbd_gt_boxes.shape == (2, 7)
scannet_indoor_rotate_data = IndoorRotateData('scannet')
scannet_results = dict()
scannet_results['point_cloud'] = 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['gt_boxes'] = np.array([[
0.55903838, 0.48201692, 0.65688646, 0.65370704, 0.60029864, 0.5163464
], [ ], [
-0.03226406, 1.70392646, 0.60348618, 0.65165804, 0.72084366, 0.64667457 0.06281816, -1.70307376, 0.60348618, 0.66448129, 0.73241497, 0.64667457
]]) ]])
scannet_results = scannet_indoor_rotate_data(scannet_results) assert np.allclose(scannet_point_cloud, expected_scannet_point_cloud)
scannet_point_cloud = scannet_results.get('point_cloud', None) assert np.allclose(scannet_gt_boxes, expected_scannet_gt_boxes)
scannet_gt_boxes = scannet_results.get('gt_boxes', None)
assert scannet_point_cloud.shape == (2, 4)
assert scannet_gt_boxes.shape == (2, 6)
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