Added new augmentation methods and the option of augmentation by addi…

…tion instead of replacement (#653)

* NOTEBOOK FOR DATA AUGMENTATION

* SMALL SAMPLE DATASET

* TRANSLACTIONS GLOBAL AND LOCAL

* LOCAL AUGMENTATIONS FINISHED

* ADJUSTED DATA FORMATS TO PCDET YAML COMPATIBILITY

* BUGFIX

* ADDED TITLES

* ADAPTED OPENPCDET TO FIT NEW AUGMENTATIONS

* SAVE IMAGES OF THE DATA AUGMENTATION EXAMPLES

* LOCAL TRANSFORMATION RANDOM VALUES VARY WITH OBJECT AND NOT SCENE

* FIXED YAML

* FIXED BUGS IN PARAMETERS

* FIX IN LOCAL BOXES

* NEW BOX FILTER

* FRUSTUM DROPOUT

* FRUSTUM DROPOUT IMPLEMENTED ON OPENPCDET

* REMOVE FIGURES FROM VERSION CONTROL

* MORE EFFICIENT LOCAL METHOD

* LOCAL TRANSFORMATION WITH NUMPY SPED UP

* BUGFIX ON AUGMENTATION FUNCTION

* GET PREDICTED LABELS

* AUGMENTATION DONE BY ADDITION

* RANDOM SELECTION OF REMAINIG FRAMES

* LOGGER WITH TRANSFORMATION DETAILS

* CHANGED COUNTER TO SAMPLE IDX SELECTION METHOD

* BUGFIX ON KITTY_DATABASE

* BUGFIX

* BUGFIX

* FIX ON GENERATOR

* DEBUG LOGS TO FILE

* FIXED LOGGER

* SAVE POINT CLOUDS TO OBJ FILE

* UPDATED NOTEBOOK

* CLEANUP

* .gitignore is now working

* fixed .gitignore

* REMOVED YAML FILES FROM .gitignore

* parse results to a csv

* yaml changed

* fixed gitignore

* FIXED BUG THAT DUPLICATED THE TEST SET

* small fixes on yaml

* fix gitignore

* updated generate_sets.py

* fixed csv_parser.py

* Delete output directory

* Delete output directory

* saving figures with better resolution

* implemented a way to chose to augment the raw pcs

* updated csv_parser.py

* fixed yamls to include the new selection system

* notebook updated

* mosaic added

* new models

* fixed gitignore

* fixed gitignore

* applied changes

* Delete data_aumentation.ipynb

* Delete pointpillar_fullaug.yaml

* updated gitignore

* bugfix

* updated training and notebook

* updated gitignore

* cleanup code

* updated yamls

* reset

* new augmentation methods

* example conf for pointpillars

* updated example config

* bugfix on the config file

* fixed get_points_in_box function

* fixed local frustum dropout

* fixed pointpillar_newaugs.yaml

* bugfixed: import math for some augmentations

* increased margin in point selection for local augmentation

* bugfixed: velocity should not be translated for translation augs. Add warning to check the correctness before using

* bugfixed: add back the aug of random_image_flip

Co-authored-by: Shaoshuai Shi <shaoshuaics@gmail.com>

pcdet/datasets/augmentor/augmentor_utils.py

@@ -1,5 +1,5 @@
-import copy
 import numpy as np
+import math
 
 from ...utils import common_utils
 
@@ -22,7 +22,6 @@ def random_flip_along_x(gt_boxes, points):
 
     return gt_boxes, points
 
-
 def random_flip_along_y(gt_boxes, points):
     """
     Args:
@@ -41,7 +40,6 @@ def random_flip_along_y(gt_boxes, points):
 
     return gt_boxes, points
 
-
 def global_rotation(gt_boxes, points, rot_range):
     """
     Args:
@@ -62,7 +60,6 @@ def global_rotation(gt_boxes, points, rot_range):
 
     return gt_boxes, points
 
-
 def global_scaling(gt_boxes, points, scale_range):
     """
     Args:
@@ -76,6 +73,7 @@ def global_scaling(gt_boxes, points, scale_range):
     noise_scale = np.random.uniform(scale_range[0], scale_range[1])
     points[:, :3] *= noise_scale
     gt_boxes[:, :6] *= noise_scale
+
     return gt_boxes, points
 
 def random_image_flip_horizontal(image, depth_map, gt_boxes, calib):
@@ -115,4 +113,358 @@ def random_image_flip_horizontal(image, depth_map, gt_boxes, calib):
         aug_depth_map = depth_map
         aug_gt_boxes = gt_boxes
 
-    return aug_image, aug_depth_map, aug_gt_boxes
+    return aug_image, aug_depth_map, aug_gt_boxes
+
+def random_translation_along_x(gt_boxes, points, offset_range):
+    """
+    Args:
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
+        points: (M, 3 + C),
+        offset_range: [min max]]
+    Returns:
+    """
+    offset = np.random.uniform(offset_range[0], offset_range[1])
+
+    points[:, 0] += offset
+    gt_boxes[:, 0] += offset
+
+    # if gt_boxes.shape[1] > 7:
+    #     gt_boxes[:, 7] += offset
+
+    return gt_boxes, points
+
+def random_translation_along_y(gt_boxes, points, offset_range):
+    """
+    Args:
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
+        points: (M, 3 + C),
+        offset_range: [min max]]
+    Returns:
+    """
+    offset = np.random.uniform(offset_range[0], offset_range[1])
+
+    points[:, 1] += offset
+    gt_boxes[:, 1] += offset
+
+    # if gt_boxes.shape[1] > 8:
+    #     gt_boxes[:, 8] += offset
+
+    return gt_boxes, points
+
+def random_translation_along_z(gt_boxes, points, offset_range):
+    """
+    Args:
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
+        points: (M, 3 + C),
+        offset_range: [min max]]
+    Returns:
+    """
+    offset = np.random.uniform(offset_range[0], offset_range[1])
+
+    points[:, 2] += offset
+    gt_boxes[:, 2] += offset
+
+    return gt_boxes, points
+
+def random_local_translation_along_x(gt_boxes, points, offset_range):
+    """
+    Args:
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
+        points: (M, 3 + C),
+        offset_range: [min max]]
+    Returns:
+    """
+    # augs = {}
+    for idx, box in enumerate(gt_boxes):
+        offset = np.random.uniform(offset_range[0], offset_range[1])
+        # augs[f'object_{idx}'] = offset
+        points_in_box, mask = get_points_in_box(points, box)
+        points[mask, 0] += offset
+
+        gt_boxes[idx, 0] += offset
+
+        # if gt_boxes.shape[1] > 7:
+        #     gt_boxes[idx, 7] += offset
+
+    return gt_boxes, points
+
+def random_local_translation_along_y(gt_boxes, points, offset_range):
+    """
+    Args:
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
+        points: (M, 3 + C),
+        offset_range: [min max]]
+    Returns:
+    """
+    # augs = {}
+    for idx, box in enumerate(gt_boxes):
+        offset = np.random.uniform(offset_range[0], offset_range[1])
+        # augs[f'object_{idx}'] = offset
+        points_in_box, mask = get_points_in_box(points, box)
+        points[mask, 1] += offset
+
+        gt_boxes[idx, 1] += offset
+
+        # if gt_boxes.shape[1] > 8:
+        #     gt_boxes[idx, 8] += offset
+
+    return gt_boxes, points
+
+def random_local_translation_along_z(gt_boxes, points, offset_range):
+    """
+    Args:
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
+        points: (M, 3 + C),
+        offset_range: [min max]]
+    Returns:
+    """
+    # augs = {}
+    for idx, box in enumerate(gt_boxes):
+        offset = np.random.uniform(offset_range[0], offset_range[1])
+        # augs[f'object_{idx}'] = offset
+        points_in_box, mask = get_points_in_box(points, box)
+        points[mask, 2] += offset
+
+        gt_boxes[idx, 2] += offset
+
+    return gt_boxes, points
+
+def global_frustum_dropout_top(gt_boxes, points, intensity_range):
+    """
+    Args:
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]],
+        points: (M, 3 + C),
+        intensity: [min, max]
+    Returns:
+    """
+    intensity = np.random.uniform(intensity_range[0], intensity_range[1])
+
+    threshold = np.max(points[:, 2]) - intensity * (np.max(points[:, 2]) - np.min(points[:, 2]))
+    points = points[points[:,2] < threshold]
+    gt_boxes = gt_boxes[gt_boxes[:,2] < threshold]
+
+    return gt_boxes, points
+
+def global_frustum_dropout_bottom(gt_boxes, points, intensity_range):
+    """
+    Args:
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]],
+        points: (M, 3 + C),
+        intensity: [min, max]
+    Returns:
+    """
+    intensity = np.random.uniform(intensity_range[0], intensity_range[1])
+
+    threshold = np.min(points[:, 2]) + intensity * (np.max(points[:, 2]) - np.min(points[:, 2]))
+    points = points[points[:,2] > threshold]
+    gt_boxes = gt_boxes[gt_boxes[:,2] > threshold]
+
+    return gt_boxes, points
+
+def global_frustum_dropout_left(gt_boxes, points, intensity_range):
+    """
+    Args:
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]],
+        points: (M, 3 + C),
+        intensity: [min, max]
+    Returns:
+    """
+    intensity = np.random.uniform(intensity_range[0], intensity_range[1])
+
+    threshold = np.max(points[:, 1]) - intensity * (np.max(points[:, 1]) - np.min(points[:, 1]))
+    points = points[points[:,1] < threshold]
+    gt_boxes = gt_boxes[gt_boxes[:,1] < threshold]
+
+    return gt_boxes, points
+
+def global_frustum_dropout_right(gt_boxes, points, intensity_range):
+    """
+    Args:
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]],
+        points: (M, 3 + C),
+        intensity: [min, max]
+    Returns:
+    """
+    intensity = np.random.uniform(intensity_range[0], intensity_range[1])
+
+    threshold = np.min(points[:, 1]) + intensity * (np.max(points[:, 1]) - np.min(points[:, 1]))
+    points = points[points[:,1] > threshold]
+    gt_boxes = gt_boxes[gt_boxes[:,1] > threshold]
+
+    return gt_boxes, points
+
+def local_scaling(gt_boxes, points, scale_range):
+    """
+    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
+
+    # augs = {}
+    for idx, box in enumerate(gt_boxes):
+        noise_scale = np.random.uniform(scale_range[0], scale_range[1])
+        # augs[f'object_{idx}'] = noise_scale
+        points_in_box, mask = get_points_in_box(points, box)
+
+        # tranlation to axis center
+        points[mask, 0] -= box[0]
+        points[mask, 1] -= box[1]
+        points[mask, 2] -= box[2]
+
+        # apply scaling
+        points[mask, :3] *= noise_scale
+
+        # tranlation back to original position
+        points[mask, 0] += box[0]
+        points[mask, 1] += box[1]
+        points[mask, 2] += box[2]
+
+        gt_boxes[idx, 3:6] *= noise_scale
+    return gt_boxes, points
+
+
+def local_rotation(gt_boxes, points, rot_range):
+    """
+    Args:
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
+        points: (M, 3 + C),
+        rot_range: [min, max]
+    Returns:
+    """
+    # augs = {}
+    for idx, box in enumerate(gt_boxes):
+        noise_rotation = np.random.uniform(rot_range[0], rot_range[1])
+        # augs[f'object_{idx}'] = noise_rotation
+        points_in_box, mask = get_points_in_box(points, box)
+
+        centroid_x = box[0]
+        centroid_y = box[1]
+        centroid_z = box[2]
+
+        # tranlation to axis center
+        points[mask, 0] -= centroid_x
+        points[mask, 1] -= centroid_y
+        points[mask, 2] -= centroid_z
+        box[0] -= centroid_x
+        box[1] -= centroid_y
+        box[2] -= centroid_z
+
+        # apply rotation
+        points[mask, :] = common_utils.rotate_points_along_z(points[np.newaxis, mask, :], np.array([noise_rotation]))[0]
+        box[0:3] = common_utils.rotate_points_along_z(box[np.newaxis, np.newaxis, 0:3], np.array([noise_rotation]))[0][0]
+
+        # tranlation back to original position
+        points[mask, 0] += centroid_x
+        points[mask, 1] += centroid_y
+        points[mask, 2] += centroid_z
+        box[0] += centroid_x
+        box[1] += centroid_y
+        box[2] += centroid_z
+
+        gt_boxes[idx, 6] += noise_rotation
+        if gt_boxes.shape[1] > 8:
+            gt_boxes[idx, 7:9] = common_utils.rotate_points_along_z(
+                np.hstack((gt_boxes[idx, 7:9], np.zeros((gt_boxes.shape[0], 1))))[np.newaxis, :, :],
+                np.array([noise_rotation])
+            )[0][:, 0:2]
+
+    return gt_boxes, points
+
+def local_frustum_dropout_top(gt_boxes, points, intensity_range):
+    """
+    Args:
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]],
+        points: (M, 3 + C),
+        intensity: [min, max]
+    Returns:
+    """
+    for idx, box in enumerate(gt_boxes):
+        x, y, z, dx, dy, dz = box[0], box[1], box[2], box[3], box[4], box[5]
+
+        intensity = np.random.uniform(intensity_range[0], intensity_range[1])
+        points_in_box, mask = get_points_in_box(points, box)
+        threshold = (z + dz/2) - intensity * dz
+
+        points = points[np.logical_not(np.logical_and(mask, points[:,2] >= threshold))]
+
+    return gt_boxes, points
+
+def local_frustum_dropout_bottom(gt_boxes, points, intensity_range):
+    """
+    Args:
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]],
+        points: (M, 3 + C),
+        intensity: [min, max]
+    Returns:
+    """
+    for idx, box in enumerate(gt_boxes):
+        x, y, z, dx, dy, dz = box[0], box[1], box[2], box[3], box[4], box[5]
+
+        intensity = np.random.uniform(intensity_range[0], intensity_range[1])
+        points_in_box, mask = get_points_in_box(points, box)
+        threshold = (z - dz/2) + intensity * dz
+
+        points = points[np.logical_not(np.logical_and(mask, points[:,2] <= threshold))]
+
+    return gt_boxes, points
+
+def local_frustum_dropout_left(gt_boxes, points, intensity_range):
+    """
+    Args:
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]],
+        points: (M, 3 + C),
+        intensity: [min, max]
+    Returns:
+    """
+    for idx, box in enumerate(gt_boxes):
+        x, y, z, dx, dy, dz = box[0], box[1], box[2], box[3], box[4], box[5]
+
+        intensity = np.random.uniform(intensity_range[0], intensity_range[1])
+        points_in_box, mask = get_points_in_box(points, box)
+        threshold = (y + dy/2) - intensity * dy
+
+        points = points[np.logical_not(np.logical_and(mask, points[:,1] >= threshold))]
+
+    return gt_boxes, points
+
+def local_frustum_dropout_right(gt_boxes, points, intensity_range):
+    """
+    Args:
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]],
+        points: (M, 3 + C),
+        intensity: [min, max]
+    Returns:
+    """
+    for idx, box in enumerate(gt_boxes):
+        x, y, z, dx, dy, dz = box[0], box[1], box[2], box[3], box[4], box[5]
+
+        intensity = np.random.uniform(intensity_range[0], intensity_range[1])
+        points_in_box, mask = get_points_in_box(points, box)
+        threshold = (y - dy/2) + intensity * dy
+
+        points = points[np.logical_not(np.logical_and(mask, points[:,1] <= threshold))]
+
+    return gt_boxes, points
+
+def get_points_in_box(points, gt_box):
+    x, y, z = points[:,0], points[:,1], points[:,2]
+    cx, cy, cz = gt_box[0], gt_box[1], gt_box[2]
+    dx, dy, dz, rz = gt_box[3], gt_box[4], gt_box[5], gt_box[6]
+    shift_x, shift_y, shift_z = x - cx, y - cy, z - cz
+
+    MARGIN = 1e-1
+    cosa, sina = math.cos(-rz), math.sin(-rz)
+    local_x = shift_x * cosa + shift_y * (-sina)
+    local_y = shift_x * sina + shift_y * cosa
+
+    mask = np.logical_and(abs(shift_z) <= dz / 2.0, \
+             np.logical_and(abs(local_x) <= dx / 2.0 + MARGIN, \
+                 abs(local_y) <= dy / 2.0 + MARGIN ))
+
+    points = points[mask]
+
+    return points, mask