Laser origins

pcdet/datasets/augmentor/augmentor_utils.py

@@ -5,109 +5,104 @@
 from ...utils import box_utils
 
 
-def random_flip_along_x(gt_boxes, points, return_flip=False, enable=None):
+def random_flip_along(dim, return_flip=False, enable=None):
  """
  Args:
- gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
- points: (M, 3 + C)
+ gt_boxes: (*, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
+ points: (*, 3 + C)
  Returns:
  """
+ assert dim in [0, 1] # corresponds to x-, y-axis respectively
+ other_dim = 1 - dim
  if enable is None:
  enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5])
+
  if enable:
- gt_boxes[:, 1] = -gt_boxes[:, 1]
- gt_boxes[:, 6] = -gt_boxes[:, 6]
- points[:, 1] = -points[:, 1]
-
- if gt_boxes.shape[1] > 7:
- gt_boxes[:, 8] = -gt_boxes[:, 8]
- if return_flip:
- return gt_boxes, points, enable
- return gt_boxes, points
+ def flip_pointlike(points):
+ points[..., other_dim] = -points[..., other_dim]
+ return points
 
+ def flip_boxlike(boxes):
+ boxes[..., other_dim] = -boxes[..., other_dim]
+ boxes[..., 6] = -(boxes[..., 6] + np.pi * dim)
 
-def random_flip_along_y(gt_boxes, points, return_flip=False, enable=None):
- """
- Args:
- gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
- points: (M, 3 + C)
- Returns:
- """
- if enable is None:
- enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5])
- if enable:
- gt_boxes[:, 0] = -gt_boxes[:, 0]
- gt_boxes[:, 6] = -(gt_boxes[:, 6] + np.pi)
- points[:, 0] = -points[:, 0]
+ if boxes.shape[-1] > 7:
+ boxes[..., 7 + other_dim] = -boxes[..., 7 + other_dim]
+
+ return boxes
+
+ tfs = dict(point=flip_pointlike, box=flip_boxlike)
+ else:
+ tfs = dict()
 
- if gt_boxes.shape[1] > 7:
- gt_boxes[:, 7] = -gt_boxes[:, 7]
  if return_flip:
- return gt_boxes, points, enable
- return gt_boxes, points
+ return tfs, enable
+ return tfs
 
 
-def global_rotation(gt_boxes, points, rot_range, return_rot=False, noise_rotation=None):
+def global_rotation(rot_range, return_rot=False, noise_rotation=None):
  """
  Args:
- gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
+ gt_boxes: (*, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
  points: (M, 3 + C),
  rot_range: [min, max]
  Returns:
  """
  if noise_rotation is None: 
  noise_rotation = np.random.uniform(rot_range[0], rot_range[1])
- points = common_utils.rotate_points_along_z(points[np.newaxis, :, :], np.array([noise_rotation]))[0]
- gt_boxes[:, 0:3] = common_utils.rotate_points_along_z(gt_boxes[np.newaxis, :, 0:3], np.array([noise_rotation]))[0]
- gt_boxes[:, 6] += noise_rotation
- if gt_boxes.shape[1] > 7:
- gt_boxes[:, 7:9] = common_utils.rotate_points_along_z(
- np.hstack((gt_boxes[:, 7:9], np.zeros((gt_boxes.shape[0], 1))))[np.newaxis, :, :],
- np.array([noise_rotation])
- )[0][:, 0:2]
+
+ def rotate_pointlike(points):
+ points = common_utils.rotate_points_along_z(points[np.newaxis, :, :], np.array([noise_rotation]))[0]
+ return points
+
+ def rotate_boxlike(boxes):
+ boxes[..., 0:3] = common_utils.rotate_points_along_z(boxes[np.newaxis, ..., 0:3], np.array([noise_rotation]))[0]
+ boxes[..., 6] += noise_rotation
+ if boxes.shape[-1] > 7:
+ boxes[..., 7:9] = common_utils.rotate_points_along_z(
+ np.concatenate((boxes[..., 7:9], np.zeros((*boxes.shape[:-1], 1))), axis=-1)[np.newaxis, ...],
+ np.array([noise_rotation])
+ )[0, ..., 0:2]
+ return boxes
+
+ tfs = dict(point=rotate_pointlike, box=rotate_boxlike)
 
  if return_rot:
- return gt_boxes, points, noise_rotation
- return gt_boxes, points
+ return tfs, noise_rotation
+ return tfs
 
 
-def global_scaling(gt_boxes, points, scale_range, return_scale=False):
+def global_scaling(scale_range, return_scale=False):
  """
  Args:
- gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading]
+ gt_boxes: (*, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
  points: (M, 3 + C),
  scale_range: [min, max]
  Returns:
  """
  if scale_range[1] - scale_range[0] < 1e-3:
- return gt_boxes, points
+ noise_scale = sum(scale_range) / len(scale_range)
+ assert noise_scale == 1.0, (noise_scale, scale_range)
  noise_scale = np.random.uniform(scale_range[0], scale_range[1])
- points[:, :3] *= noise_scale
- gt_boxes[:, :6] *= noise_scale
- if gt_boxes.shape[1] > 7:
- gt_boxes[:, 7:] *= noise_scale
-
- if return_scale:
- return gt_boxes, points, noise_scale
- return gt_boxes, points
 
-def global_scaling_with_roi_boxes(gt_boxes, roi_boxes, points, scale_range, return_scale=False):
- """
- Args:
- gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading]
- points: (M, 3 + C),
- scale_range: [min, max]
- Returns:
- """
- if scale_range[1] - scale_range[0] < 1e-3:
- return gt_boxes, points
- noise_scale = np.random.uniform(scale_range[0], scale_range[1])
- points[:, :3] *= noise_scale
- gt_boxes[:, :6] *= noise_scale
- roi_boxes[:,:, [0,1,2,3,4,5,7,8]] *= noise_scale
+ def scale_pointlike(points):
+ points[:, :3] *= noise_scale
+ return points
+
+ def scale_boxlike(boxes):
+ boxes[..., :6] *= noise_scale
+ if boxes.shape[-1] > 7:
+ boxes[..., 7:9] *= noise_scale
+ return boxes
+
+ if noise_scale != 1.0:
+ tfs = dict(point=scale_pointlike, box=scale_boxlike)
+ else:
+ tfs = {}
+
  if return_scale:
- return gt_boxes,roi_boxes, points, noise_scale
- return gt_boxes, roi_boxes, points
+ return tfs, noise_scale
+ return tfs
 
 
 def random_image_flip_horizontal(image, depth_map, gt_boxes, calib):

pcdet/datasets/augmentor/data_augmentor.py

@@ -56,22 +56,15 @@ def __setstate__(self, d):
  def random_world_flip(self, data_dict=None, config=None):
  if data_dict is None:
  return partial(self.random_world_flip, config=config)
- gt_boxes, points = data_dict['gt_boxes'], data_dict['points']
  for cur_axis in config['ALONG_AXIS_LIST']:
  assert cur_axis in ['x', 'y']
- gt_boxes, points, enable = getattr(augmentor_utils, 'random_flip_along_%s' % cur_axis)(
- gt_boxes, points, return_flip=True
+ cur_dim = ['x', 'y'].index(cur_axis)
+ tfs, enable = augmentor_utils.random_flip_along(
+ cur_dim, return_flip=True
  )
+ common_utils.apply_data_transform(data_dict, tfs)
  data_dict['flip_%s'%cur_axis] = enable
- if 'roi_boxes' in data_dict.keys():
- num_frame, num_rois,dim = data_dict['roi_boxes'].shape
- roi_boxes, _, _ = getattr(augmentor_utils, 'random_flip_along_%s' % cur_axis)(
- data_dict['roi_boxes'].reshape(-1,dim), np.zeros([1,3]), return_flip=True, enable=enable
- )
- data_dict['roi_boxes'] = roi_boxes.reshape(num_frame, num_rois,dim)
 
- data_dict['gt_boxes'] = gt_boxes
- data_dict['points'] = points
  return data_dict
 
  def random_world_rotation(self, data_dict=None, config=None):
@@ -80,36 +73,18 @@ def random_world_rotation(self, data_dict=None, config=None):
  rot_range = config['WORLD_ROT_ANGLE']
  if not isinstance(rot_range, list):
  rot_range = [-rot_range, rot_range]
- gt_boxes, points, noise_rot = augmentor_utils.global_rotation(
- data_dict['gt_boxes'], data_dict['points'], rot_range=rot_range, return_rot=True
+ tfs, noise_rot = augmentor_utils.global_rotation(
+ rot_range=rot_range, return_rot=True
  )
- if 'roi_boxes' in data_dict.keys():
- num_frame, num_rois,dim = data_dict['roi_boxes'].shape
- roi_boxes, _, _ = augmentor_utils.global_rotation(
- data_dict['roi_boxes'].reshape(-1, dim), np.zeros([1, 3]), rot_range=rot_range, return_rot=True, noise_rotation=noise_rot)
- data_dict['roi_boxes'] = roi_boxes.reshape(num_frame, num_rois,dim)
-
- data_dict['gt_boxes'] = gt_boxes
- data_dict['points'] = points
+ common_utils.apply_data_transform(data_dict, tfs)
  data_dict['noise_rot'] = noise_rot
  return data_dict
 
  def random_world_scaling(self, data_dict=None, config=None):
  if data_dict is None:
  return partial(self.random_world_scaling, config=config)
-
- if 'roi_boxes' in data_dict.keys():
- gt_boxes, roi_boxes, points, noise_scale = augmentor_utils.global_scaling_with_roi_boxes(
- data_dict['gt_boxes'], data_dict['roi_boxes'], data_dict['points'], config['WORLD_SCALE_RANGE'], return_scale=True
- )
- data_dict['roi_boxes'] = roi_boxes
- else:
- gt_boxes, points, noise_scale = augmentor_utils.global_scaling(
- data_dict['gt_boxes'], data_dict['points'], config['WORLD_SCALE_RANGE'], return_scale=True
- )
-
- data_dict['gt_boxes'] = gt_boxes
- data_dict['points'] = points
+ tfs, noise_scale = augmentor_utils.global_scaling(scale_range=config['WORLD_SCALE_RANGE'], return_scale=True)
+ common_utils.apply_data_transform(data_dict, tfs)
  data_dict['noise_scale'] = noise_scale
  return data_dict
 
@@ -143,15 +118,12 @@ def random_world_translation(self, data_dict=None, config=None):
  np.random.normal(0, noise_translate_std[2], 1),
  ], dtype=np.float32).T
 
- gt_boxes, points = data_dict['gt_boxes'], data_dict['points']
- points[:, :3] += noise_translate
- gt_boxes[:, :3] += noise_translate
-
- if 'roi_boxes' in data_dict.keys():
- data_dict['roi_boxes'][:, :3] += noise_translate
+ def translate_locationlike(locations):
+ locations[..., :3] += noise_translate
+ return locations
 
- data_dict['gt_boxes'] = gt_boxes
- data_dict['points'] = points
+ tfs = dict(point=translate_locationlike, box=translate_locationlike)
+ common_utils.apply_data_transform(data_dict, tfs)
  data_dict['noise_translate'] = noise_translate
  return data_dict
 

pcdet/datasets/dataset.py

@@ -234,7 +234,7 @@ def collate_batch(batch_list, _unused=False):
  batch_size_ratio = len(val[0])
  val = [i for item in val for i in item]
  ret[key] = np.concatenate(val, axis=0)
- elif key in ['points', 'voxel_coords']:
+ elif key in ['points', 'origins', 'voxel_coords']:
  coors = []
  if isinstance(val[0], list):
  val = [i for item in val for i in item]

pcdet/datasets/processor/data_processor.py

@@ -82,7 +82,8 @@ def mask_points_and_boxes_outside_range(self, data_dict=None, config=None):
 
  if data_dict.get('points', None) is not None:
  mask = common_utils.mask_points_by_range(data_dict['points'], self.point_cloud_range)
- data_dict['points'] = data_dict['points'][mask]
+ tfs = dict(point=lambda x: x[mask])
+ common_utils.apply_data_transform(data_dict, tfs)
 
  if data_dict.get('gt_boxes', None) is not None and config.REMOVE_OUTSIDE_BOXES and self.training:
  mask = box_utils.mask_boxes_outside_range_numpy(
@@ -97,10 +98,9 @@ def shuffle_points(self, data_dict=None, config=None):
  return partial(self.shuffle_points, config=config)
 
  if config.SHUFFLE_ENABLED[self.mode]:
- points = data_dict['points']
- shuffle_idx = np.random.permutation(points.shape[0])
- points = points[shuffle_idx]
- data_dict['points'] = points
+ shuffle_idx = np.random.permutation(data_dict['points'].shape[0])
+ tfs = dict(point=lambda x: x[shuffle_idx])
+ common_utils.apply_data_transform(data_dict, tfs)
 
  return data_dict
 
@@ -208,7 +208,8 @@ def sample_points(self, data_dict=None, config=None):
  extra_choice = np.random.choice(choice, num_points - len(points), replace=False)
  choice = np.concatenate((choice, extra_choice), axis=0)
  np.random.shuffle(choice)
- data_dict['points'] = points[choice]
+ tfs = dict(point=lambda x: x[choice])
+ common_utils.apply_data_transform(data_dict, tfs)
  return data_dict
 
  def calculate_grid_size(self, data_dict=None, config=None):