generate_spatial_img_coco.py

import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
import argparse
import copy
import json
import logging
from pathlib import Path

import colorsys
import cv2
import numpy as np
from PIL import Image
import torch
from pycocotools.coco import COCO
from tqdm import tqdm

from moge.model import MoGeModel
from perspective2d import PerspectiveFields

import utils3d


def generate_unique_colors(n):
    """
    Generate n unique colors.

    Args:
        n (int): Number of unique colors to generate.

    Returns:
        List[List[int]]: List of RGB colors.
    """
    colors = []
    for i in range(n):
        hue = i / n
        saturation = 0.7
        value = 0.9
        rgb_float = colorsys.hsv_to_rgb(hue, saturation, value)
        rgb = [int(c * 255) for c in rgb_float]
        colors.append(rgb)
    return colors


def create_rotation_matrix(
    roll: float,
    pitch: float,
    yaw: float,
    degrees: bool = False,
) -> np.ndarray:
    r"""Create rotation matrix from extrinsic parameters
    Args:
        roll (float): camera rotation about camera frame z-axis
        pitch (float): camera rotation about camera frame x-axis
        yaw (float): camera rotation about camera frame y-axis

    Returns:
        np.ndarray: rotation R_z @ R_x @ R_y
    """
    if degrees:
        roll = np.radians(roll)
        pitch = np.radians(pitch)
        yaw = np.radians(yaw)
    # calculate rotation about the x-axis
    R_x = np.array(
        [
            [1.0, 0.0, 0.0],
            [0.0, np.cos(pitch), np.sin(pitch)],
            [0.0, -np.sin(pitch), np.cos(pitch)],
        ]
    )
    # calculate rotation about the y-axis
    R_y = np.array(
        [
            [np.cos(yaw), 0.0, -np.sin(yaw)],
            [0.0, 1.0, 0.0],
            [np.sin(yaw), 0.0, np.cos(yaw)],
        ]
    )
    # calculate rotation about the z-axis
    R_z = np.array(
        [
            [np.cos(roll), np.sin(roll), 0.0],
            [-np.sin(roll), np.cos(roll), 0.0],
            [0.0, 0.0, 1.0],
        ]
    )

    return R_z @ R_x @ R_y

# Configure logging to record processing information and errors
logging.basicConfig(
    filename='log_generate_spatial_img_coco.log',        # Log file name
    filemode='a',                                        # Append mode
    format='%(asctime)s - %(levelname)s - %(message)s',  # Log format
    level=logging.INFO                                   # Logging level
)


def save_point_cloud_xyzrgb(output_ply_path, points, colors):
    """
    Save point cloud data in XYZRGB format.

    Args:
        output_ply_path (str): Path where the PLY file will be saved.
        points (np.ndarray): Array of 3D points.
        colors (np.ndarray): Array of RGB colors.
    """
    data = np.hstack([points, colors])
    header = '''ply
format ascii 1.0
element vertex {n_vertices}
property float x
property float y
property float z
property uchar red
property uchar green
property uchar blue
end_header
'''.format(n_vertices=len(data))
    with open(output_ply_path, 'w') as f:
        f.write(header)
        np.savetxt(f, data, fmt='%f %f %f %d %d %d')

def process_image(image_path: str, output_rgb_path: str, output_ply_path: str, output_depth_path: str, output_depth_no_edge_path: str, output_params_path: str,
                 coco: COCO, output_json_path: str, model: MoGeModel, depth_model, pf_model: PerspectiveFields, camera_intrinsic_model,
                 remove_edge: bool = True):
    """
    Process a single image to generate a 3D point cloud, depth map, camera parameters, and save them as PLY,
    depth image, camera parameters JSON, and category information JSON files.

    Args:
        image_path (str): Path to the input image.
        output_rgb_path (str): Path where the RGB mask image will be saved.
        output_ply_path (str): Path where the PLY file will be saved.
        output_depth_path (str): Path where the depth map will be saved.
        output_depth_no_edge_path (str): Path where the depth map without edges will be saved.
        output_params_path (str): Path where the camera parameters will be saved.
        coco (COCO): COCO API object.
        output_json_path (str): Path where the JSON file will be saved.
        model (MoGeModel): Loaded MoGeModel for depth and point cloud inference.
        depth_model: Loaded depth model for inference.
        pf_model (PerspectiveFields): Loaded PerspectiveFields model for rotation predictions.
        camera_intrinsic_model: Loaded camera intrinsic model for inference.
        remove_edge (bool): Whether to remove edges from the depth map.
    """
    # Derive scene_id from the image filename (assuming filename is the scene_id)
    scene_id = Path(image_path).stem

    # Load image info from COCO
    img_info = coco.loadImgs(int(scene_id))[0]

    # Get all annotation IDs for the image
    ann_ids = coco.getAnnIds(imgIds=img_info['id'], iscrowd=None)
    anns = coco.loadAnns(ann_ids)
    num_instances = len(anns)

    if num_instances == 0:
        logging.info(f"No annotations found for image: {image_path}. Skipping.")
        return

    unique_instance_colors = generate_unique_colors(num_instances)
    instance_color_map = {ann['id']: color for ann, color in zip(anns, unique_instance_colors)}

    # Read and convert the image using Pillow for efficiency
    color_image = Image.open(image_path).convert('RGB')
    image = np.array(color_image)
    height, width = image.shape[:2]

    # Generate RGB mask image based on category_color_map and annotations
    rgb_mask = generate_rgb_mask(image_size=(height, width), anns=anns, instance_color_map=instance_color_map, coco=coco)

    # Convert the image to a PyTorch tensor and normalize
    image_tensor = torch.tensor(image, dtype=torch.float32, device='cuda').permute(2, 0, 1) / 255.0
    image_tensor = image_tensor.unsqueeze(0)  # Add batch dimension

    with torch.no_grad():
        # Perform inference to get depth and point cloud data
        predictions = pf_model.inference(img_bgr=image)
        intrinsics, _ = camera_intrinsic_model.inference(color_image, wtassumption=False)

        input_size = (616, 1064)
        scale = min(input_size[0] / height, input_size[1] / width)
        intrinsic_matrix_scaled = intrinsics * scale
        rgb_resized = cv2.resize(image, (int(width * scale), int(height * scale)), interpolation=cv2.INTER_LINEAR)

        padding = [123.675, 116.28, 103.53]
        pad_h = input_size[0] - rgb_resized.shape[0]
        pad_w = input_size[1] - rgb_resized.shape[1]
        pad_h_half = pad_h // 2
        pad_w_half = pad_w // 2
        rgb_padded = cv2.copyMakeBorder(rgb_resized, pad_h_half, pad_h - pad_h_half, pad_w_half, pad_w - pad_w_half, cv2.BORDER_CONSTANT, value=padding)

        rgb_tensor = torch.from_numpy(rgb_padded.transpose((2, 0, 1))).float()
        rgb_tensor = ((torch.tensor([123.675, 116.28, 103.53]).float().view(3, 1, 1) - rgb_tensor) / torch.tensor([58.395, 57.12, 57.375]).float().view(3, 1, 1)).unsqueeze(0).cuda()

        pred_metric_depth, _, _ = depth_model.inference({'input': rgb_tensor})

        pred_metric_depth = pred_metric_depth.squeeze()
        pred_metric_depth = pred_metric_depth[pad_h_half: pred_metric_depth.shape[0] - (pad_h - pad_h_half), pad_w_half: pred_metric_depth.shape[1] - (pad_w - pad_w_half)]
        pred_metric_depth = torch.nn.functional.interpolate(pred_metric_depth.unsqueeze(0).unsqueeze(0), (height, width), mode='bilinear').squeeze()
        canonical_to_real_scale = intrinsic_matrix_scaled[0][0] / 1000.0
        pred_metric_depth = pred_metric_depth * canonical_to_real_scale

        # Extract rotation predictions
        pred_roll = predictions['pred_roll'].cpu().item()
        pred_pitch = predictions['pred_pitch'].cpu().item()

        # Create the rotation matrix based on predictions
        perspective_R = create_rotation_matrix(roll=pred_roll, pitch=pred_pitch, yaw=0, degrees=True)
        perspective_R[:, [1, 2]] = perspective_R[:, [2, 1]]  # Swap Y and Z axes
        perspective_T = np.zeros(3)  # Translation vector (unused)
        extrinsics = np.eye(4)
        extrinsics[:3, :3] = perspective_R
        extrinsics[:3, 3] = perspective_T

        output = model.infer(image_tensor, resolution_level=9, apply_mask=True, intrinsics=intrinsics, extrinsics=extrinsics, pred_metric_depth=pred_metric_depth)

    # Extract point cloud data from the model output
    points = output['points'].squeeze(0).cpu().numpy()    # Shape: (H, W, 3)
    depth = output['depth'].squeeze(0).cpu().numpy()      # Shape: (H, W)
    mask = output['mask'].squeeze(0).cpu().numpy()        # Shape: (H, W)
    intrinsics = output['intrinsics'].cpu().numpy()       # Shape: (3, 3)

    # Apply depth edge filtering if required
    if remove_edge:
        valid_mask = mask & ~utils3d.numpy.depth_edge(depth, mask=mask, rtol=0.02)
    else:
        valid_mask = mask

    # Flatten the points and colors
    valid_points = points[valid_mask]
    valid_colors = image[valid_mask]
    valid_colors_mask = rgb_mask[valid_mask]

    # Adjust vertex coordinates to match the coordinate system
    valid_points = valid_points * [1, -1, -1]

    # Ensure vertex colors are within [0, 255] and convert to uint8
    valid_colors = valid_colors.astype(np.uint8)

    # Validate that the number of vertex colors matches the number of vertices
    assert valid_colors.shape[0] == valid_points.shape[0], "Number of vertex colors does not match number of vertices"
    assert valid_colors.dtype == np.uint8, "vertex_colors should be of type uint8"
    
    # save rgb
    cv2.imwrite(output_rgb_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))

    # Save the point cloud in XYZRGB format
    save_point_cloud_xyzrgb(output_ply_path, valid_points, valid_colors)
    logging.info(f"PLY file saved to: {output_ply_path}")

    # Save depth map with colorization for better visualization
    depth[~mask] = 0  # Set invalid depth pixels to zero
    depth = np.nan_to_num(depth, nan=0.0, posinf=0.0, neginf=0.0)
    depth_mm = depth * 1000.0  # Convert depth to millimeters
    depth_uint16 = depth_mm.astype(np.uint16)
    cv2.imwrite(output_depth_path, depth_uint16)
    logging.info(f"Depth map saved to: {output_depth_path}")

    depth_no_edges = depth.copy()
    depth_no_edges[~valid_mask] = 0
    depth_no_edges = np.nan_to_num(depth_no_edges, nan=0.0, posinf=0.0, neginf=0.0)
    depth_no_edges_mm = depth_no_edges * 1000.0
    depth_no_edges_uint16 = depth_no_edges_mm.astype(np.uint16)
    cv2.imwrite(output_depth_no_edge_path, depth_no_edges_uint16)
    logging.info(f"Depth map with remove edge saved to: {output_depth_no_edge_path}")

    # Save camera parameters as a JSON file
    camera_params = {
        "intrinsics": intrinsics.tolist(),
        "rotation": perspective_R.tolist(),
        "translation": perspective_T.tolist()
    }
    with open(output_params_path, 'w') as params_file:
        json.dump(camera_params, params_file, indent=4)
    logging.info(f"Camera parameters saved to: {output_params_path}")

    # ========================== JSON Saving Logic ==========================

    reverse_color_map = {tuple(v): k for k, v in instance_color_map.items()}

    colors_as_tuples = [tuple(color) for color in valid_colors_mask]

    instance_ids = []
    for color in colors_as_tuples:
        instance_id = reverse_color_map.get(color, None)
        instance_ids.append(instance_id)

    instance_ids = np.array(instance_ids)

    unique_instance_ids = np.unique(instance_ids[instance_ids != None])

    # 按实例 ID 分组索引
    instance_indices = {inst_id: [] for inst_id in unique_instance_ids}
    for idx, inst_id in enumerate(instance_ids):
        if inst_id is not None:
            instance_indices[inst_id].append(int(idx))

    seg_groups = []
    for instance_id, indices in instance_indices.items():
        if not indices:
            continue
        # 获取对应的注释
        ann = coco.loadAnns(instance_id)[0]
        category_id = ann['category_id']
        cat_info = coco.loadCats(category_id)[0]
        cat_name = cat_info['name']
        points_subset = valid_points[indices]
        if len(points_subset) == 0:
            continue
        centroid = points_subset.mean(axis=0).tolist()
        seg_group = {
            "objectId": instance_id,
            "categoryId": category_id,
            "label": cat_name,
            "segments": indices,
            "centroid": centroid
        }
        seg_groups.append(seg_group)

    if seg_groups:
        json_data = {
            "sceneId": scene_id,
            "segGroups": seg_groups
        }
        with open(output_json_path, 'w') as json_file:
            json.dump(json_data, json_file, indent=4)
        logging.info(f"JSON file saved to: {output_json_path}")
    else:
        logging.info(f"No valid segments found for {image_path}. JSON file not created.")


def traverse_and_process(input_dir: str, output_dir: str, coco: COCO, model: MoGeModel, depth_model, pf_model: PerspectiveFields, camera_intrinsic_model,
                        remove_edge: bool = True):
    """
    Traverse the input directory, process each image, and save the corresponding PLY, depth map,
    camera parameters, and JSON files into their respective folders.

    Args:
        input_dir (str): Path to the input directory containing images.
        output_dir (str): Path to the output root directory where separate folders will be created.
        coco (COCO): COCO API object.
        model (MoGeModel): Loaded MoGeModel for depth and point cloud inference.
        depth_model: Loaded depth model for inference.
        pf_model (PerspectiveFields): Loaded PerspectiveFields model for rotation predictions.
        camera_intrinsic_model: Loaded camera intrinsic model for inference.
        remove_edge (bool): Whether to remove edges from the depth map.
    """
    input_path = Path(input_dir)
    output_path = Path(output_dir)

    if not input_path.is_dir():
        logging.error(f"The provided input path '{input_dir}' is not a directory.")
        raise NotADirectoryError(f"The provided input path '{input_dir}' is not a directory.")

    # Define subdirectories for different data types
    rgb_dir = output_path / "rgb"
    point_dir = output_path / "point"
    depth_dir = output_path / "depth"
    camera_params_dir = output_path / "camera_parameters"
    json_dir = output_path / "json"

    # Create subdirectories if they don't exist
    rgb_dir.mkdir(parents=True, exist_ok=True)
    point_dir.mkdir(parents=True, exist_ok=True)
    depth_dir.mkdir(parents=True, exist_ok=True)
    camera_params_dir.mkdir(parents=True, exist_ok=True)
    json_dir.mkdir(parents=True, exist_ok=True)

    # Supported image extensions
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}

    # Collect all image file paths
    tasks = []
    for root, _, files in os.walk(input_path):
        for file in files:
            file_path = Path(root) / file
            if file_path.suffix.lower() in image_extensions:
                tasks.append(file_path)

    logging.info(f"Total number of image files to process: {len(tasks)}")

    # Process each image sequentially with a progress bar
    for file_path in tqdm(tasks, desc="Processing images"):
        base_name = file_path.stem
        output_rgb_path = rgb_dir / f"{base_name}.png"
        output_ply_path = point_dir / f"{base_name}.ply"
        output_depth_path = depth_dir / f"{base_name}.png"
        output_depth_no_edge_path = depth_dir / f"{base_name}_remove_edges.png"
        output_params_path = camera_params_dir / f"{base_name}.json"
        output_json_path = json_dir / f"{base_name}.json"

        try:
            logging.info(f"Starting processing image: {file_path}")
            process_image(
                image_path=str(file_path),
                output_rgb_path=str(output_rgb_path),
                output_ply_path=str(output_ply_path),
                output_depth_path=str(output_depth_path),
                output_depth_no_edge_path=str(output_depth_no_edge_path),
                output_params_path=str(output_params_path),
                coco=coco,
                output_json_path=str(output_json_path),
                model=model,
                depth_model=depth_model,
                pf_model=pf_model,
                camera_intrinsic_model=camera_intrinsic_model,
                remove_edge=remove_edge
            )
            logging.info(f"Successfully processed image: {file_path}")
        except Exception as e:
            logging.error(f"Error processing image {file_path}: {e}", exc_info=True)
            continue

    logging.info("All images have been processed.")

def generate_rgb_mask(image_size, anns, instance_color_map, coco):
    """
    Generate an RGB mask image based on instance_color_map and COCO annotations.

    Args:
        image_size (tuple): Size of the image as (height, width).
        anns (list): List of COCO annotations for the image.
        instance_color_map (dict): Mapping from annotation_id to RGB color.
        coco (COCO): COCO API object.

    Returns:
        np.ndarray: RGB mask image of shape (H, W, 3) with dtype uint8.
    """
    height, width = image_size
    mask_image = np.zeros((height, width, 3), dtype=np.uint8)

    for ann in anns:
        instance_id = ann['id']
        color = instance_color_map.get(instance_id, [0, 0, 0])  # Default to black if category_id not found
        mask = coco.annToMask(ann)  # Binary mask of shape (H, W)

        # Assign the color to the mask regions
        mask_image[mask > 0] = color

    return mask_image

def parse_arguments():
    """
    Parse command-line arguments.

    Returns:
        argparse.Namespace: Parsed arguments.
    """
    parser = argparse.ArgumentParser(
        description="Traverse a directory of images, generate 3D point clouds, depth maps, camera parameters, "
                    "and save them as PLY, depth image, camera parameters JSON, and category information JSON files."
    )
    parser.add_argument(
        '--input_dir', '-i',
        type=str,
        default='./demo_inputs/',
        help='Path to the input directory containing images.'
    )
    parser.add_argument(
        '--output_dir', '-o',
        type=str,
        default=None,
        help='Path to the output root directory where separate folders will be created for PLY, depth maps, '
             'camera parameters, and JSON files. If not provided, the parent directory of input_dir will be used.'
    )
    parser.add_argument(
        '--ann_file', '-a',
        type=str,
        default='./instances_val2017.json',
        help='Path to the COCO annotation file (e.g., instances_train2017.json).'
    )
    parser.add_argument(
        '--moge_model_path', '-moge',
        type=str,
        default='Ruicheng/moge-vitl',
        help='Path to the pre-trained MoGeModel directory or specific model files.'
    )
    parser.add_argument(
        '--pf_model_path', '-pf',
        type=str,
        default='Paramnet-360Cities-edina-centered',
        help='Path to the pre-trained PerspectiveFields model directory or specific model files.'
    )
    parser.add_argument(
        '--camera_model_path', '-cmp',
        type=str,
        default='ShngJZ/WildCamera',
        help='Path to the pre-trained WildCamera model directory or specific model files.'
    )
    parser.add_argument(
        '--metric3d_model_path', '-mmp',
        type=str,
        default='yvanyin/metric3d',
        help='Path to the pre-trained Metric3D model directory or specific model files.'
    )
    parser.add_argument(
        '--no_remove_edge', '-nr',
        action='store_false',
        dest='remove_edge',
        help='Do not remove edges from the depth map. By default, edges are removed.'
    )
    return parser.parse_args()


def main():
    args = parse_arguments()

    input_directory = args.input_dir
    output_directory = args.output_dir
    ann_file = args.ann_file
    mo_ge_model_path = args.moge_model_path
    pf_model_path = args.pf_model_path
    remove_edge = args.remove_edge
    camera_model_path = args.camera_model_path
    metric3d_model_path = args.metric3d_model_path

    # If output directory is not specified, use the parent directory of input_dir
    if output_directory is None:
        input_path = Path(input_directory)
        output_directory = str(input_path.parent)

    # Initialize COCO API
    try:
        coco = COCO(ann_file)
        logging.info(f"Loaded COCO annotations from: {ann_file}")
    except Exception as e:
        logging.error(f"Error loading COCO annotation file: {e}", exc_info=True)
        return

    # Check if CUDA is available
    if not torch.cuda.is_available():
        logging.error("CUDA is not available. Please check your CUDA installation and GPU configuration.")
        return

    # Load MoGeModel from the specified mo_ge_model_path
    try:
        model = MoGeModel.from_pretrained(mo_ge_model_path).cuda().eval()
        logging.info("MoGeModel loaded and set to evaluation mode.")
    except Exception as e:
        logging.error(f"Error loading MoGeModel: {e}", exc_info=True)
        return

    # Load PerspectiveFields model from the specified pf_model_path
    try:
        pf_model = PerspectiveFields(pf_model_path).cuda().eval()
        logging.info("PerspectiveFields model loaded and set to evaluation mode.")
    except Exception as e:
        logging.error(f"Error loading PerspectiveFields model: {e}", exc_info=True)
        return
    
    try:
        # camera_intrinsic_model = torch.hub.load(camera_model_path, "WildCamera", pretrained=True).cuda().eval()
        camera_intrinsic_model = torch.hub.load('./WildCamera', "WildCamera", source='local', pretrained=True).cuda().eval()
        logging.info("camera_intrinsic_model model loaded and set to evaluation mode.")
    except Exception as e:
        logging.error(f"Error loading camera_intrinsic_model model: {e}", exc_info=True)
        return
    
    try:
        # depth_model = torch.hub.load(metric3d_model_path, 'metric3d_vit_giant2', pretrain=True).cuda().eval()
        depth_model = torch.hub.load('./Metric3D', 'metric3d_vit_giant2', source='local', pretrain=True).cuda().eval()
        logging.info("Metric3D loaded and set to evaluation mode.")
    except Exception as e:
        logging.error(f"Error loading Metric3D: {e}", exc_info=True)
        return

    # Start processing images
    traverse_and_process(input_directory, output_directory, coco, model, depth_model, pf_model, camera_intrinsic_model, remove_edge)


if __name__ == '__main__':
    main()