datasets_preprocess/preprocess_ase.py

#!/usr/bin/env python3
# --------------------------------------------------------
# Script to pre-process the aria-ase dataset
# Usage:
# 1. Prepare the codebase and environment for the projectaria_tools
# 2. copy this script to the project root directory
# 3. Run the script
# --------------------------------------------------------
import matplotlib.colors as colors
import matplotlib.pyplot as plt
import numpy as np
import plotly.graph_objects as go
from pathlib import Path
import os
from PIL import Image
from scipy.spatial.transform import Rotation as R
from projectaria_tools.projects import ase
from projectaria_tools.core import data_provider, calibration
from projectaria_tools.core.image import InterpolationMethod
from projects.AriaSyntheticEnvironment.tutorial.code_snippets.readers import read_trajectory_file
import cv2
from tqdm import tqdm
import os, sys, json
import open3d as o3d
import random


def save_pointcloud(points_3d_array, rgb ,pcd_name):
    # Flatten the instance values array
    rgb_values_flat = rgb

    # Check if the number of points matches the number of instance values
    assert points_3d_array.shape[0] == rgb_values_flat.shape[0], "The number of points must match the number of instance values"

    # Create an Open3D point cloud object
    pcd = o3d.geometry.PointCloud()

    # Assign the 3D points to the point cloud object
    pcd.points = o3d.utility.Vector3dVector(points_3d_array)

    # Assign the colors to the point cloud
    pcd.colors = o3d.utility.Vector3dVector(rgb_values_flat / 255.0)  # Normalize colors to [0, 1]

    # Define the file path where you want to save the point cloud
    output_file_path = pcd_name+'.pcd'

    # Save the point cloud in PCD format
    o3d.io.write_point_cloud(output_file_path, pcd)

    print(f"Point cloud saved to {output_file_path}")


def unproject(camera_params, undistorted_depth,undistorted_rgb):
    # Get the height and width of the depth image
    height, width = undistorted_depth.shape

    # Generate pixel coordinates
    y, x = np.indices((height, width))
    pixel_coords = np.stack((x, y), axis=-1).reshape(-1, 2)

    # Flatten the depth image to create a 1D array of depth values
    depth_values_flat = undistorted_depth.flatten()
    rgb_values_flat = undistorted_rgb.reshape(-1,3)

    # Initialize an array to store 3D points
    points_3d = []
    valid_rgb = []

    for pixel_coord, depth, rgb in zip(pixel_coords, depth_values_flat, rgb_values_flat):
        # Format the pixel coordinate for unproject (reshape to [2, 1])
        pixel_coord_reshaped = np.array([[pixel_coord[0]], [pixel_coord[1]]], dtype=np.float64)

        # Unproject the pixel to get the direction vector (ray)
        #   direction_vector = device.unproject(pixel_coord_reshaped)
        X = (pixel_coord_reshaped[0] - camera_params[2]) / camera_params[0] # X = (u - cx) / fx
        Y = (pixel_coord_reshaped[1] - camera_params[3]) / camera_params[1] # Y = (v - cy) / fy
        direction_vector = np.array([X[0], Y[0], 1],dtype=np.float32)
        if direction_vector is not None:
            # Replace the z-value of the direction vector with the depth value
            # Assuming the direction vector is normalized
            direction_vector_normalized = direction_vector / np.linalg.norm(direction_vector)
            point_3d = direction_vector_normalized * (depth / 1000)

            # Append the computed 3D point and the corresponding instance
            points_3d.append(point_3d.flatten())
            valid_rgb.append(rgb)

    # Convert the list of 3D points to a numpy array
    points_3d_array = np.array(points_3d)
    points_rgb = np.array(valid_rgb)
    return points_3d_array,points_rgb

def distance_to_depth(K, dist, uv=None):
    if uv is None and len(dist.shape) >= 2:
        # create mesh grid according to d
        uv = np.stack(np.meshgrid(np.arange(dist.shape[1]), np.arange(dist.shape[0])), -1)
        uv = uv.reshape(-1, 2)
        dist = dist.reshape(-1)
        if not isinstance(dist, np.ndarray):
            import torch
            uv = torch.from_numpy(uv).to(dist)
    if isinstance(dist, np.ndarray):
        # z * np.sqrt(x_temp**2+y_temp**2+z_temp**2) = dist
        uvh = np.concatenate([uv, np.ones((len(uv), 1))], -1)
        uvh = uvh.T # N, 3
        temp_point = np.linalg.inv(K) @ uvh # 3, N  
        temp_point = temp_point.T # N, 3
        z = dist / np.linalg.norm(temp_point, axis=1)
    else:
        uvh = torch.cat([uv, torch.ones(len(uv), 1).to(uv)], -1)
        temp_point = torch.inverse(K) @ uvh
        z = dist / torch.linalg.norm(temp_point, dim=1)
    return z

def transform_3d_points(transform, points):
    N = len(points)
    points_h = np.concatenate([points, np.ones((N, 1))], axis=1)
    transformed_points_h = (transform @ points_h.T).T
    transformed_points = transformed_points_h[:, :-1]
    return transformed_points


def aria_export_to_scannet(scene_id, seed):
    random.seed(int(seed + scene_id))
    src_folder = Path("ase_raw/"+str(scene_id))
    trgt_folder = Path("ase_processed/"+str(scene_id))
    trgt_folder.mkdir(parents=True, exist_ok=True)
    SCENE_ID = src_folder.stem
    print("SCENE_ID:", SCENE_ID)

    scene_max_depth = 0
    scene_min_depth = np.inf
    Path(trgt_folder, "intrinsic").mkdir(exist_ok=True)
    Path(trgt_folder, "pose").mkdir(exist_ok=True)
    Path(trgt_folder, "depth").mkdir(exist_ok=True)
    Path(trgt_folder, "color").mkdir(exist_ok=True)

    rgb_dir = src_folder / "rgb"
    depth_dir = src_folder / "depth"
    # Load camera calibration
    device = ase.get_ase_rgb_calibration()
    # Load the trajectory using read_trajectory_file() 
    trajectory_path = src_folder / "trajectory.csv"
    trajectory = read_trajectory_file(trajectory_path)
    all_points_3d = []
    all_rgb = []
    num_frames = len(list(rgb_dir.glob("*.jpg")))
    # Path('./debug').mkdir(exist_ok=True)
    for frame_idx in tqdm(range(num_frames)):   
        frame_id = str(frame_idx).zfill(7)
        rgb_path = rgb_dir / f"vignette{frame_id}.jpg"
        depth_path = depth_dir / f"depth{frame_id}.png"
        depth = Image.open(depth_path) # uint16        
        rgb = cv2.imread(str(rgb_path), cv2.IMREAD_UNCHANGED)
        depth = np.array(depth)
        scene_min_depth = min(depth.min(), scene_min_depth)
        inf_value = np.iinfo(np.array(depth).dtype).max
        depth[depth == inf_value] = 0 # consider it as invalid, inplace with 0
        T_world_from_device = trajectory["Ts_world_from_device"][frame_idx] # camera-to-world
        assert device.get_image_size()[0] == 704
        # https://facebookresearch.github.io/projectaria_tools/docs/data_utilities/advanced_code_snippets/image_utilities
        focal_length = device.get_focal_lengths()[0]
        pinhole = calibration.get_linear_camera_calibration(
            512,
            512,
            focal_length,
            "camera-rgb",
            device.get_transform_device_camera() # important to get correct transformation matrix in pinhole_cw90
            )
        # distort image
        rectified_rgb = calibration.distort_by_calibration(np.array(rgb), pinhole, device, InterpolationMethod.BILINEAR)
        # raw_image = np.array(depth) # Will not work
        depth = np.array(depth).astype(np.float32) # WILL WORK
        rectified_depth = calibration.distort_by_calibration(depth, pinhole, device)
        
        rotated_image = np.rot90(rectified_rgb, k=3)
        rotated_depth = np.rot90(rectified_depth, k=3)

        cv2.imwrite(str(Path(trgt_folder, "color", f"{frame_id}.jpg")), rotated_image)
        # # TODO: check this
        # plt.imsave(Path(f"./debug/debug_undistort_{frame_id}.png"), np.uint16(rotated_depth), cmap="plasma")
        # Get rotated image calibration
        pinhole_cw90 = calibration.rotate_camera_calib_cw90deg(pinhole)
        principal = pinhole_cw90.get_principal_point()
        cx, cy = principal[0], principal[1]
        focal_lengths = pinhole_cw90.get_focal_lengths()
        fx, fy = focal_lengths 
        K = np.array([ # camera-to-pixel
            [fx, 0, cx],
            [0, fy, cy],
            [0, 0, 1.0]])

        c2w = T_world_from_device 
        c2w_rotation = pinhole_cw90.get_transform_device_camera().to_matrix()
        c2w_final = c2w @ c2w_rotation   # right-matmul!
        cam2world = c2w_final
        
        # save depth
        rotated_depth = np.uint16(rotated_depth)
        depth_image = Image.fromarray(rotated_depth, mode='I;16')
        depth_image.save(str(Path(trgt_folder, "depth", f"{frame_id}.png")))
        # for debug; load depth and convert to pointcloud
        # depth_image = np.array(Image.open(str(Path(trgt_folder, "depth", f"{frame_id}.png"))), dtype=np.uint16)
        # points_3d_array, points_rgb = unproject((fx, fy, cx, cy), depth_image, rotated_image)
        # points_3d_world = transform_3d_points(cam2world, points_3d_array)
        # all_points_3d.append(points_3d_world)
        # all_rgb.append(points_rgb)
        # distance-to-depth
        # rotated_depth = distance_to_depth(K, rotated_depth).reshape((rotated_depth.shape[0], rotated_depth.shape[1]))#.reshape((dpt.shape[0], dpt.shape[1]))        

        Path(trgt_folder, "intrinsic", "intrinsic_color.txt").write_text(f"""{K[0][0]} {K[0][1]} {K[0][2]} 0.00\n{K[1][0]} {K[1][1]} {K[1][2]} 0.00\n{K[2][0]} {K[2][1]} {K[2][2]} 0.00\n0.00 0.00 0.00 1.00""")
        Path(trgt_folder, "pose", f"{frame_id}.txt").write_text(f"""{cam2world[0, 0]} {cam2world[0, 1]} {cam2world[0, 2]} {cam2world[0, 3]}\n{cam2world[1, 0]} {cam2world[1, 1]} {cam2world[1, 2]} {cam2world[1, 3]}\n{cam2world[2, 0]} {cam2world[2, 1]} {cam2world[2, 2]} {cam2world[2, 3]}\n0.00 0.00 0.00 1.00""")
    


if __name__ == "__main__": 
    seed = 42   
    for scene_id in tqdm(range(0, 500)):
        aria_export_to_scannet(scene_id=scene_id, seed = seed)