convertor.py

import json
import os
import pickle
from collections import defaultdict

import numpy as np
import yaml


DATASET_ROOT = "./kitti-data"
LEFT_IMAGES_DIR = os.path.join(DATASET_ROOT, "data_2d_raw/2013_05_28_drive_0009_sync/image_00/data_rect")
RIGHT_IMAGES_DIR = os.path.join(DATASET_ROOT, "data_2d_raw/2013_05_28_drive_0009_sync/image_01/data_rect")

LEFT_FISH_EYE_IMAGES_DIR = os.path.join(DATASET_ROOT, "data_2d_raw/2013_05_28_drive_0009_sync/image_02/data_rgb")
RIGHT_FISH_EYE_IMAGES_DIR = os.path.join(DATASET_ROOT, "data_2d_raw/2013_05_28_drive_0009_sync/image_03/data_rgb")

LEFT_FISH_EYE_CALIBRATION_FILE = os.path.join(DATASET_ROOT, "calibration/image_02.yaml")
RIGHT_FISH_EYE_CALIBRATION_FILE = os.path.join(DATASET_ROOT, "calibration/image_03.yaml")

POSES_FILE = os.path.join(DATASET_ROOT, "data_poses/2013_05_28_drive_0009_sync/cam0_to_world.txt")
CALIBRATION_FILE = os.path.join(DATASET_ROOT, "calibration/perspective.txt")
EXTRINSICS_FILE = os.path.join(DATASET_ROOT, "calibration/calib_cam_to_pose.txt")
POINT_CLOUD_FILE = "0000002826_0000003034_pruned.ply"
OUTPUT_FILE = "./transforms.json"

STARTING_FRAME = 2854
ENDING_FRAME = 2932
AABB_SCALE = 16

with open("config.json") as config_file:
    config = json.load(config_file)

error_config = config['error_config']

def load_intrinsics(calibration_file: str) -> dict[str, dict[str, float]]:
    intrinsics = defaultdict(dict)
    with open(calibration_file, "r") as f:
        for line in f:
            if line.startswith("P_rect_00"):
                values = [float(x) for x in line.split()[1:]]
                intrinsics["left"].update(
                    {
                        "fl_x": values[0],
                        "fl_y": values[5],
                        "cx": values[2],
                        "cy": values[6],
                    }
                )
            if line.startswith("P_rect_01"):
                values = [float(x) for x in line.split()[1:]]
                intrinsics["right"].update(
                    {
                        "fl_x": values[0],
                        "fl_y": values[5],
                        "cx": values[2],
                        "cy": values[6],
                        "tx": values[3] / values[0],  # Translation in x direction (normalized by focal length)
                    }
                )
            if line.startswith("S_rect_00"):
                values = [float(x) for x in line.split()[1:]]
                intrinsics["left"].update(
                    {
                        "w": values[0],
                        "h": values[1],
                    }
                )
            if line.startswith("S_rect_01"):
                values = [float(x) for x in line.split()[1:]]
                intrinsics["right"].update(
                    {
                        "w": values[0],
                        "h": values[1],
                    }
                )
    return dict(intrinsics)


def load_poses(poses_file: str) -> dict[int, list[list[float]]]:
    poses = {}
    with open(poses_file, "r") as f:
        for line in f:
            parts = line.strip().split()
            frame_id = int(parts[0])
            matrix = np.array([float(x) for x in parts[1:]]).reshape(4, 4)
            poses[frame_id] = matrix.tolist()
    return poses

def load_extrinsics(extrinsics_file: str) -> dict[str, list[list[float]]]:
    extrinsics = {}
    with open(extrinsics_file, "r") as f:
        for line in f:
            parts = line.strip().split(':')
            frame_id = parts[0].strip()
            matrix_values = [float(x) for x in parts[1].strip().split()]
            matrix = np.array(matrix_values).reshape(3, 4)
            matrix = np.vstack([matrix, [0, 0, 0, 1]]).astype(np.float16)
            extrinsics[frame_id] = matrix.tolist()
    return extrinsics


def load_fish_intrinsics(file_path):
    with open(file_path, 'r') as file:
        content = file.read()
    
    content = content.replace('%YAML:1.0', '').strip()
    
    data = yaml.safe_load(content)
    
    model_type = data.get('model_type')
    camera_name = data.get('camera_name')
    image_width = data.get('image_width')
    image_height = data.get('image_height')
    mirror_parameters = data.get('mirror_parameters', {})
    distortion_parameters = data.get('distortion_parameters', {})
    projection_parameters = data.get('projection_parameters', {})
    
    xi = mirror_parameters.get('xi')
    k1 = distortion_parameters.get('k1')
    k2 = distortion_parameters.get('k2')
    p1 = distortion_parameters.get('p1')
    p2 = distortion_parameters.get('p2')
    gamma1 = projection_parameters.get('gamma1')
    gamma2 = projection_parameters.get('gamma2')
    u0 = projection_parameters.get('u0')
    v0 = projection_parameters.get('v0')
    
    result = {
        "w": data.get('image_width'),
        "h": data.get('image_height'),
        "fl_x": projection_parameters.get('gamma1'),
        "fl_y": projection_parameters.get('gamma2'),
        "cx": projection_parameters.get('u0'),
        "cy": projection_parameters.get('v0'),
        "k1": distortion_parameters.get('k1', 0.0),
        "k2": distortion_parameters.get('k2', 0.0),
        "p1": distortion_parameters.get('p1', 0.0),
        "p2": distortion_parameters.get('p2', 0.0)
    }
    return result

def rotate_camera_origin(transform_matrix):
    conversion = np.array(
        [
            [1, 0, 0, 0],
            [0, -1, 0, 0],  # Flip Y
            [0, 0, -1, 0],  # Flip Z
            [0, 0, 0, 1],
        ],
        dtype=float,
    )
    new_transform = np.array(transform_matrix) @ conversion
    return new_transform.tolist()


def move_transform_matrix(transform_matrix, tx):
    shift = np.eye(4)
    shift[0, 3] = -tx
    adjusted_matrix = np.array(transform_matrix) @ shift
    return adjusted_matrix.tolist()


def calculate_transform_error(target_transform, current_transform):
    target = np.array(target_transform)
    current = np.array(current_transform)
    rotation_error = np.sum(np.abs(target[:3, :3] - current[:3, :3]))
    translation_error = np.sum(np.abs(target[:3, 3] - current[:3, 3]))
    return rotation_error, translation_error


def skew_symmetric(v):
    return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])


def rodrigues(rotation_vector):
    theta = np.linalg.norm(rotation_vector)
    if theta < 1e-8:
        return np.eye(3)
    k = rotation_vector / theta
    K = skew_symmetric(k)
    return np.eye(3) + np.sin(theta) * K + (1 - np.cos(theta)) * (K @ K)


def apply_fixed_error(transform_matrix, error_config):
    # Ensure transform_matrix is a NumPy array
    transform_matrix = np.array(transform_matrix)
    
    # Convert cm to meters and degrees to radians
    translation_error = np.array([
        error_config['translation']['x_cm'] / 100,
        error_config['translation']['y_cm'] / 100,
        error_config['translation']['z_cm'] / 100
    ])
    
    rotation_error = np.radians([
        error_config['rotation']['x_deg'],
        error_config['rotation']['y_deg'],
        error_config['rotation']['z_deg']
    ])
    
    # Create rotation matrix from Euler angles
    R_x = np.array([[1, 0, 0],
                   [0, np.cos(rotation_error[0]), -np.sin(rotation_error[0])],
                   [0, np.sin(rotation_error[0]), np.cos(rotation_error[0])]])
    
    R_y = np.array([[np.cos(rotation_error[1]), 0, np.sin(rotation_error[1])],
                   [0, 1, 0],
                   [-np.sin(rotation_error[1]), 0, np.cos(rotation_error[1])]])
    
    R_z = np.array([[np.cos(rotation_error[2]), -np.sin(rotation_error[2]), 0],
                   [np.sin(rotation_error[2]), np.cos(rotation_error[2]), 0],
                   [0, 0, 1]])
    
    rotation_matrix = R_z @ R_y @ R_x
    
    # Apply translation and rotation to transform matrix
    transform_matrix[:3, 3] += translation_error
    transform_matrix[:3, :3] = rotation_matrix @ transform_matrix[:3, :3]
    
    return transform_matrix


def generate_transforms(
    left_images_dir: str,
    right_images_dir: str,
    intrinsics: dict[str, dict[str, float]],
    poses: dict[int, list[list[float]]],
    extrinsics: dict[int, list[list[float]]],
    point_cloud_file: str,
    output_file: str,
) -> None:
    frames = []
    left_image_files = sorted(os.listdir(left_images_dir))
    right_image_files = sorted(os.listdir(right_images_dir))
    left_fish_eye_image_files = sorted(os.listdir(LEFT_FISH_EYE_IMAGES_DIR))
    right_fish_eye_image_files = sorted(os.listdir(RIGHT_FISH_EYE_IMAGES_DIR))

    processed_frames = 0
    skipped_frames = 0
    non_trainable_indices = []

    # Process left camera images first
    for idx, image_file in enumerate(left_image_files):
        frame_id = int(image_file.split(".")[0])
        if STARTING_FRAME <= frame_id <= ENDING_FRAME:
            if frame_id not in poses or idx % 2 == 0:
                skipped_frames += 1
                continue

            adjusted_transform = rotate_camera_origin(poses[frame_id])
            frame_data = {
                "file_path": os.path.join(left_images_dir, image_file),
                "transform_matrix": adjusted_transform,
                **intrinsics["left"],
                "k1": 0.0,
                "k2": 0.0,
                "p1": 0.0,
                "p2": 0.0,
            }
            frames.append(frame_data)
            non_trainable_indices.append(processed_frames)
            processed_frames += 1

    for idx, image_file in enumerate(right_image_files):
        frame_id = int(image_file.split(".")[0])
        if STARTING_FRAME <= frame_id <= ENDING_FRAME:
            if frame_id not in poses or idx % 2 == 0:
                skipped_frames += 1
                continue

            adjusted_transform = move_transform_matrix(poses[frame_id], intrinsics["right"].get("tx", 0))
            adjusted_transform = rotate_camera_origin(adjusted_transform)

            # Apply the fixed error using post-multiplication
            transform_matrix = apply_fixed_error(adjusted_transform, error_config)
            frame_data = {
                "file_path": os.path.join(right_images_dir, image_file),
                "transform_matrix": transform_matrix.tolist(),
                **intrinsics["right"],
                "k1": 0.0,
                "k2": 0.0,
                "p1": 0.0,
                "p2": 0.0,
            }
            frames.append(frame_data)
            processed_frames += 1

    for image_file in left_fish_eye_image_files:
        frame_id = int(image_file.split(".")[0])
        if STARTING_FRAME <= frame_id <= ENDING_FRAME:
            if frame_id not in poses:
                skipped_frames += 1
                continue

            adjusted_transform = rotate_camera_origin(poses[frame_id])
            left_perspective_to_fish_eye = np.matmul(np.linalg.inv(extrinsics["image_00"]), extrinsics["image_02"])
            transform = np.matmul(adjusted_transform, left_perspective_to_fish_eye)

            frame_data = {
                "file_path": os.path.join(LEFT_FISH_EYE_IMAGES_DIR, image_file),
                "transform_matrix": transform.tolist(),
                **load_fish_intrinsics(LEFT_FISH_EYE_CALIBRATION_FILE),
            }
            frames.append(frame_data)
            processed_frames += 1

    for image_file in right_fish_eye_image_files:
        frame_id = int(image_file.split(".")[0])
        if STARTING_FRAME <= frame_id <= ENDING_FRAME:
            if frame_id not in poses:
                skipped_frames += 1
                continue

            adjusted_transform = rotate_camera_origin(poses[frame_id])
            right_perspective_to_fish_eye = np.matmul(np.linalg.inv(extrinsics["image_01"]), extrinsics["image_03"])
            transform = np.matmul(adjusted_transform, right_perspective_to_fish_eye)

            frame_data = {
                "file_path": os.path.join(RIGHT_FISH_EYE_IMAGES_DIR, image_file),
                "transform_matrix": transform.tolist(),
                **load_fish_intrinsics(RIGHT_FISH_EYE_CALIBRATION_FILE),
            }
            frames.append(frame_data)
            processed_frames += 1

    print(f"Processed {processed_frames} frames, skipped {skipped_frames} frames.")
    print(f"Number of left camera frames (non-trainable): {len(non_trainable_indices)}")
    print(f"Number of right camera frames (trainable): {processed_frames - len(non_trainable_indices)}")

    data = {
        "camera_model": "OPENCV",
        "aabb_scale": AABB_SCALE,
        "frames": frames,
        "ply_file_path": point_cloud_file,
        "non_trainable_cameras": non_trainable_indices,
    }

    with open(output_file, "w") as f:
        json.dump(data, f, indent=4)


def main():
    intrinsics = load_intrinsics(CALIBRATION_FILE)
    poses = load_poses(POSES_FILE)
    extrinsics = load_extrinsics(EXTRINSICS_FILE)
    
    # Save initial errors as vectors
    initial_trans = np.array([
        config['error_config']['translation']['x_cm'],
        config['error_config']['translation']['y_cm'],
        config['error_config']['translation']['z_cm']
    ]) / 100  # Convert cm to meters
    
    initial_rot = np.radians([
        config['error_config']['rotation']['x_deg'],
        config['error_config']['rotation']['y_deg'],
        config['error_config']['rotation']['z_deg']
    ])
    
    with open('initial_errors.pkl', 'wb') as f:
        pickle.dump({
            'translation': initial_trans,
            'rotation': initial_rot
        }, f)

    generate_transforms(LEFT_IMAGES_DIR, RIGHT_IMAGES_DIR, intrinsics, poses, extrinsics, POINT_CLOUD_FILE, OUTPUT_FILE)
    print(f"Transforms.json generated at {OUTPUT_FILE}\n")


if __name__ == "__main__":
    main()