extern/CUT3R/datasets_preprocess/preprocess_irs.py

#!/usr/bin/env python3
"""
Preprocess the IRS dataset.

This script converts disparity EXR files into depth maps, copies corresponding RGB images,
and saves camera intrinsics computed from a given focal length and baseline. Processing is
done per sequence directory using parallel processing.

Usage:
    python preprocess_irs.py
       --root_dir /path/to/data_irs
       --out_dir /path/to/processed_irs
"""

import os
import shutil
import re
import glob
import time
from concurrent.futures import ProcessPoolExecutor, as_completed

import numpy as np
import OpenEXR
import Imath
import imageio
from PIL import Image
from tqdm import tqdm
import argparse

# Ensure OpenEXR support in OpenCV if needed.
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"


def exr2hdr(exrpath):
    """
    Read an OpenEXR file and return an HDR image as a NumPy array.
    """
    file = OpenEXR.InputFile(exrpath)
    pixType = Imath.PixelType(Imath.PixelType.FLOAT)
    dw = file.header()["dataWindow"]
    num_channels = len(file.header()["channels"].keys())
    if num_channels > 1:
        channels = ["R", "G", "B"]
        num_channels = 3
    else:
        channels = ["G"]

    size = (dw.max.x - dw.min.x + 1, dw.max.y - dw.min.y + 1)
    pixels = [
        np.fromstring(file.channel(c, pixType), dtype=np.float32) for c in channels
    ]
    hdr = np.zeros((size[1], size[0], num_channels), dtype=np.float32)
    if num_channels == 1:
        hdr[:, :, 0] = np.reshape(pixels[0], (size[1], size[0]))
    else:
        hdr[:, :, 0] = np.reshape(pixels[0], (size[1], size[0]))
        hdr[:, :, 1] = np.reshape(pixels[1], (size[1], size[0]))
        hdr[:, :, 2] = np.reshape(pixels[2], (size[1], size[0]))
    return hdr


def writehdr(hdrpath, hdr):
    """
    Write an HDR image to a file using the HDR format.
    If the input has one channel, duplicate it across R, G, and B.
    """
    h, w, c = hdr.shape
    if c == 1:
        hdr = np.pad(hdr, ((0, 0), (0, 0), (0, 2)), "constant")
        hdr[:, :, 1] = hdr[:, :, 0]
        hdr[:, :, 2] = hdr[:, :, 0]
    imageio.imwrite(hdrpath, hdr, format="hdr")


def load_exr(filename):
    """
    Load an EXR file and return the HDR image as a NumPy array.
    """
    hdr = exr2hdr(filename)
    h, w, c = hdr.shape
    if c == 1:
        hdr = np.squeeze(hdr)
    return hdr


def process_basename(args):
    """
    Process a single basename:
      - Load an RGB image and disparity (EXR) file.
      - Compute a depth map from disparity using: depth = (baseline * f) / disparity.
      - Copy the RGB image and save the computed depth and camera intrinsics.

    Parameters:
      args: tuple containing
            (basename, seq_dir, out_rgb_dir, out_depth_dir, out_cam_dir, f, baseline)
    Returns:
      None on success or an error string on failure.
    """
    basename, seq_dir, out_rgb_dir, out_depth_dir, out_cam_dir, f, baseline = args
    out_img_path = os.path.join(out_rgb_dir, f"{basename}.png")
    out_depth_path = os.path.join(out_depth_dir, f"{basename}.npy")
    out_cam_path = os.path.join(out_cam_dir, f"{basename}.npz")
    if os.path.exists(out_cam_path):
        return

    try:
        img_file = os.path.join(seq_dir, f"l_{basename}.png")
        disp_file = os.path.join(seq_dir, f"d_{basename}.exr")

        # Load image using PIL.
        img = Image.open(img_file)

        # Load disparity using the custom load_exr function.
        disp = load_exr(disp_file).astype(np.float32)
        H, W = disp.shape

        # Verify that the image size matches the disparity map.
        if img.size != (W, H):
            return f"Size mismatch for {basename}: Image size {img.size}, Disparity size {(W, H)}"

        # Create a simple camera intrinsics matrix.
        K = np.eye(3, dtype=np.float32)
        K[0, 0] = f
        K[1, 1] = f
        K[0, 2] = W // 2
        K[1, 2] = H // 2

        # Compute depth from disparity.
        depth = baseline * f / disp

        # Copy the RGB image.
        shutil.copyfile(img_file, out_img_path)
        # Save the depth map.
        np.save(out_depth_path, depth)
        # Save the camera intrinsics.
        np.savez(out_cam_path, intrinsics=K)

    except Exception as e:
        return f"Error processing {basename}: {e}"

    return None


def main():
    parser = argparse.ArgumentParser(
        description="Preprocess IRS dataset: convert EXR disparity to depth, "
        "copy RGB images, and save camera intrinsics."
    )
    parser.add_argument(
        "--root_dir",
        type=str,
        default="/path/to/data_raw_videos/data_irs",
        help="Root directory of the raw IRS data.",
    )
    parser.add_argument(
        "--out_dir",
        type=str,
        default="/path/to/data_raw_videos/processed_irs",
        help="Output directory for processed IRS data.",
    )
    args = parser.parse_args()

    # Example parameters (adjust as needed)
    baseline = 0.1
    f = 480

    root = args.root_dir
    out_dir = args.out_dir

    # Gather sequence directories.
    seq_dirs = []
    for d in os.listdir(root):
        if os.path.isdir(os.path.join(root, d)):
            if d == "Store":
                for sub in os.listdir(os.path.join(root, d)):
                    if os.path.isdir(os.path.join(root, d, sub)):
                        seq_dirs.append(os.path.join(d, sub))
            elif d == "IRS_small":
                for sub in os.listdir(os.path.join(root, d)):
                    if os.path.isdir(os.path.join(root, d, sub)):
                        for subsub in os.listdir(os.path.join(root, d, sub)):
                            if os.path.isdir(os.path.join(root, d, sub, subsub)):
                                seq_dirs.append(os.path.join(d, sub, subsub))
            else:
                seq_dirs.append(d)

    seq_dirs.sort()

    # Process each sequence.
    for seq in seq_dirs:
        seq_dir = os.path.join(root, seq)
        out_rgb_dir = os.path.join(out_dir, seq, "rgb")
        out_depth_dir = os.path.join(out_dir, seq, "depth")
        out_cam_dir = os.path.join(out_dir, seq, "cam")

        os.makedirs(out_rgb_dir, exist_ok=True)
        os.makedirs(out_depth_dir, exist_ok=True)
        os.makedirs(out_cam_dir, exist_ok=True)

        # Get basenames from disparity files.
        basenames = sorted([d[2:-4] for d in os.listdir(seq_dir) if d.endswith(".exr")])

        tasks = []
        for basename in basenames:
            task = (
                basename,
                seq_dir,
                out_rgb_dir,
                out_depth_dir,
                out_cam_dir,
                f,
                baseline,
            )
            tasks.append(task)

        num_workers = os.cpu_count() // 2
        with ProcessPoolExecutor(max_workers=num_workers) as executor:
            futures = {
                executor.submit(process_basename, task): task[0] for task in tasks
            }
            for future in tqdm(
                as_completed(futures), total=len(futures), desc=f"Processing {seq}"
            ):
                error = future.result()
                if error:
                    print(error)


if __name__ == "__main__":
    main()