tools/augmentation.py

import open3d as o3
import numpy as np
import copy

# Set random seed for reproducibility for all libraries
np.random.seed(42)

def apply_noise(pcd, noise_level):
    '''
    This function adds gaussian noise to the point cloud by adding random values to the x, y, and z coordinates of the points.
    Based on https://github.com/MIT-SPARK/TEASER-plusplus/blob/master/examples/teaser_python_ply/teaser_python_ply.py#L7

    Args:
        pcd: Open3D point cloud
        noise_level (float): level of noise to add

    Returns:
        pcloud: Open3D point cloud with added noise
    '''
    pcloud = copy.deepcopy(pcd)
    pcloud_points =  np.transpose(np.asarray(pcloud.points))

    N = pcloud_points.shape[1]
    noise = (np.random.rand(3, N) - 0.5) * 2 * noise_level  # gaussian noise with mean 0 and std = noise_level
    pcloud_points += noise
    
    pcloud.points = o3.utility.Vector3dVector(pcloud_points.T)

    return pcloud

def add_outliers(pcd, outlier_level, outlier_lowerbound, outlier_upperbound):
    '''
    This function adds points to a given point cloud that have no counterpart in the other point cloud (outliers).
    The outliers are generated by adding random values between the lower and upper bound to randomly chosen points in the original point cloud.
    Based on https://github.com/MIT-SPARK/TEASER-plusplus/blob/master/examples/teaser_python_ply/teaser_python_ply.py#L7

    Args:
        pcd: Open3D point cloud
        outlier_level: level of outliers to add
        outlier_lowerbound: lower bound for the random values to add to the original points
        outlier_upperbound: upper bound for the random values to add to the original points
    
    Returns:
        pcloud: Open3D point cloud with added outliers
    '''
    pcloud = copy.deepcopy(pcd)
    pcloud_points = np.asarray(pcloud.points) # (number_of_points, 3)

    n_outliers = int(outlier_level/50 * pcloud_points.shape[0])
    outliers_amounts = outlier_lowerbound + np.random.rand(n_outliers) * (outlier_upperbound - outlier_lowerbound) # (outlier_upperbound - outlier_lowerbound) =  the distance of the outliers from the original points.
    outliers_amounts = outliers_amounts[:, np.newaxis]  # reshape to (n_outliers, 1)

    original_indices = np.random.randint(0, pcloud_points.shape[0], size=n_outliers) # (n_outliers,)
    outlier_points = pcloud_points[original_indices] + outliers_amounts
    new_pcd = np.concatenate((pcloud_points, outlier_points))
    pcloud.points = o3.utility.Vector3dVector(new_pcd)

    return pcloud


def apply_occlusion(pc, radius_factor):
    '''
    This function removes points from the point cloud to simulate occlusion.
    The points are removed based on the distance from the camera, not randomly.
    Source: https://towardsdatascience.com/3d-data-processing-with-open3d-c3062aadc72e

    Args:
        pc (open3d.geometry.PointCloud): Point cloud
        radius_factor (float): Factor to determine the radius for hidden point removal
    
    Returns:
        modified_pc (open3d.geometry.PointCloud): Point cloud with occlusion
        pt_map (np.array): Indices of the points that are not hidden
    '''
    # Get the diameter of the target point cloud (max bound - min bound)
    diameter = np.linalg.norm(np.asarray(pc.get_min_bound()) - np.asarray(pc.get_max_bound()))

    # Parameters for hidden point removal
    camera = [0, 0, diameter]
    radius = diameter * radius_factor  # bigger radius removes fewer points

    _, pt_map = pc.hidden_point_removal(camera, radius) # returns a tuple of the point cloud and the indices of the points that are not hidden

    # Applly the hidden point removal to the target point cloud
    modified_pc = pc.select_by_index(pt_map)

    return modified_pc, pt_map