batch_augment.py

import open3d as o3d
import numpy as np
import os, json
from math import cos, sin
from tqdm import tqdm
import argparse

'''
# reference
http://www.open3d.org/docs/release/python_example/io/index.html
http://www.open3d.org/docs/release/tutorial/geometry/pointcloud.html#Paint-point-cloud
http://www.open3d.org/docs/latest/tutorial/Advanced/multiway_registration.html
http://www.open3d.org/docs/0.9.0/tutorial/Basic/working_with_numpy.html
'''

################################
# Check if a point in bbox or not
################################
def check_point_within_bbox(point, x_centroid, y_centroid, z_centroid,
                                length, width, height, rotation):

    # Calculate half-dimensions for convenience
    half_length = length / 2
    half_width = width / 2
    half_height = height / 2
    
    # Create rotation matrix
    rotation_matrix = np.array([[np.cos(rotation), -np.sin(rotation), 0],
                                [np.sin(rotation), np.cos(rotation), 0],
                                [0, 0, 1]])
    
    # Calculate inverse rotation matrix
    inverse_rotation_matrix = np.array([[np.cos(-rotation), -np.sin(-rotation), 0],
                                        [np.sin(-rotation), np.cos(-rotation), 0],
                                        [0, 0, 1]])
    
    # Translate the point relative to the bbox's centroid
    translated_point = point - np.array([x_centroid, y_centroid, z_centroid])
    
    # Apply inverse rotation to the translated point
    aligned_point = np.dot(translated_point, inverse_rotation_matrix.T)
    
    # Check if the aligned point is within the bbox
    is_within_bbox = (
        np.abs(aligned_point[0]) <= half_length and
        np.abs(aligned_point[1]) <= half_width and
        np.abs(aligned_point[2]) <= half_height
    )
    
    #print("Point is within bbox:", is_within_bbox)
    return is_within_bbox

################################
# Augment
################################
def rotate_point(point, rot_angle_z, rot_origin_x, rot_origin_y):
    # extract xyz of point to be rotated
    x = point[0]
    y = point[1]
    z = point[2]
    new_x = np.cos(rot_angle_z) * (x - rot_origin_x) - np.sin(rot_angle_z) * (y - rot_origin_y) + rot_origin_x
    new_y = np.sin(rot_angle_z) * (x - rot_origin_x) + np.cos(rot_angle_z) * (y - rot_origin_y) + rot_origin_y
    new_point = np.array([new_x, new_y, z])
    return new_point
    
def augment(np_pcd, obj, aug_type=None):
    '''
    all augmentation type
    -1: no augment
    0: remove obj
    1: move upward
    2: top/down flip
    3: rotate left/right

    augmentation is performed at two sections (refer AUGMENT SECTION header below)
    '''
    # no augment
    if aug_type == -1:
        return np_pcd, obj
    elif aug_type == 0:
        pass
    elif aug_type == 1:
        displacement = np.random.uniform(0.2, 0.4, 1)[0]
    elif aug_type == 2:
        pass
    elif aug_type == 3:
        rot_angle_z = np.random.uniform(0.4, 1.5, 1)[0] # radian
        if np.random.rand() > 0.5:
            rot_angle_z *= -1

    # extract bbox properties
    x_centroid = obj['centroid']['x']
    y_centroid = obj['centroid']['y']
    z_centroid = obj['centroid']['z']
    length = obj['dimensions']['length']
    width = obj['dimensions']['width']
    height = obj['dimensions']['height']
    rotation = obj['rotations']['z']

    selected_points = []
    selected_indexs = []

    ave_x, ave_y = [], [] # to calculate new centroid (for type 3: rotate left/right)
    for i in range(len(np_pcd)):
        point = np_pcd[i]
        is_within_bbox = check_point_within_bbox(point, x_centroid, y_centroid, z_centroid,
                                                    length, width, height, rotation)
        # if the point is in the bbox
        if is_within_bbox:
            selected_points.append(point)
            selected_indexs.append(i)
            
            ######################
            # AUGMENT SECTION 1
            ######################
            # 0: if delete
            if aug_type == 0: 
                # no augmentation required here, just have to delete the points in Section 2
                pass
                
            # 1: if move upward
            elif aug_type == 1: 
                new_point = point + np.array([0, 0, displacement])
                np_pcd[i] = new_point    
            
            # 2: if top/down flip
            elif aug_type == 2:
                # no augmentation required here, do in Section 2
                pass
            
            # 3: rotate left/right (< y mid)
            # find the origin for rotation
            elif aug_type == 3:
                if y_centroid < y_mid:
                    # dunno how to explain but yup, we need two part
                    rot_origin_x = x_centroid
                    rot_origin_y = (y_centroid + width / 2)
                    origin_1 = rotate_point([rot_origin_x, rot_origin_y, -1], -1 * (3.14159265 - rotation), x_centroid, y_centroid)
                    x1, y1 = origin_1[0], origin_1[1]

                    # dunno how to explain but yup, we need two part
                    rot_origin_x = x_centroid
                    rot_origin_y = (y_centroid - width / 2)
                    origin_2 = rotate_point([rot_origin_x, rot_origin_y, -1], -1 * (3.14159265 - rotation), x_centroid, y_centroid)
                    x2, y2 = origin_2[0], origin_2[1]

                    if y1 < y2:
                        rot_origin_x = x1
                        rot_origin_y = y1
                    else:
                        rot_origin_x = x2
                        rot_origin_y = y2

                else:
                    # dunno how to explain but yup, we need two part
                    rot_origin_x = x_centroid
                    rot_origin_y = (y_centroid + width / 2)
                    origin_1 = rotate_point([rot_origin_x, rot_origin_y, -1], -1 * (3.14159265 - rotation), x_centroid, y_centroid)
                    x1, y1 = origin_1[0], origin_1[1]

                    # dunno how to explain but yup, we need two part
                    rot_origin_x = x_centroid
                    rot_origin_y = (y_centroid - width / 2)
                    origin_2 = rotate_point([rot_origin_x, rot_origin_y, -1], -1 * (3.14159265 - rotation), x_centroid, y_centroid)
                    x2, y2 = origin_2[0], origin_2[1]

                    if y1 > y2:
                        rot_origin_x = x1
                        rot_origin_y = y1
                    else:
                        rot_origin_x = x2
                        rot_origin_y = y2

                # rotation angle at z            
                rot_angle_z = rot_angle_z

                # edit the point
                new_point = rotate_point(point, rot_angle_z, rot_origin_x, rot_origin_y)
                np_pcd[i] = new_point
                
                # update ave_x, ave_y to calculate new centroid
                ave_x.append(new_point[0])
                ave_y.append(new_point[1])


    ######################
    # AUGMENT SECTION 2
    ###################### 
    # 0: if delete
    if aug_type == 0:
        # augment ply
        np_pcd = np.delete(np_pcd, selected_indexs, axis=0)
        
        # augment label
        obj = None

    # 1: if move upward    
    elif aug_type == 1:   
        # augment label
        obj['centroid']['z'] = obj['centroid']['z'] + displacement

    # 2: if top/down flip
    elif aug_type == 2: 
        # augment ply
        selected_points = np.array(selected_points) 
        z_plane = z_min = np.min(selected_points[:, 2]) 
        z_coords = selected_points[:, 2]
        z_coords = z_coords - z_plane # bring down the object to z=0 plane
        z_coords = -1 * z_coords # flip the opject at z=0
        z_coords = z_coords * 0.5 # squeeze in z dimension
        z_coords = z_coords + z_plane # bring back up
        selected_points[:, 2] = z_coords 
        np_pcd[selected_indexs] = selected_points
        
        # augmetn label similarly
        obj['centroid']['z'] = (-1 * (z_centroid - z_plane)) * 0.5 + z_plane
        obj['dimensions']['height'] = obj['dimensions']['height'] * 0.5
    
    # 3: rotate left/right
    elif aug_type == 3:
        # augment label
        x_centroid = sum(ave_x) / len(ave_x)
        y_centroid = sum(ave_y) / len(ave_y)
        rotation = rotation + rot_angle_z
        
        obj['centroid']['x'] = x_centroid
        obj['centroid']['y'] = y_centroid
        obj['rotations']['z'] = rotation      

    if obj is not None: obj['name'] = 'defect'

    return np_pcd, obj

################################
# Loop ply files
################################

# get argument from user
parser = argparse.ArgumentParser()
parser.add_argument('--label_dir', type = str, required = False, default = 'custom', help="where is the directory for the labels")
parser.add_argument('--ply_dir', type = str, required = False, default = 'custom', help="where is the directory for the ply data")
args = parser.parse_args()

# start looping the data
print("Load a ply point cloud, then augment it")

ply_dir = os.path.join(os.getcwd(), args.ply_dir)
label_dir = os.path.join(os.getcwd(), args.label_dir)

filenames = os.listdir(label_dir)
for fn_idx in tqdm(range(len(filenames)), desc =f'Data Augmentation'):
    # augment 2 times for each ply
    for aug_idx in range(2):
    
        filename = filenames[fn_idx]
        ################################
        # read ply
        ################################
        ply_filename = os.path.join(ply_dir, filename).replace('.json', '.ply')
        pcd = o3d.io.read_point_cloud(ply_filename)
        temp = o3d.geometry.PointCloud()
        temp.points = o3d.utility.Vector3dVector(np.array(pcd.points)) 
        #o3d.visualization.draw_geometries([temp])
        np_pcd = np.asarray(pcd.points)
        y_mid = (np.max(np_pcd[:,1]) + np.min(np_pcd[:,1])) / 2

        ################################
        # read label
        ################################
        label_filename = os.path.join(label_dir, filename)
        with open(label_filename) as f:
            d = json.load(f)
            objs = d['objects']

        if len(objs) != 13:
            continue

        ################################
        # arrange the obj following the sequence
        ################################
        all_x_c, all_y_c = [], []
        for obj in objs:
            x_centroid = obj['centroid']['x']
            y_centroid = obj['centroid']['y']
            all_x_c.append(x_centroid)
            all_y_c.append(y_centroid)

        # get y mid
        y_mid = (max(all_y_c) + min(all_y_c)) / 2

        # arrange according to x value, in ascending order
        all_x_c = np.array(all_x_c)  
        new_objs = np.array(objs)
        idx = np.argsort(all_x_c) 
        new_objs = list(new_objs[idx])

        # split to bottom and top
        new_objs_1, new_objs_2 = [], []
        for obj in new_objs:
            y_centroid = obj['centroid']['y']
            if y_centroid <= y_mid:
                new_objs_1.append(obj)
            else:
                new_objs_2.append(obj)

        # combine part bottom and top
        objs = new_objs_1 + new_objs_2
     

        ################################
        # create augmentation config
        ################################
        aug_list = [-1 for i in range(len(objs))]
        list_of_objs_yet_to_augment = [i for i in range(len(objs))]
        np.random.shuffle(list_of_objs_yet_to_augment)

        # random rotate a few legs
        if np.random.rand() > 0.5:
            aug_list[0] = 3
            list_of_objs_yet_to_augment.remove(0)
        if np.random.rand() > 0.5:
            aug_list[3] = 3
            list_of_objs_yet_to_augment.remove(3)

        # random remove 4 legs
        for i in range(4):
            idx = list_of_objs_yet_to_augment.pop(0)
            aug_list[idx] = 0

        # at least 5 legs are normal
        for i in range(5):
            idx = list_of_objs_yet_to_augment.pop(0)
            aug_list[idx] = -1

        # random augment the rest using aug_type = 1 or 2
        for i in list_of_objs_yet_to_augment:
            aug_list[i] = np.random.choice([1, 2])

        print(aug_list)

        ################################
        # loop all legs for augmentation
        ################################
        #aug_list = [-1 for i in range(len(objs))]
        #aug_list[3] = 3
        final_objs = []
        for obj_idx, obj in enumerate(objs):
            #print(obj)

            # augment here
            rand_int = aug_list[obj_idx] #np.random.randint(0, high=3+1, size=None, dtype=int)
            np_pcd, obj = augment(np_pcd, obj, aug_type=rand_int)
            if obj is not None:
                final_objs.append(obj)

        # plot the augmented points
        augmented_points = np.array(np_pcd)
        color_grad = o3d.geometry.PointCloud()
        color_grad.points = o3d.utility.Vector3dVector(augmented_points) 
        #o3d.visualization.draw_geometries([color_grad])
        
        # save ply and label
        new_ply_filename = os.path.join(ply_dir, f'augmented_00{aug_idx} ' + filename).replace('.json', '.ply')
        o3d.io.write_point_cloud(new_ply_filename, color_grad) # save the filtered point cloud
        
        new_label_filename = os.path.join(label_dir, f'augmented_00{aug_idx} ' + filename)
        d['folder'] = os.path.join(os.getcwd(), ply_dir)
        d['filename'] = os.path.basename(new_ply_filename)
        d['path'] = new_ply_filename
        d['objects'] = final_objs
        with open(new_label_filename, 'w', encoding='utf-8') as f:
            json.dump(d, f, ensure_ascii=False, indent=4)