pointcept/datasets/preprocessing/concerto/cap3d/preprocess_cap3d.py

"""
Preprocessing Script for Cap3D

Author: Yujia Zhang (yujia.zhang.cs@gmail.com)
Please cite our work if the code is helpful to you.
"""

import os
import argparse
import torch
import numpy as np
from scipy.spatial import cKDTree
from PIL import Image
from pathlib import Path
import multiprocessing as mp
import zipfile
import glob
import json
import math

NUM_VIEWS = 4
MAX_DEPTH = 5.0


def recover_depth_from_image(depth_image_path, max_depth):
    with Image.open(depth_image_path) as depth_pil:
        depth_uint16 = np.array(depth_pil)
    invalid_depth_mask = depth_uint16 == 65535
    depth_normalized_float = depth_uint16.astype(np.float32) / 65535.0
    depth_meters = depth_normalized_float * max_depth
    return depth_meters, invalid_depth_mask


def normalize_point_cloud(point_cloud: np.ndarray) -> np.ndarray:
    bbox_min = np.min(point_cloud, axis=0)
    bbox_max = np.max(point_cloud, axis=0)
    bbox_size = bbox_max - bbox_min
    max_dim = np.max(bbox_size)
    if max_dim < 1e-8:
        return point_cloud - bbox_min

    scale = 1.0 / max_dim
    scaled_point_cloud = point_cloud * scale
    new_bbox_min = np.min(scaled_point_cloud, axis=0)
    new_bbox_max = np.max(scaled_point_cloud, axis=0)
    center = (new_bbox_min + new_bbox_max) / 2.0
    offset = -center
    normalized_point_cloud = scaled_point_cloud + offset

    return normalized_point_cloud


def get_cam_params(cam_params_file, width, height):
    with open(cam_params_file, "r") as f:
        data = json.load(f)

    x_vector = data["x"]
    y_vector = data["y"]
    z_vector = data["z"]
    origin = data["origin"]

    rotation_matrix = np.array([x_vector, y_vector, z_vector]).T

    translation_vector = np.array(origin)

    rt_matrix = np.eye(4)
    rt_matrix[:3, :3] = rotation_matrix
    rt_matrix[:3, 3] = translation_vector

    x_fov = data["x_fov"]
    y_fov = data["y_fov"]

    cx = width / 2.0
    cy = height / 2.0

    fx = (width / 2.0) / math.tan(x_fov / 2.0)
    fy = (height / 2.0) / math.tan(y_fov / 2.0)

    s = 0

    K = np.array([[fx, s, cx], [0, fy, cy], [0, 0, 1]])

    return rt_matrix, K


def correspondenceGet(depth, invalid_depth_mask, T, K, width, height):
    pixel = np.transpose(np.indices((width, height)), (2, 1, 0))
    pixel = pixel.reshape((-1, 2))
    pixel = np.hstack((pixel, np.ones((pixel.shape[0], 1))))
    depth = depth.reshape((-1, 1))
    invalid_depth_mask = invalid_depth_mask.reshape((-1, 1))
    valid = ~np.logical_or(np.isinf(depth), invalid_depth_mask).squeeze(-1)
    coord = np.zeros_like(pixel, dtype=np.float32)
    coord[valid] = depth[valid] * (np.linalg.inv(K) @ pixel[valid].T).T  # coord_camera
    coord[valid] = coord[valid] @ T[:3, :3].T + T[:3, 3]  # column then row

    pixel = pixel[valid]
    coord = coord[valid]
    if coord.shape[0] == 0:
        return None
    pixel = pixel[:, :2]
    coord_dict = {"pixel": pixel, "coord": coord}
    return coord_dict


def correspondenceSave(
    depth_paths, data_name, point_cloud_gt, cam_params_paths, output_root
):
    """
    Calculates and saves the correspondences for all views of a given mesh and point cloud.
    """
    # Build a KDTree for the ground truth point cloud for fast querying.
    tree = cKDTree(point_cloud_gt)

    # Iterate over each predefined view.
    for view_idx in range(len(depth_paths)):
        # Construct file paths.
        cam_params_path = cam_params_paths[view_idx]
        depth_path = depth_paths[view_idx]

        view_name = cam_params_path.split("/")[-1][:-5]

        # Load camera parameters and image dimensions.
        depth, invalid_depth_mask = recover_depth_from_image(depth_path, MAX_DEPTH)
        height, width = depth.shape
        T, K = get_cam_params(cam_params_path, width, height)

        coord_dict = correspondenceGet(depth, invalid_depth_mask, T, K, width, height)
        if coord_dict is None:
            correspondences = -np.ones((1, 3))
        else:
            pixels_ = coord_dict["pixel"]
            coords_ = coord_dict["coord"]
            dis, idx = tree.query(coords_, k=1)
            idx_valid = idx[dis < 0.01]
            pixels_valid = pixels_[dis < 0.01]
            correspondences = np.hstack((pixels_valid, idx_valid.reshape(-1, 1)))
        output_dir = output_root / data_name
        os.makedirs(output_dir, exist_ok=True)
        output_filename = output_dir / f"{view_name}.npy"
        np.save(output_filename, correspondences)


def handle_process(pt_path, output_root, cam_root):
    """
    The complete processing pipeline for a single GLB file.
    """
    data_name = pt_path.stem

    if not pt_path.exists():
        return
    cam_path = cam_root / f"{data_name}"
    cam_path_zip = cam_root / f"{data_name}.zip"

    # Check if files exist, just in case.
    if not cam_path.exists():
        if not cam_path_zip.exists():
            print(f"not exist {str(cam_path)} and its zip file")
            return
        with zipfile.ZipFile(cam_path_zip, "r") as zip_ref:
            zip_ref.extractall(cam_path)

    # Load the point cloud.
    point_cloud = torch.load(pt_path)
    point_cloud_np = (
        point_cloud.cpu().numpy()
        if hasattr(point_cloud, "numpy")
        else np.array(point_cloud)
    )
    point_cloud_np = point_cloud_np[:3, :].T
    depth_paths = glob.glob(os.path.join(cam_path, "*_depth.png"))
    if len(depth_paths) >= NUM_VIEWS:
        depth_paths = sorted(depth_paths)[:: len(depth_paths) // NUM_VIEWS]
    else:
        depth_paths = sorted(depth_paths)
    cam_params_paths = glob.glob(os.path.join(cam_path, "*.json"))
    if os.path.isfile(os.path.join(cam_path, "transforms_train.json")):
        cam_params_paths.remove(os.path.join(cam_path, "transforms_train.json"))
    if os.path.isfile(os.path.join(cam_path, "info.json")):
        cam_params_paths.remove(os.path.join(cam_path, "info.json"))
    if len(depth_paths) >= NUM_VIEWS:
        cam_params_paths = sorted(cam_params_paths)[
            :: len(cam_params_paths) // NUM_VIEWS
        ]
    else:
        cam_params_paths = sorted(cam_params_paths)

    valid_depth_paths = {
        os.path.basename(p).replace("_depth.png", ""): p for p in depth_paths
    }
    valid_cam_params_paths = {
        os.path.basename(p).replace(".json", ""): p for p in cam_params_paths
    }
    common_ids = set(valid_depth_paths.keys()) & set(valid_cam_params_paths.keys())
    sorted_common_ids = sorted(list(common_ids))

    valid_depth_paths_ = [valid_depth_paths[uid] for uid in sorted_common_ids]
    valid_cam_params_paths_ = [valid_cam_params_paths[uid] for uid in sorted_common_ids]
    # Calculate and save correspondences for all views.
    correspondenceSave(
        valid_depth_paths_,
        data_name,
        point_cloud_np,
        valid_cam_params_paths_,
        output_root,
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Generate 2D-3D correspondences for GLB models and point clouds."
    )
    parser.add_argument(
        "--cam_root",
        type=Path,
        required=True,
        help="Root directory containing camera parameters.",
    )
    parser.add_argument(
        "--point_cloud_root",
        type=Path,
        required=True,
        help="Directory containing .pt point cloud files.",
    )
    parser.add_argument(
        "--output_root",
        type=Path,
        required=True,
        help="Output directory to save correspondence files.",
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=mp.cpu_count(),
        help="Number of processes to use for processing.",
    )
    parser.add_argument(
        "--thread_id", type=int, default=0, help="thread id to use for processing."
    )
    config = parser.parse_args()

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

    print("Scanning for .pt files...")
    pt_files = list(config.point_cloud_root.glob("**/*.pt"))
    print(f"Found {len(pt_files)} .pt files.")

    # Load scene paths
    pt_files_list = np.array_split(pt_files, config.num_workers)
    pt_files_ = pt_files_list[config.thread_id]
    # Preprocess data.
    print("Processing scenes...")
    for pt_files_i in pt_files_:
        handle_process(
            pt_files_i,
            config.output_root,
            config.cam_root,
        )

    print("\nAll files processed.")