Skip to content

GitLab

Commit bdce5236 authored by liyinhao's avatar liyinhao
Browse files

divide augment into flip and global_rot_scale

parent 223c7ae6
Hide whitespace changes
Inline Side-by-side
Showing with 121 additions and 49 deletions
mmdet3d/datasets/pipelines/indoor_augment.py
View file @ bdce5236
......@@ -20,18 +20,57 @@ def _rotz(t):
return rot_mat
# TODO: merge outdoor indoor transform
@PIPELINES.register_module()
class IndoorAugment(object):
"""Indoor Augment.
class IndoorFlipData(object):
"""Indoor Flip Data
Augment sunrgbd and scannet data.
Flip point_cloud and groundtruth boxes.
Args:
seed (int): Numpy random seed.
"""
def __init__(self, seed=None):
if seed is not None:
np.random.seed(seed)
def __call__(self, results):
point_cloud = results.get('point_cloud', None)
gt_boxes = results.get('gt_boxes', None)
name = 'scannet' if gt_boxes.shape[1] == 6 else 'sunrgbd'
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 name == 'sunrgbd':
gt_boxes[:, 6] = np.pi - gt_boxes[:, 6]
results['gt_boxes'] = gt_boxes
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]
results['gt_boxes'] = gt_boxes
results['point_cloud'] = point_cloud
return results
def __repr__(self):
repr_str = self.__class__.__name__
return repr_str
# TODO: merge outdoor indoor transform.
# TODO: try transform noise.
@PIPELINES.register_module()
class IndoorGlobalRotScale(object):
"""Indoor Global Rotate Scale.
Augment sunrgbd and scannet data with global rotating and scaling.
Args:
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.
......@@ -41,27 +80,23 @@ class IndoorAugment(object):
def __init__(self,
seed=None,
use_flip=True,
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]):
color_mean=[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
self.color_mean = color_mean
def _rotate_aligned_boxes(self, input_boxes, rot_mat):
"""Rotate Aligned Boxes.
......@@ -101,19 +136,6 @@ class IndoorAugment(object):
gt_boxes = results.get('gt_boxes', None)
name = 'scannet' if gt_boxes.shape[1] == 6 else 'sunrgbd'
if self.use_flip:
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 name == 'sunrgbd':
gt_boxes[:, 6] = np.pi - gt_boxes[:, 6]
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
......@@ -128,12 +150,9 @@ class IndoorAugment(object):
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 = point_cloud[:, 3:6] + self.color_mean
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
......@@ -145,7 +164,7 @@ class IndoorAugment(object):
# 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
point_cloud[:, 3:6] = rgb_color - self.color_mean
if self.use_scale:
# Augment point cloud scale: 0.85x-1.15x
......@@ -165,8 +184,6 @@ class IndoorAugment(object):
def __repr__(self):
repr_str = self.__class__.__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)
......
tests/test_indoor_augment.py
View file @ bdce5236
import numpy as np
from mmdet3d.datasets.pipelines.indoor_augment import IndoorAugment
from mmdet3d.datasets.pipelines.indoor_augment import (IndoorFlipData,
IndoorGlobalRotScale)
def test_indoor_augment():
sunrgbd_augment = IndoorAugment(0, True, True, False, False, True, True)
def test_indoor_flip_data():
sunrgbd_indoor_flip_data = IndoorFlipData(0)
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_flip_data(sunrgbd_results)
sunrgbd_point_cloud = sunrgbd_results.get('point_cloud', None)
sunrgbd_gt_boxes = sunrgbd_results.get('gt_boxes', None)
expected_sunrgbd_point_cloud = np.array(
[[-1.02828765, 3.65790772, 0.1972947, 1.61959505],
[0.39597902, 1.05465031, -0.74920434, 0.673096]])
expected_sunrgbd_gt_boxes = np.array([[
-0.213684, 1.036364, -0.982323, 0.61541, 0.572574, 0.872728, 0.07130739
], [
0.449953, 1.395455, -1.027778, 1.500956, 1.637298, 0.636364, 4.72401624
]])
assert np.allclose(sunrgbd_point_cloud, expected_sunrgbd_point_cloud)
assert np.allclose(sunrgbd_gt_boxes, expected_sunrgbd_gt_boxes)
scannet_indoor_flip_data = IndoorFlipData(0)
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
]])
scannet_results = scannet_indoor_flip_data(scannet_results)
scannet_point_cloud = scannet_results.get('point_cloud', None)
scannet_gt_boxes = scannet_results.get('gt_boxes', None)
expected_scannet_point_cloud = np.array(
[[-1.6110241, 0.16903955, 0.5811581, 0.5989725],
[-1.3978075, -0.42035791, 0.38729519, 0.40510958]])
expected_scannet_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
]])
assert np.allclose(scannet_point_cloud, expected_scannet_point_cloud)
assert np.allclose(scannet_gt_boxes, expected_scannet_gt_boxes)
def test_global_rot_scale():
sunrgbd_augment = IndoorGlobalRotScale(0, True, False, True, True)
sunrgbd_results = dict()
sunrgbd_results['point_cloud'] = np.array(
[[1.02828765e+00, 3.65790772e+00, 1.97294697e-01, 1.61959505e+00],
......@@ -22,23 +77,23 @@ def test_indoor_augment():
sunrgbd_point_cloud = sunrgbd_results.get('point_cloud', None)
sunrgbd_gt_boxes = sunrgbd_results.get('gt_boxes', None)
expected_sunrgbd_point_cloud = np.array(
[[-1.87572197, 3.4384955, 0.2033771, 1.66952557],
[0.15494677, 1.15088388, -0.77230157, 0.69384689]])
[[0.89427376, 3.94489646, 0.21003141, 1.72415094],
[-0.47835783, 1.09972989, -0.79757058, 0.71654893]])
expected_sunrgbd_gt_boxes = np.array([[
-0.45341026, 0.99208554, -1.01260705, 0.63438248, 0.59022589,
0.89963334, -0.15403838
0.17080999, 1.11345031, -1.04573864, 0.65513891, 0.60953755,
0.92906854, 3.01916788
],
[
0.13067981, 1.50574494,
-1.05946338, 1.54722899,
1.68777428, 0.65598247,
4.49867047
-0.55427876, 1.45912611,
-1.09412807, 1.59785293,
1.74299674, 0.67744563,
-1.63354097
]])
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_augment = IndoorGlobalRotScale(
0, True, False, False, True, rot_range=1 / 18)
scannet_results = dict()
scannet_results['point_cloud'] = np.array(
[[1.6110241e+00, -1.6903955e-01, 5.8115810e-01, 5.9897250e-01],
......@@ -52,12 +107,12 @@ def test_indoor_augment():
scannet_point_cloud = scannet_results.get('point_cloud', None)
scannet_gt_boxes = scannet_results.get('gt_boxes', None)
expected_scannet_point_cloud = np.array(
[[-1.61379665, 0.14011924, 0.5811581, 0.5989725],
[-1.39004371, -0.44535946, 0.38729519, 0.40510958]])
[[1.61240576, -0.15530836, 0.5811581, 0.5989725],
[1.39417555, 0.43225122, 0.38729519, 0.40510958]])
expected_scannet_gt_boxes = np.array([[
-0.55030367, -0.49196554, 0.65688646, 0.66436803, 0.61192608, 0.5163464
0.55491157, 0.48676213, 0.65688646, 0.65879754, 0.60584609, 0.5163464
], [
0.06281816, -1.70307376, 0.60348618, 0.66448129, 0.73241497, 0.64667457
-0.04677942, 1.70358975, 0.60348618, 0.65777559, 0.72636927, 0.64667457
]])
assert np.allclose(scannet_point_cloud, expected_scannet_point_cloud)
assert np.allclose(scannet_gt_boxes, expected_scannet_gt_boxes)
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