extern/CUT3R/surfel_inference.py

#!/usr/bin/env python3
"""
3D Point Cloud Inference and Visualization Script

This script performs inference using the ARCroco3DStereo model and visualizes the
resulting 3D point clouds with the PointCloudViewer. Use the command-line arguments
to adjust parameters such as the model checkpoint path, image sequence directory,
image size, device, etc.

Usage:
    python demo_ga.py [--model_path MODEL_PATH] [--seq_path SEQ_PATH] [--size IMG_SIZE]
                            [--device DEVICE] [--vis_threshold VIS_THRESHOLD] [--output_dir OUT_DIR]

Example:
    python demo_ga.py --model_path src/cut3r_512_dpt_4_64.pth \
        --seq_path examples/001 --device cuda --size 512
"""

import os
import numpy as np
import torch
import time
import glob
import random
import cv2
import argparse
import tempfile
import shutil
from copy import deepcopy
from add_ckpt_path import add_path_to_dust3r
import imageio.v2 as iio

# Set random seed for reproducibility.
random.seed(42)



def listify(elems):
    return [x for e in elems for x in e]


def collate_with_cat(whatever, lists=False):
    if isinstance(whatever, dict):
        return {k: collate_with_cat(vals, lists=lists) for k, vals in whatever.items()}

    elif isinstance(whatever, (tuple, list)):
        if len(whatever) == 0:
            return whatever
        elem = whatever[0]
        T = type(whatever)

        if elem is None:
            return None
        if isinstance(elem, (bool, float, int, str)):
            return whatever
        if isinstance(elem, tuple):
            return T(collate_with_cat(x, lists=lists) for x in zip(*whatever))
        if isinstance(elem, dict):
            return {
                k: collate_with_cat([e[k] for e in whatever], lists=lists) for k in elem
            }

        if isinstance(elem, torch.Tensor):
            return listify(whatever) if lists else torch.cat(whatever)
        if isinstance(elem, np.ndarray):
            return (
                listify(whatever)
                if lists
                else torch.cat([torch.from_numpy(x) for x in whatever])
            )

        # otherwise, we just chain lists
        return sum(whatever, T())


def parse_args():
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        description="Run 3D point cloud inference and visualization using ARCroco3DStereo."
    )
    parser.add_argument(
        "--model_path",
        type=str,
        default="src/cut3r_512_dpt_4_64.pth",
        help="Path to the pretrained model checkpoint.",
    )
    parser.add_argument(
        "--seq_path",
        type=str,
        default="",
        help="Path to the directory containing the image sequence.",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda",
        help="Device to run inference on (e.g., 'cuda' or 'cpu').",
    )
    parser.add_argument(
        "--size",
        type=int,
        default="512",
        help="Shape that input images will be rescaled to; if using 224+linear model, choose 224 otherwise 512",
    )
    parser.add_argument(
        "--vis_threshold",
        type=float,
        default=1.5,
        help="Visualization threshold for the point cloud viewer. Ranging from 1 to INF",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="./demo_tmp",
        help="value for tempfile.tempdir",
    )

    return parser.parse_args()


def prepare_input(
    img_paths, img_mask, size, raymaps=None, raymap_mask=None, revisit=1, update=True
):
    """
    Prepare input views for inference from a list of image paths.

    Args:
        img_paths (list): List of image file paths.
        img_mask (list of bool): Flags indicating valid images.
        size (int): Target image size.
        raymaps (list, optional): List of ray maps.
        raymap_mask (list, optional): Flags indicating valid ray maps.
        revisit (int): How many times to revisit each view.
        update (bool): Whether to update the state on revisits.

    Returns:
        list: A list of view dictionaries.
    """
    # Import image loader (delayed import needed after adding ckpt path).
    from src.dust3r.utils.image import load_images

    images = load_images(img_paths, size=size)
    num_views = len(images)
    views = []

    for i in range(num_views):
        
        view = {
        "img": images[i]["img"],
        "ray_map": torch.full(
                (
                    images[i]["img"].shape[0],
                    6,
                    images[i]["img"].shape[-2],
                    images[i]["img"].shape[-1],
                ),
                torch.nan,
            ),
            "true_shape": torch.from_numpy(images[i]["true_shape"]),
            "idx": i,
            "instance": str(i),
            "camera_pose": torch.from_numpy(np.eye(4).astype(np.float32)).unsqueeze(
                0
            ),
            "img_mask": torch.tensor(True).unsqueeze(0),
            "ray_mask": torch.tensor(False).unsqueeze(0),
            "update": torch.tensor(True).unsqueeze(0),
            "reset": torch.tensor(False).unsqueeze(0),
        }
        views.append(view)
    return views


def prepare_output(output, poses, depths, lr, niter, outdir, device, save_flag=False):
    from cloud_opt.dust3r_opt import global_aligner, GlobalAlignerMode

    with torch.enable_grad():
        mode = GlobalAlignerMode.PointCloudOptimizer
        
        scene = global_aligner(
            output,
            device=device,
            mode=mode,
            verbose=True,
        )
        if depths is not None:
            scene.preset_depth(depths)
        if poses is not None:
            scene.preset_pose(poses)
        
        loss = scene.compute_global_alignment(
            init="mst",
            niter=niter,
            schedule="linear",
            lr=lr,
        )
    scene.clean_pointcloud()
    pts3d = scene.get_pts3d()
    depths = scene.get_depthmaps()
    poses = scene.get_im_poses()
    focals = scene.get_focals()
    pps = scene.get_principal_points()
    confs = scene.get_conf(mode="none")

    pts3ds_other = [pts.detach().cpu().unsqueeze(0) for pts in pts3d]
    depths = [d.detach().cpu().unsqueeze(0) for d in depths]
    colors = [torch.from_numpy(img).unsqueeze(0) for img in scene.imgs]
    confs = [conf.detach().cpu().unsqueeze(0) for conf in confs]
    cam_dict = {
        "focal": focals.detach().cpu().numpy(),
        "pp": pps.detach().cpu().numpy(),
        "R": poses.detach().cpu().numpy()[..., :3, :3],
        "t": poses.detach().cpu().numpy()[..., :3, 3],
    }
    if save_flag:
        depths_tosave = torch.cat(depths)  # B, H, W
        pts3ds_other_tosave = torch.cat(pts3ds_other)  # B, H, W, 3
        conf_self_tosave = torch.cat(confs)  # B, H, W
        colors_tosave = torch.cat(colors)  # [B, H, W, 3]
        cam2world_tosave = poses.detach().cpu()  # B, 4, 4
        intrinsics_tosave = (
            torch.eye(3).unsqueeze(0).repeat(cam2world_tosave.shape[0], 1, 1)
        )  # B, 3, 3
        intrinsics_tosave[:, 0, 0] = focals[:, 0].detach().cpu()
        intrinsics_tosave[:, 1, 1] = focals[:, 0].detach().cpu()
        intrinsics_tosave[:, 0, 2] = pps[:, 0].detach().cpu()
        intrinsics_tosave[:, 1, 2] = pps[:, 1].detach().cpu()

    
        os.makedirs(os.path.join(outdir, "depth"), exist_ok=True)
        os.makedirs(os.path.join(outdir, "conf"), exist_ok=True)
        os.makedirs(os.path.join(outdir, "color"), exist_ok=True)
        os.makedirs(os.path.join(outdir, "camera"), exist_ok=True)
        
        for f_id in range(len(depths_tosave)):
            depth = depths_tosave[f_id].cpu().numpy()
            conf = conf_self_tosave[f_id].cpu().numpy()
            color = colors_tosave[f_id].cpu().numpy()
            c2w = cam2world_tosave[f_id].cpu().numpy()
            intrins = intrinsics_tosave[f_id].cpu().numpy()
            np.save(os.path.join(outdir, "depth", f"{f_id:06d}.npy"), depth)
            np.save(os.path.join(outdir, "conf", f"{f_id:06d}.npy"), conf)
            iio.imwrite(
                os.path.join(outdir, "color", f"{f_id:06d}.png"),
                (color * 255).astype(np.uint8),
            )
            np.savez(
                os.path.join(outdir, "camera", f"{f_id:06d}.npz"),
                pose=c2w,
                intrinsics=intrins,
            )

    return pts3ds_other, colors, depths, confs, cam_dict


def parse_seq_path(p):
    if os.path.isdir(p):
        img_paths = sorted(glob.glob(f"{p}/*"))
        tmpdirname = None
    else:
        cap = cv2.VideoCapture(p)
        if not cap.isOpened():
            raise ValueError(f"Error opening video file {p}")
        video_fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        if video_fps == 0:
            cap.release()
            raise ValueError(f"Error: Video FPS is 0 for {p}")
        frame_interval = 1
        frame_indices = list(range(0, total_frames, frame_interval))
        print(
            f" - Video FPS: {video_fps}, Frame Interval: {frame_interval}, Total Frames to Read: {len(frame_indices)}"
        )
        img_paths = []
        tmpdirname = tempfile.mkdtemp()
        for i in frame_indices:
            cap.set(cv2.CAP_PROP_POS_FRAMES, i)
            ret, frame = cap.read()
            if not ret:
                break
            frame_path = os.path.join(tmpdirname, f"frame_{i}.jpg")
            cv2.imwrite(frame_path, frame)
            img_paths.append(frame_path)
        cap.release()
    return img_paths, tmpdirname




def run_inference_from_pil(
    pil_images, 
    model,
    poses=None,
    depths=None,
    lr = 0.01,
    niter = 300,
    device="cuda", 
    size=512, 
    output_dir="./demo_tmp", 
    visualize=False, 
    vis_threshold=1.5,
    save_flag=False
):
    """
    Run 3D reconstruction from a list of PIL images.
    
    Args:
        pil_images (list): List of PIL image objects.
        poses (list): List of camera poses.
        model_path (str): Path to the pretrained model checkpoint.
        device (str): Device to run inference on ('cuda' or 'cpu').
        size (int): Target image size for processing.
        output_dir (str): Directory to save outputs.
        visualize (bool): Whether to launch the point cloud viewer.
        vis_threshold (float): Visualization threshold for point cloud viewer.
        
    Returns:
        dict: A dictionary containing the reconstruction results:
            - point_clouds: List of point cloud tensors
            - colors: List of color tensors
            - confidences: List of confidence tensors
            - camera_info: Camera parameters dictionary
    """
    # Set up the computation device
    if device == "cuda" and not torch.cuda.is_available():
        print("CUDA not available. Switching to CPU.")
        device = "cpu"

    # Add the checkpoint path (required for model imports in the dust3r package)
    

    # Import model and inference functions after adding the ckpt path
    from src.dust3r.inference import inference, inference_recurrent
 
    
    # Prepare input views directly from PIL images
    print(f"Processing {len(pil_images)} images...")
    views = prepare_input_from_pil(
        pil_images=pil_images,
        size=size,
        revisit=1,
        update=True,
    )

    # Run inference
    print("Running inference...")
    start_time = time.time()
    output = {
        "view1": [],
        "view2": [],
        "pred1": [],
        "pred2": [],
    }
    edges = []

    outputs, state_args = inference(views, model, device)
    for view_id in range(1, len(outputs["views"])):
        output["view1"].append(outputs["views"][0])
        output["view2"].append(outputs["views"][view_id])
        output["pred1"].append(outputs["pred"][0])
        output["pred2"].append(outputs["pred"][view_id])

        edges.append((outputs["views"][0]["idx"], outputs["views"][view_id]["idx"]))

    list_of_tuples = edges
    sorted_indices = sorted(
        range(len(list_of_tuples)),
        key=lambda x: (
            list_of_tuples[x][0] > list_of_tuples[x][1],  # Grouping condition
            (
                list_of_tuples[x][1]
                if list_of_tuples[x][0] > list_of_tuples[x][1]
                else list_of_tuples[x][0]
            ),  # First sort key
            (
                list_of_tuples[x][0]
                if list_of_tuples[x][0] > list_of_tuples[x][1]
                else list_of_tuples[x][1]
            ),  # Second sort key
        ),
    )
    
    new_output = {
        "view1": [],
        "view2": [],
        "pred1": [],
        "pred2": [],
    }
    for i in sorted_indices:
        new_output["view1"].append(output["view1"][i])
        new_output["view2"].append(output["view2"][i])
        new_output["pred1"].append(output["pred1"][i])
        new_output["pred2"].append(output["pred2"][i])
        
    output["view1"] = collate_with_cat(new_output["view1"])
    output["view2"] = collate_with_cat(new_output["view2"])
    output["pred1"] = collate_with_cat(new_output["pred1"])
    output["pred2"] = collate_with_cat(new_output["pred2"])

    total_time = time.time() - start_time
    per_frame_time = total_time / len(views)
    print(f"Inference completed in {total_time:.2f} seconds (average {per_frame_time:.2f} s per frame).")

    # Process outputs
    print("Processing reconstruction output...")
    pts3ds_other, colors, depths, conf, cam_dict = prepare_output(output, poses, depths, lr, niter, output_dir, device, save_flag)
    
    # Create result dictionary
    result = {
        "point_clouds": pts3ds_other,
        "colors": colors,
        "depths": depths,
        "confidences": conf,
        "camera_info": cam_dict
    }

    # Visualize if requested
    if visualize:
        from viser_utils import PointCloudViewer
        
        # Convert tensors to numpy arrays for visualization
        pts3ds_to_vis = [p.cpu().numpy() for p in pts3ds_other]
        colors_to_vis = [c.cpu().numpy() for c in colors]
        edge_colors = [None] * len(pts3ds_to_vis)

        # Create and run the point cloud viewer
        print("Launching point cloud viewer...")
        viewer = PointCloudViewer(
            model,
            state_args,
            pts3ds_to_vis,
            colors_to_vis,
            conf,
            cam_dict,
            device=device,
            edge_color_list=edge_colors,
            show_camera=True,
            vis_threshold=vis_threshold,
            size=size,
        )
        viewer.run()
    
    return result


def prepare_input_from_pil(
    pil_images, size, square_ok=False, raymaps=None, raymap_mask=None, revisit=1, update=True
):
    """
    Prepare input views for inference from a list of PIL images.

    Args:
        pil_images (list): List of PIL image objects.
        size (int): Target image size.
        raymaps (list, optional): List of ray maps.
        raymap_mask (list, optional): Flags indicating valid ray maps.
        revisit (int): How many times to revisit each view.
        update (bool): Whether to update the state on revisits.

    Returns:
        list: A list of view dictionaries.
    """
    # Import needed utilities (delayed import needed after adding ckpt path)
    from src.dust3r.utils.image import _resize_pil_image, ImgNorm, exif_transpose
    import PIL
    
    # Process PIL images to have the same format as the load_images output
    imgs = []
    for i, img in enumerate(pil_images):
        # Convert to RGB to ensure consistency
        img = exif_transpose(img).convert("RGB")
        W1, H1 = img.size
        if size == 224:

            img = _resize_pil_image(img, round(size * max(W1 / H1, H1 / W1)))
        else:

            img = _resize_pil_image(img, size)
        W, H = img.size
        cx, cy = W // 2, H // 2
        if size == 224:
            half = min(cx, cy)
            img = img.crop((cx - half, cy - half, cx + half, cy + half))
        else:
            halfw, halfh = ((2 * cx) // 16) * 8, ((2 * cy) // 16) * 8
            if not (square_ok) and W == H:
                halfh = 3 * halfw / 4
            img = img.crop((cx - halfw, cy - halfh, cx + halfw, cy + halfh))

        
        # Create dictionary with the same structure as in load_images
        imgs.append({
            "img": ImgNorm(img)[None],  # Using ImgNorm for normalization
            "true_shape": np.int32([img.size[::-1]]),
            "idx": i,
            "instance": str(i),
        })
    
    # Prepare views similar to prepare_input
    views = []
    num_views = len(imgs)
    
    for i in range(num_views):
        view = {
            "img": imgs[i]["img"],
            "ray_map": torch.full(
                (
                    imgs[i]["img"].shape[0],
                    6,
                    imgs[i]["img"].shape[-2],
                    imgs[i]["img"].shape[-1],
                ),
                torch.nan,
            ),
            "true_shape": torch.from_numpy(imgs[i]["true_shape"]),
            "idx": i,
            "instance": str(i),
            "camera_pose": torch.from_numpy(np.eye(4).astype(np.float32)).unsqueeze(0),
            "img_mask": torch.tensor(True).unsqueeze(0),
            "ray_mask": torch.tensor(False).unsqueeze(0),
            "update": torch.tensor(True).unsqueeze(0),
            "reset": torch.tensor(False).unsqueeze(0),
        }
        views.append(view)
    
    return views