bugfixed: flip and rot velocity when doing augmentation

pcdet/datasets/augmentor/augmentor_utils.py

@@ -5,7 +5,7 @@ from ...utils import common_utils
 def random_flip_along_x(gt_boxes, points):
     """
     Args:
-        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading]
+        gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
         points: (M, 3 + C)
     Returns:
     """
@@ -14,13 +14,17 @@ def random_flip_along_x(gt_boxes, points):
         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]
+
     return gt_boxes, points
 
 
 def random_flip_along_y(gt_boxes, points):
     """
     Args:
-        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading]
+        gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
         points: (M, 3 + C)
     Returns:
     """
@@ -29,13 +33,17 @@ def random_flip_along_y(gt_boxes, points):
         gt_boxes[:, 0] = -gt_boxes[:, 0]
         gt_boxes[:, 6] = -(gt_boxes[:, 6] + np.pi)
         points[:, 0] = -points[:, 0]
+
+        if gt_boxes.shape[1] > 7:
+            gt_boxes[:, 7] = -gt_boxes[:, 7]
+
     return gt_boxes, points
 
 
 def global_rotation(gt_boxes, points, rot_range):
     """
     Args:
-        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading]
+        gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
         points: (M, 3 + C),
         rot_range: [min, max]
     Returns:
@@ -44,6 +52,12 @@ def global_rotation(gt_boxes, points, rot_range):
     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]
+
     return gt_boxes, points