extern/CUT3R/datasets_preprocess/preprocess_synscapes.py

#!/usr/bin/env python3
"""
Preprocess Synscapes Data

This script processes Synscapes data by:
  1. Copying the RGB images.
  2. Reading the EXR depth data and saving it as .npy.
  3. Generating a sky mask using the class labels.
  4. Extracting camera intrinsics from the meta file.

The directory structure is expected to be:
    synscapes_dir/
        img/
            rgb/
            depth/
            class/
        meta/
    Each file shares the same base name, e.g. 000000.png/exr in corresponding folders.

Usage:
    python preprocess_synscapes.py \
        --synscapes_dir /path/to/Synscapes/Synscapes \
        --output_dir /path/to/processed_synscapes
"""

import os
import json
import shutil
import argparse
import numpy as np
import cv2
import OpenEXR
from tqdm import tqdm

# Enable EXR support in OpenCV if desired:
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"


def process_basename(
    basename,
    rgb_dir,
    depth_dir,
    class_dir,
    meta_dir,
    out_rgb_dir,
    out_depth_dir,
    out_mask_dir,
    out_cam_dir,
    sky_id=23,
):
    """
    Process a single sample of the Synscapes dataset:
      1. Reads an RGB .png and depth .exr file.
      2. Reads a class label .png, generating a sky mask.
      3. Reads camera intrinsics from the meta .json file.
      4. Saves the resulting data to the specified output folders.

    Args:
        basename (str): The base filename (without extension).
        rgb_dir (str): Directory containing RGB .png files.
        depth_dir (str): Directory containing depth .exr files.
        class_dir (str): Directory containing class .png files.
        meta_dir (str): Directory containing camera metadata .json files.
        out_rgb_dir (str): Output directory for RGB files.
        out_depth_dir (str): Output directory for depth .npy files.
        out_mask_dir (str): Output directory for sky masks.
        out_cam_dir (str): Output directory for camera intrinsics (.npz).
        sky_id (int): Class ID for sky pixels in the class label images.

    Returns:
        None or str:
            If an error occurs, returns an error message (str). Otherwise, returns None.
    """
    try:
        # Input file paths
        rgb_file = os.path.join(rgb_dir, f"{basename}.png")
        depth_file = os.path.join(depth_dir, f"{basename}.exr")
        class_file = os.path.join(class_dir, f"{basename}.png")
        meta_file = os.path.join(meta_dir, f"{basename}.json")

        # Output file paths
        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_mask_path = os.path.join(out_mask_dir, f"{basename}.png")
        out_cam_path = os.path.join(out_cam_dir, f"{basename}.npz")

        # --- Read Depth Data ---
        # If you want to use OpenEXR directly (matching your code), do so here:
        exr_file = OpenEXR.InputFile(depth_file)
        # e.g. reading "Z" channel. Adjust channel name as needed.
        # It's possible that the data is stored in multiple channels (R/G/B or separate "Z").
        # Check your file structure to match the correct channel name.
        # The snippet below is just an example approach using .parts and .channels.
        # If your EXR file is a single-part file with a standard channel, you'd do something like:
        #   depth = np.frombuffer(exr_file.channel('Z', Imath.PixelType(Imath.PixelType.FLOAT)), dtype=np.float32)
        # The way you've shown "parts[0].channels['Z'].pixels" may or may not be valid for your version of PyOpenEXR.

        # This example code is approximate and may need to be adapted:
        # If your version of OpenEXR has a different interface, change accordingly.
        # The snippet below won't work unless you install a specific PyOpenEXR wrapper that supports .parts, .channels, etc.
        #
        # For demonstration, let's assume a single-part EXR with channel 'Z':
        # depth_data = exr_file.channel('Z')  # returns raw bytes
        # depth = np.frombuffer(depth_data, dtype=np.float32).reshape((height, width))  # you need to know (height, width) or read header

        # As you mentioned "np.array(OpenEXR.File(depth_file).parts[0].channels['Z'].pixels)",
        # let's keep it consistent with your original snippet:
        depth = np.array(OpenEXR.InputFile(depth_file).parts[0].channels["Z"].pixels)

        # --- Read Class Image (for Sky Mask) ---
        class_img = cv2.imread(class_file, cv2.IMREAD_UNCHANGED)
        # Create sky mask
        sky_mask = (class_img == sky_id).astype(np.uint8) * 255

        # --- Read Meta Data (for Camera Intrinsics) ---
        with open(meta_file, "r") as f:
            cam_info = json.load(f)["camera"]
            intrinsic = cam_info["intrinsic"]
            fx, fy, cx, cy = (
                intrinsic["fx"],
                intrinsic["fy"],
                intrinsic["u0"],
                intrinsic["v0"],
            )

        K = np.eye(3, dtype=np.float32)
        K[0, 0] = fx
        K[1, 1] = fy
        K[0, 2] = cx
        K[1, 2] = cy

        # --- Copy RGB ---
        shutil.copy(rgb_file, out_img_path)

        # --- Save Depth, Mask, and Intrinsics ---
        np.save(out_depth_path, depth)
        cv2.imwrite(out_mask_path, sky_mask)
        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 Synscapes data.")
    parser.add_argument(
        "--synscapes_dir",
        required=True,
        help="Path to the main Synscapes directory (contains 'img' and 'meta' folders).",
    )
    parser.add_argument(
        "--output_dir",
        required=True,
        help="Path to the output directory for processed data.",
    )
    parser.add_argument(
        "--sky_id",
        type=int,
        default=23,
        help="Class ID for sky pixels in class .png. Default is 23.",
    )
    args = parser.parse_args()

    synscapes_dir = os.path.abspath(args.synscapes_dir)
    output_dir = os.path.abspath(args.output_dir)
    os.makedirs(output_dir, exist_ok=True)

    # Define input subdirectories
    rgb_dir = os.path.join(synscapes_dir, "img", "rgb")
    depth_dir = os.path.join(synscapes_dir, "img", "depth")
    class_dir = os.path.join(synscapes_dir, "img", "class")
    meta_dir = os.path.join(synscapes_dir, "meta")

    # Define output subdirectories
    out_rgb_dir = os.path.join(output_dir, "rgb")
    out_depth_dir = os.path.join(output_dir, "depth")
    out_mask_dir = os.path.join(output_dir, "sky_mask")
    out_cam_dir = os.path.join(output_dir, "cam")
    for d in [out_rgb_dir, out_depth_dir, out_mask_dir, out_cam_dir]:
        os.makedirs(d, exist_ok=True)

    # Collect all EXR depth filenames (excluding extension)
    basenames = sorted(
        [
            os.path.splitext(fname)[0]
            for fname in os.listdir(depth_dir)
            if fname.endswith(".exr")
        ]
    )

    # Parallel processing
    from concurrent.futures import ProcessPoolExecutor, as_completed

    num_workers = max(1, os.cpu_count() // 2)

    with ProcessPoolExecutor(max_workers=num_workers) as executor:
        future_to_basename = {
            executor.submit(
                process_basename,
                bname,
                rgb_dir,
                depth_dir,
                class_dir,
                meta_dir,
                out_rgb_dir,
                out_depth_dir,
                out_mask_dir,
                out_cam_dir,
                args.sky_id,
            ): bname
            for bname in basenames
        }

        for future in tqdm(
            as_completed(future_to_basename),
            total=len(future_to_basename),
            desc="Processing Synscapes",
        ):
            basename = future_to_basename[future]
            error = future.result()
            if error:
                print(error)


if __name__ == "__main__":
    main()