extern/CUT3R/datasets_preprocess/preprocess_omniobject3d.py

#!/usr/bin/env python3
"""
This script processes scene data by reading images, depth maps, and camera poses,
computing camera intrinsics, and saving the results in a structured format.

Usage:
    python preprocess_omniobject3d.py --input_dir /path/to/input_root --output_dir /path/to/output_root
"""

import os
import os.path as osp
import json
import argparse
from concurrent.futures import ProcessPoolExecutor

import numpy as np
import cv2
import imageio.v2 as imageio
from tqdm import tqdm
import math

# Enable OpenEXR support in OpenCV
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"


def prepare_scene_args(scene, input_root, output_root):
    """
    Prepare processing arguments for a given scene.

    Args:
        scene (str): Scene directory name.
        input_root (str): Root directory for input data.
        output_root (str): Root directory for output data.

    Returns:
        list or None: A list of arguments for each frame in the scene or None if preparation fails.
    """
    seq_dir = osp.join(input_root, scene, "render")
    rgb_dir = osp.join(seq_dir, "images")
    depth_dir = osp.join(seq_dir, "depths")
    pose_file = osp.join(seq_dir, "transforms.json")
    out_seq_dir = osp.join(output_root, scene)

    # Check if the necessary file exists
    if not osp.exists(pose_file):
        print(f"Pose file not found: {pose_file}")
        return None

    # Load metadata from JSON
    with open(pose_file, "r") as fp:
        meta = json.load(fp)

    camera_angle_x = float(meta["camera_angle_x"])

    # Create output directories for this scene
    os.makedirs(osp.join(out_seq_dir, "rgb"), exist_ok=True)
    os.makedirs(osp.join(out_seq_dir, "depth"), exist_ok=True)
    os.makedirs(osp.join(out_seq_dir, "cam"), exist_ok=True)

    # Prepare a list of frame processing arguments
    frame_args = [
        (frame, camera_angle_x, rgb_dir, depth_dir, out_seq_dir)
        for frame in meta.get("frames", [])
    ]

    return frame_args


def process_frame(args):
    """
    Process a single frame:
      - Reads the image and depth data.
      - Handles alpha channels by compositing over a white background.
      - Computes the camera intrinsics.
      - Saves the processed RGB image, depth map, and camera parameters.

    Args:
        args (tuple): A tuple containing:
            - frame (dict): Frame metadata.
            - camera_angle_x (float): Camera field-of-view.
            - rgb_dir (str): Directory containing RGB images.
            - depth_dir (str): Directory containing depth maps.
            - out_seq_dir (str): Output directory for the processed scene.
    """
    frame, camera_angle_x, rgb_dir, depth_dir, out_seq_dir = args

    # Derive the base name from the frame's file path
    frame_name = osp.basename(frame["file_path"])

    # Define file paths for input and output
    image_path = osp.join(rgb_dir, frame_name + ".png")
    depth_path = osp.join(depth_dir, frame_name + "_depth.exr")
    out_img_path = osp.join(out_seq_dir, "rgb", frame_name + ".png")
    out_depth_path = osp.join(out_seq_dir, "depth", frame_name + ".npy")
    out_cam_path = osp.join(out_seq_dir, "cam", frame_name + ".npz")

    # Skip processing if outputs already exist
    if (
        osp.exists(out_img_path)
        and osp.exists(out_depth_path)
        and osp.exists(out_cam_path)
    ):
        return

    # Read image using imageio
    img = imageio.imread(image_path)

    # If image has an alpha channel, composite it over a white background
    if img.shape[-1] == 4:
        alpha_channel = img[..., 3]
        rgb_channels = img[..., :3]
        white_background = np.full_like(rgb_channels, 255)
        img = np.where(alpha_channel[..., None] == 0, white_background, rgb_channels)
    else:
        img = img[..., :3]

    H, W, _ = img.shape

    # Process the camera pose
    pose = np.array(frame["transform_matrix"], dtype=np.float32)
    pose[:, 1:3] *= -1  # Invert Y and Z axes if necessary

    # Compute camera intrinsics using the provided camera angle
    focal = 0.5 * W / np.tan(0.5 * camera_angle_x)
    intrinsics = np.array(
        [[focal, 0, W / 2], [0, focal, H / 2], [0, 0, 1]], dtype=np.float32
    )

    # Read depth data using OpenCV (which supports OpenEXR)
    depth = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED)
    if depth is None:
        print(f"Warning: Depth file not found or failed to read: {depth_path}")
        return

    # Use the last channel of the depth data and convert to float32
    depth = depth[..., -1].astype(np.float32)
    depth[depth >= 65504.0] = 0.0  # Set invalid depth values to 0

    # Save the processed outputs
    imageio.imwrite(out_img_path, img.astype(np.uint8))
    np.save(out_depth_path, depth)
    np.savez_compressed(out_cam_path, intrinsics=intrinsics, pose=pose)


def process_scene(frame_args):
    """
    Process all frames within a single scene.

    Args:
        frame_args (list): List of frame arguments for the scene.
    """
    if frame_args is None:
        return

    for args in frame_args:
        process_frame(args)


def main():
    parser = argparse.ArgumentParser(
        description="Preprocess scene data by extracting RGB images, depth maps, and camera parameters."
    )
    parser.add_argument(
        "--input_dir",
        type=str,
        required=True,
        help="Path to the root input directory containing scene data.",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        required=True,
        help="Path to the directory where processed data will be saved.",
    )
    parser.add_argument(
        "--max_workers",
        type=int,
        default=None,
        help="Maximum number of worker processes. Defaults to the number of CPU cores.",
    )
    args = parser.parse_args()

    input_root = args.input_dir
    output_root = args.output_dir

    # Ensure the output root directory exists
    os.makedirs(output_root, exist_ok=True)

    # List all scene directories in the input root
    scenes = sorted(
        [d for d in os.listdir(input_root) if osp.isdir(osp.join(input_root, d))]
    )

    # Determine the number of workers to use
    max_workers = (
        args.max_workers if args.max_workers is not None else os.cpu_count() or 1
    )

    # Prepare processing arguments for each scene in parallel
    with ProcessPoolExecutor(max_workers=max_workers) as executor:
        scene_args_list = list(
            tqdm(
                executor.map(
                    lambda s: prepare_scene_args(s, input_root, output_root), scenes
                ),
                total=len(scenes),
                desc="Preparing scenes",
            )
        )

    # Filter out scenes where preparation failed
    scene_frame_args = [fa for fa in scene_args_list if fa is not None]

    # Process each scene in parallel
    with ProcessPoolExecutor(max_workers=max_workers) as executor:
        list(
            tqdm(
                executor.map(process_scene, scene_frame_args),
                total=len(scene_frame_args),
                desc="Processing scenes",
            )
        )


if __name__ == "__main__":
    main()