Add initial project structure with core files, configurations, and sample images

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+test_samples/open_door.jpg filter=lfs diff=lfs merge=lfs -text
+test_samples/oxford.jpeg filter=lfs diff=lfs merge=lfs -text
+test_samples/changi.jpg filter=lfs diff=lfs merge=lfs -text
+test_samples/friends.jpg filter=lfs diff=lfs merge=lfs -text
+test_samples/jesus.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+assets/*
+pycache/*
+__pycache__/*
+.DS_Store

README.md

@@ -1,13 +1,76 @@
----
-title: Vmem
-emoji: 👁
-colorFrom: yellow
-colorTo: gray
-sdk: gradio
-sdk_version: 5.33.2
-app_file: app.py
-pinned: false
-license: apache-2.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+<div align="center">
+<h1>VMem: Consistent Video Scene Generation with Surfel-Indexed View Memory</h1>
+
+<a href="https://v-mem.github.io/"><img src="https://img.shields.io/badge/%F0%9F%8F%A0%20Project%20Page-gray.svg"></a>
+<a href="http://arxiv.org/abs/2503.14489"><img src="https://img.shields.io/badge/%F0%9F%93%84%20arXiv-2503.14489-B31B1B.svg"></a>
+<a href="https://huggingface.co/liguang0115/vmem"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Model_Card-Huggingface-orange"></a>
+<a href="https://huggingface.co/spaces/stabilityai/stable-virtual-camera"><img src="https://img.shields.io/badge/%F0%9F%9A%80%20Gradio%20Demo-Huggingface-orange"></a>
+
+[Runjia Li](https://runjiali-rl.github.io/), [Philip Torr](https://www.robots.ox.ac.uk/~phst/), [Andrea Vedaldi](https://www.robots.ox.ac.uk/~vedaldi/), [Tomas Jakab](https://www.robots.ox.ac.uk/~tomj/)
+<br>
+<br>
+[University of Oxford](https://www.robots.ox.ac.uk/~vgg/)
+</div>
+
+<p align="center">
+  <img src="assets/demo_teaser.gif" width="100%" alt="Teaser" style="border-radius:10px;"/>
+</p>
+
+<!-- <p align="center" border-radius="10px">
+  <img src="assets/benchmark.png" width="100%" alt="teaser_page1"/>
+</p> -->
+
+# Overview
+
+`VMem` is a plug-and-play memory mechanism of image-set models for consistent scene generation.
+Existing methods either rely on inpainting with explicit geometry estimation, which suffers from inaccuracies, or use limited context windows in video-based approaches, leading to poor long-term coherence. To overcome these issues, we introduce Surfel Memory of Views (VMem), which anchors past views to surface elements (surfels) they observed. This enables conditioning novel view generation on the most relevant past views rather than just the most recent ones, enhancing long-term scene consistency while reducing computational cost.
+
+
+# :wrench: Installation
+
+```bash
+conda create -n vmem python=3.10
+conda activate vmem
+pip install -r requirements.txt
+```
+
+
+# :rocket: Usage
+
+You need to properly authenticate with Hugging Face to download our model weights. Once set up, our code will handle it automatically at your first run. You can authenticate by running
+
+```bash
+# This will prompt you to enter your Hugging Face credentials.
+huggingface-cli login
+```
+
+Once authenticated, go to our model card [here](https://huggingface.co/stabilityai/stable-virtual-camera) and enter your information for access.
+
+We provide a demo for you to interact with `VMem`. Simply run
+
+```bash
+python app.py
+```
+
+
+## :heart: Acknowledgement
+This work is built on top of [CUT3R](https://github.com/CUT3R/CUT3R), [DUSt3R](https://github.com/naver/dust3r) and [Stable Virtual Camera](https://github.com/stability-ai/stable-virtual-camera). We thank them for their great works.
+
+
+
+
+
+# :books: Citing
+
+If you find this repository useful, please consider giving a star :star: and citation.
+
+```
+@article{zhou2025stable,
+    title={Stable Virtual Camera: Generative View Synthesis with Diffusion Models},
+    author={Jensen (Jinghao) Zhou and Hang Gao and Vikram Voleti and Aaryaman Vasishta and Chun-Han Yao and Mark Boss and
+    Philip Torr and Christian Rupprecht and Varun Jampani
+    },
+    journal={arXiv preprint arXiv:2503.14489},
+    year={2025}
+}
+```

app.py

@@ -0,0 +1,933 @@
+from typing import List, Literal
+from pathlib import Path
+from functools import partial
+import spaces
+import gradio as gr
+import numpy as np
+import torch
+from torchvision.datasets.utils import download_and_extract_archive
+from einops import repeat
+from omegaconf import OmegaConf
+from modeling.pipeline import VMemPipeline
+from diffusers.utils import export_to_video, export_to_gif
+from scipy.spatial.transform import Rotation, Slerp
+from navigation import Navigator
+from PIL import Image
+from utils import tensor_to_pil, encode_vae_image, encode_image, get_default_intrinsics, load_img_and_K, transform_img_and_K
+import os
+import glob
+
+
+CONFIG_PATH = "configs/inference/inference.yaml"
+CONFIG = OmegaConf.load(CONFIG_PATH)
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+MODEL = VMemPipeline(CONFIG, DEVICE)
+NAVIGATORS = []
+
+
+NAVIGATION_FPS = 3
+WIDTH = 576
+HEIGHT = 576
+
+
+IMAGE_PATHS = ['test_samples/changi.jpg', 'test_samples/oxford.jpeg', 'test_samples/open_door.jpg', 'test_samples/jesus.jpg', 'test_samples/friends.jpg']
+
+# for asset_dir in ASSET_DIRS:
+#     if os.path.exists(asset_dir):
+#         for ext in ["*.jpg", "*.jpeg", "*.png"]:
+#             IMAGE_PATHS.extend(glob.glob(os.path.join(asset_dir, ext)))
+
+# If no images found, create placeholders
+if not IMAGE_PATHS:
+    def create_placeholder_images(num_samples=5, height=HEIGHT, width=WIDTH):
+        """Create placeholder images for the demo"""
+        images = []
+        for i in range(num_samples):
+            # Create a gradient image as placeholder
+            img = np.zeros((height, width, 3), dtype=np.uint8)
+            for h in range(height):
+                for w in range(width):
+                    img[h, w, 0] = int(255 * h / height)  # Red gradient
+                    img[h, w, 1] = int(255 * w / width)   # Green gradient
+                    img[h, w, 2] = int(255 * (i+1) / num_samples)  # Blue varies by image
+            images.append(img)
+        return images
+
+    # Create placeholder video frames and poses
+    def create_placeholder_video_and_poses(num_samples=5, num_frames=1, height=HEIGHT, width=WIDTH):
+        """Create placeholder videos and poses for the demo"""
+        videos = []
+        poses = []
+        
+        for i in range(num_samples):
+            # Create a simple video (just one frame initially for each sample)
+            frames = []
+            for j in range(num_frames):
+                # Create a gradient frame
+                img = np.zeros((height, width, 3), dtype=np.uint8)
+                for h in range(height):
+                    for w in range(width):
+                        img[h, w, 0] = int(255 * h / height)  # Red gradient
+                        img[h, w, 1] = int(255 * w / width)   # Green gradient
+                        img[h, w, 2] = int(255 * (i+1) / num_samples)  # Blue varies by video
+                
+                # Convert to torch tensor [C, H, W] with normalized values
+                frame = torch.from_numpy(img.transpose(2, 0, 1)).float() / 255.0
+                frames.append(frame)
+            
+            video = torch.stack(frames)
+            videos.append(video)
+            
+            # Create placeholder poses (identity matrices flattened)
+            # This creates a 4x4 identity matrix flattened to match expected format
+            # pose = torch.eye(4).flatten()[:-4]  # Remove last row of 4x4 matrix
+            poses.append(torch.eye(4).unsqueeze(0).repeat(num_frames, 1, 1))
+        
+        return videos, poses
+
+    first_frame_list = create_placeholder_images(num_samples=5)
+    video_list, poses_list = create_placeholder_video_and_poses(num_samples=5)
+
+# Function to load image from path
+def load_image_for_navigation(image_path):
+    """Load image from path and prepare for navigation"""
+    # Load image and get default intrinsics
+    image, _ = load_img_and_K(image_path, None, K=None, device=DEVICE)
+    
+    # Transform image to the target size
+    config = OmegaConf.load(CONFIG_PATH)
+    image, _ = transform_img_and_K(image, (config.model.height, config.model.width), mode="crop", K=None)
+    
+    # Create initial video with single frame and pose
+    video = image
+    pose = torch.eye(4).unsqueeze(0)  # [1, 4, 4]
+    
+    return {
+        "image": tensor_to_pil(image),
+        "video": video,
+        "pose": pose
+    }
+
+
+class CustomProgressBar:
+    def __init__(self, pbar):
+        self.pbar = pbar
+
+    def set_postfix(self, **kwargs):
+        pass
+
+    def __getattr__(self, attr):
+        return getattr(self.pbar, attr)
+
+def get_duration_navigate_video(video: torch.Tensor,
+    poses: torch.Tensor,
+    x_angle: float,
+    y_angle: float,
+    distance: float
+):
+    # Estimate processing time based on navigation complexity and number of frames
+    base_duration = 15  # Base duration in seconds
+    
+    # Add time for more complex navigation operations
+    if abs(x_angle) > 20 or abs(y_angle) > 30:
+        base_duration += 10  # More time for sharp turns
+    
+    if distance > 100:
+        base_duration += 10  # More time for longer distances
+    
+    # Add time proportional to existing video length (more frames = more processing)
+    base_duration += min(10, len(video))
+    
+    return base_duration
+
+@spaces.GPU(duration=get_duration_navigate_video)
+@torch.autocast("cuda")
+@torch.no_grad()
+def navigate_video(
+    video: torch.Tensor,
+    poses: torch.Tensor,
+    x_angle: float,
+    y_angle: float,
+    distance: float,
+):
+    """
+    Generate new video frames by navigating in the 3D scene.
+    This function uses the Navigator class from navigation.py to handle movement:
+    - y_angle parameter controls left/right turning (turn_left/turn_right methods)
+    - distance parameter controls forward movement (move_forward method)
+    - x_angle parameter controls vertical angle (not directly implemented in Navigator)
+    
+    Each Navigator instance is stored based on the video session to maintain state.
+    """
+    try:
+        # Convert first frame to PIL Image for navigator
+        initial_frame = tensor_to_pil(video[0])
+        
+        # Initialize the navigator for this session if not already done
+        if len(NAVIGATORS) == 0:
+            # Create a new navigator instance
+            NAVIGATORS.append(Navigator(MODEL, step_size=0.1, num_interpolation_frames=4))
+            
+            # Get the initial pose and convert to numpy
+            initial_pose = poses[0].cpu().numpy().reshape(4, 4)
+            
+            # Default camera intrinsics if not available
+            initial_K = np.array(get_default_intrinsics()[0])
+            
+            # Initialize the navigator
+            NAVIGATORS[0].initialize(initial_frame, initial_pose, initial_K)
+
+        navigator = NAVIGATORS[0]
+        
+        # Generate new frames based on navigation commands
+        new_frames = []
+        
+        # First handle any x-angle (vertical angle) adjustments
+        # Note: This is approximated as Navigator doesn't directly support this
+        if abs(x_angle) > 0:
+            # Implementation for x-angle could be added here
+            # For now, we'll skip this as it's not directly supported
+            pass
+        
+        # Next handle y-angle (turning left/right)
+        if abs(y_angle) > 0:
+            # Use Navigator's turn methods
+            if y_angle > 0:
+                new_frames = navigator.turn_left(abs(y_angle//2))
+            else:
+                new_frames = navigator.turn_right(abs(y_angle//2))
+        # Finally handle distance (moving forward)
+        elif distance > 0:
+            # Calculate number of steps based on distance
+            steps = max(1, int(distance / 10))
+            new_frames = navigator.move_forward(steps)
+        elif distance < 0:
+            # Handle moving backward if needed
+            steps = max(1, int(abs(distance) / 10))
+            new_frames = navigator.move_backward(steps)
+        
+        if not new_frames:
+            # If no new frames were generated, return the current state
+            return video, poses, tensor_to_pil(video[-1]), export_to_video([tensor_to_pil(video[i]) for i in range(len(video))], fps=NAVIGATION_FPS), [(tensor_to_pil(video[i]), f"t={i}") for i in range(len(video))]
+        
+        # Convert PIL images to tensors
+        new_frame_tensors = []
+        for frame in new_frames:
+            # Convert PIL Image to tensor [C, H, W]
+            frame_np = np.array(frame) / 255.0
+            # Convert to [-1, 1] range to match the expected format
+            frame_tensor = torch.from_numpy(frame_np.transpose(2, 0, 1)).float() * 2.0 - 1.0
+            new_frame_tensors.append(frame_tensor)
+        
+        new_frames_tensor = torch.stack(new_frame_tensors)
+        
+        # Get the updated camera poses from the navigator
+        current_pose = navigator.current_pose
+        new_poses = torch.from_numpy(current_pose).float().unsqueeze(0).repeat(len(new_frames), 1, 1)
+        
+        # Reshape the poses to match the expected format
+        new_poses = new_poses.view(len(new_frames), 4, 4)
+        
+        # Concatenate new frames and poses with existing ones
+        updated_video = torch.cat([video.cpu(), new_frames_tensor], dim=0)
+        updated_poses = torch.cat([poses.cpu(), new_poses], dim=0)
+        
+        # Create output images for gallery
+        all_images = [(tensor_to_pil(updated_video[i]), f"t={i}") for i in range(len(updated_video))]
+        updated_video_pil = [tensor_to_pil(updated_video[i]) for i in range(len(updated_video))]
+        
+        return (
+            updated_video,
+            updated_poses,
+            tensor_to_pil(updated_video[-1]),  # Current view
+            export_to_video(updated_video_pil, fps=NAVIGATION_FPS),  # Video
+            all_images,  # Gallery
+        )
+    except Exception as e:
+        print(f"Error in navigate_video: {e}")
+        gr.Warning(f"Navigation error: {e}")
+        # Return the original inputs to avoid crashes
+        current_frame = tensor_to_pil(video[-1]) if len(video) > 0 else None
+        all_frames = [(tensor_to_pil(video[i]), f"t={i}") for i in range(len(video))]
+        video_frames = [tensor_to_pil(video[i]) for i in range(len(video))]
+        video_output = export_to_video(video_frames, fps=NAVIGATION_FPS) if video_frames else None
+        return video, poses, current_frame, video_output, all_frames
+
+
+def undo_navigation(
+    video: torch.Tensor,
+    poses: torch.Tensor,
+):
+    """
+    Undo the last navigation step by removing the last set of frames.
+    Uses the Navigator's undo method which in turn uses the pipeline's undo_latest_move
+    to properly handle surfels and state management.
+    """
+    if len(NAVIGATORS) > 0:
+        navigator = NAVIGATORS[0]
+        
+        # Call the Navigator's undo method to handle the operation
+        success = navigator.undo()
+        
+        if success:
+            # Since the navigator has handled the frame removal internally,
+            # we need to update our video and poses tensors to match
+            updated_video = video[:len(navigator.frames)]
+            updated_poses = poses[:len(navigator.frames)]
+            
+            # Create gallery images
+            all_images = [(tensor_to_pil(updated_video[i]), f"t={i}") for i in range(len(updated_video))]
+            
+            return (
+                updated_video,
+                updated_poses,
+                tensor_to_pil(updated_video[-1]),
+                export_to_video([tensor_to_pil(updated_video[i]) for i in range(len(updated_video))], fps=NAVIGATION_FPS),
+                all_images,
+            )
+        else:
+            gr.Warning("You have no moves left to undo!")
+    else:
+        gr.Warning("No navigation session available!")
+    
+    # If undo wasn't successful or no navigator exists, return original state
+    all_images = [(tensor_to_pil(video[i]), f"t={i}") for i in range(len(video))]
+    
+    return (
+        video,
+        poses,
+        tensor_to_pil(video[-1]),
+        export_to_video([tensor_to_pil(video[i]) for i in range(len(video))], fps=NAVIGATION_FPS),
+        all_images,
+    )
+
+
+
+
+
+def render_demo3(
+    s: Literal["Selection", "Generation"],
+    idx: int,
+    demo3_stage: gr.State,
+    demo3_selected_index: gr.State,
+    demo3_current_video: gr.State,
+    demo3_current_poses: gr.State
+):
+    gr.Markdown(
+        """
+        ## Single Image → Consistent Scene Navigation
+        > #### _Select an image and navigate through the scene by controlling camera movements._
+    """,
+    elem_classes=["task-title"]
+    )
+    match s:
+        case "Selection":
+            with gr.Group():
+                # Add upload functionality
+                with gr.Group(elem_classes=["gradio-box"]):
+                    gr.Markdown("### Upload Your Own Image")
+                    gr.Markdown("_Upload an image to navigate through its 3D scene_")
+                    with gr.Row():
+                        with gr.Column(scale=3):
+                            upload_image = gr.Image(
+                                label="Upload an image",
+                                type="filepath",
+                                height=300,
+                                elem_id="upload-image"
+                            )
+                        with gr.Column(scale=1):
+                            gr.Markdown("#### Instructions:")
+                            gr.Markdown("1. Upload a clear, high-quality image")
+                            gr.Markdown("2. Images with distinct visual features work best")
+                            gr.Markdown("3. Landscape or architectural scenes are ideal")
+                            upload_btn = gr.Button("Start Navigation", variant="primary", size="lg")
+                    
+                    def process_uploaded_image(image_path):
+                        if image_path is None:
+                            gr.Warning("Please upload an image first")
+                            return "Selection", None, None, None
+                        try:
+                            # Load image and prepare for navigation
+                            result = load_image_for_navigation(image_path)
+                            
+                            # Clear any existing navigators
+                            global NAVIGATORS
+                            NAVIGATORS = []
+                            
+                            return (
+                                "Generation",
+                                None,  # No predefined index for uploaded images
+                                result["video"],
+                                result["pose"],
+                            )
+                        except Exception as e:
+                            print(f"Error in process_uploaded_image: {e}")
+                            gr.Warning(f"Error processing uploaded image: {e}")
+                            return "Selection", None, None, None
+                    
+                    upload_btn.click(
+                        fn=process_uploaded_image,
+                        inputs=[upload_image],
+                        outputs=[demo3_stage, demo3_selected_index, demo3_current_video, demo3_current_poses]
+                    )
+                
+                gr.Markdown("### Or Choose From Our Examples")
+                # Define image captions
+                image_captions = {
+                    'test_samples/changi.jpg': 'Changi Airport',
+                    'test_samples/oxford.jpeg': 'Oxford University',
+                    'test_samples/open_door.jpg': 'Bedroom Interior',
+                    'test_samples/jesus.jpg': 'Jesus College',
+                    'test_samples/friends.jpg': 'Friends Café'
+                }
+                
+                # Load all images for the gallery with captions
+                gallery_images = []
+                for img_path in IMAGE_PATHS:
+                    try:
+                        # Get caption or default to basename
+                        caption = image_captions.get(img_path, os.path.basename(img_path))
+                        gallery_images.append((img_path, caption))
+                    except Exception as e:
+                        print(f"Error loading image {img_path}: {e}")
+                
+                # Show image gallery for selection
+                demo3_image_gallery = gr.Gallery(
+                    value=gallery_images,
+                    label="Select an Image to Start Navigation",
+                    columns=len(gallery_images),
+                    height=400,
+                    allow_preview=True,
+                    preview=False,
+                    elem_id="navigation-gallery"
+                )
+                
+                gr.Markdown("_Click on an image to begin navigation_")
+                
+                def start_navigation(evt: gr.SelectData):
+                    try:
+                        # Get the selected image path
+                        selected_path = IMAGE_PATHS[evt.index]
+                        
+                        # Load image and prepare for navigation
+                        result = load_image_for_navigation(selected_path)
+                        
+                        # Clear any existing navigators
+                        global NAVIGATORS
+                        NAVIGATORS = []
+                        
+                        return (
+                            "Generation",
+                            evt.index,
+                            result["video"],
+                            result["pose"],
+                        )
+                    except Exception as e:
+                        print(f"Error in start_navigation: {e}")
+                        gr.Warning(f"Error starting navigation: {e}")
+                        return "Selection", None, None, None
+                
+                demo3_image_gallery.select(
+                    fn=start_navigation,
+                    inputs=None, 
+                    outputs=[demo3_stage, demo3_selected_index, demo3_current_video, demo3_current_poses]
+                )
+
+        case "Generation":
+            with gr.Row():
+                with gr.Column(scale=3):
+                    with gr.Row():
+                        demo3_current_view = gr.Image(
+                            label="Current View",
+                            width=256,
+                            height=256,
+                        )
+                        demo3_video = gr.Video(
+                            label="Generated Video",
+                            width=256,
+                            height=256,
+                            autoplay=True,
+                            loop=True,
+                            show_share_button=True,
+                            show_download_button=True,
+                        )
+
+                    demo3_generated_gallery = gr.Gallery(
+                        value=[],
+                        label="Generated Frames",
+                        columns=[6],
+                    )
+                    
+                    # Initialize the current view with the selected image if available
+                    if idx is not None:
+                        try:
+                            selected_path = IMAGE_PATHS[idx]
+                            result = load_image_for_navigation(selected_path)
+                            demo3_current_view.value = result["image"]
+                        except Exception as e:
+                            print(f"Error initializing current view: {e}")
+
+                with gr.Column():
+                    gr.Markdown("### Navigation Controls ↓")
+                    with gr.Accordion("Instructions", open=False):
+                        gr.Markdown("""
+                            - **The model will predict the next few frames based on your camera movements. Repeat the process to continue navigating through the scene.**
+                            - **Use the navigation controls to move forward/backward and turn left/right.**
+                            - **At the end of your navigation, you can save your camera path for later use.**
+                           
+                        """)
+                    # with gr.Tab("Basic", elem_id="basic-controls-tab"):
+                    with gr.Group():
+                        gr.Markdown("_**Select a direction to move:**_")
+                        # First row: Turn left/right
+                        with gr.Row(elem_id="basic-controls"):
+                            gr.Button(
+                                "↰20°\nVeer",
+                                size="sm",
+                                min_width=0,
+                                variant="primary",
+                            ).click(
+                                fn=partial(
+                                    navigate_video,
+                                    x_angle=0,
+                                    y_angle=20,
+                                    distance=0,
+                                ),
+                                inputs=[
+                                    demo3_current_video,
+                                    demo3_current_poses,
+                                ],
+                                outputs=[
+                                    demo3_current_video,
+                                    demo3_current_poses,
+                                    demo3_current_view,
+                                    demo3_video,
+                                    demo3_generated_gallery,
+                                ],
+                            )
+
+                            gr.Button(
+                                "↖10°\nTurn",
+                                size="sm",
+                                min_width=0,
+                                variant="primary",
+                            ).click(
+                                fn=partial(
+                                    navigate_video,
+                                    x_angle=0,
+                                    y_angle=10,
+                                    distance=0,
+                                ),
+                                inputs=[
+                                    demo3_current_video,
+                                    demo3_current_poses,
+                                ],
+                                outputs=[
+                                    demo3_current_video,
+                                    demo3_current_poses,
+                                    demo3_current_view,
+                                    demo3_video,
+                                    demo3_generated_gallery,
+                                ],
+                            )
+
+                            # gr.Button(
+                            #     "↑0°\nAhead",
+                            #     size="sm",
+                            #     min_width=0,
+                            #     variant="primary",
+                            # ).click(
+                            #     fn=partial(
+                            #         navigate_video,
+                            #         x_angle=0,
+                            #         y_angle=0,
+                            #         distance=10,
+                            #     ),
+                            #     inputs=[
+                            #         demo3_current_video,
+                            #         demo3_current_poses,
+                            #     ],
+                            #     outputs=[
+                            #         demo3_current_video,
+                            #         demo3_current_poses,
+                            #         demo3_current_view,
+                            #         demo3_video,
+                            #         demo3_generated_gallery,
+                            #     ],
+                            # )
+                            gr.Button(
+                                "↗10°\nTurn",
+                                size="sm",
+                                min_width=0,
+                                variant="primary",
+                            ).click(
+                                fn=partial(
+                                    navigate_video,
+                                    x_angle=0,
+                                    y_angle=-10,
+                                    distance=0,
+                                ),
+                                inputs=[
+                                    demo3_current_video,
+                                    demo3_current_poses,
+                                ],
+                                outputs=[
+                                    demo3_current_video,
+                                    demo3_current_poses,
+                                    demo3_current_view,
+                                    demo3_video,
+                                    demo3_generated_gallery,
+                                ],
+                            )
+                            gr.Button(
+                                "↱\n20° Veer",
+                                size="sm",
+                                min_width=0,
+                                variant="primary",
+                            ).click(
+                                fn=partial(
+                                    navigate_video,
+                                    x_angle=0,
+                                    y_angle=-20,
+                                    distance=0,
+                                ),
+                                inputs=[
+                                    demo3_current_video,
+                                    demo3_current_poses,
+                                ],
+                                outputs=[
+                                    demo3_current_video,
+                                    demo3_current_poses,
+                                    demo3_current_view,
+                                    demo3_video,
+                                    demo3_generated_gallery,
+                                ],
+                            )
+                        
+                        # Second row: Forward/Backward movement
+                        with gr.Row(elem_id="forward-backward-controls"):
+                            gr.Button(
+                                "↓\nBackward",
+                                size="sm",
+                                min_width=0,
+                                variant="secondary",
+                            ).click(
+                                fn=partial(
+                                    navigate_video,
+                                    x_angle=0,
+                                    y_angle=0,
+                                    distance=-10,
+                                ),
+                                inputs=[
+                                    demo3_current_video,
+                                    demo3_current_poses,
+                                ],
+                                outputs=[
+                                    demo3_current_video,
+                                    demo3_current_poses,
+                                    demo3_current_view,
+                                    demo3_video,
+                                    demo3_generated_gallery,
+                                ],
+                            )
+                            
+                            gr.Button(
+                                "↑\nForward",
+                                size="sm",
+                                min_width=0,
+                                variant="secondary",
+                            ).click(
+                                fn=partial(
+                                    navigate_video,
+                                    x_angle=0,
+                                    y_angle=0,
+                                    distance=10,
+                                ),
+                                inputs=[
+                                    demo3_current_video,
+                                    demo3_current_poses,
+                                ],
+                                outputs=[
+                                    demo3_current_video,
+                                    demo3_current_poses,
+                                    demo3_current_view,
+                                    demo3_video,
+                                    demo3_generated_gallery,
+                                ],
+                            )
+                    # with gr.Tab("Advanced", elem_id="advanced-controls-tab"):
+                    #     with gr.Group():
+                    #         gr.Markdown("_**Select angles and distance:**_")
+
+                    #         demo3_y_angle = gr.Slider(
+                    #             minimum=-90,
+                    #             maximum=90,
+                    #             value=0,
+                    #             step=10,
+                    #             label="Horizontal Angle",
+                    #             interactive=True,
+                    #         )
+                    #         demo3_x_angle = gr.Slider(
+                    #             minimum=-40,
+                    #             maximum=40,
+                    #             value=0,
+                    #             step=10,
+                    #             label="Vertical Angle",
+                    #             interactive=True,
+                    #         )
+                    #         demo3_distance = gr.Slider(
+                    #             minimum=-200,
+                    #             maximum=200,
+                    #             value=100,
+                    #             step=10,
+                    #             label="Distance (negative = backward)",
+                    #             interactive=True,
+                    #         )
+
+                    #         gr.Button(
+                    #             "Generate Next Move", variant="primary"
+                    #         ).click(
+                    #             fn=navigate_video,
+                    #             inputs=[
+                    #                 demo3_current_video,
+                    #                 demo3_current_poses,
+                    #                 demo3_x_angle,
+                    #                 demo3_y_angle,
+                    #                 demo3_distance,
+                    #             ],
+                    #             outputs=[
+                    #                 demo3_current_video,
+                    #                 demo3_current_poses,
+                    #                 demo3_current_view,
+                    #                 demo3_video,
+                    #                 demo3_generated_gallery,
+                    #             ],
+                    #         )
+                    gr.Markdown("---")
+                    with gr.Group():
+                        gr.Markdown("_**Navigation controls:**_")
+                        with gr.Row():
+                            gr.Button("Undo Last Move", variant="huggingface").click(
+                                fn=undo_navigation,
+                                inputs=[demo3_current_video, demo3_current_poses],
+                                outputs=[
+                                    demo3_current_video,
+                                    demo3_current_poses,
+                                    demo3_current_view,
+                                    demo3_video,
+                                    demo3_generated_gallery,
+                                ],
+                            )
+                            
+                            # Add a function to save camera poses
+                            def save_camera_poses(video, poses):
+                                if len(NAVIGATORS) > 0:
+                                    navigator = NAVIGATORS[0]
+                                    # Create a directory for saved poses
+                                    os.makedirs("./visualization", exist_ok=True)
+                                    save_path = f"./visualization/transforms_{len(navigator.frames)}_frames.json"
+                                    navigator.save_camera_poses(save_path)
+                                    return gr.Info(f"Camera poses saved to {save_path}")
+                                return gr.Warning("No navigation instance found")
+                            
+                            gr.Button("Save Camera", variant="huggingface").click(
+                                fn=save_camera_poses,
+                                inputs=[demo3_current_video, demo3_current_poses],
+                                outputs=[]
+                            )
+                            
+                            # Add a button to return to image selection
+                            def reset_navigation():
+                                # Clear current navigator
+                                global NAVIGATORS
+                                NAVIGATORS = []
+                                return "Selection", None, None, None
+                            
+                            gr.Button("Choose New Image", variant="secondary").click(
+                                fn=reset_navigation,
+                                inputs=[],
+                                outputs=[demo3_stage, demo3_selected_index, demo3_current_video, demo3_current_poses]
+                            )
+
+
+# Create the Gradio Blocks
+with gr.Blocks(theme=gr.themes.Base(primary_hue="teal")) as demo:
+    gr.HTML(
+        """
+    <style>
+    [data-tab-id="task-1"], [data-tab-id="task-2"], [data-tab-id="task-3"] {
+        font-size: 16px !important;
+        font-weight: bold;
+    }
+    #page-title h1 {
+        color: #0D9488 !important;
+    }
+    .task-title h2 {
+        color: #F59E0C !important;
+    }
+    .header-button-row {
+        gap: 4px !important;
+    }
+    .header-button-row div {
+        width: 131.0px !important;
+    }
+    .header-button-column {
+        width: 131.0px !important;
+        gap: 5px !important;
+    }
+    .header-button a {
+        border: 1px solid #e4e4e7;
+    }
+    .header-button .button-icon {
+        margin-right: 8px;
+    }
+    .demo-button-column .gap {
+        gap: 5px !important;
+    }
+    #basic-controls {
+        column-gap: 0px;
+    }
+    #basic-controls-tab {
+        padding: 0px;
+    }
+    #advanced-controls-tab {
+        padding: 0px;
+    }
+    #forward-backward-controls {
+        column-gap: 0px;
+        justify-content: center;
+        margin-top: 8px;
+    }
+    #selected-demo-button {
+        color: #F59E0C;
+        text-decoration: underline;
+    }
+    .demo-button {
+        text-align: left !important;
+        display: block !important;
+    }
+    #navigation-gallery {
+        margin-bottom: 15px;
+    }
+    #navigation-gallery .gallery-item {
+        cursor: pointer;
+        border-radius: 6px;
+        transition: transform 0.2s, box-shadow 0.2s;
+    }
+    #navigation-gallery .gallery-item:hover {
+        transform: scale(1.02);
+        box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
+    }
+    #navigation-gallery .gallery-item.selected {
+        border: 3px solid #0D9488;
+    }
+    /* Upload image styling */
+    #upload-image {
+        border-radius: 8px;
+        border: 2px dashed #0D9488;
+        padding: 10px;
+        transition: all 0.3s ease;
+    }
+    #upload-image:hover {
+        border-color: #F59E0C;
+        box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
+    }
+    /* Box styling */
+    .gradio-box {
+        border-radius: 10px;
+        margin-bottom: 20px;
+        padding: 15px;
+        background-color: #f8f9fa;
+        border: 1px solid #e9ecef;
+    }
+    </style>
+    """
+    )
+
+    demo_idx = gr.State(value=3)
+
+    with gr.Sidebar():
+        gr.Markdown("# VMem: Consistent Scene Generation with Surfel Memory of Views", elem_id="page-title")
+        gr.Markdown(
+            "### Official Interactive Demo for [_VMem_](https://arxiv.org/abs/2502.06764)"
+        )
+        gr.Markdown("---")
+        gr.Markdown("#### Links ↓")
+        with gr.Row(elem_classes=["header-button-row"]):
+            with gr.Column(elem_classes=["header-button-column"], min_width=0):
+                gr.Button(
+                    value="Website",
+                    link="https://v-mem.github.io/",
+                    icon="https://simpleicons.org/icons/googlechrome.svg",
+                    elem_classes=["header-button"],
+                    size="md",
+                    min_width=0,
+                )
+                gr.Button(
+                    value="Paper",
+                    link="https://arxiv.org/abs/2502.06764",
+                    icon="https://simpleicons.org/icons/arxiv.svg",
+                    elem_classes=["header-button"],
+                    size="md",
+                    min_width=0,
+                )
+            with gr.Column(elem_classes=["header-button-column"], min_width=0):
+                gr.Button(
+                    value="Code",
+                    link="https://github.com/kwsong0113/diffusion-forcing-transformer",
+                    icon="https://simpleicons.org/icons/github.svg",
+                    elem_classes=["header-button"],
+                    size="md",
+                    min_width=0,
+                )
+                gr.Button(
+                    value="Weights",
+                    link="https://huggingface.co/liguang0115/vmem",
+                    icon="https://simpleicons.org/icons/huggingface.svg",
+                    elem_classes=["header-button"],
+                    size="md",
+                    min_width=0,
+                )
+        gr.Markdown("---")
+        gr.Markdown("#### Choose a Demo ↓")
+        with gr.Column(elem_classes=["demo-button-column"]):
+            @gr.render(inputs=[demo_idx])
+            def render_demo_tabs(idx):
+                demo_tab_button3 = gr.Button(
+                    "Navigate Image",
+                    size="md", elem_classes=["demo-button"],  **{"elem_id": "selected-demo-button"} if idx == 3 else {}
+                ).click(
+                    fn=lambda: 3,
+                    outputs=demo_idx
+                )
+        gr.Markdown("---")
+        gr.Markdown("#### Troubleshooting ↓")
+        with gr.Group():
+            with gr.Accordion("Error or Unexpected Results?", open=False):
+                gr.Markdown("Please try again after refreshing the page and ensure you do not click the same button multiple times.")
+            with gr.Accordion("Too Slow or No GPU Allocation?", open=False):
+                gr.Markdown(
+                    "Consider running the demo locally (click the dots in the top-right corner). Alternatively, you can subscribe to Hugging Face Pro for an increased GPU quota."
+                )
+
+
+    demo3_stage = gr.State(value="Selection")
+    demo3_selected_index = gr.State(value=None)
+    demo3_current_video = gr.State(value=None)
+    demo3_current_poses = gr.State(value=None)
+
+    @gr.render(inputs=[demo_idx, demo3_stage, demo3_selected_index])
+    def render_demo(
+        _demo_idx, _demo3_stage, _demo3_selected_index
+    ):
+        match _demo_idx:
+            case 3:
+                render_demo3(_demo3_stage, _demo3_selected_index, demo3_stage, demo3_selected_index, demo3_current_video, demo3_current_poses)
+                
+
+if __name__ == "__main__":
+    demo.launch(debug=True,
+                share=True,
+                max_threads=1,  # Limit concurrent processing
+                show_error=True,  # Show detailed error messages
+                )

configs/inference/inference.yaml

@@ -0,0 +1,69 @@
+
+model:
+    height: 576
+    width: 576
+    original_height: 288
+    original_width: 512
+    cache_dir: "/homes/55/runjia/storage/svd_weights"
+    # pretrained_model_path: "stabilityai/stable-diffusion-2-1"
+    # pretrained_video_model_path: "stabilityai/stable-video-diffusion-img2vid"
+    
+    context_num_frames: 4
+    target_num_frames: 4
+    num_frames: 8
+    vae_spatial_scale: 8
+    latent_channels: 4
+    # num_ray_blocks: 2
+    vae_scale_factor: 8
+    inference_mode: false
+
+    temporal_only: false
+    use_non_maximum_suppression: true
+    translation_distance_weight: 0.1
+
+    camera_scale: 2.0
+    inference_num_steps: 50
+    cfg_min: 1.2
+    cfg: 3.0
+    guider_types: 1
+    
+    samples_dir: "./visualization"
+    save_flag: false
+    use_wandb: false
+
+    
+
+    # model_path: "/homes/55/runjia/storage/simview_weights/2025-04-30_12-08-55/checkpoint_230000.pth"
+    model_path: "liguang0115/vmem"
+
+
+surfel:
+    use_surfel: true
+    shrink_factor: 0.05
+    radius_scale: 0.5
+    conf_thresh: 1
+    merge_position_threshold: 0.2
+    merge_normal_threshold: 0.6
+    lr: 0.01
+    niter: 1000
+    model_path: "./extern/CUT3R/src/cut3r_512_dpt_4_64.pth"
+    width: 512
+    height: 288
+
+inference:
+    visualize: true
+    visualize_pointcloud: false
+    visualize_surfel: false
+    save_surfels: false
+    image_dir: "/homes/55/runjia/storage/realestate10k/video_data/test"
+    meta_info_dir: "/homes/55/runjia/storage/realestate10k/RealEstate10K/test"
+
+
+
+    
+
+
+
+
+visualization_dir: "./visualization"
+seed: 42

extern/CUT3R/.gitignore

@@ -0,0 +1,55 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+
+# C extensions
+*.so
+
+# Distribution / packaging
+bin/
+build/
+develop-eggs/
+dist/
+eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+.tox/
+.coverage
+.cache
+nosetests.xml
+coverage.xml
+
+# Translations
+*.mo
+
+# Mr Developer
+.mr.developer.cfg
+.project
+.pydevproject
+
+# Rope
+.ropeproject
+
+# Django stuff:
+*.log
+*.pot
+
+# Sphinx documentation
+docs/_build/
+
+# Ignore data and ckpts
+*.pth
+data
+src/checkpoints

extern/CUT3R/LICENSE

@@ -0,0 +1,6 @@
+Copyright [2025–present] 
+
+CUT3R is licensed under the Creative Commons Attribution-NonCommercial-ShareAlike 4.0 License.  
+
+To view a copy of the CC BY-NC-SA 4.0, visit:  
+https://creativecommons.org/licenses/by-nc-sa/4.0/  

extern/CUT3R/README.md

@@ -0,0 +1,208 @@
+# Continuous 3D Perception Model with Persistent State
+<div align="center">
+  <img src="./assets/factory-ezgif.com-video-speed.gif"  alt="CUT3R" />
+</div>
+
+<hr>
+
+<br>
+Official implementation of <strong>Continuous 3D Perception Model with Persistent State</strong>, CVPR 2025 (Oral)
+
+[*QianqianWang**](https://qianqianwang68.github.io/),
+[*Yifei Zhang**](https://forrest-110.github.io/),
+[*Aleksander Holynski*](https://holynski.org/),
+[*Alexei A Efros*](https://people.eecs.berkeley.edu/~efros/),
+[*Angjoo Kanazawa*](https://people.eecs.berkeley.edu/~kanazawa/)
+
+
+(*: equal contribution)
+
+<div style="line-height: 1;">
+  <a href="https://cut3r.github.io/" target="_blank" style="margin: 2px;">
+    <img alt="Website" src="https://img.shields.io/badge/Website-CUT3R-536af5?color=536af5&logoColor=white" style="display: inline-block; vertical-align: middle;"/>
+  </a>
+  <a href="https://arxiv.org/pdf/2501.12387" target="_blank" style="margin: 2px;">
+    <img alt="Arxiv" src="https://img.shields.io/badge/Arxiv-CUT3R-red?logo=%23B31B1B" style="display: inline-block; vertical-align: middle;"/>
+  </a>
+</div>
+
+
+![Example of capabilities](assets/ezgif.com-video-to-gif-converter.gif)
+
+## Table of Contents
+- [TODO](#todo)
+- [Get Started](#getting-started)
+  - [Installation](#installation)
+  - [Checkpoints](#download-checkpoints)
+  - [Inference](#inference)
+- [Datasets](#datasets)
+- [Evaluation](#evaluation)
+  - [Datasets](#datasets-1)
+  - [Evaluation Scripts](#evaluation-scripts)
+- [Training and Fine-tuning](#training-and-fine-tuning)
+- [Acknowledgements](#acknowledgements)
+- [Citation](#citation)
+
+## TODO
+- [x] Release multi-view stereo results of DL3DV dataset.
+- [ ] Online demo integrated with WebCam
+
+## Getting Started
+
+### Installation
+
+1. Clone CUT3R.
+```bash
+git clone https://github.com/CUT3R/CUT3R.git
+cd CUT3R
+```
+
+2. Create the environment.
+```bash
+conda create -n cut3r python=3.11 cmake=3.14.0
+conda activate cut3r
+conda install pytorch torchvision pytorch-cuda=12.1 -c pytorch -c nvidia  # use the correct version of cuda for your system
+pip install -r requirements.txt
+# issues with pytorch dataloader, see https://github.com/pytorch/pytorch/issues/99625
+conda install 'llvm-openmp<16'
+# for training logging
+pip install git+https://github.com/nerfstudio-project/gsplat.git
+# for evaluation
+pip install evo
+pip install open3d
+```
+
+3. Compile the cuda kernels for RoPE (as in CroCo v2).
+```bash
+cd src/croco/models/curope/
+python setup.py build_ext --inplace
+cd ../../../../
+```
+
+### Download Checkpoints
+
+We currently provide checkpoints on Google Drive:
+
+| Modelname   | Training resolutions | #Views| Head |
+|-------------|----------------------|-------|------|
+| [`cut3r_224_linear_4.pth`](https://drive.google.com/file/d/11dAgFkWHpaOHsR6iuitlB_v4NFFBrWjy/view?usp=drive_link) | 224x224 | 16 | Linear |
+| [`cut3r_512_dpt_4_64.pth`](https://drive.google.com/file/d/1Asz-ZB3FfpzZYwunhQvNPZEUA8XUNAYD/view?usp=drive_link) | 512x384, 512x336, 512x288, 512x256, 512x160, 384x512, 336x512, 288x512, 256x512, 160x512 | 4-64 | DPT |
+
+> `cut3r_224_linear_4.pth` is our intermediate checkpoint and `cut3r_512_dpt_4_64.pth` is our final checkpoint.
+
+To download the weights, run the following commands:
+```bash
+cd src
+# for 224 linear ckpt
+gdown --fuzzy https://drive.google.com/file/d/11dAgFkWHpaOHsR6iuitlB_v4NFFBrWjy/view?usp=drive_link 
+# for 512 dpt ckpt
+gdown --fuzzy https://drive.google.com/file/d/1Asz-ZB3FfpzZYwunhQvNPZEUA8XUNAYD/view?usp=drive_link
+cd ..
+```
+
+### Inference
+
+To run the inference code, you can use the following command:
+```bash
+# the following script will run inference offline and visualize the output with viser on port 8080
+python demo.py --model_path MODEL_PATH --seq_path SEQ_PATH --size SIZE --vis_threshold VIS_THRESHOLD --output_dir OUT_DIR  # input can be a folder or a video
+# Example:
+#     python demo.py --model_path src/cut3r_512_dpt_4_64.pth --size 512 \
+#         --seq_path examples/001 --vis_threshold 1.5 --output_dir tmp
+#
+#     python demo.py --model_path src/cut3r_224_linear_4.pth --size 224 \
+#         --seq_path examples/001 --vis_threshold 1.5 --output_dir tmp
+
+# the following script will run inference with global alignment and visualize the output with viser on port 8080
+python demo_ga.py --model_path MODEL_PATH --seq_path SEQ_PATH --size SIZE --vis_threshold VIS_THRESHOLD --output_dir OUT_DIR
+```
+Output results will be saved to `output_dir`.
+
+> Currently, we accelerate the feedforward process by processing inputs in parallel within the encoder, which results in linear memory consumption as the number of frames increases.
+
+## Datasets
+Our training data includes 32 datasets listed below. We provide processing scripts for all of them. Please download the datasets from their official sources, and refer to [preprocess.md](docs/preprocess.md) for processing scripts and more information about the datasets.
+
+  - [ARKitScenes](https://github.com/apple/ARKitScenes) 
+  - [BlendedMVS](https://github.com/YoYo000/BlendedMVS)
+  - [CO3Dv2](https://github.com/facebookresearch/co3d)
+  - [MegaDepth](https://www.cs.cornell.edu/projects/megadepth/)
+  - [ScanNet++](https://kaldir.vc.in.tum.de/scannetpp/) 
+  - [ScanNet](http://www.scan-net.org/ScanNet/)
+  - [WayMo Open dataset](https://github.com/waymo-research/waymo-open-dataset)
+  - [WildRGB-D](https://github.com/wildrgbd/wildrgbd/)
+  - [Map-free](https://research.nianticlabs.com/mapfree-reloc-benchmark/dataset)
+  - [TartanAir](https://theairlab.org/tartanair-dataset/)
+  - [UnrealStereo4K](https://github.com/fabiotosi92/SMD-Nets) 
+  - [Virtual KITTI 2](https://europe.naverlabs.com/research/computer-vision/proxy-virtual-worlds-vkitti-2/)
+  - [3D Ken Burns](https://github.com/sniklaus/3d-ken-burns.git)
+  - [BEDLAM](https://bedlam.is.tue.mpg.de/)
+  - [COP3D](https://github.com/facebookresearch/cop3d)
+  - [DL3DV](https://github.com/DL3DV-10K/Dataset)
+  - [Dynamic Replica](https://github.com/facebookresearch/dynamic_stereo)
+  - [EDEN](https://lhoangan.github.io/eden/)
+  - [Hypersim](https://github.com/apple/ml-hypersim)
+  - [IRS](https://github.com/HKBU-HPML/IRS)
+  - [Matterport3D](https://niessner.github.io/Matterport/)
+  - [MVImgNet](https://github.com/GAP-LAB-CUHK-SZ/MVImgNet)
+  - [MVS-Synth](https://phuang17.github.io/DeepMVS/mvs-synth.html)
+  - [OmniObject3D](https://omniobject3d.github.io/)
+  - [PointOdyssey](https://pointodyssey.com/)
+  - [RealEstate10K](https://google.github.io/realestate10k/)
+  - [SmartPortraits](https://mobileroboticsskoltech.github.io/SmartPortraits/)
+  - [Spring](https://spring-benchmark.org/)
+  - [Synscapes](https://synscapes.on.liu.se/)
+  - [UASOL](https://osf.io/64532/)
+  - [UrbanSyn](https://www.urbansyn.org/)
+  - [HOI4D](https://hoi4d.github.io/)
+
+
+## Evaluation
+
+### Datasets
+Please follow [MonST3R](https://github.com/Junyi42/monst3r/blob/main/data/evaluation_script.md) and [Spann3R](https://github.com/HengyiWang/spann3r/blob/main/docs/data_preprocess.md) to prepare **Sintel**, **Bonn**, **KITTI**, **NYU-v2**, **TUM-dynamics**, **ScanNet**, **7scenes** and **Neural-RGBD** datasets.
+
+The datasets should be organized as follows:
+```
+data/
+├── 7scenes
+├── bonn
+├── kitti
+├── neural_rgbd
+├── nyu-v2
+├── scannetv2
+├── sintel
+└── tum
+```
+
+### Evaluation Scripts
+Please refer to the [eval.md](docs/eval.md) for more details.
+
+## Training and Fine-tuning
+Please refer to the [train.md](docs/train.md) for more details.
+
+## Acknowledgements
+Our code is based on the following awesome repositories:
+
+- [DUSt3R](https://github.com/naver/dust3r)
+- [MonST3R](https://github.com/Junyi42/monst3r.git)
+- [Spann3R](https://github.com/HengyiWang/spann3r.git)
+- [Viser](https://github.com/nerfstudio-project/viser)
+
+We thank the authors for releasing their code!
+
+
+
+## Citation
+
+If you find our work useful, please cite:
+
+```bibtex
+@article{wang2025continuous,
+  title={Continuous 3D Perception Model with Persistent State},
+  author={Wang, Qianqian and Zhang, Yifei and Holynski, Aleksander and Efros, Alexei A and Kanazawa, Angjoo},
+  journal={arXiv preprint arXiv:2501.12387},
+  year={2025}
+}
+```
+

extern/CUT3R/add_ckpt_path.py

@@ -0,0 +1,9 @@
+import sys
+import os
+import os.path as path
+
+
+def add_path_to_dust3r(ckpt):
+    HERE_PATH = os.path.dirname(os.path.abspath(ckpt))
+    # workaround for sibling import
+    sys.path.insert(0, HERE_PATH)

extern/CUT3R/cloud_opt/base_opt.py

@@ -0,0 +1,301 @@
+from copy import deepcopy
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import roma
+from copy import deepcopy
+import tqdm
+import os
+import matplotlib.pyplot as plt
+
+from cloud_opt.utils import *
+from cloud_opt.utils import _check_edges, _compute_img_conf
+import cloud_opt.init_all as init_fun
+
+
+class BaseOptimizer(nn.Module):
+    """Optimize a global scene, given a graph-organized observations.
+    Graph node: images
+    Graph edges: observations = (pred1, pred2), pred2 is in pred1's coordinate
+    """
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def _init_from_views(
+        self,
+        view1s,
+        view2s,
+        pred1s,
+        pred2s,  # whatever predictions, they should be organized into pairwise for graph optimization
+        dist="l1",
+        conf="log",
+        min_conf_thr=3,
+        thr_for_init_conf=False,
+        base_scale=0.5,
+        allow_pw_adaptors=False,
+        pw_break=20,
+        rand_pose=torch.randn,
+        empty_cache=False,
+        verbose=True,
+    ):
+        super().__init__()
+        self.edges = [
+            (int(view1["idx"]), int(view2["idx"]))
+            for view1, view2 in zip(view1s, view2s)
+        ]
+        self.dist = ALL_DISTS[dist]
+        self.n_imgs = _check_edges(self.edges)
+
+        self.edge2pts_i = NoGradParamDict(
+            {ij: pred1s[n]["pts3d_is_self_view"] for n, ij in enumerate(self.str_edges)}
+        )  # ij: the name of the edge
+        self.edge2pts_j = NoGradParamDict(
+            {
+                ij: pred2s[n]["pts3d_in_other_view"]
+                for n, ij in enumerate(self.str_edges)
+            }
+        )
+        self.edge2conf_i = NoGradParamDict(
+            {ij: pred1s[n]["conf_self"] for n, ij in enumerate(self.str_edges)}
+        )
+        self.edge2conf_j = NoGradParamDict(
+            {ij: pred2s[n]["conf"] for n, ij in enumerate(self.str_edges)}
+        )
+
+        self.imshapes = get_imshapes(self.edges, pred1s, pred2s)
+        self.min_conf_thr = min_conf_thr
+        self.thr_for_init_conf = thr_for_init_conf
+        self.conf_trf = get_conf_trf(conf)
+
+        self.im_conf = _compute_img_conf(
+            self.imshapes, self.device, self.edges, self.edge2conf_i, self.edge2conf_j
+        )
+        for i in range(len(self.im_conf)):
+            self.im_conf[i].requires_grad = False
+
+        self.init_conf_maps = [c.clone() for c in self.im_conf]
+
+        self.base_scale = base_scale
+        self.norm_pw_scale = True
+        self.pw_break = pw_break
+        self.POSE_DIM = 7
+        self.pw_poses = nn.Parameter(
+            rand_pose((self.n_edges, 1 + self.POSE_DIM))
+        )  # pairwise poses
+        self.pw_adaptors = nn.Parameter(
+            torch.zeros((self.n_edges, 2))
+        )  # slight xy/z adaptation
+        self.pw_adaptors.requires_grad_(allow_pw_adaptors)
+        self.has_im_poses = False
+        self.rand_pose = rand_pose
+
+    def get_known_poses(self):
+        if self.has_im_poses:
+            known_poses_msk = torch.tensor(
+                [not (p.requires_grad) for p in self.im_poses]
+            )
+            known_poses = self.get_im_poses()
+            return known_poses_msk.sum(), known_poses_msk, known_poses
+        else:
+            return 0, None, None
+
+    def get_pw_norm_scale_factor(self):
+        if self.norm_pw_scale:
+            # normalize scales so that things cannot go south
+            # we want that exp(scale) ~= self.base_scale
+            return (np.log(self.base_scale) - self.pw_poses[:, -1].mean()).exp()
+        else:
+            return 1  # don't norm scale for known poses
+
+    def _set_pose(self, poses, idx, R, T=None, scale=None, force=False):
+        # all poses == cam-to-world
+        pose = poses[idx]
+        if not (pose.requires_grad or force):
+            return pose
+
+        if R.shape == (4, 4):
+            assert T is None
+            T = R[:3, 3]
+            R = R[:3, :3]
+
+        if R is not None:
+            pose.data[0:4] = roma.rotmat_to_unitquat(R)
+        if T is not None:
+            pose.data[4:7] = signed_log1p(
+                T / (scale or 1)
+            )  # translation is function of scale
+
+        if scale is not None:
+            assert poses.shape[-1] in (8, 13)
+            pose.data[-1] = np.log(float(scale))
+        return pose
+
+    def forward(self, ret_details=False):
+        pw_poses = self.get_pw_poses()  # cam-to-world
+        pw_adapt = self.get_adaptors()
+        proj_pts3d = self.get_pts3d()
+        # pre-compute pixel weights
+        weight_i = {i_j: self.conf_trf(c) for i_j, c in self.conf_i.items()}
+        weight_j = {i_j: self.conf_trf(c) for i_j, c in self.conf_j.items()}
+
+        loss = 0
+        if ret_details:
+            details = -torch.ones((self.n_imgs, self.n_imgs))
+
+        for e, (i, j) in enumerate(self.edges):
+            i_j = edge_str(i, j)
+            # distance in image i and j
+            aligned_pred_i = geotrf(pw_poses[e], pw_adapt[e] * self.pred_i[i_j])
+            aligned_pred_j = geotrf(pw_poses[e], pw_adapt[e] * self.pred_j[i_j])
+            li = self.dist(proj_pts3d[i], aligned_pred_i, weight=weight_i[i_j]).mean()
+            lj = self.dist(proj_pts3d[j], aligned_pred_j, weight=weight_j[i_j]).mean()
+            loss = loss + li + lj
+
+            if ret_details:
+                details[i, j] = li + lj
+        loss /= self.n_edges  # average over all pairs
+
+        if ret_details:
+            return loss, details
+        return loss
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def compute_global_alignment(self, init=None, niter_PnP=10, **kw):
+        if init is None:
+            pass
+        elif init == "msp" or init == "mst":
+            init_fun.init_minimum_spanning_tree(self, niter_PnP=niter_PnP)
+        elif init == "known_poses":
+            raise NotImplementedError
+            self.preset_pose(known_poses=self.camera_poses, requires_grad=True)
+            init_fun.init_from_known_poses(
+                self, min_conf_thr=self.min_conf_thr, niter_PnP=niter_PnP
+            )
+        else:
+            raise ValueError(f"bad value for {init=}")
+
+        return global_alignment_loop(self, **kw)
+
+    @property
+    def str_edges(self):
+        return [edge_str(i, j) for i, j in self.edges]
+
+    @property
+    def n_edges(self):
+        return len(self.edges)
+
+
+def global_alignment_loop(
+    net,
+    lr=0.01,
+    niter=300,
+    schedule="cosine",
+    lr_min=1e-3,
+    temporal_smoothing_weight=0,
+    depth_map_save_dir=None,
+):
+    params = [p for p in net.parameters() if p.requires_grad]
+    if not params:
+        return net
+
+    verbose = net.verbose
+    if verbose:
+        print("Global alignement - optimizing for:")
+        print([name for name, value in net.named_parameters() if value.requires_grad])
+
+    lr_base = lr
+    optimizer = torch.optim.Adam(params, lr=lr, betas=(0.9, 0.9))
+
+    loss = float("inf")
+    if verbose:
+        with tqdm.tqdm(total=niter) as bar:
+            while bar.n < bar.total:
+                if bar.n % 500 == 0 and depth_map_save_dir is not None:
+                    if not os.path.exists(depth_map_save_dir):
+                        os.makedirs(depth_map_save_dir)
+                    # visualize the depthmaps
+                    depth_maps = net.get_depthmaps()
+                    for i, depth_map in enumerate(depth_maps):
+                        depth_map_save_path = os.path.join(
+                            depth_map_save_dir, f"depthmaps_{i}_iter_{bar.n}.png"
+                        )
+                        plt.imsave(
+                            depth_map_save_path,
+                            depth_map.detach().cpu().numpy(),
+                            cmap="jet",
+                        )
+                    print(
+                        f"Saved depthmaps at iteration {bar.n} to {depth_map_save_dir}"
+                    )
+                loss, lr = global_alignment_iter(
+                    net,
+                    bar.n,
+                    niter,
+                    lr_base,
+                    lr_min,
+                    optimizer,
+                    schedule,
+                    temporal_smoothing_weight=temporal_smoothing_weight,
+                )
+                bar.set_postfix_str(f"{lr=:g} loss={loss:g}")
+                bar.update()
+    else:
+        for n in range(niter):
+            loss, _ = global_alignment_iter(
+                net,
+                n,
+                niter,
+                lr_base,
+                lr_min,
+                optimizer,
+                schedule,
+                temporal_smoothing_weight=temporal_smoothing_weight,
+            )
+    return loss
+
+
+def global_alignment_iter(
+    net,
+    cur_iter,
+    niter,
+    lr_base,
+    lr_min,
+    optimizer,
+    schedule,
+    temporal_smoothing_weight=0,
+):
+    t = cur_iter / niter
+    if schedule == "cosine":
+        lr = cosine_schedule(t, lr_base, lr_min)
+    elif schedule == "linear":
+        lr = linear_schedule(t, lr_base, lr_min)
+    elif schedule.startswith("cycle"):
+        try:
+            num_cycles = int(schedule[5:])
+        except ValueError:
+            num_cycles = 2
+        lr = cycled_linear_schedule(t, lr_base, lr_min, num_cycles=num_cycles)
+    else:
+        raise ValueError(f"bad lr {schedule=}")
+
+    adjust_learning_rate_by_lr(optimizer, lr)
+    optimizer.zero_grad()
+
+    if net.empty_cache:
+        torch.cuda.empty_cache()
+
+    loss = net(epoch=cur_iter)
+
+    if net.empty_cache:
+        torch.cuda.empty_cache()
+
+    loss.backward()
+
+    if net.empty_cache:
+        torch.cuda.empty_cache()
+
+    optimizer.step()
+
+    return float(loss), lr

extern/CUT3R/cloud_opt/commons.py

@@ -0,0 +1,102 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# utility functions for global alignment
+# --------------------------------------------------------
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+def edge_str(i, j):
+    return f"{i}_{j}"
+
+
+def i_j_ij(ij):
+    return edge_str(*ij), ij
+
+
+def edge_conf(conf_i, conf_j, edge):
+    return float(conf_i[edge].mean() * conf_j[edge].mean())
+
+
+def compute_edge_scores(edges, conf_i, conf_j):
+    return {(i, j): edge_conf(conf_i, conf_j, e) for e, (i, j) in edges}
+
+
+def NoGradParamDict(x):
+    assert isinstance(x, dict)
+    return nn.ParameterDict(x).requires_grad_(False)
+
+
+def get_imshapes(edges, pred_i, pred_j):
+    n_imgs = max(max(e) for e in edges) + 1
+    imshapes = [None] * n_imgs
+    for e, (i, j) in enumerate(edges):
+        shape_i = tuple(pred_i[e].shape[0:2])
+        shape_j = tuple(pred_j[e].shape[0:2])
+        if imshapes[i]:
+            assert imshapes[i] == shape_i, f"incorrect shape for image {i}"
+        if imshapes[j]:
+            assert imshapes[j] == shape_j, f"incorrect shape for image {j}"
+        imshapes[i] = shape_i
+        imshapes[j] = shape_j
+    return imshapes
+
+
+def get_conf_trf(mode):
+    if mode == "log":
+
+        def conf_trf(x):
+            return x.log()
+
+    elif mode == "sqrt":
+
+        def conf_trf(x):
+            return x.sqrt()
+
+    elif mode == "m1":
+
+        def conf_trf(x):
+            return x - 1
+
+    elif mode in ("id", "none"):
+
+        def conf_trf(x):
+            return x
+
+    else:
+        raise ValueError(f"bad mode for {mode=}")
+    return conf_trf
+
+
+def l2_dist(a, b, weight):
+    return (a - b).square().sum(dim=-1) * weight
+
+
+def l1_dist(a, b, weight):
+    return (a - b).norm(dim=-1) * weight
+
+
+ALL_DISTS = dict(l1=l1_dist, l2=l2_dist)
+
+
+def signed_log1p(x):
+    sign = torch.sign(x)
+    return sign * torch.log1p(torch.abs(x))
+
+
+def signed_expm1(x):
+    sign = torch.sign(x)
+    return sign * torch.expm1(torch.abs(x))
+
+
+def cosine_schedule(t, lr_start, lr_end):
+    assert 0 <= t <= 1
+    return lr_end + (lr_start - lr_end) * (1 + np.cos(t * np.pi)) / 2
+
+
+def linear_schedule(t, lr_start, lr_end):
+    assert 0 <= t <= 1
+    return lr_start + (lr_end - lr_start) * t

extern/CUT3R/cloud_opt/dust3r_opt/__init__.py

@@ -0,0 +1,31 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# global alignment optimization wrapper function
+# --------------------------------------------------------
+from enum import Enum
+
+from .optimizer import PointCloudOptimizer
+
+
+class GlobalAlignerMode(Enum):
+    PointCloudOptimizer = "PointCloudOptimizer"
+    ModularPointCloudOptimizer = "ModularPointCloudOptimizer"
+    PairViewer = "PairViewer"
+
+
+def global_aligner(
+    dust3r_output, device, mode=GlobalAlignerMode.PointCloudOptimizer, **optim_kw
+):
+    # extract all inputs
+    view1, view2, pred1, pred2 = [
+        dust3r_output[k] for k in "view1 view2 pred1 pred2".split()
+    ]
+    # build the optimizer
+    if mode == GlobalAlignerMode.PointCloudOptimizer:
+        net = PointCloudOptimizer(view1, view2, pred1, pred2, **optim_kw).to(device)
+    else:
+        raise NotImplementedError(f"Unknown mode {mode}")
+
+    return net

extern/CUT3R/cloud_opt/dust3r_opt/base_opt.py

@@ -0,0 +1,620 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Base class for the global alignement procedure
+# --------------------------------------------------------
+from copy import deepcopy
+
+import numpy as np
+import torch
+import torch.nn as nn
+import roma
+from copy import deepcopy
+import tqdm
+import cv2
+from PIL import Image
+from dust3r.utils.geometry import inv, geotrf
+from dust3r.utils.device import to_numpy
+from dust3r.utils.image import rgb
+from dust3r.viz import SceneViz, segment_sky, auto_cam_size
+
+from cloud_opt.dust3r_opt.commons import (
+    edge_str,
+    ALL_DISTS,
+    NoGradParamDict,
+    get_imshapes,
+    signed_expm1,
+    signed_log1p,
+    cosine_schedule,
+    linear_schedule,
+    get_conf_trf,
+)
+import cloud_opt.dust3r_opt.init_im_poses as init_fun
+from pathlib import Path
+from scipy.spatial.transform import Rotation
+from evo.core.trajectory import PosePath3D, PoseTrajectory3D
+
+
+def adjust_learning_rate_by_lr(optimizer, lr):
+    for param_group in optimizer.param_groups:
+        if "lr_scale" in param_group:
+            param_group["lr"] = lr * param_group["lr_scale"]
+        else:
+            param_group["lr"] = lr
+
+
+def make_traj(args) -> PoseTrajectory3D:
+    if isinstance(args, tuple) or isinstance(args, list):
+        traj, tstamps = args
+        return PoseTrajectory3D(
+            positions_xyz=traj[:, :3],
+            orientations_quat_wxyz=traj[:, 3:],
+            timestamps=tstamps,
+        )
+    assert isinstance(args, PoseTrajectory3D), type(args)
+    return deepcopy(args)
+
+
+def save_trajectory_tum_format(traj, filename):
+    traj = make_traj(traj)
+    tostr = lambda a: " ".join(map(str, a))
+    with Path(filename).open("w") as f:
+        for i in range(traj.num_poses):
+            f.write(
+                f"{traj.timestamps[i]} {tostr(traj.positions_xyz[i])} {tostr(traj.orientations_quat_wxyz[i][[0,1,2,3]])}\n"
+            )
+    print(f"Saved trajectory to {filename}")
+
+
+def c2w_to_tumpose(c2w):
+    """
+    Convert a camera-to-world matrix to a tuple of translation and rotation
+
+    input: c2w: 4x4 matrix
+    output: tuple of translation and rotation (x y z qw qx qy qz)
+    """
+    # convert input to numpy
+    c2w = to_numpy(c2w)
+    xyz = c2w[:3, -1]
+    rot = Rotation.from_matrix(c2w[:3, :3])
+    qx, qy, qz, qw = rot.as_quat()
+    tum_pose = np.concatenate([xyz, [qw, qx, qy, qz]])
+    return tum_pose
+
+
+class BasePCOptimizer(nn.Module):
+    """Optimize a global scene, given a list of pairwise observations.
+    Graph node: images
+    Graph edges: observations = (pred1, pred2)
+    """
+
+    def __init__(self, *args, **kwargs):
+        if len(args) == 1 and len(kwargs) == 0:
+            other = deepcopy(args[0])
+            attrs = """edges is_symmetrized dist n_imgs pred_i pred_j imshapes 
+                        min_conf_thr conf_thr conf_i conf_j im_conf
+                        base_scale norm_pw_scale POSE_DIM pw_poses 
+                        pw_adaptors pw_adaptors has_im_poses rand_pose imgs verbose""".split()
+            self.__dict__.update({k: other[k] for k in attrs})
+        else:
+            self._init_from_views(*args, **kwargs)
+
+    def _init_from_views(
+        self,
+        view1,
+        view2,
+        pred1,
+        pred2,
+        dist="l1",
+        conf="log",
+        min_conf_thr=3,
+        base_scale=0.5,
+        allow_pw_adaptors=False,
+        pw_break=20,
+        rand_pose=torch.randn,
+        iterationsCount=None,
+        verbose=True,
+    ):
+        super().__init__()
+        if not isinstance(view1["idx"], list):
+            view1["idx"] = view1["idx"].tolist()
+        if not isinstance(view2["idx"], list):
+            view2["idx"] = view2["idx"].tolist()
+        self.edges = [(int(i), int(j)) for i, j in zip(view1["idx"], view2["idx"])]
+        self.is_symmetrized = set(self.edges) == {(j, i) for i, j in self.edges}
+        self.dist = ALL_DISTS[dist]
+        self.verbose = verbose
+
+        self.n_imgs = self._check_edges()
+
+        # input data
+        pred1_pts = pred1["pts3d_in_self_view"]
+        pred2_pts = pred2["pts3d_in_other_view"]
+        self.pred_i = NoGradParamDict(
+            {ij: pred1_pts[n] for n, ij in enumerate(self.str_edges)}
+        )
+        self.pred_j = NoGradParamDict(
+            {ij: pred2_pts[n] for n, ij in enumerate(self.str_edges)}
+        )
+        self.imshapes = get_imshapes(self.edges, pred1_pts, pred2_pts)
+
+        # work in log-scale with conf
+        pred1_conf = pred1["conf_self"]
+        pred2_conf = pred2["conf"]
+        self.min_conf_thr = min_conf_thr
+        self.conf_trf = get_conf_trf(conf)
+
+        self.conf_i = NoGradParamDict(
+            {ij: pred1_conf[n] for n, ij in enumerate(self.str_edges)}
+        )
+        self.conf_j = NoGradParamDict(
+            {ij: pred2_conf[n] for n, ij in enumerate(self.str_edges)}
+        )
+        self.im_conf = self._compute_img_conf(pred1_conf, pred2_conf)
+        for i in range(len(self.im_conf)):
+            self.im_conf[i].requires_grad = False
+
+        # pairwise pose parameters
+        self.base_scale = base_scale
+        self.norm_pw_scale = True
+        self.pw_break = pw_break
+        self.POSE_DIM = 7
+        self.pw_poses = nn.Parameter(
+            rand_pose((self.n_edges, 1 + self.POSE_DIM))
+        )  # pairwise poses
+        self.pw_adaptors = nn.Parameter(
+            torch.zeros((self.n_edges, 2))
+        )  # slight xy/z adaptation
+        self.pw_adaptors.requires_grad_(allow_pw_adaptors)
+        self.has_im_poses = False
+        self.rand_pose = rand_pose
+
+        # possibly store images for show_pointcloud
+        self.imgs = None
+        if "img" in view1 and "img" in view2:
+            imgs = [torch.zeros((3,) + hw) for hw in self.imshapes]
+            for v in range(len(self.edges)):
+                idx = view1["idx"][v]
+                imgs[idx] = view1["img"][v]
+                idx = view2["idx"][v]
+                imgs[idx] = view2["img"][v]
+            self.imgs = rgb(imgs)
+
+    @property
+    def n_edges(self):
+        return len(self.edges)
+
+    @property
+    def str_edges(self):
+        return [edge_str(i, j) for i, j in self.edges]
+
+    @property
+    def imsizes(self):
+        return [(w, h) for h, w in self.imshapes]
+
+    @property
+    def device(self):
+        return next(iter(self.parameters())).device
+
+    def state_dict(self, trainable=True):
+        all_params = super().state_dict()
+        return {
+            k: v
+            for k, v in all_params.items()
+            if k.startswith(("_", "pred_i.", "pred_j.", "conf_i.", "conf_j."))
+            != trainable
+        }
+
+    def load_state_dict(self, data):
+        return super().load_state_dict(self.state_dict(trainable=False) | data)
+
+    def _check_edges(self):
+        indices = sorted({i for edge in self.edges for i in edge})
+        assert indices == list(range(len(indices))), "bad pair indices: missing values "
+        return len(indices)
+
+    @torch.no_grad()
+    def _compute_img_conf(self, pred1_conf, pred2_conf):
+        im_conf = nn.ParameterList(
+            [torch.zeros(hw, device=self.device) for hw in self.imshapes]
+        )
+        for e, (i, j) in enumerate(self.edges):
+            im_conf[i] = torch.maximum(im_conf[i], pred1_conf[e])
+            im_conf[j] = torch.maximum(im_conf[j], pred2_conf[e])
+        return im_conf
+
+    def get_adaptors(self):
+        adapt = self.pw_adaptors
+        adapt = torch.cat(
+            (adapt[:, 0:1], adapt), dim=-1
+        )  # (scale_xy, scale_xy, scale_z)
+        if self.norm_pw_scale:  # normalize so that the product == 1
+            adapt = adapt - adapt.mean(dim=1, keepdim=True)
+        return (adapt / self.pw_break).exp()
+
+    def _get_poses(self, poses):
+        # normalize rotation
+        Q = poses[:, :4]
+        T = signed_expm1(poses[:, 4:7])
+        RT = roma.RigidUnitQuat(Q, T).normalize().to_homogeneous()
+        return RT
+
+    def _set_pose(self, poses, idx, R, T=None, scale=None, force=False):
+        # all poses == cam-to-world
+        pose = poses[idx]
+        if not (pose.requires_grad or force):
+            return pose
+
+        if R.shape == (4, 4):
+            assert T is None
+            T = R[:3, 3]
+            R = R[:3, :3]
+
+        if R is not None:
+            pose.data[0:4] = roma.rotmat_to_unitquat(R)
+        if T is not None:
+            pose.data[4:7] = signed_log1p(
+                T / (scale or 1)
+            )  # translation is function of scale
+
+        if scale is not None:
+            assert poses.shape[-1] in (8, 13)
+            pose.data[-1] = np.log(float(scale))
+        return pose
+
+    def get_pw_norm_scale_factor(self):
+        if self.norm_pw_scale:
+            # normalize scales so that things cannot go south
+            # we want that exp(scale) ~= self.base_scale
+            return (np.log(self.base_scale) - self.pw_poses[:, -1].mean()).exp()
+        else:
+            return 1  # don't norm scale for known poses
+
+    def get_pw_scale(self):
+        scale = self.pw_poses[:, -1].exp()  # (n_edges,)
+        scale = scale * self.get_pw_norm_scale_factor()
+        return scale
+
+    def get_pw_poses(self):  # cam to world
+        RT = self._get_poses(self.pw_poses)
+        scaled_RT = RT.clone()
+        scaled_RT[:, :3] *= self.get_pw_scale().view(
+            -1, 1, 1
+        )  # scale the rotation AND translation
+        return scaled_RT
+
+    def get_masks(self):
+        return [(conf > self.min_conf_thr) for conf in self.im_conf]
+
+    def depth_to_pts3d(self):
+        raise NotImplementedError()
+
+    def get_pts3d(self, raw=False):
+        res = self.depth_to_pts3d()
+        if not raw:
+            res = [dm[: h * w].view(h, w, 3) for dm, (h, w) in zip(res, self.imshapes)]
+        return res
+
+    def _set_focal(self, idx, focal, force=False):
+        raise NotImplementedError()
+
+    def get_focals(self):
+        raise NotImplementedError()
+
+    def get_known_focal_mask(self):
+        raise NotImplementedError()
+
+    def get_principal_points(self):
+        raise NotImplementedError()
+
+    def get_conf(self, mode=None):
+        trf = self.conf_trf if mode is None else get_conf_trf(mode)
+        return [trf(c) for c in self.im_conf]
+
+    def get_im_poses(self):
+        raise NotImplementedError()
+
+    def _set_depthmap(self, idx, depth, force=False):
+        raise NotImplementedError()
+
+    def get_depthmaps(self, raw=False):
+        raise NotImplementedError()
+
+    def save_depth_maps(self, path):
+        depth_maps = self.get_depthmaps()
+        images = []
+
+        for i, depth_map in enumerate(depth_maps):
+            # Apply color map to depth map
+            depth_map_colored = cv2.applyColorMap(
+                (depth_map * 255).detach().cpu().numpy().astype(np.uint8),
+                cv2.COLORMAP_JET,
+            )
+            img_path = f"{path}/frame_{(i):04d}.png"
+            cv2.imwrite(img_path, depth_map_colored)
+            images.append(Image.open(img_path))
+            np.save(f"{path}/frame_{(i):04d}.npy", depth_map.detach().cpu().numpy())
+
+        images[0].save(
+            f"{path}/_depth_maps.gif",
+            save_all=True,
+            append_images=images[1:],
+            duration=100,
+            loop=0,
+        )
+
+        return depth_maps
+
+    def clean_pointcloud(self, **kw):
+        cams = inv(self.get_im_poses())
+        K = self.get_intrinsics()
+        depthmaps = self.get_depthmaps()
+        all_pts3d = self.get_pts3d()
+
+        new_im_confs = clean_pointcloud(
+            self.im_conf, K, cams, depthmaps, all_pts3d, **kw
+        )
+        for i, new_conf in enumerate(new_im_confs):
+            self.im_conf[i].data[:] = new_conf
+        return self
+
+    def get_tum_poses(self):
+        poses = self.get_im_poses()
+        tt = np.arange(len(poses)).astype(float)
+        tum_poses = [c2w_to_tumpose(p) for p in poses]
+        tum_poses = np.stack(tum_poses, 0)
+        return [tum_poses, tt]
+
+    def save_tum_poses(self, path):
+        traj = self.get_tum_poses()
+        save_trajectory_tum_format(traj, path)
+        return traj[0]  # return the poses
+
+    def save_focals(self, path):
+        # convert focal to txt
+        focals = self.get_focals()
+        np.savetxt(path, focals.detach().cpu().numpy(), fmt="%.6f")
+        return focals
+
+    def save_intrinsics(self, path):
+        K_raw = self.get_intrinsics()
+        K = K_raw.reshape(-1, 9)
+        np.savetxt(path, K.detach().cpu().numpy(), fmt="%.6f")
+        return K_raw
+
+    def save_conf_maps(self, path):
+        conf = self.get_conf()
+        for i, c in enumerate(conf):
+            np.save(f"{path}/conf_{i}.npy", c.detach().cpu().numpy())
+        return conf
+
+    def save_init_conf_maps(self, path):
+        conf = self.get_init_conf()
+        for i, c in enumerate(conf):
+            np.save(f"{path}/init_conf_{i}.npy", c.detach().cpu().numpy())
+        return conf
+
+    def save_rgb_imgs(self, path):
+        imgs = self.imgs
+        for i, img in enumerate(imgs):
+            # convert from rgb to bgr
+            img = img[..., ::-1]
+            cv2.imwrite(f"{path}/frame_{i:04d}.png", img * 255)
+        return imgs
+
+    def save_dynamic_masks(self, path):
+        dynamic_masks = (
+            self.dynamic_masks
+            if getattr(self, "sam2_dynamic_masks", None) is None
+            else self.sam2_dynamic_masks
+        )
+        for i, dynamic_mask in enumerate(dynamic_masks):
+            cv2.imwrite(
+                f"{path}/dynamic_mask_{i}.png",
+                (dynamic_mask * 255).detach().cpu().numpy().astype(np.uint8),
+            )
+        return dynamic_masks
+
+    def save_depth_maps(self, path):
+        depth_maps = self.get_depthmaps()
+        images = []
+
+        for i, depth_map in enumerate(depth_maps):
+            # Apply color map to depth map
+            depth_map_colored = cv2.applyColorMap(
+                (depth_map * 255).detach().cpu().numpy().astype(np.uint8),
+                cv2.COLORMAP_JET,
+            )
+            img_path = f"{path}/frame_{(i):04d}.png"
+            cv2.imwrite(img_path, depth_map_colored)
+            images.append(Image.open(img_path))
+            np.save(f"{path}/frame_{(i):04d}.npy", depth_map.detach().cpu().numpy())
+
+        images[0].save(
+            f"{path}/_depth_maps.gif",
+            save_all=True,
+            append_images=images[1:],
+            duration=100,
+            loop=0,
+        )
+
+        return depth_maps
+
+    def forward(self, ret_details=False):
+        pw_poses = self.get_pw_poses()  # cam-to-world
+        pw_adapt = self.get_adaptors()
+        proj_pts3d = self.get_pts3d()
+        # pre-compute pixel weights
+        weight_i = {i_j: self.conf_trf(c) for i_j, c in self.conf_i.items()}
+        weight_j = {i_j: self.conf_trf(c) for i_j, c in self.conf_j.items()}
+
+        loss = 0
+        if ret_details:
+            details = -torch.ones((self.n_imgs, self.n_imgs))
+
+        for e, (i, j) in enumerate(self.edges):
+            i_j = edge_str(i, j)
+            # distance in image i and j
+            aligned_pred_i = geotrf(pw_poses[e], pw_adapt[e] * self.pred_i[i_j])
+            aligned_pred_j = geotrf(pw_poses[e], pw_adapt[e] * self.pred_j[i_j])
+            li = self.dist(proj_pts3d[i], aligned_pred_i, weight=weight_i[i_j]).mean()
+            lj = self.dist(proj_pts3d[j], aligned_pred_j, weight=weight_j[i_j]).mean()
+            loss = loss + li + lj
+
+            if ret_details:
+                details[i, j] = li + lj
+        loss /= self.n_edges  # average over all pairs
+
+        if ret_details:
+            return loss, details
+        return loss
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def compute_global_alignment(self, init=None, niter_PnP=10, **kw):
+        if init is None:
+            pass
+        elif init == "msp" or init == "mst":
+            init_fun.init_minimum_spanning_tree(self, niter_PnP=niter_PnP)
+        elif init == "known_poses":
+            init_fun.init_from_known_poses(
+                self, min_conf_thr=self.min_conf_thr, niter_PnP=niter_PnP
+            )
+        else:
+            raise ValueError(f"bad value for {init=}")
+        return global_alignment_loop(self, **kw)
+
+    @torch.no_grad()
+    def mask_sky(self):
+        res = deepcopy(self)
+        for i in range(self.n_imgs):
+            sky = segment_sky(self.imgs[i])
+            res.im_conf[i][sky] = 0
+        return res
+
+    def show(self, show_pw_cams=False, show_pw_pts3d=False, cam_size=None, **kw):
+        viz = SceneViz()
+        if self.imgs is None:
+            colors = np.random.randint(0, 256, size=(self.n_imgs, 3))
+            colors = list(map(tuple, colors.tolist()))
+            for n in range(self.n_imgs):
+                viz.add_pointcloud(self.get_pts3d()[n], colors[n], self.get_masks()[n])
+        else:
+            viz.add_pointcloud(self.get_pts3d(), self.imgs, self.get_masks())
+            colors = np.random.randint(256, size=(self.n_imgs, 3))
+
+        # camera poses
+        im_poses = to_numpy(self.get_im_poses())
+        if cam_size is None:
+            cam_size = auto_cam_size(im_poses)
+        viz.add_cameras(
+            im_poses,
+            self.get_focals(),
+            colors=colors,
+            images=self.imgs,
+            imsizes=self.imsizes,
+            cam_size=cam_size,
+        )
+        if show_pw_cams:
+            pw_poses = self.get_pw_poses()
+            viz.add_cameras(pw_poses, color=(192, 0, 192), cam_size=cam_size)
+
+            if show_pw_pts3d:
+                pts = [
+                    geotrf(pw_poses[e], self.pred_i[edge_str(i, j)])
+                    for e, (i, j) in enumerate(self.edges)
+                ]
+                viz.add_pointcloud(pts, (128, 0, 128))
+
+        viz.show(**kw)
+        return viz
+
+
+def global_alignment_loop(net, lr=0.01, niter=300, schedule="cosine", lr_min=1e-6):
+    params = [p for p in net.parameters() if p.requires_grad]
+    if not params:
+        return net
+
+    verbose = net.verbose
+    if verbose:
+        print("Global alignement - optimizing for:")
+        print([name for name, value in net.named_parameters() if value.requires_grad])
+
+    lr_base = lr
+    optimizer = torch.optim.Adam(params, lr=lr, betas=(0.9, 0.9))
+
+    loss = float("inf")
+    if verbose:
+        with tqdm.tqdm(total=niter) as bar:
+            while bar.n < bar.total:
+                loss, lr = global_alignment_iter(
+                    net, bar.n, niter, lr_base, lr_min, optimizer, schedule
+                )
+                bar.set_postfix_str(f"{lr=:g} loss={loss:g}")
+                bar.update()
+    else:
+        for n in range(niter):
+            loss, _ = global_alignment_iter(
+                net, n, niter, lr_base, lr_min, optimizer, schedule
+            )
+    return loss
+
+
+def global_alignment_iter(net, cur_iter, niter, lr_base, lr_min, optimizer, schedule):
+    t = cur_iter / niter
+    if schedule == "cosine":
+        lr = cosine_schedule(t, lr_base, lr_min)
+    elif schedule == "linear":
+        lr = linear_schedule(t, lr_base, lr_min)
+    else:
+        raise ValueError(f"bad lr {schedule=}")
+    adjust_learning_rate_by_lr(optimizer, lr)
+    optimizer.zero_grad()
+    loss = net()
+    loss.backward()
+    optimizer.step()
+
+    return float(loss), lr
+
+
+@torch.no_grad()
+def clean_pointcloud(
+    im_confs, K, cams, depthmaps, all_pts3d, tol=0.001, bad_conf=0, dbg=()
+):
+    """Method:
+    1) express all 3d points in each camera coordinate frame
+    2) if they're in front of a depthmap --> then lower their confidence
+    """
+    assert len(im_confs) == len(cams) == len(K) == len(depthmaps) == len(all_pts3d)
+    assert 0 <= tol < 1
+    res = [c.clone() for c in im_confs]
+
+    # reshape appropriately
+    all_pts3d = [p.view(*c.shape, 3) for p, c in zip(all_pts3d, im_confs)]
+    depthmaps = [d.view(*c.shape) for d, c in zip(depthmaps, im_confs)]
+
+    for i, pts3d in enumerate(all_pts3d):
+        for j in range(len(all_pts3d)):
+            if i == j:
+                continue
+
+            # project 3dpts in other view
+            proj = geotrf(cams[j], pts3d)
+            proj_depth = proj[:, :, 2]
+            u, v = geotrf(K[j], proj, norm=1, ncol=2).round().long().unbind(-1)
+
+            # check which points are actually in the visible cone
+            H, W = im_confs[j].shape
+            msk_i = (proj_depth > 0) & (0 <= u) & (u < W) & (0 <= v) & (v < H)
+            msk_j = v[msk_i], u[msk_i]
+
+            # find bad points = those in front but less confident
+            bad_points = (proj_depth[msk_i] < (1 - tol) * depthmaps[j][msk_j]) & (
+                res[i][msk_i] < res[j][msk_j]
+            )
+
+            bad_msk_i = msk_i.clone()
+            bad_msk_i[msk_i] = bad_points
+            res[i][bad_msk_i] = res[i][bad_msk_i].clip_(max=bad_conf)
+
+    return res

extern/CUT3R/cloud_opt/dust3r_opt/commons.py

@@ -0,0 +1,102 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# utility functions for global alignment
+# --------------------------------------------------------
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+def edge_str(i, j):
+    return f"{i}_{j}"
+
+
+def i_j_ij(ij):
+    return edge_str(*ij), ij
+
+
+def edge_conf(conf_i, conf_j, edge):
+    return float(conf_i[edge].mean() * conf_j[edge].mean())
+
+
+def compute_edge_scores(edges, conf_i, conf_j):
+    return {(i, j): edge_conf(conf_i, conf_j, e) for e, (i, j) in edges}
+
+
+def NoGradParamDict(x):
+    assert isinstance(x, dict)
+    return nn.ParameterDict(x).requires_grad_(False)
+
+
+def get_imshapes(edges, pred_i, pred_j):
+    n_imgs = max(max(e) for e in edges) + 1
+    imshapes = [None] * n_imgs
+    for e, (i, j) in enumerate(edges):
+        shape_i = tuple(pred_i[e].shape[0:2])
+        shape_j = tuple(pred_j[e].shape[0:2])
+        if imshapes[i]:
+            assert imshapes[i] == shape_i, f"incorrect shape for image {i}"
+        if imshapes[j]:
+            assert imshapes[j] == shape_j, f"incorrect shape for image {j}"
+        imshapes[i] = shape_i
+        imshapes[j] = shape_j
+    return imshapes
+
+
+def get_conf_trf(mode):
+    if mode == "log":
+
+        def conf_trf(x):
+            return x.log()
+
+    elif mode == "sqrt":
+
+        def conf_trf(x):
+            return x.sqrt()
+
+    elif mode == "m1":
+
+        def conf_trf(x):
+            return x - 1
+
+    elif mode in ("id", "none"):
+
+        def conf_trf(x):
+            return x
+
+    else:
+        raise ValueError(f"bad mode for {mode=}")
+    return conf_trf
+
+
+def l2_dist(a, b, weight):
+    return (a - b).square().sum(dim=-1) * weight
+
+
+def l1_dist(a, b, weight):
+    return (a - b).norm(dim=-1) * weight
+
+
+ALL_DISTS = dict(l1=l1_dist, l2=l2_dist)
+
+
+def signed_log1p(x):
+    sign = torch.sign(x)
+    return sign * torch.log1p(torch.abs(x))
+
+
+def signed_expm1(x):
+    sign = torch.sign(x)
+    return sign * torch.expm1(torch.abs(x))
+
+
+def cosine_schedule(t, lr_start, lr_end):
+    assert 0 <= t <= 1
+    return lr_end + (lr_start - lr_end) * (1 + np.cos(t * np.pi)) / 2
+
+
+def linear_schedule(t, lr_start, lr_end):
+    assert 0 <= t <= 1
+    return lr_start + (lr_end - lr_start) * t

extern/CUT3R/cloud_opt/dust3r_opt/init_im_poses.py

@@ -0,0 +1,382 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Initialization functions for global alignment
+# --------------------------------------------------------
+from functools import cache
+
+import numpy as np
+import scipy.sparse as sp
+import torch
+import cv2
+import roma
+from tqdm import tqdm
+
+from dust3r.utils.geometry import geotrf, inv, get_med_dist_between_poses
+from dust3r.post_process import estimate_focal_knowing_depth
+from dust3r.viz import to_numpy
+
+from cloud_opt.commons import edge_str, i_j_ij, compute_edge_scores
+
+
+@torch.no_grad()
+def init_from_known_poses(self, niter_PnP=10, min_conf_thr=3):
+    device = self.device
+
+    # indices of known poses
+    nkp, known_poses_msk, known_poses = get_known_poses(self)
+    assert nkp == self.n_imgs, "not all poses are known"
+
+    # get all focals
+    nkf, _, im_focals = get_known_focals(self)
+    assert nkf == self.n_imgs
+    im_pp = self.get_principal_points()
+
+    best_depthmaps = {}
+    # init all pairwise poses
+    for e, (i, j) in enumerate(tqdm(self.edges, disable=not self.verbose)):
+        i_j = edge_str(i, j)
+
+        # find relative pose for this pair
+        P1 = torch.eye(4, device=device)
+        msk = self.conf_i[i_j] > min(min_conf_thr, self.conf_i[i_j].min() - 0.1)
+        _, P2 = fast_pnp(
+            self.pred_j[i_j],
+            float(im_focals[i].mean()),
+            pp=im_pp[i],
+            msk=msk,
+            device=device,
+            niter_PnP=niter_PnP,
+        )
+
+        # align the two predicted camera with the two gt cameras
+        s, R, T = align_multiple_poses(torch.stack((P1, P2)), known_poses[[i, j]])
+        # normally we have known_poses[i] ~= sRT_to_4x4(s,R,T,device) @ P1
+        # and geotrf(sRT_to_4x4(1,R,T,device), s*P2[:3,3])
+        self._set_pose(self.pw_poses, e, R, T, scale=s)
+
+        # remember if this is a good depthmap
+        score = float(self.conf_i[i_j].mean())
+        if score > best_depthmaps.get(i, (0,))[0]:
+            best_depthmaps[i] = score, i_j, s
+
+    # init all image poses
+    for n in range(self.n_imgs):
+        assert known_poses_msk[n]
+        _, i_j, scale = best_depthmaps[n]
+        depth = self.pred_i[i_j][:, :, 2]
+        self._set_depthmap(n, depth * scale)
+
+
+@torch.no_grad()
+def init_minimum_spanning_tree(self, **kw):
+    """Init all camera poses (image-wise and pairwise poses) given
+    an initial set of pairwise estimations.
+    """
+    device = self.device
+    pts3d, _, im_focals, im_poses = minimum_spanning_tree(
+        self.imshapes,
+        self.edges,
+        self.pred_i,
+        self.pred_j,
+        self.conf_i,
+        self.conf_j,
+        self.im_conf,
+        self.min_conf_thr,
+        device,
+        has_im_poses=self.has_im_poses,
+        verbose=self.verbose,
+        **kw,
+    )
+
+    return init_from_pts3d(self, pts3d, im_focals, im_poses)
+
+
+def init_from_pts3d(self, pts3d, im_focals, im_poses):
+    # init poses
+    nkp, known_poses_msk, known_poses = get_known_poses(self)
+    if nkp == 1:
+        raise NotImplementedError(
+            "Would be simpler to just align everything afterwards on the single known pose"
+        )
+    elif nkp > 1:
+        # global rigid SE3 alignment
+        s, R, T = align_multiple_poses(
+            im_poses[known_poses_msk], known_poses[known_poses_msk]
+        )
+        trf = sRT_to_4x4(s, R, T, device=known_poses.device)
+
+        # rotate everything
+        im_poses = trf @ im_poses
+        im_poses[:, :3, :3] /= s  # undo scaling on the rotation part
+        for img_pts3d in pts3d:
+            img_pts3d[:] = geotrf(trf, img_pts3d)
+
+    # set all pairwise poses
+    for e, (i, j) in enumerate(self.edges):
+        i_j = edge_str(i, j)
+        # compute transform that goes from cam to world
+        s, R, T = rigid_points_registration(
+            self.pred_i[i_j], pts3d[i], conf=self.conf_i[i_j]
+        )
+        self._set_pose(self.pw_poses, e, R, T, scale=s)
+
+    # take into account the scale normalization
+    s_factor = self.get_pw_norm_scale_factor()
+    im_poses[:, :3, 3] *= s_factor  # apply downscaling factor
+    for img_pts3d in pts3d:
+        img_pts3d *= s_factor
+
+    # init all image poses
+    if self.has_im_poses:
+        for i in range(self.n_imgs):
+            cam2world = im_poses[i]
+            depth = geotrf(inv(cam2world), pts3d[i])[..., 2]
+            self._set_depthmap(i, depth)
+            self._set_pose(self.im_poses, i, cam2world)
+            if im_focals[i] is not None:
+                self._set_focal(i, im_focals[i])
+
+    if self.verbose:
+        pass
+        # print(' init loss =', float(self()))
+
+
+def minimum_spanning_tree(
+    imshapes,
+    edges,
+    pred_i,
+    pred_j,
+    conf_i,
+    conf_j,
+    im_conf,
+    min_conf_thr,
+    device,
+    has_im_poses=True,
+    niter_PnP=10,
+    verbose=True,
+):
+    n_imgs = len(imshapes)
+    sparse_graph = -dict_to_sparse_graph(
+        compute_edge_scores(map(i_j_ij, edges), conf_i, conf_j)
+    )
+    print(sparse_graph)
+    msp = sp.csgraph.minimum_spanning_tree(sparse_graph).tocoo()
+
+    # temp variable to store 3d points
+    pts3d = [None] * len(imshapes)
+
+    todo = sorted(zip(-msp.data, msp.row, msp.col))  # sorted edges
+    im_poses = [None] * n_imgs
+    im_focals = [None] * n_imgs
+
+    # init with strongest edge
+    score, i, j = todo.pop()
+    if verbose:
+        print(f" init edge ({i}*,{j}*) {score=}")
+    i_j = edge_str(i, j)
+    pts3d[i] = pred_i[i_j].clone()
+    pts3d[j] = pred_j[i_j].clone()
+    done = {i, j}
+    if has_im_poses:
+        im_poses[i] = torch.eye(4, device=device)
+        im_focals[i] = estimate_focal(pred_i[i_j])
+
+    # set initial pointcloud based on pairwise graph
+    msp_edges = [(i, j)]
+    while todo:
+        # each time, predict the next one
+        score, i, j = todo.pop()
+
+        if im_focals[i] is None:
+            im_focals[i] = estimate_focal(pred_i[i_j])
+
+        if i in done:
+            if verbose:
+                print(f" init edge ({i},{j}*) {score=}")
+            assert j not in done
+            # align pred[i] with pts3d[i], and then set j accordingly
+            i_j = edge_str(i, j)
+            s, R, T = rigid_points_registration(pred_i[i_j], pts3d[i], conf=conf_i[i_j])
+            trf = sRT_to_4x4(s, R, T, device)
+            pts3d[j] = geotrf(trf, pred_j[i_j])
+            done.add(j)
+            msp_edges.append((i, j))
+
+            if has_im_poses and im_poses[i] is None:
+                im_poses[i] = sRT_to_4x4(1, R, T, device)
+
+        elif j in done:
+            if verbose:
+                print(f" init edge ({i}*,{j}) {score=}")
+            assert i not in done
+            i_j = edge_str(i, j)
+            s, R, T = rigid_points_registration(pred_j[i_j], pts3d[j], conf=conf_j[i_j])
+            trf = sRT_to_4x4(s, R, T, device)
+            pts3d[i] = geotrf(trf, pred_i[i_j])
+            done.add(i)
+            msp_edges.append((i, j))
+
+            if has_im_poses and im_poses[i] is None:
+                im_poses[i] = sRT_to_4x4(1, R, T, device)
+        else:
+            # let's try again later
+            todo.insert(0, (score, i, j))
+
+    if has_im_poses:
+        # complete all missing informations
+        pair_scores = list(
+            sparse_graph.values()
+        )  # already negative scores: less is best
+        edges_from_best_to_worse = np.array(list(sparse_graph.keys()))[
+            np.argsort(pair_scores)
+        ]
+        for i, j in edges_from_best_to_worse.tolist():
+            if im_focals[i] is None:
+                im_focals[i] = estimate_focal(pred_i[edge_str(i, j)])
+
+        for i in range(n_imgs):
+            if im_poses[i] is None:
+                msk = im_conf[i] > min_conf_thr
+                res = fast_pnp(
+                    pts3d[i], im_focals[i], msk=msk, device=device, niter_PnP=niter_PnP
+                )
+                if res:
+                    im_focals[i], im_poses[i] = res
+            if im_poses[i] is None:
+                im_poses[i] = torch.eye(4, device=device)
+        im_poses = torch.stack(im_poses)
+    else:
+        im_poses = im_focals = None
+
+    return pts3d, msp_edges, im_focals, im_poses
+
+
+def dict_to_sparse_graph(dic):
+    n_imgs = max(max(e) for e in dic) + 1
+    res = sp.dok_array((n_imgs, n_imgs))
+    for edge, value in dic.items():
+        res[edge] = value
+    return res
+
+
+def rigid_points_registration(pts1, pts2, conf):
+    R, T, s = roma.rigid_points_registration(
+        pts1.reshape(-1, 3),
+        pts2.reshape(-1, 3),
+        weights=conf.ravel(),
+        compute_scaling=True,
+    )
+    return s, R, T  # return un-scaled (R, T)
+
+
+def sRT_to_4x4(scale, R, T, device):
+    trf = torch.eye(4, device=device)
+    trf[:3, :3] = R * scale
+    trf[:3, 3] = T.ravel()  # doesn't need scaling
+    return trf
+
+
+def estimate_focal(pts3d_i, pp=None):
+    if pp is None:
+        H, W, THREE = pts3d_i.shape
+        assert THREE == 3
+        pp = torch.tensor((W / 2, H / 2), device=pts3d_i.device)
+    focal = estimate_focal_knowing_depth(
+        pts3d_i.unsqueeze(0), pp.unsqueeze(0), focal_mode="weiszfeld"
+    ).ravel()
+    return float(focal)
+
+
+@cache
+def pixel_grid(H, W):
+    return np.mgrid[:W, :H].T.astype(np.float32)
+
+
+def fast_pnp(pts3d, focal, msk, device, pp=None, niter_PnP=10):
+    # extract camera poses and focals with RANSAC-PnP
+    if msk.sum() < 4:
+        return None  # we need at least 4 points for PnP
+    pts3d, msk = map(to_numpy, (pts3d, msk))
+
+    H, W, THREE = pts3d.shape
+    assert THREE == 3
+    pixels = pixel_grid(H, W)
+
+    if focal is None:
+        S = max(W, H)
+        tentative_focals = np.geomspace(S / 2, S * 3, 21)
+    else:
+        tentative_focals = [focal]
+
+    if pp is None:
+        pp = (W / 2, H / 2)
+    else:
+        pp = to_numpy(pp)
+
+    best = (0,)
+    for focal in tentative_focals:
+        K = np.float32([(focal, 0, pp[0]), (0, focal, pp[1]), (0, 0, 1)])
+        try:
+            success, R, T, inliers = cv2.solvePnPRansac(
+                pts3d[msk],
+                pixels[msk],
+                K,
+                None,
+                iterationsCount=niter_PnP,
+                reprojectionError=5,
+                flags=cv2.SOLVEPNP_SQPNP,
+            )
+            if not success:
+                continue
+        except:
+            continue
+
+        score = len(inliers)
+        if success and score > best[0]:
+            best = score, R, T, focal
+
+    if not best[0]:
+        return None
+
+    _, R, T, best_focal = best
+    R = cv2.Rodrigues(R)[0]  # world to cam
+    R, T = map(torch.from_numpy, (R, T))
+    return best_focal, inv(sRT_to_4x4(1, R, T, device))  # cam to world
+
+
+def get_known_poses(self):
+    if self.has_im_poses:
+        known_poses_msk = torch.tensor([not (p.requires_grad) for p in self.im_poses])
+        known_poses = self.get_im_poses()
+        return known_poses_msk.sum(), known_poses_msk, known_poses
+    else:
+        return 0, None, None
+
+
+def get_known_focals(self):
+    if self.has_im_poses:
+        known_focal_msk = self.get_known_focal_mask()
+        known_focals = self.get_focals()
+        return known_focal_msk.sum(), known_focal_msk, known_focals
+    else:
+        return 0, None, None
+
+
+def align_multiple_poses(src_poses, target_poses):
+    N = len(src_poses)
+    assert src_poses.shape == target_poses.shape == (N, 4, 4)
+
+    def center_and_z(poses):
+        # Add small epsilon to prevent division by zero when all poses are at origin
+        eps = max(get_med_dist_between_poses(poses) / 100, 1e-6)
+        return torch.cat((poses[:, :3, 3], poses[:, :3, 3] + eps * poses[:, :3, 2]))
+
+    R, T, s = roma.rigid_points_registration(
+        center_and_z(src_poses), center_and_z(target_poses), compute_scaling=True
+    )
+    # If scale is too small (near zero), set it to 1 to prevent numerical issues
+    if abs(s) < 1e-6:
+        s = 1.0
+    return s, R, T

extern/CUT3R/cloud_opt/dust3r_opt/optimizer.py

@@ -0,0 +1,341 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Main class for the implementation of the global alignment
+# --------------------------------------------------------
+import numpy as np
+import torch
+import torch.nn as nn
+
+from cloud_opt.dust3r_opt.base_opt import BasePCOptimizer
+from dust3r.utils.geometry import xy_grid, geotrf
+from dust3r.utils.device import to_cpu, to_numpy
+
+
+class PointCloudOptimizer(BasePCOptimizer):
+    """Optimize a global scene, given a list of pairwise observations.
+    Graph node: images
+    Graph edges: observations = (pred1, pred2)
+    """
+
+    def __init__(self, *args, optimize_pp=False, focal_break=20, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.has_im_poses = True  # by definition of this class
+        self.focal_break = focal_break
+
+        # adding thing to optimize
+        self.im_depthmaps = nn.ParameterList(
+            torch.randn(H, W) / 10 - 3 for H, W in self.imshapes
+        )  # log(depth)
+        self.im_poses = nn.ParameterList(
+            self.rand_pose(self.POSE_DIM) for _ in range(self.n_imgs)
+        )  # camera poses
+        self.im_focals = nn.ParameterList(
+            torch.FloatTensor([self.focal_break * np.log(max(H, W))])
+            for H, W in self.imshapes
+        )  # camera intrinsics
+        self.im_pp = nn.ParameterList(
+            torch.zeros((2,)) for _ in range(self.n_imgs)
+        )  # camera intrinsics
+        self.im_pp.requires_grad_(optimize_pp)
+
+        self.imshape = self.imshapes[0]
+        im_areas = [h * w for h, w in self.imshapes]
+        self.max_area = max(im_areas)
+
+        # adding thing to optimize
+        # self.im_depthmaps = ParameterStack(
+        #     self.im_depthmaps, is_param=True, fill=self.max_area
+        # )
+
+        self.im_poses = ParameterStack(self.im_poses, is_param=True)
+        self.im_focals = ParameterStack(self.im_focals, is_param=True)
+        self.im_pp = ParameterStack(self.im_pp, is_param=True)
+        self.register_buffer(
+            "_pp", torch.tensor([(w / 2, h / 2) for h, w in self.imshapes])
+        )
+        self.register_buffer(
+            "_grid",
+            ParameterStack(
+                [xy_grid(W, H, device=self.device) for H, W in self.imshapes],
+                fill=self.max_area,
+            ),
+        )
+
+        # pre-compute pixel weights
+        self.register_buffer(
+            "_weight_i",
+            ParameterStack(
+                [self.conf_trf(self.conf_i[i_j]) for i_j in self.str_edges],
+                fill=self.max_area,
+            ),
+        )
+        self.register_buffer(
+            "_weight_j",
+            ParameterStack(
+                [self.conf_trf(self.conf_j[i_j]) for i_j in self.str_edges],
+                fill=self.max_area,
+            ),
+        )
+
+        # precompute aa
+        self.register_buffer(
+            "_stacked_pred_i",
+            ParameterStack(self.pred_i, self.str_edges, fill=self.max_area),
+        )
+        self.register_buffer(
+            "_stacked_pred_j",
+            ParameterStack(self.pred_j, self.str_edges, fill=self.max_area),
+        )
+        self.register_buffer("_ei", torch.tensor([i for i, j in self.edges]))
+        self.register_buffer("_ej", torch.tensor([j for i, j in self.edges]))
+        self.total_area_i = sum([im_areas[i] for i, j in self.edges])
+        self.total_area_j = sum([im_areas[j] for i, j in self.edges])
+
+    def _check_all_imgs_are_selected(self, msk):
+        assert np.all(
+            self._get_msk_indices(msk) == np.arange(self.n_imgs)
+        ), "incomplete mask!"
+
+    def preset_pose(self, known_poses, pose_msk=None):  # cam-to-world
+        self._check_all_imgs_are_selected(pose_msk)
+
+        if isinstance(known_poses, torch.Tensor) and known_poses.ndim == 2:
+            known_poses = [known_poses]
+        for idx, pose in zip(self._get_msk_indices(pose_msk), known_poses):
+            if self.verbose:
+                print(f" (setting pose #{idx} = {pose[:3,3]})")
+            self._no_grad(self._set_pose(self.im_poses, idx, torch.tensor(pose)))
+
+        # normalize scale if there's less than 1 known pose
+        self.im_poses.requires_grad_(False)
+        for p in self.im_poses:
+            print(p.requires_grad)
+            print(p.data)
+        n_known_poses = sum((p.requires_grad is False) for p in self.im_poses)
+        self.norm_pw_scale = n_known_poses <= 1
+
+        
+        self.norm_pw_scale = False
+
+    def preset_focal(self, known_focals, msk=None):
+        self._check_all_imgs_are_selected(msk)
+
+        for idx, focal in zip(self._get_msk_indices(msk), known_focals):
+            if self.verbose:
+                print(f" (setting focal #{idx} = {focal})")
+            self._no_grad(self._set_focal(idx, focal))
+
+        self.im_focals.requires_grad_(False)
+
+    def preset_principal_point(self, known_pp, msk=None):
+        self._check_all_imgs_are_selected(msk)
+
+        for idx, pp in zip(self._get_msk_indices(msk), known_pp):
+            if self.verbose:
+                print(f" (setting principal point #{idx} = {pp})")
+            self._no_grad(self._set_principal_point(idx, pp))
+
+        self.im_pp.requires_grad_(False)
+
+
+        
+
+    def _get_msk_indices(self, msk):
+        if msk is None:
+            return range(self.n_imgs)
+        elif isinstance(msk, int):
+            return [msk]
+        elif isinstance(msk, (tuple, list)):
+            return self._get_msk_indices(np.array(msk))
+        elif msk.dtype in (bool, torch.bool, np.bool_):
+            assert len(msk) == self.n_imgs
+            return np.where(msk)[0]
+        elif np.issubdtype(msk.dtype, np.integer):
+            return msk
+        else:
+            raise ValueError(f"bad {msk=}")
+
+    def _no_grad(self, tensor):
+        assert (
+            tensor.requires_grad
+        ), "it must be True at this point, otherwise no modification occurs"
+
+    def _set_focal(self, idx, focal, force=False):
+        param = self.im_focals[idx]
+        if (
+            param.requires_grad or force
+        ):  # can only init a parameter not already initialized
+            param.data[:] = self.focal_break * np.log(focal)
+        return param
+
+    def get_focals(self):
+        log_focals = torch.stack(list(self.im_focals), dim=0)
+        return (log_focals / self.focal_break).exp()
+
+    def get_known_focal_mask(self):
+        return torch.tensor([not (p.requires_grad) for p in self.im_focals])
+
+    def _set_principal_point(self, idx, pp, force=False):
+        param = self.im_pp[idx]
+        H, W = self.imshapes[idx]
+        if (
+            param.requires_grad or force
+        ):  # can only init a parameter not already initialized
+            param.data[:] = to_cpu(to_numpy(pp) - (W / 2, H / 2)) / 10
+        return param
+
+    def get_principal_points(self):
+        return self._pp + 10 * self.im_pp
+
+    def get_intrinsics(self):
+        K = torch.zeros((self.n_imgs, 3, 3), device=self.device)
+        focals = self.get_focals().flatten()
+        K[:, 0, 0] = K[:, 1, 1] = focals
+        K[:, :2, 2] = self.get_principal_points()
+        K[:, 2, 2] = 1
+        return K
+
+    def get_im_poses(self):  # cam to world
+        cam2world = self._get_poses(self.im_poses)
+        return cam2world
+
+
+    def preset_depth(self, known_depths, msk=None):
+        """Preset known depth maps for specified images.
+        
+        Args:
+            known_depths: List of depth maps or single depth map (should be in normal depth space, not log space)
+            msk: Mask or indices indicating which images to preset. If None, applies to all images.
+        """
+        self._check_all_imgs_are_selected(msk)
+
+        if isinstance(known_depths, (torch.Tensor, np.ndarray)) and known_depths.ndim == 2:
+            known_depths = [known_depths]
+            
+        for idx, depth in zip(self._get_msk_indices(msk), known_depths):
+            if self.verbose:
+                print(f" (setting depth #{idx})")
+            # No need to take log here since _set_depthmap already expects depths in normal space
+            depth = _ravel_hw(depth, self.max_area).view(self.imshapes[idx])
+            self._no_grad(self._set_depthmap(idx, torch.tensor(depth)))
+            self.im_depthmaps[idx].requires_grad_(False)
+
+
+    def _set_depthmap(self, idx, depth, force=False):
+        """Set a depth map for an image.
+        
+        Args:
+            idx: Image index
+            depth: Depth map in normal space (not log space)
+            force: Whether to force setting even if already initialized
+        """
+        depth = _ravel_hw(depth, self.max_area)
+        depth = depth.view(self.imshapes[idx])
+        depth = depth.nan_to_num(neginf=0)
+        param = self.im_depthmaps[idx]
+        if (
+            param.requires_grad or force
+        ):  # can only init a parameter not already initialized
+            param.data[:] = depth.log().nan_to_num(neginf=0)  # Store in log space
+        return param
+
+    def get_depthmaps(self, raw=False):
+        res = ParameterStack(self.im_depthmaps, is_param=False).exp()
+        if not raw:
+            res = [dm[: h * w].view(h, w) for dm, (h, w) in zip(res, self.imshapes)]
+        return res
+
+    def depth_to_pts3d(self):
+        # Get depths and  projection params if not provided
+        focals = self.get_focals()
+        pp = self.get_principal_points()
+        im_poses = self.get_im_poses()
+        depth = self.get_depthmaps(raw=True)
+
+        # get pointmaps in camera frame
+        rel_ptmaps = _fast_depthmap_to_pts3d(depth, self._grid, focals, pp=pp)
+        # project to world frame
+        return geotrf(im_poses, rel_ptmaps)
+
+    def get_pts3d(self, raw=False):
+        res = self.depth_to_pts3d()
+        if not raw:
+            res = [dm[: h * w].view(h, w, 3) for dm, (h, w) in zip(res, self.imshapes)]
+        return res
+
+    def forward(self):
+        pw_poses = self.get_pw_poses()  # cam-to-world
+        pw_adapt = self.get_adaptors().unsqueeze(1)
+        proj_pts3d = self.get_pts3d(raw=True)
+
+        # rotate pairwise prediction according to pw_poses
+        aligned_pred_i = geotrf(pw_poses, pw_adapt * self._stacked_pred_i)
+        aligned_pred_j = geotrf(pw_poses, pw_adapt * self._stacked_pred_j)
+
+        # compute the less
+        li = (
+            self.dist(proj_pts3d[self._ei], aligned_pred_i, weight=self._weight_i).sum()
+            / self.total_area_i
+        )
+        lj = (
+            self.dist(proj_pts3d[self._ej], aligned_pred_j, weight=self._weight_j).sum()
+            / self.total_area_j
+        )
+
+        return li + lj
+
+
+def _fast_depthmap_to_pts3d(depth, pixel_grid, focal, pp):
+    pp = pp.unsqueeze(1)
+    focal = focal.unsqueeze(1)
+    if depth.ndim == 3:
+        depth = depth.view(depth.shape[0], -1)
+    assert focal.shape == (len(depth), 1, 1)
+    assert pp.shape == (len(depth), 1, 2)
+    assert pixel_grid.shape == depth.shape + (2,)
+    depth = depth.unsqueeze(-1)
+    return torch.cat((depth * (pixel_grid - pp) / focal, depth), dim=-1)
+
+
+def ParameterStack(params, keys=None, is_param=None, fill=0):
+    if keys is not None:
+        params = [params[k] for k in keys]
+
+    if fill > 0:
+        params = [_ravel_hw(p, fill) for p in params]
+
+    requires_grad = params[0].requires_grad
+    assert all(p.requires_grad == requires_grad for p in params) if is_param else True
+
+    params = torch.stack(list(params)).float().detach()
+    if is_param or requires_grad:
+        params = nn.Parameter(params)
+        params.requires_grad_(requires_grad)
+    return params
+
+
+def _ravel_hw(tensor, fill=0):
+    # ravel H,W
+    tensor = tensor.view((tensor.shape[0] * tensor.shape[1],) + tensor.shape[2:])
+
+    if len(tensor) < fill:
+        tensor = torch.cat(
+            (tensor, tensor.new_zeros((fill - len(tensor),) + tensor.shape[1:]))
+        )
+    return tensor
+
+
+def acceptable_focal_range(H, W, minf=0.5, maxf=3.5):
+    focal_base = max(H, W) / (
+        2 * np.tan(np.deg2rad(60) / 2)
+    )  # size / 1.1547005383792515
+    return minf * focal_base, maxf * focal_base
+
+
+def apply_mask(img, msk):
+    img = img.copy()
+    img[msk] = 0
+    return img

extern/CUT3R/cloud_opt/init_all.py

@@ -0,0 +1,222 @@
+from functools import cache
+import numpy as np
+import scipy.sparse as sp
+import torch
+import cv2
+import roma
+from tqdm import tqdm
+
+from cloud_opt.utils import *
+
+
+def compute_edge_scores(edges, edge2conf_i, edge2conf_j):
+    """
+    edges: 'i_j', (i,j)
+    """
+    score_dict = {
+        (i, j): edge_conf(edge2conf_i[e], edge2conf_j[e]) for e, (i, j) in edges
+    }
+    return score_dict
+
+
+def dict_to_sparse_graph(dic):
+    n_imgs = max(max(e) for e in dic) + 1
+    res = sp.dok_array((n_imgs, n_imgs))
+    for edge, value in dic.items():
+        res[edge] = value
+    return res
+
+
+@torch.no_grad()
+def init_minimum_spanning_tree(self, **kw):
+    """Init all camera poses (image-wise and pairwise poses) given
+    an initial set of pairwise estimations.
+    """
+    device = self.device
+    pts3d, _, im_focals, im_poses = minimum_spanning_tree(
+        self.imshapes,
+        self.edges,
+        self.edge2pts_i,
+        self.edge2pts_j,
+        self.edge2conf_i,
+        self.edge2conf_j,
+        self.im_conf,
+        self.min_conf_thr,
+        device,
+        has_im_poses=self.has_im_poses,
+        verbose=self.verbose,
+        **kw,
+    )
+
+    return init_from_pts3d(self, pts3d, im_focals, im_poses)
+
+
+def minimum_spanning_tree(
+    imshapes,
+    edges,
+    edge2pred_i,
+    edge2pred_j,
+    edge2conf_i,
+    edge2conf_j,
+    im_conf,
+    min_conf_thr,
+    device,
+    has_im_poses=True,
+    niter_PnP=10,
+    verbose=True,
+    save_score_path=None,
+):
+    n_imgs = len(imshapes)
+    eadge_and_scores = compute_edge_scores(map(i_j_ij, edges), edge2conf_i, edge2conf_j)
+    sparse_graph = -dict_to_sparse_graph(eadge_and_scores)
+    msp = sp.csgraph.minimum_spanning_tree(sparse_graph).tocoo()
+
+    # temp variable to store 3d points
+    pts3d = [None] * len(imshapes)
+
+    todo = sorted(zip(-msp.data, msp.row, msp.col))  # sorted edges
+    im_poses = [None] * n_imgs
+    im_focals = [None] * n_imgs
+
+    # init with strongest edge
+    score, i, j = todo.pop()
+    if verbose:
+        print(f" init edge ({i}*,{j}*) {score=}")
+    i_j = edge_str(i, j)
+
+    pts3d[i] = edge2pred_i[i_j].clone()
+    pts3d[j] = edge2pred_j[i_j].clone()
+    done = {i, j}
+    if has_im_poses:
+        im_poses[i] = torch.eye(4, device=device)
+        im_focals[i] = estimate_focal(edge2pred_i[i_j])
+
+    # set initial pointcloud based on pairwise graph
+    msp_edges = [(i, j)]
+    while todo:
+        # each time, predict the next one
+        score, i, j = todo.pop()
+
+        if im_focals[i] is None:
+            im_focals[i] = estimate_focal(edge2pred_i[i_j])
+
+        if i in done:
+            if verbose:
+                print(f" init edge ({i},{j}*) {score=}")
+            assert j not in done
+            # align pred[i] with pts3d[i], and then set j accordingly
+            i_j = edge_str(i, j)
+            s, R, T = rigid_points_registration(
+                edge2pred_i[i_j], pts3d[i], conf=edge2conf_i[i_j]
+            )
+            trf = sRT_to_4x4(s, R, T, device)
+            pts3d[j] = geotrf(trf, edge2pred_j[i_j])
+            done.add(j)
+            msp_edges.append((i, j))
+
+            if has_im_poses and im_poses[i] is None:
+                im_poses[i] = sRT_to_4x4(1, R, T, device)
+
+        elif j in done:
+            if verbose:
+                print(f" init edge ({i}*,{j}) {score=}")
+            assert i not in done
+            i_j = edge_str(i, j)
+            s, R, T = rigid_points_registration(
+                edge2pred_j[i_j], pts3d[j], conf=edge2conf_j[i_j]
+            )
+            trf = sRT_to_4x4(s, R, T, device)
+            pts3d[i] = geotrf(trf, edge2pred_i[i_j])
+            done.add(i)
+            msp_edges.append((i, j))
+
+            if has_im_poses and im_poses[i] is None:
+                im_poses[i] = sRT_to_4x4(1, R, T, device)
+        else:
+            # let's try again later
+            todo.insert(0, (score, i, j))
+
+    if has_im_poses:
+        # complete all missing informations
+        pair_scores = list(
+            sparse_graph.values()
+        )  # already negative scores: less is best
+        edges_from_best_to_worse = np.array(list(sparse_graph.keys()))[
+            np.argsort(pair_scores)
+        ]
+        for i, j in edges_from_best_to_worse.tolist():
+            if im_focals[i] is None:
+                im_focals[i] = estimate_focal(edge2pred_i[edge_str(i, j)])
+
+        for i in range(n_imgs):
+            if im_poses[i] is None:
+                msk = im_conf[i] > min_conf_thr
+                res = fast_pnp(
+                    pts3d[i], im_focals[i], msk=msk, device=device, niter_PnP=niter_PnP
+                )
+                if res:
+                    im_focals[i], im_poses[i] = res
+            if im_poses[i] is None:
+                im_poses[i] = torch.eye(4, device=device)
+        im_poses = torch.stack(im_poses)
+    else:
+        im_poses = im_focals = None
+
+    return pts3d, msp_edges, im_focals, im_poses
+
+
+def init_from_pts3d(self, pts3d, im_focals, im_poses):
+    # init poses
+    nkp, known_poses_msk, known_poses = self.get_known_poses()
+    if nkp == 1:
+        raise NotImplementedError(
+            "Would be simpler to just align everything afterwards on the single known pose"
+        )
+    elif nkp > 1:
+        # global rigid SE3 alignment
+        s, R, T = align_multiple_poses(
+            im_poses[known_poses_msk], known_poses[known_poses_msk]
+        )
+        trf = sRT_to_4x4(s, R, T, device=known_poses.device)
+
+        # rotate everything
+        im_poses = trf @ im_poses
+        im_poses[:, :3, :3] /= s  # undo scaling on the rotation part
+        for img_pts3d in pts3d:
+            img_pts3d[:] = geotrf(trf, img_pts3d)
+    else:
+        pass  # no known poses
+
+    # set all pairwise poses
+    for e, (i, j) in enumerate(self.edges):
+        i_j = edge_str(i, j)
+        # compute transform that goes from cam to world
+        s, R, T = rigid_points_registration(
+            self.pred_i[i_j], pts3d[i], conf=self.conf_i[i_j]
+        )
+        self._set_pose(self.pw_poses, e, R, T, scale=s)
+
+    # take into account the scale normalization
+    s_factor = self.get_pw_norm_scale_factor()
+    im_poses[:, :3, 3] *= s_factor  # apply downscaling factor
+    for img_pts3d in pts3d:
+        img_pts3d *= s_factor
+
+    # init all image poses
+    if self.has_im_poses:
+        for i in range(self.n_imgs):
+            cam2world = im_poses[i]
+            depth = geotrf(inv(cam2world), pts3d[i])[..., 2]
+            self._set_depthmap(i, depth)
+            self._set_pose(self.im_poses, i, cam2world)
+            if im_focals[i] is not None:
+                if not self.shared_focal:
+                    self._set_focal(i, im_focals[i])
+        if self.shared_focal:
+            self._set_focal(0, sum(im_focals) / self.n_imgs)
+        if self.n_imgs > 2:
+            self._set_init_depthmap()
+
+    if self.verbose:
+        with torch.no_grad():
+            print(" init loss =", float(self()))

extern/CUT3R/cloud_opt/utils.py

@@ -0,0 +1,443 @@
+import torch.nn as nn
+import torch
+import roma
+import numpy as np
+import cv2
+from functools import cache
+
+
+def todevice(batch, device, callback=None, non_blocking=False):
+    """Transfer some variables to another device (i.e. GPU, CPU:torch, CPU:numpy).
+
+    batch: list, tuple, dict of tensors or other things
+    device: pytorch device or 'numpy'
+    callback: function that would be called on every sub-elements.
+    """
+    if callback:
+        batch = callback(batch)
+
+    if isinstance(batch, dict):
+        return {k: todevice(v, device) for k, v in batch.items()}
+
+    if isinstance(batch, (tuple, list)):
+        return type(batch)(todevice(x, device) for x in batch)
+
+    x = batch
+    if device == "numpy":
+        if isinstance(x, torch.Tensor):
+            x = x.detach().cpu().numpy()
+    elif x is not None:
+        if isinstance(x, np.ndarray):
+            x = torch.from_numpy(x)
+        if torch.is_tensor(x):
+            x = x.to(device, non_blocking=non_blocking)
+    return x
+
+
+to_device = todevice  # alias
+
+
+def to_numpy(x):
+    return todevice(x, "numpy")
+
+
+def to_cpu(x):
+    return todevice(x, "cpu")
+
+
+def to_cuda(x):
+    return todevice(x, "cuda")
+
+
+def signed_log1p(x):
+    sign = torch.sign(x)
+    return sign * torch.log1p(torch.abs(x))
+
+
+def l2_dist(a, b, weight):
+    return (a - b).square().sum(dim=-1) * weight
+
+
+def l1_dist(a, b, weight):
+    return (a - b).norm(dim=-1) * weight
+
+
+ALL_DISTS = dict(l1=l1_dist, l2=l2_dist)
+
+
+def _check_edges(edges):
+    indices = sorted({i for edge in edges for i in edge})
+    assert indices == list(range(len(indices))), "bad pair indices: missing values "
+    return len(indices)
+
+
+def NoGradParamDict(x):
+    assert isinstance(x, dict)
+    return nn.ParameterDict(x).requires_grad_(False)
+
+
+def edge_str(i, j):
+    return f"{i}_{j}"
+
+
+def i_j_ij(ij):
+    # inputs are (i, j)
+    return edge_str(*ij), ij
+
+
+def edge_conf(conf_i, conf_j):
+    score = float(conf_i.mean() * conf_j.mean())
+    return score
+
+
+def get_imshapes(edges, pred_i, pred_j):
+    n_imgs = max(max(e) for e in edges) + 1
+    imshapes = [None] * n_imgs
+    for e, (i, j) in enumerate(edges):
+        shape_i = tuple(pred_i[e]["pts3d_is_self_view"].shape[0:2])
+        shape_j = tuple(pred_j[e]["pts3d_in_other_view"].shape[0:2])
+        if imshapes[i]:
+            assert imshapes[i] == shape_i, f"incorrect shape for image {i}"
+        if imshapes[j]:
+            assert imshapes[j] == shape_j, f"incorrect shape for image {j}"
+        imshapes[i] = shape_i
+        imshapes[j] = shape_j
+    return imshapes
+
+
+def get_conf_trf(mode):
+    if mode == "log":
+
+        def conf_trf(x):
+            return x.log()
+
+    elif mode == "sqrt":
+
+        def conf_trf(x):
+            return x.sqrt()
+
+    elif mode == "m1":
+
+        def conf_trf(x):
+            return x - 1
+
+    elif mode in ("id", "none"):
+
+        def conf_trf(x):
+            return x
+
+    else:
+        raise ValueError(f"bad mode for {mode=}")
+    return conf_trf
+
+
+@torch.no_grad()
+def _compute_img_conf(imshapes, device, edges, edge2conf_i, edge2conf_j):
+    im_conf = nn.ParameterList([torch.zeros(hw, device=device) for hw in imshapes])
+    for e, (i, j) in enumerate(edges):
+        im_conf[i] = torch.maximum(im_conf[i], edge2conf_i[edge_str(i, j)])
+        im_conf[j] = torch.maximum(im_conf[j], edge2conf_j[edge_str(i, j)])
+    return im_conf
+
+
+def xy_grid(
+    W,
+    H,
+    device=None,
+    origin=(0, 0),
+    unsqueeze=None,
+    cat_dim=-1,
+    homogeneous=False,
+    **arange_kw,
+):
+    """Output a (H,W,2) array of int32
+    with output[j,i,0] = i + origin[0]
+         output[j,i,1] = j + origin[1]
+    """
+    if device is None:
+        # numpy
+        arange, meshgrid, stack, ones = np.arange, np.meshgrid, np.stack, np.ones
+    else:
+        # torch
+        arange = lambda *a, **kw: torch.arange(*a, device=device, **kw)
+        meshgrid, stack = torch.meshgrid, torch.stack
+        ones = lambda *a: torch.ones(*a, device=device)
+
+    tw, th = [arange(o, o + s, **arange_kw) for s, o in zip((W, H), origin)]
+    grid = meshgrid(tw, th, indexing="xy")
+    if homogeneous:
+        grid = grid + (ones((H, W)),)
+    if unsqueeze is not None:
+        grid = (grid[0].unsqueeze(unsqueeze), grid[1].unsqueeze(unsqueeze))
+    if cat_dim is not None:
+        grid = stack(grid, cat_dim)
+    return grid
+
+
+def estimate_focal_knowing_depth(
+    pts3d, pp, focal_mode="median", min_focal=0.0, max_focal=np.inf
+):
+    """Reprojection method, for when the absolute depth is known:
+    1) estimate the camera focal using a robust estimator
+    2) reproject points onto true rays, minimizing a certain error
+    """
+    B, H, W, THREE = pts3d.shape
+    assert THREE == 3
+
+    # centered pixel grid
+    pixels = xy_grid(W, H, device=pts3d.device).view(1, -1, 2) - pp.view(
+        -1, 1, 2
+    )  # B,HW,2
+    pts3d = pts3d.flatten(1, 2)  # (B, HW, 3)
+
+    if focal_mode == "median":
+        with torch.no_grad():
+            # direct estimation of focal
+            u, v = pixels.unbind(dim=-1)
+            x, y, z = pts3d.unbind(dim=-1)
+            fx_votes = (u * z) / x
+            fy_votes = (v * z) / y
+
+            # assume square pixels, hence same focal for X and Y
+            f_votes = torch.cat((fx_votes.view(B, -1), fy_votes.view(B, -1)), dim=-1)
+            focal = torch.nanmedian(f_votes, dim=-1).values
+
+    elif focal_mode == "weiszfeld":
+        # init focal with l2 closed form
+        # we try to find focal = argmin Sum | pixel - focal * (x,y)/z|
+        xy_over_z = (pts3d[..., :2] / pts3d[..., 2:3]).nan_to_num(
+            posinf=0, neginf=0
+        )  # homogeneous (x,y,1)
+
+        dot_xy_px = (xy_over_z * pixels).sum(dim=-1)
+        dot_xy_xy = xy_over_z.square().sum(dim=-1)
+
+        focal = dot_xy_px.mean(dim=1) / dot_xy_xy.mean(dim=1)
+
+        # iterative re-weighted least-squares
+        for iter in range(10):
+            # re-weighting by inverse of distance
+            dis = (pixels - focal.view(-1, 1, 1) * xy_over_z).norm(dim=-1)
+            # print(dis.nanmean(-1))
+            w = dis.clip(min=1e-8).reciprocal()
+            # update the scaling with the new weights
+            focal = (w * dot_xy_px).mean(dim=1) / (w * dot_xy_xy).mean(dim=1)
+    else:
+        raise ValueError(f"bad {focal_mode=}")
+
+    focal_base = max(H, W) / (
+        2 * np.tan(np.deg2rad(60) / 2)
+    )  # size / 1.1547005383792515
+    focal = focal.clip(min=min_focal * focal_base, max=max_focal * focal_base)
+    # print(focal)
+    return focal
+
+
+def estimate_focal(pts3d_i, pp=None):
+    if pp is None:
+        H, W, THREE = pts3d_i.shape
+        assert THREE == 3
+        pp = torch.tensor((W / 2, H / 2), device=pts3d_i.device)
+    focal = estimate_focal_knowing_depth(
+        pts3d_i.unsqueeze(0), pp.unsqueeze(0), focal_mode="weiszfeld"
+    ).ravel()
+    return float(focal)
+
+
+def rigid_points_registration(pts1, pts2, conf):
+    R, T, s = roma.rigid_points_registration(
+        pts1.reshape(-1, 3),
+        pts2.reshape(-1, 3),
+        weights=conf.ravel(),
+        compute_scaling=True,
+    )
+    return s, R, T  # return un-scaled (R, T)
+
+
+def sRT_to_4x4(scale, R, T, device):
+    trf = torch.eye(4, device=device)
+    trf[:3, :3] = R * scale
+    trf[:3, 3] = T.ravel()  # doesn't need scaling
+    return trf
+
+
+def geotrf(Trf, pts, ncol=None, norm=False):
+    """Apply a geometric transformation to a list of 3-D points.
+
+    H: 3x3 or 4x4 projection matrix (typically a Homography)
+    p: numpy/torch/tuple of coordinates. Shape must be (...,2) or (...,3)
+
+    ncol: int. number of columns of the result (2 or 3)
+    norm: float. if != 0, the resut is projected on the z=norm plane.
+
+    Returns an array of projected 2d points.
+    """
+    assert Trf.ndim >= 2
+    if isinstance(Trf, np.ndarray):
+        pts = np.asarray(pts)
+    elif isinstance(Trf, torch.Tensor):
+        pts = torch.as_tensor(pts, dtype=Trf.dtype)
+
+    # adapt shape if necessary
+    output_reshape = pts.shape[:-1]
+    ncol = ncol or pts.shape[-1]
+
+    # optimized code
+    if (
+        isinstance(Trf, torch.Tensor)
+        and isinstance(pts, torch.Tensor)
+        and Trf.ndim == 3
+        and pts.ndim == 4
+    ):
+        d = pts.shape[3]
+        if Trf.shape[-1] == d:
+            pts = torch.einsum("bij, bhwj -> bhwi", Trf, pts)
+        elif Trf.shape[-1] == d + 1:
+            pts = (
+                torch.einsum("bij, bhwj -> bhwi", Trf[:, :d, :d], pts)
+                + Trf[:, None, None, :d, d]
+            )
+        else:
+            raise ValueError(f"bad shape, not ending with 3 or 4, for {pts.shape=}")
+    else:
+        if Trf.ndim >= 3:
+            n = Trf.ndim - 2
+            assert Trf.shape[:n] == pts.shape[:n], "batch size does not match"
+            Trf = Trf.reshape(-1, Trf.shape[-2], Trf.shape[-1])
+
+            if pts.ndim > Trf.ndim:
+                # Trf == (B,d,d) & pts == (B,H,W,d) --> (B, H*W, d)
+                pts = pts.reshape(Trf.shape[0], -1, pts.shape[-1])
+            elif pts.ndim == 2:
+                # Trf == (B,d,d) & pts == (B,d) --> (B, 1, d)
+                pts = pts[:, None, :]
+
+        if pts.shape[-1] + 1 == Trf.shape[-1]:
+            Trf = Trf.swapaxes(-1, -2)  # transpose Trf
+            pts = pts @ Trf[..., :-1, :] + Trf[..., -1:, :]
+        elif pts.shape[-1] == Trf.shape[-1]:
+            Trf = Trf.swapaxes(-1, -2)  # transpose Trf
+            pts = pts @ Trf
+        else:
+            pts = Trf @ pts.T
+            if pts.ndim >= 2:
+                pts = pts.swapaxes(-1, -2)
+
+    if norm:
+        pts = pts / pts[..., -1:]  # DONT DO /= BECAUSE OF WEIRD PYTORCH BUG
+        if norm != 1:
+            pts *= norm
+
+    res = pts[..., :ncol].reshape(*output_reshape, ncol)
+    return res
+
+
+def inv(mat):
+    """Invert a torch or numpy matrix"""
+    if isinstance(mat, torch.Tensor):
+        return torch.linalg.inv(mat)
+    if isinstance(mat, np.ndarray):
+        return np.linalg.inv(mat)
+    raise ValueError(f"bad matrix type = {type(mat)}")
+
+
+@cache
+def pixel_grid(H, W):
+    return np.mgrid[:W, :H].T.astype(np.float32)
+
+
+def fast_pnp(pts3d, focal, msk, device, pp=None, niter_PnP=10):
+    # extract camera poses and focals with RANSAC-PnP
+    if msk.sum() < 4:
+        return None  # we need at least 4 points for PnP
+    pts3d, msk = map(to_numpy, (pts3d, msk))
+
+    H, W, THREE = pts3d.shape
+    assert THREE == 3
+    pixels = pixel_grid(H, W)
+
+    if focal is None:
+        S = max(W, H)
+        tentative_focals = np.geomspace(S / 2, S * 3, 21)
+    else:
+        tentative_focals = [focal]
+
+    if pp is None:
+        pp = (W / 2, H / 2)
+    else:
+        pp = to_numpy(pp)
+
+    best = (0,)
+    for focal in tentative_focals:
+        K = np.float32([(focal, 0, pp[0]), (0, focal, pp[1]), (0, 0, 1)])
+
+        success, R, T, inliers = cv2.solvePnPRansac(
+            pts3d[msk],
+            pixels[msk],
+            K,
+            None,
+            iterationsCount=niter_PnP,
+            reprojectionError=5,
+            flags=cv2.SOLVEPNP_SQPNP,
+        )
+        if not success:
+            continue
+
+        score = len(inliers)
+        if success and score > best[0]:
+            best = score, R, T, focal
+
+    if not best[0]:
+        return None
+
+    _, R, T, best_focal = best
+    R = cv2.Rodrigues(R)[0]  # world to cam
+    R, T = map(torch.from_numpy, (R, T))
+    return best_focal, inv(sRT_to_4x4(1, R, T, device))  # cam to world
+
+
+def get_med_dist_between_poses(poses):
+    from scipy.spatial.distance import pdist
+
+    return np.median(pdist([to_numpy(p[:3, 3]) for p in poses]))
+
+
+def align_multiple_poses(src_poses, target_poses):
+    N = len(src_poses)
+    assert src_poses.shape == target_poses.shape == (N, 4, 4)
+
+    def center_and_z(poses):
+        eps = get_med_dist_between_poses(poses) / 100
+        return torch.cat((poses[:, :3, 3], poses[:, :3, 3] + eps * poses[:, :3, 2]))
+
+    R, T, s = roma.rigid_points_registration(
+        center_and_z(src_poses), center_and_z(target_poses), compute_scaling=True
+    )
+    return s, R, T
+
+
+def cosine_schedule(t, lr_start, lr_end):
+    assert 0 <= t <= 1
+    return lr_end + (lr_start - lr_end) * (1 + np.cos(t * np.pi)) / 2
+
+
+def linear_schedule(t, lr_start, lr_end):
+    assert 0 <= t <= 1
+    return lr_start + (lr_end - lr_start) * t
+
+
+def cycled_linear_schedule(t, lr_start, lr_end, num_cycles=2):
+    assert 0 <= t <= 1
+    cycle_t = t * num_cycles
+    cycle_t = cycle_t - int(cycle_t)
+    if t == 1:
+        cycle_t = 1
+    return linear_schedule(cycle_t, lr_start, lr_end)
+
+
+def adjust_learning_rate_by_lr(optimizer, lr):
+    for param_group in optimizer.param_groups:
+        if "lr_scale" in param_group:
+            param_group["lr"] = lr * param_group["lr_scale"]
+        else:
+            param_group["lr"] = lr

extern/CUT3R/config/dpt_512_vary_4_64.yaml

@@ -0,0 +1,103 @@
+model: "ARCroco3DStereo(ARCroco3DStereoConfig(freeze='encoder', state_size=768, state_pe='2d', pos_embed='RoPE100', rgb_head=True, pose_head=True, patch_embed_cls='ManyAR_PatchEmbed', img_size=(512, 512), head_type='dpt', output_mode='pts3d+pose', depth_mode=('exp', -inf, inf), conf_mode=('exp', 1, inf), pose_mode=('exp', -inf, inf), enc_embed_dim=1024, enc_depth=24, enc_num_heads=16, dec_embed_dim=768, dec_depth=12, dec_num_heads=12, landscape_only=False))"
+pretrained: cut3r_512_dpt_4_64.pth
+load_only_encoder: False
+long_context: True
+fixed_length: False
+resume: null
+benchmark: False
+num_views : 64
+num_test_views : 4
+n_corres_train: 0
+n_corres_test: 0
+
+train_criterion: ConfLoss(Regr3DPoseBatchList(L21, norm_mode='?avg_dis'), alpha=0.2) + RGBLoss(MSE)
+test_criterion: Regr3DPose(L21, norm_mode='?avg_dis', gt_scale=True, sky_loss_value=0) + Regr3DPose_ScaleInv(L21, norm_mode='?avg_dis', gt_scale=True, sky_loss_value=0) + RGBLoss(L21)
+
+resolution: [(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)]
+
+allow_repeat: True
+dataset1: Co3d_Multi(allow_repeat=${allow_repeat}, split='train', ROOT='../../data/dust3r_data/processed_co3d/', aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset2: WildRGBD_Multi(allow_repeat=${allow_repeat}, split='train', ROOT="../../data/dust3r_data/processed_wildrgbd", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset3: ARKitScenes_Multi(allow_repeat=${allow_repeat}, split='train', ROOT='../../data/dust3r_data/processed_arkitscenes/', aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset4: ARKitScenesHighRes_Multi(allow_repeat=${allow_repeat}, split='train', ROOT="../../data/dust3r_data/processed_arkitscenes_highres", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset5: ScanNetpp_Multi(allow_repeat=${allow_repeat}, split='train', ROOT="../../data/dust3r_data/processed_scannetpp/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset6: ScanNet_Multi(allow_repeat=${allow_repeat}, split='train', ROOT="../../data/dust3r_data/processed_scannet/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset7: HyperSim_Multi(allow_repeat=${allow_repeat}, split='train', ROOT="../../data/custom_data/processed_hypersim", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset8: BlendedMVS_Multi(allow_repeat=${allow_repeat}, split='train', ROOT="../../data/dust3r_data/processed_blendedmvs/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset9: MegaDepth_Multi(allow_repeat=${allow_repeat}, split="train", ROOT="../../data/dust3r_data/processed_megadepth", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset10: MapFree_Multi(allow_repeat=${allow_repeat}, split=None, ROOT="../../data/mast3r_data/processed_mapfree/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset11: Waymo_Multi(allow_repeat=${allow_repeat}, split=None, ROOT="../../data/dust3r_data/processed_waymo/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset12: VirtualKITTI2_Multi(allow_repeat=${allow_repeat}, split=None, ROOT="../../data/mast3r_data/processed_vkitti", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset13: UnReal4K_Multi(allow_repeat=${allow_repeat}, split=None, ROOT="../../data/mast3r_data/processed_unreal4k/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset14: TartanAir_Multi(allow_repeat=${allow_repeat}, split=None, ROOT="../../data/mast3r_data/processed_tartanair/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset15: DL3DV_Multi(allow_repeat=${allow_repeat}, split='train', ROOT="../../data/custom_data/processed_dl3dv", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset16: Cop3D_Multi(allow_repeat=${allow_repeat}, split='train', ROOT="../../data/custom_data/processed_cop3d/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset17: MVImgNet_Multi(allow_repeat=${allow_repeat}, split='train', ROOT="../../data/custom_data/processed_mvimgnet/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset18: RE10K_Multi(allow_repeat=${allow_repeat}, split=None, ROOT="../../data/custom_data/processed_re10k/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset19: OmniObject3D_Multi(allow_repeat=${allow_repeat}, split=None, ROOT="../../data/custom_data/processed_omniobject3d/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset20: ThreeDKenBurns(allow_repeat=${allow_repeat}, split=None, ROOT="../../data/custom_data/processed_3dkb/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset21: IRS(allow_repeat=${allow_repeat}, split=None, ROOT="../../data/custom_data/processed_irs/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset22: SynScapes(allow_repeat=${allow_repeat}, split=None, ROOT="../../data/custom_data/processed_synscapes/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset23: UrbanSyn(allow_repeat=${allow_repeat}, split=None, ROOT="../../data/custom_data/processed_urbansyn/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset24: EDEN_Multi(allow_repeat=${allow_repeat}, split='train', ROOT="../../data/custom_data/processed_eden", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset25: SmartPortraits_Multi(allow_repeat=${allow_repeat}, split='train', ROOT="../../data/custom_data/processed_smartportraits", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset26: DynamicReplica(allow_repeat=${allow_repeat}, split='train', ROOT="../../data/custom_data/processed_dynamic_replica/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset27: Spring(allow_repeat=${allow_repeat}, split=None, ROOT="../../data/custom_data/processed_spring/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset28: BEDLAM_Multi(allow_repeat=${allow_repeat}, split='train', ROOT="../../data/custom_data/processed_bedlam", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset29: MVS_Synth_Multi(allow_repeat=${allow_repeat}, split='train', ROOT="../../data/custom_data/processed_mvs_synth", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset30: PointOdyssey_Multi(allow_repeat=${allow_repeat}, split='train', ROOT="../../data/custom_data/processed_point_odyssey", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset31: UASOL_Multi(allow_repeat=${allow_repeat}, split='train', ROOT="../../data/custom_data/processed_uasol", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset32: MP3D_Multi(allow_repeat=${allow_repeat}, split=None, ROOT="../../data/custom_data/processed_mp3d/", aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208), (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+train_dataset: 44800 @ ${dataset1} + 56000 @ ${dataset2} + 56000 @ ${dataset3} + 22400 @ ${dataset4} 
+  + 16800 @ ${dataset5} + 22400 @ ${dataset6} + 11200 @ ${dataset7} 
+  + 22400 @ ${dataset8} + 22400 @ ${dataset9} + 84000 @ ${dataset10} + 56000 @ ${dataset11}
+  + 5600 @ ${dataset12} + 168 @ ${dataset13} + 56000 @ ${dataset14} + 84000 @ ${dataset15}
+  + 480 @ ${dataset16} + 19200 @ ${dataset17} + 4800 @ ${dataset18} + 38400 @ ${dataset19}
+  + 26400 @ ${dataset26} + 1200 @ ${dataset27} + 36000 @ ${dataset28} + 2400 @ ${dataset29}
+  + 24000 @ ${dataset30} + 14400 @ ${dataset31} + 28800 @ ${dataset32}
+test_dataset: 1000 @ ARKitScenes_Multi(split='test', ROOT='../../data/dust3r_data/processed_arkitscenes/', resolution=(512, 384), num_views=${num_test_views}, seed=42, n_corres=${n_corres_test})
+
+seed: 0
+batch_size: 4
+accum_iter: 4
+gradient_checkpointing: True
+epochs: 10
+start_epoch: 0
+weight_decay: 0.05
+lr: 1e-6
+min_lr: 1e-7
+warmup_epochs: 0.5
+amp: 1
+
+num_workers: 4
+world_size: 1
+local-rank: -1
+dist_url: 'env://'
+rank: 0
+gpu: 0
+distributed: False
+dist_backend: 'nccl'
+
+eval_freq: 1
+save_freq: 0.1
+keep_freq: 1
+print_freq: 10
+print_img_freq: 50000000
+num_imgs_vis: 4
+save_dir: 'checkpoints'
+exp_name: 'dpt_512_vary_4_64'
+task: 'cut3r'
+logdir: ./${save_dir}/${exp_name}/logs
+output_dir: ./${save_dir}/${exp_name}/
+hydra:
+  verbose: True
+  run:
+    dir: ./${save_dir}/${exp_name}

extern/CUT3R/config/linear_224_fixed_16.yaml

@@ -0,0 +1,99 @@
+model: "ARCroco3DStereo(ARCroco3DStereoConfig(freeze='encoder', state_size=768, state_pe='2d', pos_embed='RoPE100', rgb_head=True, pose_head=True,  img_size=(224, 224), head_type='linear', output_mode='pts3d+pose', depth_mode=('exp', -inf, inf), conf_mode=('exp', 1, inf), pose_mode=('exp', -inf, inf), enc_embed_dim=1024, enc_depth=24, enc_num_heads=16, dec_embed_dim=768, dec_depth=12, dec_num_heads=12))"
+pretrained: cut3r_224_linear_4.pth
+load_only_encoder: False
+long_context: False
+fixed_length: True
+resume: null
+benchmark: True
+num_views : 16
+num_test_views : 4
+n_corres_train: 0
+n_corres_test: 0
+
+train_criterion: ConfLoss(Regr3DPoseBatchList(L21, norm_mode='?avg_dis'), alpha=0.2) + RGBLoss(MSE)
+test_criterion: Regr3DPose(L21, norm_mode='?avg_dis', gt_scale=True, sky_loss_value=0) + Regr3DPose_ScaleInv(L21, norm_mode='?avg_dis', gt_scale=True, sky_loss_value=0) + RGBLoss(L21)
+
+
+dataset1: Co3d_Multi(allow_repeat=False, split='train', ROOT='../../data/dust3r_data/processed_co3d/', aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset2: WildRGBD_Multi(allow_repeat=False, split='train', ROOT="../../data/dust3r_data/processed_wildrgbd", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset3: ARKitScenes_Multi(allow_repeat=False, split='train', ROOT='../../data/dust3r_data/processed_arkitscenes/', aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset4: ARKitScenesHighRes_Multi(allow_repeat=False, split='train', ROOT="../../data/dust3r_data/processed_arkitscenes_highres", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset5: ScanNetpp_Multi(allow_repeat=False, split='train', ROOT="../../data/dust3r_data/processed_scannetpp/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset6: ScanNet_Multi(allow_repeat=False, split='train', ROOT="../../data/dust3r_data/processed_scannet/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset7: HyperSim_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_hypersim", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset8: BlendedMVS_Multi(allow_repeat=False, split='train', ROOT="../../data/dust3r_data/processed_blendedmvs/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset9: MegaDepth_Multi(allow_repeat=False, split="train", ROOT="../../data/dust3r_data/processed_megadepth", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset10: MapFree_Multi(allow_repeat=False, split=None, ROOT="../../data/mast3r_data/processed_mapfree/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset11: Waymo_Multi(allow_repeat=False, split=None, ROOT="../../data/dust3r_data/processed_waymo/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset12: VirtualKITTI2_Multi(allow_repeat=False, split=None, ROOT="../../data/mast3r_data/processed_vkitti", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset13: UnReal4K_Multi(allow_repeat=False, split=None, ROOT="../../data/mast3r_data/processed_unreal4k/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset14: TartanAir_Multi(allow_repeat=False, split=None, ROOT="../../data/mast3r_data/processed_tartanair/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset15: DL3DV_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_dl3dv", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset16: Cop3D_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_cop3d/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset17: MVImgNet_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_mvimgnet/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset18: RE10K_Multi(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_re10k/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset19: OmniObject3D_Multi(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_omniobject3d/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset20: ThreeDKenBurns(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_3dkb/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset21: IRS(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_irs/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset22: SynScapes(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_synscapes/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset23: UrbanSyn(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_urbansyn/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset24: EDEN_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_eden", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset25: SmartPortraits_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_smartportraits", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset26: DynamicReplica(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_dynamic_replica/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset27: Spring(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_spring/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset28: BEDLAM_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_bedlam", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset29: MVS_Synth_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_mvs_synth", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset30: PointOdyssey_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_point_odyssey", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset31: UASOL_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_uasol", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset32: MP3D_Multi(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_mp3d/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset33: HOI4D_Multi(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_hoi4d/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+train_dataset: 44800 @ ${dataset1} + 56000 @ ${dataset2} + 56000 @ ${dataset3} + 5600 @ ${dataset4} + 5600 @ ${dataset5} + 140000 @ ${dataset6} + 5600 @ ${dataset7} + 22400 @ ${dataset8} + 16800 @ ${dataset9} + 56000 @ ${dataset10}  + 42000 @ ${dataset11} + 5600 @ ${dataset12} + 168 @ ${dataset13} + 84000 @ ${dataset14} + 84000 @ ${dataset15} + 7200 @ ${dataset16} + 19200 @ ${dataset17} + 9600 @ ${dataset18} + 24000 @ ${dataset19} + 33600 @ ${dataset26} + 2400 @ ${dataset27} + 9600 @ ${dataset28} + 4800 @ ${dataset29} + 28800 @ ${dataset30} + 14400 @ ${dataset31} + 19200 @ ${dataset32}
+
+
+test_dataset: 1000 @ ARKitScenes_Multi(split='test', ROOT='../../data/dust3r_data/processed_arkitscenes/', resolution=224, num_views=${num_test_views}, seed=42, n_corres=${n_corres_test})
+
+seed: 0
+batch_size: 6
+accum_iter: 2
+gradient_checkpointing: False
+epochs: 10
+start_epoch: 0
+weight_decay: 0.05
+lr: 1e-6
+min_lr: 1e-7
+warmup_epochs: 0.5
+amp: 1
+
+num_workers: 16
+world_size: 1
+local-rank: -1
+dist_url: 'env://'
+rank: 0
+gpu: 0
+distributed: False
+dist_backend: 'nccl'
+
+eval_freq: 1
+save_freq: 0.1
+keep_freq: 1
+print_freq: 10
+print_img_freq: 50000000
+num_imgs_vis: 4
+save_dir: 'checkpoints'
+exp_name: 'linear_224_fixed_16'
+task: 'cut3r'
+logdir: ./${save_dir}/${exp_name}/logs
+output_dir: ./${save_dir}/${exp_name}/
+hydra:
+  verbose: True
+  run:
+    dir: ./${save_dir}/${exp_name}

extern/CUT3R/config/stage1.yaml

@@ -0,0 +1,74 @@
+model: "ARCroco3DStereo(ARCroco3DStereoConfig(state_size=768, state_pe='2d', pos_embed='RoPE100', rgb_head=True, pose_head=True,  img_size=(224, 224), head_type='linear', output_mode='pts3d+pose', depth_mode=('exp', -inf, inf), conf_mode=('exp', 1, inf), pose_mode=('exp', -inf, inf), enc_embed_dim=1024, enc_depth=24, enc_num_heads=16, dec_embed_dim=768, dec_depth=12, dec_num_heads=12))"
+pretrained: null
+load_only_encoder: False
+long_context: False
+fixed_length: True
+resume: null
+benchmark: True
+num_views : 4
+num_test_views : 4
+n_corres_train: 0
+n_corres_test: 0
+
+train_criterion: ConfLoss(Regr3DPose(L21, norm_mode='?avg_dis'), alpha=0.2) + RGBLoss(MSE)
+test_criterion: Regr3DPose(L21, norm_mode='?avg_dis', gt_scale=True, sky_loss_value=0) + Regr3DPose_ScaleInv(L21, norm_mode='?avg_dis', gt_scale=True, sky_loss_value=0) + RGBLoss(L21)
+
+dataset1: Co3d_Multi(split='train', ROOT='../../data/dust3r_data/processed_co3d/', aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset2: WildRGBD_Multi(split='train', ROOT="../../data/dust3r_data/processed_wildrgbd", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset3: ARKitScenes_Multi(split='train', ROOT='../../data/dust3r_data/processed_arkitscenes/', aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset4: ARKitScenesHighRes_Multi(split='train', ROOT="../../data/dust3r_data/processed_arkitscenes_highres", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset5: ScanNetpp_Multi(split='train', ROOT="../../data/dust3r_data/processed_scannetpp/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset6: ScanNet_Multi(split='train', ROOT="../../data/dust3r_data/processed_scannet/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset7: HyperSim_Multi(split='train', ROOT="../../data/custom_data/processed_hypersim", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset8: BlendedMVS_Multi(split='train', ROOT="../../data/dust3r_data/processed_blendedmvs/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset9: MegaDepth_Multi(split="train", ROOT="../../data/dust3r_data/processed_megadepth", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset10: MapFree_Multi(split=None, ROOT="../../data/mast3r_data/processed_mapfree/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset11: Waymo_Multi(split=None, ROOT="../../data/dust3r_data/processed_waymo/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset12: VirtualKITTI2_Multi(split=None, ROOT="../../data/mast3r_data/processed_vkitti", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset13: UnReal4K_Multi(split=None, ROOT="../../data/mast3r_data/processed_unreal4k/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset14: TartanAir_Multi(split=None, ROOT="../../data/mast3r_data/processed_tartanair/", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+dataset15: DL3DV_Multi(split='train', ROOT="../../data/custom_data/processed_dl3dv", aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+
+train_dataset: 32000 @ ${dataset1} + 48000 @ ${dataset2} + 100800 @ ${dataset3} + 56000 @ ${dataset4} + 33600 @ ${dataset5} + 56000 @ ${dataset6} + 33600 @ ${dataset7} + 33600 @ ${dataset8} + 33600 @ ${dataset9} + 100800 @ ${dataset10} + 78400 @ ${dataset11} + 5000 @ ${dataset12} + 1000 @ ${dataset13} + 33600 @ ${dataset14} + 160000 @ ${dataset15}
+test_dataset: 1000 @ ARKitScenes_Multi(split='test', ROOT='../../data/dust3r_data/processed_arkitscenes/', resolution=224, num_views=4, seed=42, n_corres=0)
+
+
+seed: 0
+batch_size: 16
+accum_iter: 1
+gradient_checkpointing: False
+epochs: 100
+start_epoch: 0
+weight_decay: 0.05
+lr: 1e-4
+min_lr: 1e-6
+warmup_epochs: 10
+amp: 1
+
+num_workers: 8
+world_size: 1
+local-rank: -1
+dist_url: 'env://'
+rank: 0
+gpu: 0
+distributed: False
+dist_backend: 'nccl'
+
+eval_freq: 1
+save_freq: 1
+keep_freq: 10
+print_freq: 10
+print_img_freq: 500
+num_imgs_vis: 4
+save_dir: 'checkpoints'
+exp_name: 'train_first_stage'
+task: 'cut3r'
+logdir: ./${save_dir}/${exp_name}/logs
+output_dir: ./${save_dir}/${exp_name}/
+hydra:
+  verbose: True
+  run:
+    dir: ./${save_dir}/${exp_name}

extern/CUT3R/config/stage2.yaml

@@ -0,0 +1,132 @@
+model: ARCroco3DStereo(ARCroco3DStereoConfig(state_size=768, state_pe='2d', pos_embed='RoPE100',
+  rgb_head=True, pose_head=True,  img_size=(224, 224), head_type='linear', output_mode='pts3d+pose',
+  depth_mode=('exp', -inf, inf), conf_mode=('exp', 1, inf), pose_mode=('exp', -inf,
+  inf), enc_embed_dim=1024, enc_depth=24, enc_num_heads=16, dec_embed_dim=768, dec_depth=12,
+  dec_num_heads=12))
+pretrained: checkpoints/train_first_stage/checkpoint-final.pth
+load_only_encoder: False
+long_context: False
+fixed_length: True
+resume: null
+benchmark: True
+num_views : 4
+num_test_views : 4
+n_corres_train: 0
+n_corres_test: 0
+
+
+train_criterion: ConfLoss(Regr3DPoseBatchList(L21, norm_mode='?avg_dis'), alpha=0.2) + RGBLoss(MSE)
+test_criterion: Regr3DPose(L21, norm_mode='?avg_dis', gt_scale=True, sky_loss_value=0) + Regr3DPose_ScaleInv(L21, norm_mode='?avg_dis', gt_scale=True, sky_loss_value=0) + RGBLoss(L21)
+
+
+dataset1: Co3d_Multi(split='train', ROOT='../../data/dust3r_data/processed_co3d/',
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset2: WildRGBD_Multi(split='train', ROOT="../../data/dust3r_data/processed_wildrgbd",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset3: ARKitScenes_Multi(split='train', ROOT='../../data/dust3r_data/processed_arkitscenes/',
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset4: ARKitScenesHighRes_Multi(split='train', ROOT="../../data/dust3r_data/processed_arkitscenes_highres",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset5: ScanNetpp_Multi(split='train', ROOT="../../data/dust3r_data/processed_scannetpp/",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset6: ScanNet_Multi(split='train', ROOT="../../data/dust3r_data/processed_scannet/",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset7: HyperSim_Multi(split='train', ROOT="../../data/custom_data/processed_hypersim",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset8: BlendedMVS_Multi(split='train', ROOT="../../data/dust3r_data/processed_blendedmvs/",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset9: MegaDepth_Multi(split="train", ROOT="../../data/dust3r_data/processed_megadepth",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset10: MapFree_Multi(split=None, ROOT="../../data/mast3r_data/processed_mapfree/",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset11: Waymo_Multi(split=None, ROOT="../../data/dust3r_data/processed_waymo/",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset12: VirtualKITTI2_Multi(split=None, ROOT="../../data/mast3r_data/processed_vkitti",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset13: UnReal4K_Multi(split=None, ROOT="../../data/mast3r_data/processed_unreal4k/",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset14: TartanAir_Multi(split=None, ROOT="../../data/mast3r_data/processed_tartanair/",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset15: DL3DV_Multi(split='train', ROOT="../../data/custom_data/processed_dl3dv",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset16: Cop3D_Multi(split='train', ROOT="../../data/custom_data/processed_cop3d/",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset17: MVImgNet_Multi(split='train', ROOT="../../data/custom_data/processed_mvimgnet/",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset18: RE10K_Multi(split=None, ROOT="../../data/custom_data/processed_re10k/",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset19: OmniObject3D_Multi(split=None, ROOT="../../data/custom_data/processed_omniobject3d/",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset20: ThreeDKenBurns(split=None, ROOT="../../data/custom_data/processed_3dkb/",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset21: IRS(split=None, ROOT="../../data/custom_data/processed_irs/", aug_crop=16,
+  resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset22: SynScapes(split=None, ROOT="../../data/custom_data/processed_synscapes/",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset23: UrbanSyn(split=None, ROOT="../../data/custom_data/processed_urbansyn/",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset24: EDEN_Multi(split='train', ROOT="../../data/custom_data/processed_eden",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset25: SmartPortraits_Multi(split='train', ROOT="../../data/custom_data/processed_smartportraits",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset26: DynamicReplica(split='train', ROOT="../../data/custom_data/processed_dynamic_replica/",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset27: Spring(split=None, ROOT="../../data/custom_data/processed_spring/", aug_crop=16,
+  resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset28: BEDLAM_Multi(split='train', ROOT="../../data/custom_data/processed_bedlam",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset29: MVS_Synth_Multi(split='train', ROOT="../../data/custom_data/processed_mvs_synth",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset30: PointOdyssey_Multi(split='train', ROOT="../../data/custom_data/processed_point_odyssey",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset31: UASOL_Multi(split='train', ROOT="../../data/custom_data/processed_uasol",
+  aug_crop=16, resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+dataset32: MP3D_Multi(split=None, ROOT="../../data/custom_data/processed_mp3d/", aug_crop=16,
+  resolution=224, transform=SeqColorJitter, num_views=${num_views}, n_corres=${n_corres_train})
+train_dataset: 48000 @ ${dataset1} + 60000 @ ${dataset2} + 54000 @ ${dataset3} + 18000
+  @ ${dataset4} + 6000 @ ${dataset5} + 42000 @ ${dataset6} + 12000 @ ${dataset7} +
+  6000 @ ${dataset8} + 6000 @ ${dataset9} + 60000 @ ${dataset10} + 48000 @ ${dataset11}
+  + 2400 @ ${dataset12} + 180 @ ${dataset13} + 18000 @ ${dataset14} + 222000 @ ${dataset15}
+  + 400 @ ${dataset16} + 16000 @ ${dataset17} + 4000 @ ${dataset18} + 32000 @ ${dataset19}
+  + 4000 @ ${dataset20} + 2000 @ ${dataset21} + 2000 @ ${dataset22} + 500 @ ${dataset23}
+  + 12000 @ ${dataset24} + 16000 @ ${dataset25} + 20000 @ ${dataset26} + 400 @ ${dataset27}
+  + 32000 @ ${dataset28} + 2000 @ ${dataset29} + 20000 @ ${dataset30} + 12000 @ ${dataset31}
+  + 24000 @ ${dataset32}
+test_dataset: 1000 @ ARKitScenes_Multi(split='test', ROOT='../../data/dust3r_data/processed_arkitscenes/', resolution=224, num_views=4, seed=42, n_corres=0)
+
+seed: 0
+batch_size: 16
+accum_iter: 1
+gradient_checkpointing: false
+epochs: 35
+start_epoch: 0
+weight_decay: 0.05
+lr: 5.0e-06
+min_lr: 1.0e-06
+warmup_epochs: 1
+amp: 1
+
+num_workers: 8
+world_size: 1
+local-rank: -1
+dist_url: 'env://'
+rank: 0
+gpu: 0
+distributed: False
+dist_backend: 'nccl'
+
+eval_freq: 1
+save_freq: 1
+keep_freq: 10
+print_freq: 10
+print_img_freq: 500
+num_imgs_vis: 4
+save_dir: 'checkpoints'
+exp_name: 'train_second_stage'
+task: 'cut3r'
+logdir: ./${save_dir}/${exp_name}/logs
+output_dir: ./${save_dir}/${exp_name}/
+hydra:
+  verbose: True
+  run:
+    dir: ./${save_dir}/${exp_name}

extern/CUT3R/config/stage3.yaml

@@ -0,0 +1,219 @@
+model: ARCroco3DStereo(ARCroco3DStereoConfig(state_size=768, state_pe='2d', pos_embed='RoPE100',
+  rgb_head=True, pose_head=True, patch_embed_cls='ManyAR_PatchEmbed', img_size=(512,
+  512), head_type='dpt', output_mode='pts3d+pose', depth_mode=('exp', -inf, inf),
+  conf_mode=('exp', 1, inf), pose_mode=('exp', -inf, inf), enc_embed_dim=1024, enc_depth=24,
+  enc_num_heads=16, dec_embed_dim=768, dec_depth=12, dec_num_heads=12, landscape_only=False))
+pretrained: checkpoints/train_second_stage/checkpoint-final.pth
+load_only_encoder: False
+long_context: False
+fixed_length: True
+resume: null
+benchmark: True
+num_views : 4
+num_test_views : 4
+n_corres_train: 0
+n_corres_test: 0
+
+
+train_criterion: ConfLoss(Regr3DPoseBatchList(L21, norm_mode='?avg_dis'), alpha=0.2) + RGBLoss(MSE)
+test_criterion: Regr3DPose(L21, norm_mode='?avg_dis', gt_scale=True, sky_loss_value=0) + Regr3DPose_ScaleInv(L21, norm_mode='?avg_dis', gt_scale=True, sky_loss_value=0) + RGBLoss(L21)
+
+resolution:
+- (512
+- 384)
+- (512
+- 336)
+- (512
+- 288)
+- (512
+- 256)
+- (512
+- 208)
+- (512
+- 144)
+- (384
+- 512)
+- (336
+- 512)
+- (288
+- 512)
+- (256
+- 512)
+dataset1: Co3d_Multi(allow_repeat=True, split='train', ROOT='../../data/dust3r_data/processed_co3d/',
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset2: WildRGBD_Multi(allow_repeat=True, split='train', ROOT="../../data/dust3r_data/processed_wildrgbd",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset3: ARKitScenes_Multi(allow_repeat=True, split='train', ROOT='../../data/dust3r_data/processed_arkitscenes/',
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset4: ARKitScenesHighRes_Multi(allow_repeat=True, split='train', ROOT="../../data/dust3r_data/processed_arkitscenes_highres",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset5: ScanNetpp_Multi(allow_repeat=True, split='train', ROOT="../../data/dust3r_data/processed_scannetpp/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset6: ScanNet_Multi(allow_repeat=True, split='train', ROOT="../../data/dust3r_data/processed_scannet/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset7: HyperSim_Multi(allow_repeat=True, split='train', ROOT="../../data/custom_data/processed_hypersim",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset8: BlendedMVS_Multi(allow_repeat=True, split='train', ROOT="../../data/dust3r_data/processed_blendedmvs/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset9: MegaDepth_Multi(allow_repeat=True, split="train", ROOT="../../data/dust3r_data/processed_megadepth",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset10: MapFree_Multi(allow_repeat=True, split=None, ROOT="../../data/mast3r_data/processed_mapfree/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset11: Waymo_Multi(allow_repeat=True, split=None, ROOT="../../data/dust3r_data/processed_waymo/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset12: VirtualKITTI2_Multi(allow_repeat=True, split=None, ROOT="../../data/mast3r_data/processed_vkitti",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset13: UnReal4K_Multi(allow_repeat=True, split=None, ROOT="../../data/mast3r_data/processed_unreal4k/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset14: TartanAir_Multi(allow_repeat=True, split=None, ROOT="../../data/mast3r_data/processed_tartanair/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset15: DL3DV_Multi(allow_repeat=True, split='train', ROOT="../../data/custom_data/processed_dl3dv",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset16: Cop3D_Multi(allow_repeat=True, split='train', ROOT="../../data/custom_data/processed_cop3d/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset17: MVImgNet_Multi(allow_repeat=True, split='train', ROOT="../../data/custom_data/processed_mvimgnet/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset18: RE10K_Multi(allow_repeat=True, split=None, ROOT="../../data/custom_data/processed_re10k/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset19: OmniObject3D_Multi(allow_repeat=True, split=None, ROOT="../../data/custom_data/processed_omniobject3d/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset20: ThreeDKenBurns(allow_repeat=True, split=None, ROOT="../../data/custom_data/processed_3dkb/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset21: IRS(allow_repeat=True, split=None, ROOT="../../data/custom_data/processed_irs/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset22: SynScapes(allow_repeat=True, split=None, ROOT="../../data/custom_data/processed_synscapes/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset23: UrbanSyn(allow_repeat=True, split=None, ROOT="../../data/custom_data/processed_urbansyn/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset24: EDEN_Multi(allow_repeat=True, split='train', ROOT="../../data/custom_data/processed_eden",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset25: SmartPortraits_Multi(allow_repeat=True, split='train', ROOT="../../data/custom_data/processed_smartportraits",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset26: DynamicReplica(allow_repeat=True, split='train', ROOT="../../data/custom_data/processed_dynamic_replica/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset27: Spring(allow_repeat=True, split=None, ROOT="../../data/custom_data/processed_spring/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset28: BEDLAM_Multi(allow_repeat=True, split='train', ROOT="../../data/custom_data/processed_bedlam",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset29: MVS_Synth_Multi(allow_repeat=True, split='train', ROOT="../../data/custom_data/processed_mvs_synth",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset30: PointOdyssey_Multi(allow_repeat=True, split='train', ROOT="../../data/custom_data/processed_point_odyssey",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset31: UASOL_Multi(allow_repeat=True, split='train', ROOT="../../data/custom_data/processed_uasol",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset32: MP3D_Multi(allow_repeat=True, split=None, ROOT="../../data/custom_data/processed_mp3d/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset33: HOI4D_Multi(allow_repeat=True, split=None, ROOT="../../data/custom_data/processed_hoi4d/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+train_dataset: 44800 @ ${dataset1} + 56000 @ ${dataset2} + 56000 @ ${dataset3} + 22400
+  @ ${dataset4} + 16800 @ ${dataset5} + 38400 @ ${dataset6} + 11200 @ ${dataset7}
+  + 22400 @ ${dataset8} + 22400 @ ${dataset9} + 84000 @ ${dataset10}  + 20000 @ ${dataset11}
+  + 5600 @ ${dataset12} + 168 @ ${dataset13} + 56000 @ ${dataset14} + 74000 @ ${dataset15}
+  + 480 @ ${dataset16} + 19200 @ ${dataset17} + 4800 @ ${dataset18} + 4800 @ ${dataset20}
+  + 2400 @ ${dataset21} + 2400 @ ${dataset22} + 600 @ ${dataset23} + 19200 @ ${dataset25}
+  + 36000 @ ${dataset26} + 9400 @ ${dataset27} + 36000 @ ${dataset28} + 1400 @ ${dataset29}
+  + 7200 @ ${dataset30} + 14400 @ ${dataset31} + 28800 @ ${dataset32} + 12000 @ ${dataset33}
+test_dataset: 1000 @ ARKitScenes_Multi(split='test', ROOT='../../data/dust3r_data/processed_arkitscenes/', resolution=224, num_views=4, seed=42, n_corres=0)
+
+seed: 0
+batch_size: 16
+accum_iter: 1
+gradient_checkpointing: true
+epochs: 40
+start_epoch: 0
+weight_decay: 0.05
+lr: 1.0e-05
+min_lr: 1.0e-06
+warmup_epochs: 2
+amp: 1
+
+num_workers: 8
+world_size: 1
+local-rank: -1
+dist_url: 'env://'
+rank: 0
+gpu: 0
+distributed: False
+dist_backend: 'nccl'
+
+eval_freq: 1
+save_freq: 1
+keep_freq: 10
+print_freq: 10
+print_img_freq: 500
+num_imgs_vis: 4
+save_dir: 'checkpoints'
+exp_name: 'train_third_stage'
+task: 'cut3r'
+logdir: ./${save_dir}/${exp_name}/logs
+output_dir: ./${save_dir}/${exp_name}/
+hydra:
+  verbose: True
+  run:
+    dir: ./${save_dir}/${exp_name}

extern/CUT3R/config/stage4.yaml

@@ -0,0 +1,219 @@
+model: ARCroco3DStereo(ARCroco3DStereoConfig(freeze='encoder', state_size=768, state_pe='2d',
+  pos_embed='RoPE100', rgb_head=True, pose_head=True, patch_embed_cls='ManyAR_PatchEmbed',
+  img_size=(512, 512), head_type='dpt', output_mode='pts3d+pose', depth_mode=('exp',
+  -inf, inf), conf_mode=('exp', 1, inf), pose_mode=('exp', -inf, inf), enc_embed_dim=1024,
+  enc_depth=24, enc_num_heads=16, dec_embed_dim=768, dec_depth=12, dec_num_heads=12,
+  landscape_only=False))
+pretrained: checkpoints/train_third_stage/checkpoint-final.pth
+load_only_encoder: False
+long_context: True
+fixed_length: True
+resume: null
+benchmark: True
+num_views : 32
+num_test_views : 4
+n_corres_train: 0
+n_corres_test: 0
+
+train_criterion: ConfLoss(Regr3DPose(L21, norm_mode='?avg_dis'), alpha=0.2) + RGBLoss(MSE)
+test_criterion: Regr3DPose(L21, norm_mode='?avg_dis', gt_scale=True, sky_loss_value=0)
+  + Regr3DPose_ScaleInv(L21, norm_mode='?avg_dis', gt_scale=True, sky_loss_value=0)
+  + RGBLoss(L21)
+resolution:
+- (512
+- 384)
+- (512
+- 336)
+- (512
+- 288)
+- (512
+- 256)
+- (512
+- 208)
+- (512
+- 144)
+- (384
+- 512)
+- (336
+- 512)
+- (288
+- 512)
+- (256
+- 512)
+dataset1: Co3d_Multi(allow_repeat=False, split='train', ROOT='../../data/dust3r_data/processed_co3d/',
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset2: WildRGBD_Multi(allow_repeat=False, split='train', ROOT="../../data/dust3r_data/processed_wildrgbd",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset3: ARKitScenes_Multi(allow_repeat=False, split='train', ROOT='../../data/dust3r_data/processed_arkitscenes/',
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset4: ARKitScenesHighRes_Multi(allow_repeat=False, split='train', ROOT="../../data/dust3r_data/processed_arkitscenes_highres",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset5: ScanNetpp_Multi(allow_repeat=False, split='train', ROOT="../../data/dust3r_data/processed_scannetpp/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset6: ScanNet_Multi(allow_repeat=False, split='train', ROOT="../../data/dust3r_data/processed_scannet/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset7: HyperSim_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_hypersim",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset8: BlendedMVS_Multi(allow_repeat=False, split='train', ROOT="../../data/dust3r_data/processed_blendedmvs/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset9: MegaDepth_Multi(allow_repeat=False, split="train", ROOT="../../data/dust3r_data/processed_megadepth",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset10: MapFree_Multi(allow_repeat=False, split=None, ROOT="../../data/mast3r_data/processed_mapfree/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset11: Waymo_Multi(allow_repeat=False, split=None, ROOT="../../data/dust3r_data/processed_waymo/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset12: VirtualKITTI2_Multi(allow_repeat=False, split=None, ROOT="../../data/mast3r_data/processed_vkitti",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset13: UnReal4K_Multi(allow_repeat=False, split=None, ROOT="../../data/mast3r_data/processed_unreal4k/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset14: TartanAir_Multi(allow_repeat=False, split=None, ROOT="../../data/mast3r_data/processed_tartanair/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset15: DL3DV_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_dl3dv",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset16: Cop3D_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_cop3d/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset17: MVImgNet_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_mvimgnet/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset18: RE10K_Multi(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_re10k/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset19: OmniObject3D_Multi(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_omniobject3d/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset20: ThreeDKenBurns(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_3dkb/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset21: IRS(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_irs/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset22: SynScapes(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_synscapes/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset23: UrbanSyn(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_urbansyn/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset24: EDEN_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_eden",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset25: SmartPortraits_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_smartportraits",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset26: DynamicReplica(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_dynamic_replica/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset27: Spring(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_spring/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset28: BEDLAM_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_bedlam",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset29: MVS_Synth_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_mvs_synth",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset30: PointOdyssey_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_point_odyssey",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset31: UASOL_Multi(allow_repeat=False, split='train', ROOT="../../data/custom_data/processed_uasol",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset32: MP3D_Multi(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_mp3d/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+dataset33: HOI4D_Multi(allow_repeat=False, split=None, ROOT="../../data/custom_data/processed_hoi4d/",
+  aug_crop=16, resolution=[(512, 384), (512, 336), (512, 288), (512, 256), (512, 208),
+  (512, 144), (384, 512), (336, 512), (288, 512), (256, 512)], transform=SeqColorJitter,
+  num_views=${num_views}, n_corres=${n_corres_train})
+train_dataset: 22400 @ ${dataset1} + 28000 @ ${dataset2} + 28000 @ ${dataset3} + 2800
+  @ ${dataset4} + 2800 @ ${dataset5} + 70000 @ ${dataset6} + 2800 @ ${dataset7} +
+  11200 @ ${dataset8} + 8400 @ ${dataset9} + 28000 @ ${dataset10} + 21000 @ ${dataset11}
+  + 2800 @ ${dataset12} + 84 @ ${dataset13} + 42000 @ ${dataset14} + 42000 @ ${dataset15}
+  + 3600 @ ${dataset16} + 9600 @ ${dataset17} + 4800 @ ${dataset18} + 12000 @ ${dataset19}
+  + 16800 @ ${dataset26} + 1200 @ ${dataset27} + 4800 @ ${dataset28} + 2400 @ ${dataset29}
+  + 14400 @ ${dataset30} + 7200 @ ${dataset31} + 9600 @ ${dataset32}
+test_dataset: 1000 @ ARKitScenes_Multi(split='test', ROOT='../../data/dust3r_data/processed_arkitscenes/', resolution=224, num_views=4, seed=42, n_corres=0)
+
+seed: 0
+batch_size: 16
+accum_iter: 1
+gradient_checkpointing: true
+epochs: 10
+start_epoch: 0
+weight_decay: 0.05
+lr: 1.0e-06
+min_lr: 1.0e-07
+warmup_epochs: 0.5
+amp: 1
+
+num_workers: 8
+world_size: 1
+local-rank: -1
+dist_url: 'env://'
+rank: 0
+gpu: 0
+distributed: False
+dist_backend: 'nccl'
+
+eval_freq: 1
+save_freq: 1
+keep_freq: 10
+print_freq: 10
+print_img_freq: 500
+num_imgs_vis: 4
+save_dir: 'checkpoints'
+exp_name: 'train_final_stage'
+task: 'cut3r'
+logdir: ./${save_dir}/${exp_name}/logs
+output_dir: ./${save_dir}/${exp_name}/
+hydra:
+  verbose: True
+  run:
+    dir: ./${save_dir}/${exp_name}

extern/CUT3R/datasets_preprocess/custom_convert2TUM.py

@@ -0,0 +1,262 @@
+import os
+import json
+import shutil
+import numpy as np
+import cv2 as cv
+import imageio
+from tqdm import tqdm
+from concurrent.futures import ProcessPoolExecutor, as_completed
+import open3d as o3d
+import scipy.ndimage
+import pickle
+
+# Set environment variable to limit OpenBLAS threads
+os.environ["OPENBLAS_NUM_THREADS"] = "1"
+
+DEPTH_SCALE_FACTOR = 5000
+
+
+# Point cloud from depth
+def pointcloudify_depth(depth, intrinsics, dist_coeff, undistort=True):
+    shape = depth.shape[::-1]
+
+    if undistort:
+        undist_intrinsics, _ = cv.getOptimalNewCameraMatrix(
+            intrinsics, dist_coeff, shape, 1, shape
+        )
+        inv_undist_intrinsics = np.linalg.inv(undist_intrinsics)
+
+        map_x, map_y = cv.initUndistortRectifyMap(
+            intrinsics, dist_coeff, None, undist_intrinsics, shape, cv.CV_32FC1
+        )
+        undist_depth = cv.remap(depth, map_x, map_y, cv.INTER_NEAREST)
+    else:
+        inv_undist_intrinsics = np.linalg.inv(intrinsics)
+        undist_depth = depth
+
+    # Generate x,y grid for H x W image
+    grid_x, grid_y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]))
+    grid = np.stack((grid_x, grid_y, np.ones_like(grid_x)), axis=-1)
+
+    # Reshape and compute local grid
+    grid_flat = grid.reshape(-1, 3).T
+    local_grid = inv_undist_intrinsics @ grid_flat
+
+    # Multiply by depth
+    local_grid = local_grid.T * undist_depth.reshape(-1, 1)
+
+    return local_grid.astype(np.float32)
+
+
+def project_pcd_to_depth(pcd, undist_intrinsics, img_size, config):
+    h, w = img_size
+    points = np.asarray(pcd.points)
+    d = points[:, 2]
+    normalized_points = points / points[:, 2][:, np.newaxis]
+    proj_pcd = np.round((undist_intrinsics @ normalized_points.T).T).astype(np.int64)
+    proj_mask = (
+        (proj_pcd[:, 0] >= 0)
+        & (proj_pcd[:, 0] < w)
+        & (proj_pcd[:, 1] >= 0)
+        & (proj_pcd[:, 1] < h)
+    )
+    proj_pcd = proj_pcd[proj_mask]
+    d = d[proj_mask]
+    pcd_image = np.zeros((config["res_h"], config["res_w"]), dtype=np.float32)
+    pcd_image[proj_pcd[:, 1], proj_pcd[:, 0]] = d
+    return pcd_image
+
+
+def smooth_depth(depth):
+    MAX_DEPTH_VAL = 1e5
+    KERNEL_SIZE = 11
+    depth = depth.copy()
+    depth[depth == 0] = MAX_DEPTH_VAL
+    smoothed_depth = scipy.ndimage.minimum_filter(depth, KERNEL_SIZE)
+    smoothed_depth[smoothed_depth == MAX_DEPTH_VAL] = 0
+    return smoothed_depth
+
+
+def align_rgb_depth(rgb, depth, roi, config, rgb_cnf, config_dict, T):
+    # Undistort rgb image
+    undist_rgb = cv.undistort(
+        rgb,
+        rgb_cnf["intrinsics"],
+        rgb_cnf["dist_coeff"],
+        None,
+        rgb_cnf["undist_intrinsics"],
+    )
+
+    # Create point cloud from depth
+    pcd = o3d.geometry.PointCloud()
+    points = pointcloudify_depth(
+        depth, config_dict["depth"]["dist_mtx"], config_dict["depth"]["dist_coef"]
+    )
+    pcd.points = o3d.utility.Vector3dVector(points)
+    # Align point cloud with depth reference frame
+    pcd.transform(T)
+
+    # Project aligned point cloud to rgb
+    aligned_depth = project_pcd_to_depth(
+        pcd, rgb_cnf["undist_intrinsics"], rgb.shape[:2], config
+    )
+
+    smoothed_aligned_depth = smooth_depth(aligned_depth)
+    x, y, w, h = roi
+
+    depth_res = smoothed_aligned_depth[y : y + h, x : x + w]
+    rgb_res = undist_rgb[y : y + h, x : x + w]
+    return rgb_res, depth_res, rgb_cnf["undist_intrinsics"]
+
+
+def process_pair(args):
+    (
+        pair,
+        smartphone_folder,
+        azure_depth_folder,
+        final_folder,
+        config,
+        rgb_cnf,
+        config_dict,
+        T,
+    ) = args
+    try:
+        rgb_image = cv.imread(os.path.join(smartphone_folder, f"{pair[0]}.png"))
+        depth_array = np.load(
+            os.path.join(azure_depth_folder, f"{pair[1]}.npy"), allow_pickle=True
+        )
+
+        rgb_image_aligned, depth_array_aligned, intrinsics = align_rgb_depth(
+            rgb_image,
+            depth_array,
+            (0, 0, config["res_w"], config["res_h"]),
+            config,
+            rgb_cnf,
+            config_dict,
+            T,
+        )
+        # Save rgb as 8-bit png
+        cv.imwrite(
+            os.path.join(final_folder, "rgb", f"{pair[0]}.png"), rgb_image_aligned
+        )
+
+        # # Save depth as 16-bit unsigned int with scale factor
+        # depth_array_aligned = (depth_array_aligned *
+        #                        DEPTH_SCALE_FACTOR).astype(np.uint16)
+        # imageio.imwrite(os.path.join(final_folder, 'depth', f"{pair[1]}.png"), depth_array_aligned)
+        np.save(
+            os.path.join(final_folder, "depth", f"{pair[0]}.npy"), depth_array_aligned
+        )
+        np.savez(
+            os.path.join(final_folder, "cam", f"{pair[0]}.npz"), intrinsics=intrinsics
+        )
+    except Exception as e:
+        return f"Error processing pair {pair}: {e}"
+    return None
+
+
+def main():
+    DATA_DIR_ = "data_smartportraits/SmartPortraits"  # REPLACE WITH YOUR OWN DATA PATH!
+    DATA_DIR = DATA_DIR_.rstrip("/")
+    print(f"{DATA_DIR_} {DATA_DIR}/")
+
+    # Folder where the data in TUM format will be put
+    curr_dir = os.path.dirname(os.path.abspath(__file__))
+    with open(os.path.join(curr_dir, "config.json")) as conf_f:
+        config = json.load(conf_f)
+
+    # Pre-load shared data
+    with open(os.path.join(curr_dir, config["depth_conf"]), "rb") as config_f:
+        config_dict = pickle.load(config_f)
+
+    rgb_cnf = np.load(
+        os.path.join(curr_dir, config["rgb_intristics"]), allow_pickle=True
+    ).item()
+
+    T = np.load(os.path.join(curr_dir, config["transform_intristics"]))
+
+    final_root = "processed_smartportraits1"  # REPLACE WITH YOUR OWN DATA PATH!
+
+    seqs = []
+    for scene in os.listdir(DATA_DIR):
+        scene_path = os.path.join(DATA_DIR, scene)
+        if not os.path.isdir(scene_path):
+            continue
+        for s in os.listdir(scene_path):
+            s_path = os.path.join(scene_path, s)
+            if not os.path.isdir(s_path):
+                continue
+            for date in os.listdir(s_path):
+                date_path = os.path.join(s_path, date)
+                if os.path.isdir(date_path):
+                    seqs.append((scene, s, date))
+
+    for seq in tqdm(seqs):
+        scene, s, date = seq
+        dataset_path = os.path.join(DATA_DIR, scene, s, date)
+        final_folder = os.path.join(final_root, "_".join([scene, s, date]))
+
+        azure_depth_folder = os.path.join(dataset_path, "_azure_depth_image_raw")
+        smartphone_folder = os.path.join(dataset_path, "smartphone_video_frames")
+
+        depth_files = [
+            file for file in os.listdir(azure_depth_folder) if file.endswith(".npy")
+        ]
+        depth_ts = np.array([int(file.split(".")[0]) for file in depth_files])
+        depth_ts.sort()
+
+        rgb_files = [
+            file for file in os.listdir(smartphone_folder) if file.endswith(".png")
+        ]
+        rgb_ts = np.array([int(file.split(".")[0]) for file in rgb_files])
+        rgb_ts.sort()
+
+        print(
+            f"Depth timestamps from {depth_ts[0]} to {depth_ts[-1]} (cnt {len(depth_ts)})"
+        )
+        print(f"RGB timestamps from {rgb_ts[0]} to {rgb_ts[-1]} (cnt {len(rgb_ts)})")
+
+        # Build correspondences between depth and rgb by nearest neighbour algorithm
+        rgbd_pairs = []
+        for depth_t in depth_ts:
+            idx = np.argmin(np.abs(rgb_ts - depth_t))
+            closest_rgb_t = rgb_ts[idx]
+            rgbd_pairs.append((closest_rgb_t, depth_t))
+
+        # Prepare folder infrastructure
+        if os.path.exists(final_folder):
+            shutil.rmtree(final_folder)
+        os.makedirs(os.path.join(final_folder, "depth"), exist_ok=True)
+        os.makedirs(os.path.join(final_folder, "rgb"), exist_ok=True)
+        os.makedirs(os.path.join(final_folder, "cam"), exist_ok=True)
+
+        # Prepare arguments for processing
+        tasks = [
+            (
+                pair,
+                smartphone_folder,
+                azure_depth_folder,
+                final_folder,
+                config,
+                rgb_cnf,
+                config_dict,
+                T,
+            )
+            for pair in rgbd_pairs
+        ]
+
+        num_workers = os.cpu_count()
+        with ProcessPoolExecutor(max_workers=num_workers) as executor:
+            futures = {executor.submit(process_pair, task): task[0] for task in tasks}
+            for future in tqdm(
+                as_completed(futures),
+                total=len(futures),
+                desc=f"Processing {scene}_{s}_{date}",
+            ):
+                error = future.result()
+                if error:
+                    print(error)
+
+
+if __name__ == "__main__":
+    main()

extern/CUT3R/datasets_preprocess/flow_IO.py

@@ -0,0 +1,476 @@
+import struct
+import numpy as np
+import png
+import re
+import sys
+import csv
+from PIL import Image
+import h5py
+
+
+FLO_TAG_FLOAT = (
+    202021.25  # first 4 bytes in flo file; check for this when READING the file
+)
+FLO_TAG_STRING = "PIEH"  # first 4 bytes in flo file; use this when WRITING the file
+FLO_UNKNOWN_FLOW_THRESH = 1e9  # flo format threshold for unknown values
+FLO_UNKNOWN_FLOW = 1e10  # value to use to represent unknown flow in flo file format
+
+
+def readFlowFile(filepath):
+    """read flow files in several formats. The resulting flow has shape height x width x 2.
+    For positions where there is no groundtruth available, the flow is set to np.nan.
+    Supports flo (Sintel), png (KITTI), npy (numpy), pfm (FlyingThings3D) and flo5 (Spring) file format.
+    filepath: path to the flow file
+    returns: flow with shape height x width x 2
+    """
+    if filepath.endswith(".flo"):
+        return readFloFlow(filepath)
+    elif filepath.endswith(".png"):
+        return readPngFlow(filepath)
+    elif filepath.endswith(".npy"):
+        return readNpyFlow(filepath)
+    elif filepath.endswith(".pfm"):
+        return readPfmFlow(filepath)
+    elif filepath.endswith(".flo5"):
+        return readFlo5Flow(filepath)
+    else:
+        raise ValueError(f"readFlowFile: Unknown file format for {filepath}")
+
+
+def writeFlowFile(flow, filepath):
+    """write optical flow to file. Supports flo (Sintel), png (KITTI) and npy (numpy) file format.
+    flow: optical flow with shape height x width x 2. Invalid values should be represented as np.nan
+    filepath: file path where to write the flow
+    """
+    if not filepath:
+        raise ValueError("writeFlowFile: empty filepath")
+
+    if len(flow.shape) != 3 or flow.shape[2] != 2:
+        raise IOError(
+            f"writeFlowFile {filepath}: expected shape height x width x 2 but received {flow.shape}"
+        )
+
+    if flow.shape[0] > flow.shape[1]:
+        print(
+            f"write flo file {filepath}: Warning: Are you writing an upright image? Expected shape height x width x 2, got {flow.shape}"
+        )
+
+    if filepath.endswith(".flo"):
+        return writeFloFlow(flow, filepath)
+    elif filepath.endswith(".png"):
+        return writePngFlow(flow, filepath)
+    elif filepath.endswith(".npy"):
+        return writeNpyFile(flow, filepath)
+    elif filepath.endswith(".flo5"):
+        return writeFlo5File(flow, filepath)
+    else:
+        raise ValueError(f"writeFlowFile: Unknown file format for {filepath}")
+
+
+def readFloFlow(filepath):
+    """read optical flow from file stored in .flo file format as used in the Sintel dataset (Butler et al., 2012)
+    filepath: path to file where to read from
+    returns: flow as a numpy array with shape height x width x 2
+    ---
+    ".flo" file format used for optical flow evaluation
+
+    Stores 2-band float image for horizontal (u) and vertical (v) flow components.
+    Floats are stored in little-endian order.
+    A flow value is considered "unknown" if either |u| or |v| is greater than 1e9.
+
+    bytes  contents
+
+    0-3     tag: "PIEH" in ASCII, which in little endian happens to be the float 202021.25
+            (just a sanity check that floats are represented correctly)
+    4-7     width as an integer
+    8-11    height as an integer
+    12-end  data (width*height*2*4 bytes total)
+            the float values for u and v, interleaved, in row order, i.e.,
+            u[row0,col0], v[row0,col0], u[row0,col1], v[row0,col1], ...
+    """
+    if filepath is None:
+        raise IOError("read flo file: empty filename")
+
+    if not filepath.endswith(".flo"):
+        raise IOError(f"read flo file ({filepath}): extension .flo expected")
+
+    with open(filepath, "rb") as stream:
+        tag = struct.unpack("f", stream.read(4))[0]
+        width = struct.unpack("i", stream.read(4))[0]
+        height = struct.unpack("i", stream.read(4))[0]
+
+        if tag != FLO_TAG_FLOAT:  # simple test for correct endian-ness
+            raise IOError(
+                f"read flo file({filepath}): wrong tag (possibly due to big-endian machine?)"
+            )
+
+        # another sanity check to see that integers were read correctly (99999 should do the trick...)
+        if width < 1 or width > 99999:
+            raise IOError(f"read flo file({filepath}): illegal width {width}")
+
+        if height < 1 or height > 99999:
+            raise IOError(f"read flo file({filepath}): illegal height {height}")
+
+        nBands = 2
+        flow = []
+
+        n = nBands * width
+        for _ in range(height):
+            data = stream.read(n * 4)
+            if data is None:
+                raise IOError(f"read flo file({filepath}): file is too short")
+            data = np.asarray(struct.unpack(f"{n}f", data))
+            data = data.reshape((width, nBands))
+            flow.append(data)
+
+        if stream.read(1) != b"":
+            raise IOError(f"read flo file({filepath}): file is too long")
+
+        flow = np.asarray(flow)
+        # unknown values are set to nan
+        flow[np.abs(flow) > FLO_UNKNOWN_FLOW_THRESH] = np.nan
+
+        return flow
+
+
+def writeFloFlow(flow, filepath):
+    """
+    write optical flow in .flo format to file as used in the Sintel dataset (Butler et al., 2012)
+    flow: optical flow with shape height x width x 2
+    filepath: optical flow file path to be saved
+    ---
+    ".flo" file format used for optical flow evaluation
+
+    Stores 2-band float image for horizontal (u) and vertical (v) flow components.
+    Floats are stored in little-endian order.
+    A flow value is considered "unknown" if either |u| or |v| is greater than 1e9.
+
+    bytes  contents
+
+    0-3     tag: "PIEH" in ASCII, which in little endian happens to be the float 202021.25
+            (just a sanity check that floats are represented correctly)
+    4-7     width as an integer
+    8-11    height as an integer
+    12-end  data (width*height*2*4 bytes total)
+            the float values for u and v, interleaved, in row order, i.e.,
+            u[row0,col0], v[row0,col0], u[row0,col1], v[row0,col1], ...
+    """
+
+    height, width, nBands = flow.shape
+
+    with open(filepath, "wb") as f:
+        if f is None:
+            raise IOError(f"write flo file {filepath}: file could not be opened")
+
+        # write header
+        result = f.write(FLO_TAG_STRING.encode("ascii"))
+        result += f.write(struct.pack("i", width))
+        result += f.write(struct.pack("i", height))
+        if result != 12:
+            raise IOError(f"write flo file {filepath}: problem writing header")
+
+        # write content
+        n = nBands * width
+        for i in range(height):
+            data = flow[i, :, :].flatten()
+            data[np.isnan(data)] = FLO_UNKNOWN_FLOW
+            result = f.write(struct.pack(f"{n}f", *data))
+            if result != n * 4:
+                raise IOError(f"write flo file {filepath}: problem writing row {i}")
+
+
+def readPngFlow(filepath):
+    """read optical flow from file stored in png file format as used in the KITTI 12 (Geiger et al., 2012) and KITTI 15 (Menze et al., 2015) dataset.
+    filepath: path to file where to read from
+    returns: flow as a numpy array with shape height x width x 2. Invalid values are represented as np.nan
+    """
+    # adapted from https://github.com/liruoteng/OpticalFlowToolkit
+    flow_object = png.Reader(filename=filepath)
+    flow_direct = flow_object.asDirect()
+    flow_data = list(flow_direct[2])
+    (w, h) = flow_direct[3]["size"]
+    flow = np.zeros((h, w, 3), dtype=np.float64)
+    for i in range(len(flow_data)):
+        flow[i, :, 0] = flow_data[i][0::3]
+        flow[i, :, 1] = flow_data[i][1::3]
+        flow[i, :, 2] = flow_data[i][2::3]
+
+    invalid_idx = flow[:, :, 2] == 0
+    flow[:, :, 0:2] = (flow[:, :, 0:2] - 2**15) / 64.0
+    flow[invalid_idx, 0] = np.nan
+    flow[invalid_idx, 1] = np.nan
+    return flow[:, :, :2]
+
+
+def writePngFlow(flow, filename):
+    """write optical flow to file png file format as used in the KITTI 12 (Geiger et al., 2012) and KITTI 15 (Menze et al., 2015) dataset.
+    flow: optical flow in shape height x width x 2, invalid values should be represented as np.nan
+    filepath: path to file where to write to
+    """
+    flow = 64.0 * flow + 2**15
+    width = flow.shape[1]
+    height = flow.shape[0]
+    valid_map = np.ones([flow.shape[0], flow.shape[1], 1])
+    valid_map[np.isnan(flow[:, :, 0]) | np.isnan(flow[:, :, 1])] = 0
+    flow = np.nan_to_num(flow)
+    flow = np.concatenate([flow, valid_map], axis=-1)
+    flow = np.clip(flow, 0, 2**16 - 1)
+    flow = flow.astype(np.uint16)
+    flow = np.reshape(flow, (-1, width * 3))
+    with open(filename, "wb") as f:
+        writer = png.Writer(width=width, height=height, bitdepth=16, greyscale=False)
+        writer.write(f, flow)
+
+
+def readNpyFlow(filepath):
+    """read numpy array from file.
+    filepath: file to read from
+    returns: numpy array
+    """
+    return np.load(filepath)
+
+
+def writeNpyFile(arr, filepath):
+    """write numpy array to file.
+    arr: numpy array to write
+    filepath: file to write to
+    """
+    np.save(filepath, arr)
+
+
+def writeFlo5File(flow, filename):
+    with h5py.File(filename, "w") as f:
+        f.create_dataset("flow", data=flow, compression="gzip", compression_opts=5)
+
+
+def readFlo5Flow(filename):
+    with h5py.File(filename, "r") as f:
+        if "flow" not in f.keys():
+            raise IOError(
+                f"File {filename} does not have a 'flow' key. Is this a valid flo5 file?"
+            )
+        return f["flow"][()]
+
+
+def readPfmFlow(filepath):
+    """read optical flow from file stored in pfm file format as used in the FlyingThings3D (Mayer et al., 2016) dataset.
+    filepath: path to file where to read from
+    returns: flow as a numpy array with shape height x width x 2.
+    """
+    flow = readPfmFile(filepath)
+    if len(flow.shape) != 3:
+        raise IOError(
+            f"read pfm flow: PFM file has wrong shape (assumed to be w x h x 3): {flow.shape}"
+        )
+    if flow.shape[2] != 3:
+        raise IOError(
+            f"read pfm flow: PFM file has wrong shape (assumed to be w x h x 3): {flow.shape}"
+        )
+    # remove third channel -> is all zeros
+    return flow[:, :, :2]
+
+
+def readPfmFile(filepath):
+    """
+    adapted from https://lmb.informatik.uni-freiburg.de/resources/datasets/SceneFlowDatasets.en.html
+    """
+    file = open(filepath, "rb")
+
+    color = None
+    width = None
+    height = None
+    scale = None
+    endian = None
+
+    header = file.readline().rstrip()
+    if header.decode("ascii") == "PF":
+        color = True
+    elif header.decode("ascii") == "Pf":
+        color = False
+    else:
+        raise Exception("Not a PFM file.")
+
+    dim_match = re.match(r"^(\d+)\s(\d+)\s$", file.readline().decode("ascii"))
+    if dim_match:
+        width, height = list(map(int, dim_match.groups()))
+    else:
+        raise Exception("Malformed PFM header.")
+
+    scale = float(file.readline().decode("ascii").rstrip())
+    if scale < 0:  # little-endian
+        endian = "<"
+        scale = -scale
+    else:
+        endian = ">"  # big-endian
+
+    data = np.fromfile(file, endian + "f")
+    shape = (height, width, 3) if color else (height, width)
+
+    data = np.reshape(data, shape)
+    data = np.flipud(data)
+    return data  # , scale
+
+
+def writePfmFile(image, filepath):
+    """
+    adapted from https://lmb.informatik.uni-freiburg.de/resources/datasets/SceneFlowDatasets.en.html
+    """
+    scale = 1
+    file = open(filepath, "wb")
+
+    color = None
+
+    if image.dtype.name != "float32":
+        raise Exception("Image dtype must be float32.")
+
+    image = np.flipud(image)
+
+    if len(image.shape) == 3 and image.shape[2] == 3:  # color image
+        color = True
+    elif (
+        len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1
+    ):  # greyscale
+        color = False
+    else:
+        raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")
+
+    file.write("PF\n" if color else "Pf\n".encode())
+    file.write("%d %d\n".encode() % (image.shape[1], image.shape[0]))
+
+    endian = image.dtype.byteorder
+
+    if endian == "<" or endian == "=" and sys.byteorder == "little":
+        scale = -scale
+
+    file.write("%f\n".encode() % scale)
+
+    image.tofile(file)
+
+
+def readDispFile(filepath):
+    """read disparity (or disparity change) from file. The resulting numpy array has shape height x width.
+    For positions where there is no groundtruth available, the value is set to np.nan.
+    Supports png (KITTI), npy (numpy) and pfm (FlyingThings3D) file format.
+    filepath: path to the flow file
+    returns: disparity with shape height x width
+    """
+    if filepath.endswith(".png"):
+        return readPngDisp(filepath)
+    elif filepath.endswith(".npy"):
+        return readNpyFlow(filepath)
+    elif filepath.endswith(".pfm"):
+        return readPfmDisp(filepath)
+    elif filepath.endswith(".dsp5"):
+        return readDsp5Disp(filepath)
+    else:
+        raise ValueError(f"readDispFile: Unknown file format for {filepath}")
+
+
+def readPngDisp(filepath):
+    """read disparity from file stored in png file format as used in the KITTI 12 (Geiger et al., 2012) and KITTI 15 (Menze et al., 2015) dataset.
+    filepath: path to file where to read from
+    returns: disparity as a numpy array with shape height x width. Invalid values are represented as np.nan
+    """
+    # adapted from https://github.com/liruoteng/OpticalFlowToolkit
+    image_object = png.Reader(filename=filepath)
+    image_direct = image_object.asDirect()
+    image_data = list(image_direct[2])
+    (w, h) = image_direct[3]["size"]
+    channel = len(image_data[0]) // w
+    if channel != 1:
+        raise IOError("read png disp: assumed channels to be 1!")
+    disp = np.zeros((h, w), dtype=np.float64)
+    for i in range(len(image_data)):
+        disp[i, :] = image_data[i][:]
+    disp[disp == 0] = np.nan
+    return disp[:, :] / 256.0
+
+
+def readPfmDisp(filepath):
+    """read disparity or disparity change from file stored in pfm file format as used in the FlyingThings3D (Mayer et al., 2016) dataset.
+    filepath: path to file where to read from
+    returns: disparity as a numpy array with shape height x width. Invalid values are represented as np.nan
+    """
+    disp = readPfmFile(filepath)
+    if len(disp.shape) != 2:
+        raise IOError(
+            f"read pfm disp: PFM file has wrong shape (assumed to be w x h): {disp.shape}"
+        )
+    return disp
+
+
+def writePngDisp(disp, filepath):
+    """write disparity to png file format as used in the KITTI 12 (Geiger et al., 2012) and KITTI 15 (Menze et al., 2015) dataset.
+    disp: disparity in shape height x width, invalid values should be represented as np.nan
+    filepath: path to file where to write to
+    """
+    disp = 256 * disp
+    width = disp.shape[1]
+    height = disp.shape[0]
+    disp = np.clip(disp, 0, 2**16 - 1)
+    disp = np.nan_to_num(disp).astype(np.uint16)
+    disp = np.reshape(disp, (-1, width))
+    with open(filepath, "wb") as f:
+        writer = png.Writer(width=width, height=height, bitdepth=16, greyscale=True)
+        writer.write(f, disp)
+
+
+def writeDsp5File(disp, filename):
+    with h5py.File(filename, "w") as f:
+        f.create_dataset("disparity", data=disp, compression="gzip", compression_opts=5)
+
+
+def readDsp5Disp(filename):
+    with h5py.File(filename, "r") as f:
+        if "disparity" not in f.keys():
+            raise IOError(
+                f"File {filename} does not have a 'disparity' key. Is this a valid dsp5 file?"
+            )
+        return f["disparity"][()]
+
+
+def writeDispFile(disp, filepath):
+    """write disparity to file. Supports png (KITTI) and npy (numpy) file format.
+    disp: disparity with shape height x width. Invalid values should be represented as np.nan
+    filepath: file path where to write the flow
+    """
+    if not filepath:
+        raise ValueError("writeDispFile: empty filepath")
+
+    if len(disp.shape) != 2:
+        raise IOError(
+            f"writeDispFile {filepath}: expected shape height x width but received {disp.shape}"
+        )
+
+    if disp.shape[0] > disp.shape[1]:
+        print(
+            f"writeDispFile {filepath}: Warning: Are you writing an upright image? Expected shape height x width, got {disp.shape}"
+        )
+
+    if filepath.endswith(".png"):
+        writePngDisp(disp, filepath)
+    elif filepath.endswith(".npy"):
+        writeNpyFile(disp, filepath)
+    elif filepath.endswith(".dsp5"):
+        writeDsp5File(disp, filepath)
+
+
+def readKITTIObjMap(filepath):
+    assert filepath.endswith(".png")
+    return np.asarray(Image.open(filepath)) > 0
+
+
+def readKITTIIntrinsics(filepath, image=2):
+    assert filepath.endswith(".txt")
+
+    with open(filepath) as f:
+        reader = csv.reader(f, delimiter=" ")
+        for row in reader:
+            if row[0] == f"K_{image:02d}:":
+                K = np.array(row[1:], dtype=np.float32).reshape(3, 3)
+                kvec = np.array([K[0, 0], K[1, 1], K[0, 2], K[1, 2]])
+                return kvec
+
+
+def writePngMapFile(map_, filename):
+    Image.fromarray(map_).save(filename)

extern/CUT3R/datasets_preprocess/generate_set_arkitscenes.py

@@ -0,0 +1,159 @@
+#!/usr/bin/env python3
+"""
+Preprocess scenes by sorting images and generating image/video collections.
+
+This script processes scenes in parallel using a thread pool, updating metadata
+with sorted images, trajectories, intrinsics, and generating pair, image collection,
+and video collection data. The processed metadata is saved to a new file in each scene directory.
+
+Usage:
+    python generate_set_arkitscenes.py --root /path/to/data --splits Training Test --max_interval 5.0 --num_workers 8
+"""
+
+import os
+import os.path as osp
+import argparse
+import numpy as np
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from tqdm import tqdm
+
+
+def get_timestamp(img_name):
+    """
+    Extract the timestamp from an image filename.
+    Assumes the timestamp is the last underscore-separated token in the name (before the file extension).
+
+    Args:
+        img_name (str): The image filename.
+
+    Returns:
+        float: The extracted timestamp.
+    """
+    return float(img_name[:-4].split("_")[-1])
+
+
+def process_scene(root, split, scene, max_interval):
+    """
+    Process a single scene by sorting its images by timestamp, updating trajectories,
+    intrinsics, and pairings, and generating image/video collections.
+
+    Args:
+        root (str): Root directory of the dataset.
+        split (str): The dataset split (e.g., 'Training', 'Test').
+        scene (str): The scene identifier.
+        max_interval (float): Maximum allowed time interval (in seconds) between images to consider them in the same video collection.
+    """
+    scene_dir = osp.join(root, split, scene)
+    metadata_path = osp.join(scene_dir, "scene_metadata.npz")
+
+    # Load the scene metadata
+    with np.load(metadata_path) as data:
+        images = data["images"]
+        trajectories = data["trajectories"]
+        intrinsics = data["intrinsics"]
+        pairs = data["pairs"]
+
+    # Sort images by timestep
+    imgs_with_indices = sorted(enumerate(images), key=lambda x: x[1])
+    indices, images = zip(*imgs_with_indices)
+    indices = np.array(indices)
+    index2sorted = {index: i for i, index in enumerate(indices)}
+
+    # Reorder trajectories and intrinsics based on the new image order
+    trajectories = trajectories[indices]
+    intrinsics = intrinsics[indices]
+
+    # Update pair indices (each pair is (id1, id2, score))
+    pairs = [(index2sorted[id1], index2sorted[id2], score) for id1, id2, score in pairs]
+
+    # Form image_collection: mapping from an image id to a list of (other image id, score)
+    image_collection = {}
+    for id1, id2, score in pairs:
+        image_collection.setdefault(id1, []).append((id2, score))
+
+    # Form video_collection: for each image, collect subsequent images within the max_interval time window
+    video_collection = {}
+    for i, image in enumerate(images):
+        j = i + 1
+        for j in range(i + 1, len(images)):
+            if get_timestamp(images[j]) - get_timestamp(image) > max_interval:
+                break
+        video_collection[i] = list(range(i + 1, j))
+
+    # Save the new metadata
+    output_path = osp.join(scene_dir, "new_scene_metadata.npz")
+    np.savez(
+        output_path,
+        images=images,
+        trajectories=trajectories,
+        intrinsics=intrinsics,
+        pairs=pairs,
+        image_collection=image_collection,
+        video_collection=video_collection,
+    )
+    print(f"Processed scene: {scene}")
+
+
+def main(args):
+    """
+    Main function to process scenes across specified dataset splits in parallel.
+    """
+    root = args.root
+    splits = args.splits
+    max_interval = args.max_interval
+    num_workers = args.num_workers
+
+    futures = []
+
+    # Create a ThreadPoolExecutor for parallel processing
+    with ThreadPoolExecutor(max_workers=num_workers) as executor:
+        for split in splits:
+            all_meta_path = osp.join(root, split, "all_metadata.npz")
+            with np.load(all_meta_path) as data:
+                scenes = data["scenes"]
+
+            # Submit processing tasks for each scene in the current split
+            for scene in scenes:
+                futures.append(
+                    executor.submit(process_scene, root, split, scene, max_interval)
+                )
+
+        # Use tqdm to display a progress bar as futures complete
+        for future in tqdm(
+            as_completed(futures), total=len(futures), desc="Processing scenes"
+        ):
+            # This will raise any exceptions caught during scene processing.
+            future.result()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Preprocess scene data to update metadata with sorted images and collections."
+    )
+    parser.add_argument(
+        "--root",
+        type=str,
+        default="",
+        help="Root directory containing the dataset splits.",
+    )
+    parser.add_argument(
+        "--splits",
+        type=str,
+        nargs="+",
+        default=["Training", "Test"],
+        help="List of dataset splits to process (e.g., Training Test).",
+    )
+    parser.add_argument(
+        "--max_interval",
+        type=float,
+        default=5.0,
+        help="Maximum time interval (in seconds) between images to consider them in the same video sequence.",
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=8,
+        help="Number of worker threads for parallel processing.",
+    )
+    args = parser.parse_args()
+    main(args)

extern/CUT3R/datasets_preprocess/generate_set_scannet.py

@@ -0,0 +1,132 @@
+#!/usr/bin/env python3
+"""
+Preprocess ScanNet scenes to generate video collections.
+
+This script processes each scene in specified splits by reading the image filenames
+from the "color" folder, grouping images into video sequences based on a maximum
+timestamp interval, and then saving the per-scene metadata as a NumPy .npz file.
+
+Usage:
+    python generate_set_scannet.py --root /path/to/processed_scannet \
+        --splits scans_test scans_train --max_interval 150 --num_workers 8
+"""
+
+import os
+import os.path as osp
+import argparse
+import numpy as np
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from tqdm import tqdm
+
+
+def get_timestamp(img_name):
+    """
+    Convert an image basename to an integer timestamp.
+
+    For ScanNet data, it is assumed that the basename is an integer string.
+
+    Args:
+        img_name (str): Image basename (without extension).
+
+    Returns:
+        int: The timestamp as an integer.
+    """
+    return int(img_name)
+
+
+def process_scene(root, split, scene, max_interval):
+    """
+    Process a single scene: group images into video sequences and save metadata.
+
+    Args:
+        root (str): Root directory for the processed data.
+        split (str): Name of the split (e.g., 'scans_test', 'scans_train').
+        scene (str): Name of the scene directory.
+        max_interval (int): Maximum allowed difference in timestamps for grouping images.
+    """
+    scene_dir = osp.join(root, split, scene)
+    color_dir = osp.join(scene_dir, "color")
+    # depth_dir and camera_dir are defined in case you need them in future modifications.
+    # depth_dir = osp.join(scene_dir, 'depth')
+    # camera_dir = osp.join(scene_dir, 'cam')
+
+    # Get all image basenames from the color folder (without file extension)
+    basenames = sorted(
+        [f.split(".")[0] for f in os.listdir(color_dir) if f.endswith(".jpg")],
+        key=lambda x: get_timestamp(x),
+    )
+
+    video_collection = {}
+    for i, image in enumerate(basenames):
+        video_collection[i] = []
+        for j in range(i + 1, len(basenames)):
+            # Group images that fall within max_interval seconds of the reference image.
+            if get_timestamp(basenames[j]) - get_timestamp(image) > max_interval:
+                break
+            video_collection[i].append(j)
+
+    # Save the scene metadata (list of basenames and the video collection) to an NPZ file.
+    out_path = osp.join(scene_dir, "new_scene_metadata.npz")
+    np.savez(out_path, images=basenames, video_collection=video_collection)
+    print(f"Processed scene: {scene} (split: {split})")
+
+
+def main(args):
+    root = args.root
+    splits = args.splits
+    max_interval = args.max_interval
+    num_workers = args.num_workers
+
+    futures = []
+    with ThreadPoolExecutor(max_workers=num_workers) as executor:
+        for split in splits:
+            split_dir = osp.join(root, split)
+            if not osp.isdir(split_dir):
+                print(
+                    f"Warning: Split directory '{split_dir}' does not exist; skipping."
+                )
+                continue
+            scenes = os.listdir(split_dir)
+            for scene in scenes:
+                futures.append(
+                    executor.submit(process_scene, root, split, scene, max_interval)
+                )
+        # Use tqdm to display progress as futures complete.
+        for future in tqdm(
+            as_completed(futures), total=len(futures), desc="Processing scenes"
+        ):
+            # This will re-raise any exceptions from process_scene.
+            future.result()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Preprocess ScanNet scenes to create video collections based on image timestamps."
+    )
+    parser.add_argument(
+        "--root",
+        type=str,
+        default="",
+        help="Root directory containing the processed ScanNet splits.",
+    )
+    parser.add_argument(
+        "--splits",
+        type=str,
+        nargs="+",
+        default=["scans_test", "scans_train"],
+        help="List of split directories to process (e.g., scans_test scans_train).",
+    )
+    parser.add_argument(
+        "--max_interval",
+        type=int,
+        default=150,
+        help="Maximum allowed timestamp difference (in integer units) for grouping images.",
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=8,
+        help="Number of worker threads for parallel processing.",
+    )
+    args = parser.parse_args()
+    main(args)

extern/CUT3R/datasets_preprocess/generate_set_scannetpp.py

@@ -0,0 +1,169 @@
+#!/usr/bin/env python3
+"""
+Preprocess processed_scannetpp scenes to update scene metadata.
+
+This script reads each scene's "scene_metadata.npz", sorts images by timestamp,
+updates trajectories, intrinsics, and pair indices, and builds two collections:
+  - image_collection: For each image, stores pairs (other image index, score)
+  - video_collection: For each image, groups subsequent images whose timestamps
+                      differ by at most a given max_interval (and share the same
+                      first character in the image name).
+
+The new metadata is saved as "new_scene_metadata.npz" in each scene folder.
+
+Usage:
+    python generate_set_scannetpp.py --root /path/to/processed_scannetpp \
+        --max_interval 150 --num_workers 8
+"""
+
+import os
+import os.path as osp
+import argparse
+import numpy as np
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from tqdm import tqdm
+
+
+def get_timestamp(img_name):
+    """
+    Convert an image name to a timestamp (integer).
+
+    If the image name starts with 'DSC', the timestamp is the integer part after 'DSC'.
+    Otherwise, it is assumed the image name has an underscore, and the second element is used.
+
+    Args:
+        img_name (str): The image basename (without extension).
+
+    Returns:
+        int: The extracted timestamp.
+    """
+    if img_name.startswith("DSC"):
+        return int(img_name[3:])
+    else:
+        return int(img_name.split("_")[1])
+
+
+def process_scene(root, scene, max_interval):
+    """
+    Process a single scene: sort images, update trajectories/intrinsics/pairs, and
+    form image and video collections. Save the updated metadata.
+
+    Args:
+        root (str): Root directory containing scene folders.
+        scene (str): Scene folder name.
+        max_interval (int): Maximum allowed difference (in timestamp units) for video grouping.
+    """
+    scene_dir = osp.join(root, scene)
+    metadata_path = osp.join(scene_dir, "scene_metadata.npz")
+    with np.load(metadata_path, allow_pickle=True) as data:
+        images = data["images"]
+        trajectories = data["trajectories"]
+        intrinsics = data["intrinsics"]
+        pairs = data["pairs"]
+
+    # Sort images by timestamp.
+    imgs_with_indices = sorted(enumerate(images), key=lambda x: x[1])
+    indices, images = zip(*imgs_with_indices)
+    indices = np.array(indices)
+    index2sorted = {index: i for i, index in enumerate(indices)}
+
+    # Update trajectories and intrinsics arrays according to the new order.
+    trajectories = trajectories[indices]
+    intrinsics = intrinsics[indices]
+
+    # Update pairs (each pair is (id1, id2, score)) with new indices.
+    pairs = [(index2sorted[id1], index2sorted[id2], score) for id1, id2, score in pairs]
+
+    # Build image_collection: for each pair, verify that both image files exist.
+    image_collection = {}
+    for id1, id2, score in pairs:
+        img1 = images[id1]
+        img2 = images[id2]
+        img1_path = osp.join(scene_dir, "images", img1 + ".jpg")
+        img2_path = osp.join(scene_dir, "images", img2 + ".jpg")
+        if not (osp.exists(img1_path) and osp.exists(img2_path)):
+            continue
+        if id1 not in image_collection:
+            image_collection[id1] = []
+        image_collection[id1].append((id2, score))
+
+    # Build video_collection: for each image, group subsequent images if:
+    #  1. Their timestamp difference is at most max_interval.
+    #  2. Their name's first character is the same as the current image.
+    video_collection = {}
+    for i, image in enumerate(images):
+        img_path = osp.join(scene_dir, "images", image + ".jpg")
+        if not osp.exists(img_path):
+            continue
+        video_collection[i] = []
+        for j in range(i + 1, len(images)):
+            next_img_path = osp.join(scene_dir, "images", images[j] + ".jpg")
+            if not osp.exists(next_img_path):
+                continue
+            if (
+                get_timestamp(images[j]) - get_timestamp(image) > max_interval
+                or images[j][0] != image[0]
+            ):
+                break
+            video_collection[i].append(j)
+
+    # Save the updated metadata to a new file.
+    out_path = osp.join(scene_dir, "new_scene_metadata.npz")
+    np.savez(
+        out_path,
+        images=images,
+        trajectories=trajectories,
+        intrinsics=intrinsics,
+        pairs=pairs,
+        image_collection=image_collection,
+        video_collection=video_collection,
+    )
+    print(f"Processed scene: {scene}")
+
+
+def main(args):
+    root = args.root
+    max_interval = args.max_interval
+    num_workers = args.num_workers
+
+    # Load the list of scenes from the 'all_metadata.npz' file.
+    all_metadata_path = osp.join(root, "all_metadata.npz")
+    with np.load(all_metadata_path, allow_pickle=True) as data:
+        scenes = data["scenes"]
+
+    # Process scenes in parallel.
+    futures = []
+    with ThreadPoolExecutor(max_workers=num_workers) as executor:
+        for scene in scenes:
+            futures.append(executor.submit(process_scene, root, scene, max_interval))
+        for future in tqdm(
+            as_completed(futures), total=len(futures), desc="Processing scenes"
+        ):
+            # This will raise any exceptions from process_scene.
+            future.result()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Preprocess processed_scannetpp scenes to update scene metadata."
+    )
+    parser.add_argument(
+        "--root",
+        type=str,
+        required=True,
+        help="Root directory containing processed_scannetpp scene folders.",
+    )
+    parser.add_argument(
+        "--max_interval",
+        type=int,
+        default=150,
+        help="Maximum timestamp interval for grouping images (default: 150).",
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=8,
+        help="Number of worker threads for parallel processing (default: 8).",
+    )
+    args = parser.parse_args()
+    main(args)

extern/CUT3R/datasets_preprocess/merge_dl3dv.py

@@ -0,0 +1,85 @@
+import os
+import shutil
+from tqdm import tqdm
+
+# Set these paths to your original and moved locations.
+src_base = "/path/to/processed_dl3dv"  # original location
+dst_base = "processed_dl3dv_ours"  # current (moved) location
+
+# Set dry_run to True for testing (no changes made), and False to perform the actions.
+dry_run = False
+
+def merge_directories(source_dir, destination_dir, dry_run=False):
+    """
+    Merge all contents from source_dir into destination_dir.
+    If an item already exists in destination_dir:
+      - For files: remove the destination file and move the source file.
+      - For directories: merge them recursively.
+    After moving items, empty directories are removed.
+    """
+    for item in os.listdir(source_dir):
+        source_item = os.path.join(source_dir, item)
+        dest_item = os.path.join(destination_dir, item)
+        if os.path.isdir(source_item):
+            if os.path.exists(dest_item):
+                # Recursively merge subdirectories.
+                merge_directories(source_item, dest_item, dry_run=dry_run)
+                # Remove the source subdirectory if empty.
+                if not os.listdir(source_item):
+                    if dry_run:
+                        print(f"[Dry-run] Would remove empty directory: {source_item}")
+                    else:
+                        os.rmdir(source_item)
+            else:
+                if dry_run:
+                    print(f"[Dry-run] Would move directory: {source_item} -> {dest_item}")
+                else:
+                    shutil.move(source_item, dest_item)
+        else:
+            # For files: if a file already exists at the destination, remove it.
+            if os.path.exists(dest_item):
+                if dry_run:
+                    print(f"[Dry-run] Would remove existing file: {dest_item}")
+                else:
+                    os.remove(dest_item)
+            if dry_run:
+                print(f"[Dry-run] Would move file: {source_item} -> {dest_item}")
+            else:
+                shutil.move(source_item, dest_item)
+
+# Build a list of relative folder paths in dst_base.
+# This assumes the structure is: dst_base/f1/f2
+all_folders = []
+for f1 in os.listdir(dst_base):
+    f1_path = os.path.join(dst_base, f1)
+    if not os.path.isdir(f1_path):
+        continue
+    for f2 in os.listdir(f1_path):
+        all_folders.append(os.path.join(f1, f2))
+
+# Process each folder and move/merge it back to the original location.
+for folder in tqdm(all_folders, desc="Moving folders back"):
+    original_folder = os.path.join(src_base, folder)  # target location in the original path
+    moved_folder = os.path.join(dst_base, folder)       # current location
+
+    # Ensure the parent directory of the original folder exists.
+    parent_dir = os.path.dirname(original_folder)
+    if dry_run:
+        if not os.path.exists(parent_dir):
+            print(f"[Dry-run] Would create directory: {parent_dir}")
+    else:
+        os.makedirs(parent_dir, exist_ok=True)
+
+    if not os.path.exists(original_folder):
+        if dry_run:
+            print(f"[Dry-run] Would move folder: {moved_folder} -> {original_folder}")
+        else:
+            shutil.move(moved_folder, original_folder)
+    else:
+        merge_directories(moved_folder, original_folder, dry_run=dry_run)
+        # Remove the moved folder if it becomes empty.
+        if not os.listdir(moved_folder):
+            if dry_run:
+                print(f"[Dry-run] Would remove empty directory: {moved_folder}")
+            else:
+                os.rmdir(moved_folder)

extern/CUT3R/datasets_preprocess/path_to_root.py

@@ -0,0 +1,14 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# DUSt3R repo root import
+# --------------------------------------------------------
+
+import sys
+import os.path as path
+
+HERE_PATH = path.normpath(path.dirname(__file__))
+DUST3R_REPO_PATH = path.normpath(path.join(HERE_PATH, "../"))
+# workaround for sibling import
+sys.path.insert(0, DUST3R_REPO_PATH)

extern/CUT3R/datasets_preprocess/preprocess_3dkb.py

@@ -0,0 +1,220 @@
+#!/usr/bin/env python3
+"""
+Process 3D Ken Burns data by selecting random view types, copying images and depth files,
+and computing camera intrinsics from a field-of-view value. The output files are stored in an
+organized folder structure.
+
+Usage:
+    python preprocess_3dkb.py --root /path/to/data_3d_ken_burns \
+                           --out_dir /path/to/processed_3dkb \
+                           [--num_workers 4] [--seed 42]
+"""
+
+import os
+import json
+import random
+import shutil
+from functools import partial
+from pathlib import Path
+import argparse
+
+import cv2  # noqa: F401; cv2 is imported to ensure OpenEXR support.
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+from concurrent.futures import ProcessPoolExecutor, as_completed
+
+# Ensure OpenCV can read OpenEXR files.
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+
+
+def fov_to_intrinsic_matrix(width, height, fov_deg, fov_type="horizontal"):
+    """
+    Converts field of view (FOV) in degrees to a camera intrinsic matrix.
+
+    Args:
+        width (int): Image width in pixels.
+        height (int): Image height in pixels.
+        fov_deg (float): Field of view in degrees.
+        fov_type (str): 'horizontal' or 'vertical'; determines which FOV is used.
+
+    Returns:
+        np.ndarray: A 3x3 camera intrinsic matrix.
+
+    Raises:
+        ValueError: If width or height is non-positive or if fov_deg is not in (0, 180).
+    """
+    if width <= 0 or height <= 0:
+        raise ValueError("Image width and height must be positive numbers.")
+    if not (0 < fov_deg < 180):
+        raise ValueError("FOV must be between 0 and 180 degrees (non-inclusive).")
+    if fov_type not in ["horizontal", "vertical"]:
+        raise ValueError("fov_type must be either 'horizontal' or 'vertical'.")
+
+    fov_rad = np.deg2rad(fov_deg)
+
+    if fov_type == "horizontal":
+        f_x = width / (2 * np.tan(fov_rad / 2))
+        aspect_ratio = height / width
+        f_y = f_x * aspect_ratio
+    else:
+        f_y = height / (2 * np.tan(fov_rad / 2))
+        aspect_ratio = width / height
+        f_x = f_y * aspect_ratio
+
+    c_x = width / 2
+    c_y = height / 2
+    K = np.array([[f_x, 0, c_x], [0, f_y, c_y], [0, 0, 1]])
+    return K
+
+
+def process_basename(root, seq, basename, view_types, out_dir):
+    """
+    Processes a single basename: selects a random view type, copies the corresponding
+    image and depth file, and computes the camera intrinsics from the JSON metadata.
+
+    Args:
+        root (str): Root directory of the raw data.
+        seq (str): Sequence directory name.
+        basename (str): Basename (common identifier) for the files.
+        view_types (list): List of view types to choose from (e.g. ['bl', 'br', 'tl', 'tr']).
+        out_dir (str): Output directory where processed data will be saved.
+
+    Returns:
+        str or None: Returns an error message string on failure; otherwise, returns None.
+    """
+    # Select a random view type.
+    view_type = random.choice(view_types)
+
+    imgname = f"{basename}-{view_type}-image.png"
+    depthname = f"{basename}-{view_type}-depth.exr"
+
+    img_path = os.path.join(root, seq, imgname)
+    cam_path = os.path.join(root, seq, f"{basename}-meta.json")
+    depth_path = os.path.join(root, f"{seq}-depth", depthname)
+
+    # Prepare output directories.
+    out_seq_dir = os.path.join(out_dir, seq)
+    out_rgb_dir = os.path.join(out_seq_dir, "rgb")
+    out_depth_dir = os.path.join(out_seq_dir, "depth")
+    out_cam_dir = os.path.join(out_seq_dir, "cam")
+
+    # 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}.exr")
+    out_cam_path = os.path.join(out_cam_dir, f"{basename}.npz")
+
+    try:
+        # Load image using PIL and save as PNG.
+        with Image.open(img_path) as img:
+            W, H = img.size
+            img.save(out_img_path, format="PNG")
+
+        # Load camera JSON metadata.
+        with open(cam_path, "r") as f:
+            cam = json.load(f)
+        fov = cam["fltFov"]
+        K = fov_to_intrinsic_matrix(W, H, fov)
+
+        # Copy depth file.
+        shutil.copy(depth_path, out_depth_path)
+
+        # Save camera intrinsics.
+        np.savez(out_cam_path, intrinsics=K)
+
+    except Exception as e:
+        return f"Error processing {seq}/{basename}: {e}"
+
+    return None  # Success indicator
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Process raw 3D Ken Burns video data and generate processed images, depth maps, and camera intrinsics."
+    )
+    parser.add_argument(
+        "--root", type=str, required=True, help="Root directory of the raw data."
+    )
+    parser.add_argument(
+        "--out_dir",
+        type=str,
+        required=True,
+        help="Output directory for processed data.",
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=None,
+        help="Number of worker processes to use (default: half of available CPUs).",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed for reproducibility (default: 42).",
+    )
+    parser.add_argument(
+        "--view_types",
+        type=str,
+        nargs="+",
+        default=["bl", "br", "tl", "tr"],
+        help="List of view types to choose from (default: bl br tl tr).",
+    )
+    args = parser.parse_args()
+
+    # Set the random seed.
+    random.seed(args.seed)
+
+    root = args.root
+    out_dir = args.out_dir
+    view_types = args.view_types
+
+    # Determine number of worker processes.
+    num_workers = (
+        args.num_workers if args.num_workers is not None else (os.cpu_count() or 4) // 2
+    )
+
+    # Collect all sequence directories from root.
+    seq_dirs = [
+        d
+        for d in os.listdir(root)
+        if os.path.isdir(os.path.join(root, d)) and not d.endswith("-depth")
+    ]
+
+    # Pre-create output directory structure.
+    for seq in seq_dirs:
+        for subfolder in ["rgb", "depth", "cam"]:
+            (Path(out_dir) / seq / subfolder).mkdir(parents=True, exist_ok=True)
+
+    # Prepare list of tasks.
+    tasks = []
+    for seq in seq_dirs:
+        seq_path = os.path.join(root, seq)
+        # Assume JSON files contain metadata and have a name ending with "-meta.json".
+        json_files = [f for f in os.listdir(seq_path) if f.endswith(".json")]
+        # Remove the trailing "-meta.json" (10 characters) to get the basename.
+        basenames = sorted([f[:-10] for f in json_files])
+        for basename in basenames:
+            tasks.append((seq, basename))
+
+    # Define a partial function with fixed root, view_types, and out_dir.
+    process_func = partial(
+        process_basename, root, view_types=view_types, out_dir=out_dir
+    )
+
+    # Process tasks in parallel using ProcessPoolExecutor.
+    with ProcessPoolExecutor(max_workers=num_workers) as executor:
+        futures = {
+            executor.submit(process_func, seq, basename): (seq, basename)
+            for seq, basename in tasks
+        }
+        for future in tqdm(
+            as_completed(futures), total=len(futures), desc="Processing"
+        ):
+            error = future.result()
+            if error:
+                print(error)
+
+
+if __name__ == "__main__":
+    main()

extern/CUT3R/datasets_preprocess/preprocess_arkitscenes.py

@@ -0,0 +1,445 @@
+import os
+import json
+import os.path as osp
+import decimal
+import argparse
+import math
+from bisect import bisect_left
+from PIL import Image
+import numpy as np
+import quaternion
+from scipy import interpolate
+import cv2
+from tqdm import tqdm
+
+
+def get_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--arkitscenes_dir",
+        default="data/dust3r_data/data_arkitscenes/raw",
+    )
+    parser.add_argument(
+        "--precomputed_pairs",
+        default="data/dust3r_data/data_arkitscenes/arkitscenes_pairs",
+    )
+    parser.add_argument(
+        "--output_dir",
+        default="data/dust3r_data/processed_arkitscenes",
+    )
+    return parser
+
+
+def value_to_decimal(value, decimal_places):
+    decimal.getcontext().rounding = decimal.ROUND_HALF_UP  # define rounding method
+    return decimal.Decimal(str(float(value))).quantize(
+        decimal.Decimal("1e-{}".format(decimal_places))
+    )
+
+
+def closest(value, sorted_list):
+    index = bisect_left(sorted_list, value)
+    if index == 0:
+        return sorted_list[0]
+    elif index == len(sorted_list):
+        return sorted_list[-1]
+    else:
+        value_before = sorted_list[index - 1]
+        value_after = sorted_list[index]
+        if value_after - value < value - value_before:
+            return value_after
+        else:
+            return value_before
+
+
+def get_up_vectors(pose_device_to_world):
+    return np.matmul(pose_device_to_world, np.array([[0.0], [-1.0], [0.0], [0.0]]))
+
+
+def get_right_vectors(pose_device_to_world):
+    return np.matmul(pose_device_to_world, np.array([[1.0], [0.0], [0.0], [0.0]]))
+
+
+def read_traj(traj_path):
+    quaternions = []
+    poses = []
+    timestamps = []
+    poses_p_to_w = []
+    with open(traj_path) as f:
+        traj_lines = f.readlines()
+        for line in traj_lines:
+            tokens = line.split()
+            assert len(tokens) == 7
+            traj_timestamp = float(tokens[0])
+
+            timestamps_decimal_value = value_to_decimal(traj_timestamp, 3)
+            timestamps.append(
+                float(timestamps_decimal_value)
+            )  # for spline interpolation
+
+            angle_axis = [float(tokens[1]), float(tokens[2]), float(tokens[3])]
+            r_w_to_p, _ = cv2.Rodrigues(np.asarray(angle_axis))
+            t_w_to_p = np.asarray(
+                [float(tokens[4]), float(tokens[5]), float(tokens[6])]
+            )
+
+            pose_w_to_p = np.eye(4)
+            pose_w_to_p[:3, :3] = r_w_to_p
+            pose_w_to_p[:3, 3] = t_w_to_p
+
+            pose_p_to_w = np.linalg.inv(pose_w_to_p)
+
+            r_p_to_w_as_quat = quaternion.from_rotation_matrix(pose_p_to_w[:3, :3])
+            t_p_to_w = pose_p_to_w[:3, 3]
+            poses_p_to_w.append(pose_p_to_w)
+            poses.append(t_p_to_w)
+            quaternions.append(r_p_to_w_as_quat)
+    return timestamps, poses, quaternions, poses_p_to_w
+
+
+def main(rootdir, pairsdir, outdir):
+    os.makedirs(outdir, exist_ok=True)
+
+    subdirs = ["Test", "Training"]
+    for subdir in subdirs:
+        # STEP 1: list all scenes
+        outsubdir = osp.join(outdir, subdir)
+        os.makedirs(outsubdir, exist_ok=True)
+        listfile = osp.join(pairsdir, subdir, "scene_list.json")
+        with open(listfile, "r") as f:
+            scene_dirs = json.load(f)
+
+        valid_scenes = []
+        for scene_subdir in tqdm(scene_dirs):
+            if not os.path.isdir(osp.join(rootdir, "Test", scene_subdir)):
+                if not os.path.isdir(osp.join(rootdir, "Training", scene_subdir)):
+                    continue
+                else:
+                    root_subdir = "Training"
+            else:
+                root_subdir = "Test"
+            out_scene_subdir = osp.join(outsubdir, scene_subdir)
+            os.makedirs(out_scene_subdir, exist_ok=True)
+
+            scene_dir = osp.join(rootdir, root_subdir, scene_subdir)
+            depth_dir = osp.join(scene_dir, "lowres_depth")
+            rgb_dir = osp.join(scene_dir, "vga_wide")
+            intrinsics_dir = osp.join(scene_dir, "vga_wide_intrinsics")
+            traj_path = osp.join(scene_dir, "lowres_wide.traj")
+
+            # STEP 2: read selected_pairs.npz
+            selected_pairs_path = osp.join(
+                pairsdir, subdir, scene_subdir, "selected_pairs.npz"
+            )
+            selected_npz = np.load(selected_pairs_path)
+            selection, pairs = selected_npz["selection"], selected_npz["pairs"]
+            selected_sky_direction_scene = str(selected_npz["sky_direction_scene"][0])
+            if len(selection) == 0 or len(pairs) == 0:
+                # not a valid scene
+                continue
+            valid_scenes.append(scene_subdir)
+
+            # STEP 3: parse the scene and export the list of valid (K, pose, rgb, depth) and convert images
+            scene_metadata_path = osp.join(out_scene_subdir, "scene_metadata.npz")
+            if osp.isfile(scene_metadata_path):
+                continue
+            else:
+                print(f"parsing {scene_subdir}")
+                # loads traj
+                timestamps, poses, quaternions, poses_cam_to_world = read_traj(
+                    traj_path
+                )
+
+                poses = np.array(poses)
+                quaternions = np.array(quaternions, dtype=np.quaternion)
+                quaternions = quaternion.unflip_rotors(quaternions)
+                timestamps = np.array(timestamps)
+
+                selected_images = [
+                    (basename, basename.split(".png")[0].split("_")[1])
+                    for basename in selection
+                ]
+                timestamps_selected = [
+                    float(frame_id) for _, frame_id in selected_images
+                ]
+
+                sky_direction_scene, trajectories, intrinsics, images = (
+                    convert_scene_metadata(
+                        scene_subdir,
+                        intrinsics_dir,
+                        timestamps,
+                        quaternions,
+                        poses,
+                        poses_cam_to_world,
+                        selected_images,
+                        timestamps_selected,
+                    )
+                )
+                assert selected_sky_direction_scene == sky_direction_scene
+
+                os.makedirs(os.path.join(out_scene_subdir, "vga_wide"), exist_ok=True)
+                os.makedirs(
+                    os.path.join(out_scene_subdir, "lowres_depth"), exist_ok=True
+                )
+                assert isinstance(sky_direction_scene, str)
+                all_exist = True
+                for basename in images:
+                    vga_wide_path = osp.join(rgb_dir, basename)
+                    depth_path = osp.join(depth_dir, basename)
+                    if not osp.isfile(vga_wide_path) or not osp.isfile(depth_path):
+                        all_exist = False
+                        break
+                if not all_exist:
+                    continue
+
+                for basename in images:
+                    img_out = os.path.join(
+                        out_scene_subdir, "vga_wide", basename.replace(".png", ".jpg")
+                    )
+                    depth_out = os.path.join(out_scene_subdir, "lowres_depth", basename)
+                    if osp.isfile(img_out) and osp.isfile(depth_out):
+                        continue
+
+                    vga_wide_path = osp.join(rgb_dir, basename)
+                    depth_path = osp.join(depth_dir, basename)
+
+                    img = Image.open(vga_wide_path)
+                    depth = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED)
+
+                    # rotate the image
+                    if sky_direction_scene == "RIGHT":
+                        try:
+                            img = img.transpose(Image.Transpose.ROTATE_90)
+                        except Exception:
+                            img = img.transpose(Image.ROTATE_90)
+                        depth = cv2.rotate(depth, cv2.ROTATE_90_COUNTERCLOCKWISE)
+                    elif sky_direction_scene == "LEFT":
+                        try:
+                            img = img.transpose(Image.Transpose.ROTATE_270)
+                        except Exception:
+                            img = img.transpose(Image.ROTATE_270)
+                        depth = cv2.rotate(depth, cv2.ROTATE_90_CLOCKWISE)
+                    elif sky_direction_scene == "DOWN":
+                        try:
+                            img = img.transpose(Image.Transpose.ROTATE_180)
+                        except Exception:
+                            img = img.transpose(Image.ROTATE_180)
+                        depth = cv2.rotate(depth, cv2.ROTATE_180)
+
+                    W, H = img.size
+                    if not osp.isfile(img_out):
+                        img.save(img_out)
+
+                    depth = cv2.resize(
+                        depth, (W, H), interpolation=cv2.INTER_NEAREST_EXACT
+                    )
+                    if not osp.isfile(
+                        depth_out
+                    ):  # avoid destroying the base dataset when you mess up the paths
+                        cv2.imwrite(depth_out, depth)
+
+                # save at the end
+                np.savez(
+                    scene_metadata_path,
+                    trajectories=trajectories,
+                    intrinsics=intrinsics,
+                    images=images,
+                    pairs=pairs,
+                )
+
+        outlistfile = osp.join(outsubdir, "scene_list.json")
+        for scene_subdir in valid_scenes:
+            scene_metadata_path = osp.join(
+                outsubdir, scene_subdir, "scene_metadata.npz"
+            )
+            if not osp.isfile(scene_metadata_path):
+                valid_scenes.remove(scene_subdir)
+        with open(outlistfile, "w") as f:
+            json.dump(valid_scenes, f)
+
+        # STEP 5: concat all scene_metadata.npz into a single file
+        scene_data = {}
+        for scene_subdir in valid_scenes:
+            scene_metadata_path = osp.join(
+                outsubdir, scene_subdir, "scene_metadata.npz"
+            )
+            with np.load(scene_metadata_path) as data:
+                trajectories = data["trajectories"]
+                intrinsics = data["intrinsics"]
+                images = data["images"]
+                pairs = data["pairs"]
+            scene_data[scene_subdir] = {
+                "trajectories": trajectories,
+                "intrinsics": intrinsics,
+                "images": images,
+                "pairs": pairs,
+            }
+        offset = 0
+        counts = []
+        scenes = []
+        sceneids = []
+        images = []
+        intrinsics = []
+        trajectories = []
+        pairs = []
+        for scene_idx, (scene_subdir, data) in enumerate(scene_data.items()):
+            num_imgs = data["images"].shape[0]
+            img_pairs = data["pairs"]
+
+            scenes.append(scene_subdir)
+            sceneids.extend([scene_idx] * num_imgs)
+
+            images.append(data["images"])
+
+            K = np.expand_dims(np.eye(3), 0).repeat(num_imgs, 0)
+            K[:, 0, 0] = [fx for _, _, fx, _, _, _ in data["intrinsics"]]
+            K[:, 1, 1] = [fy for _, _, _, fy, _, _ in data["intrinsics"]]
+            K[:, 0, 2] = [hw for _, _, _, _, hw, _ in data["intrinsics"]]
+            K[:, 1, 2] = [hh for _, _, _, _, _, hh in data["intrinsics"]]
+
+            intrinsics.append(K)
+            trajectories.append(data["trajectories"])
+
+            # offset pairs
+            img_pairs[:, 0:2] += offset
+            pairs.append(img_pairs)
+            counts.append(offset)
+
+            offset += num_imgs
+
+        images = np.concatenate(images, axis=0)
+        intrinsics = np.concatenate(intrinsics, axis=0)
+        trajectories = np.concatenate(trajectories, axis=0)
+        pairs = np.concatenate(pairs, axis=0)
+        np.savez(
+            osp.join(outsubdir, "all_metadata.npz"),
+            counts=counts,
+            scenes=scenes,
+            sceneids=sceneids,
+            images=images,
+            intrinsics=intrinsics,
+            trajectories=trajectories,
+            pairs=pairs,
+        )
+
+
+def convert_scene_metadata(
+    scene_subdir,
+    intrinsics_dir,
+    timestamps,
+    quaternions,
+    poses,
+    poses_cam_to_world,
+    selected_images,
+    timestamps_selected,
+):
+    # find scene orientation
+    sky_direction_scene, rotated_to_cam = find_scene_orientation(poses_cam_to_world)
+
+    # find/compute pose for selected timestamps
+    # most images have a valid timestamp / exact pose associated
+    timestamps_selected = np.array(timestamps_selected)
+    spline = interpolate.interp1d(timestamps, poses, kind="linear", axis=0)
+    interpolated_rotations = quaternion.squad(
+        quaternions, timestamps, timestamps_selected
+    )
+    interpolated_positions = spline(timestamps_selected)
+
+    trajectories = []
+    intrinsics = []
+    images = []
+    for i, (basename, frame_id) in enumerate(selected_images):
+        intrinsic_fn = osp.join(intrinsics_dir, f"{scene_subdir}_{frame_id}.pincam")
+        if not osp.exists(intrinsic_fn):
+            intrinsic_fn = osp.join(
+                intrinsics_dir, f"{scene_subdir}_{float(frame_id) - 0.001:.3f}.pincam"
+            )
+        if not osp.exists(intrinsic_fn):
+            intrinsic_fn = osp.join(
+                intrinsics_dir, f"{scene_subdir}_{float(frame_id) + 0.001:.3f}.pincam"
+            )
+        assert osp.exists(intrinsic_fn)
+        w, h, fx, fy, hw, hh = np.loadtxt(intrinsic_fn)  # PINHOLE
+
+        pose = np.eye(4)
+        pose[:3, :3] = quaternion.as_rotation_matrix(interpolated_rotations[i])
+        pose[:3, 3] = interpolated_positions[i]
+
+        images.append(basename)
+        if sky_direction_scene == "RIGHT" or sky_direction_scene == "LEFT":
+            intrinsics.append([h, w, fy, fx, hh, hw])  # swapped intrinsics
+        else:
+            intrinsics.append([w, h, fx, fy, hw, hh])
+        trajectories.append(
+            pose @ rotated_to_cam
+        )  # pose_cam_to_world @ rotated_to_cam = rotated(cam) to world
+
+    return sky_direction_scene, trajectories, intrinsics, images
+
+
+def find_scene_orientation(poses_cam_to_world):
+    if len(poses_cam_to_world) > 0:
+        up_vector = sum(get_up_vectors(p) for p in poses_cam_to_world) / len(
+            poses_cam_to_world
+        )
+        right_vector = sum(get_right_vectors(p) for p in poses_cam_to_world) / len(
+            poses_cam_to_world
+        )
+        up_world = np.array([[0.0], [0.0], [1.0], [0.0]])
+    else:
+        up_vector = np.array([[0.0], [-1.0], [0.0], [0.0]])
+        right_vector = np.array([[1.0], [0.0], [0.0], [0.0]])
+        up_world = np.array([[0.0], [0.0], [1.0], [0.0]])
+
+    # value between 0, 180
+    device_up_to_world_up_angle = (
+        np.arccos(np.clip(np.dot(np.transpose(up_world), up_vector), -1.0, 1.0)).item()
+        * 180.0
+        / np.pi
+    )
+    device_right_to_world_up_angle = (
+        np.arccos(
+            np.clip(np.dot(np.transpose(up_world), right_vector), -1.0, 1.0)
+        ).item()
+        * 180.0
+        / np.pi
+    )
+
+    up_closest_to_90 = abs(device_up_to_world_up_angle - 90.0) < abs(
+        device_right_to_world_up_angle - 90.0
+    )
+    if up_closest_to_90:
+        assert abs(device_up_to_world_up_angle - 90.0) < 45.0
+        # LEFT
+        if device_right_to_world_up_angle > 90.0:
+            sky_direction_scene = "LEFT"
+            cam_to_rotated_q = quaternion.from_rotation_vector(
+                [0.0, 0.0, math.pi / 2.0]
+            )
+        else:
+            # note that in metadata.csv RIGHT does not exist, but again it's not accurate...
+            # well, turns out there are scenes oriented like this
+            # for example Training/41124801
+            sky_direction_scene = "RIGHT"
+            cam_to_rotated_q = quaternion.from_rotation_vector(
+                [0.0, 0.0, -math.pi / 2.0]
+            )
+    else:
+        # right is close to 90
+        assert abs(device_right_to_world_up_angle - 90.0) < 45.0
+        if device_up_to_world_up_angle > 90.0:
+            sky_direction_scene = "DOWN"
+            cam_to_rotated_q = quaternion.from_rotation_vector([0.0, 0.0, math.pi])
+        else:
+            sky_direction_scene = "UP"
+            cam_to_rotated_q = quaternion.quaternion(1, 0, 0, 0)
+    cam_to_rotated = np.eye(4)
+    cam_to_rotated[:3, :3] = quaternion.as_rotation_matrix(cam_to_rotated_q)
+    rotated_to_cam = np.linalg.inv(cam_to_rotated)
+    return sky_direction_scene, rotated_to_cam
+
+
+if __name__ == "__main__":
+    parser = get_parser()
+    args = parser.parse_args()
+    main(args.arkitscenes_dir, args.precomputed_pairs, args.output_dir)

extern/CUT3R/datasets_preprocess/preprocess_arkitscenes_highres.py

@@ -0,0 +1,409 @@
+import os
+import json
+import os.path as osp
+import decimal
+import argparse
+import math
+from bisect import bisect_left
+from PIL import Image
+import numpy as np
+import quaternion
+from scipy import interpolate
+import cv2
+from tqdm import tqdm
+from multiprocessing import Pool
+
+
+def get_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--arkitscenes_dir",
+        default="",
+    )
+    parser.add_argument(
+        "--output_dir",
+        default="data/dust3r_data/processed_arkitscenes_highres",
+    )
+    return parser
+
+
+def value_to_decimal(value, decimal_places):
+    decimal.getcontext().rounding = decimal.ROUND_HALF_UP  # define rounding method
+    return decimal.Decimal(str(float(value))).quantize(
+        decimal.Decimal("1e-{}".format(decimal_places))
+    )
+
+
+def closest(value, sorted_list):
+    index = bisect_left(sorted_list, value)
+    if index == 0:
+        return sorted_list[0]
+    elif index == len(sorted_list):
+        return sorted_list[-1]
+    else:
+        value_before = sorted_list[index - 1]
+        value_after = sorted_list[index]
+        if value_after - value < value - value_before:
+            return value_after
+        else:
+            return value_before
+
+
+def get_up_vectors(pose_device_to_world):
+    return np.matmul(pose_device_to_world, np.array([[0.0], [-1.0], [0.0], [0.0]]))
+
+
+def get_right_vectors(pose_device_to_world):
+    return np.matmul(pose_device_to_world, np.array([[1.0], [0.0], [0.0], [0.0]]))
+
+
+def read_traj(traj_path):
+    quaternions = []
+    poses = []
+    timestamps = []
+    poses_p_to_w = []
+    with open(traj_path) as f:
+        traj_lines = f.readlines()
+        for line in traj_lines:
+            tokens = line.split()
+            assert len(tokens) == 7
+            traj_timestamp = float(tokens[0])
+
+            timestamps_decimal_value = value_to_decimal(traj_timestamp, 3)
+            timestamps.append(
+                float(timestamps_decimal_value)
+            )  # for spline interpolation
+
+            angle_axis = [float(tokens[1]), float(tokens[2]), float(tokens[3])]
+            r_w_to_p, _ = cv2.Rodrigues(np.asarray(angle_axis))
+            t_w_to_p = np.asarray(
+                [float(tokens[4]), float(tokens[5]), float(tokens[6])]
+            )
+
+            pose_w_to_p = np.eye(4)
+            pose_w_to_p[:3, :3] = r_w_to_p
+            pose_w_to_p[:3, 3] = t_w_to_p
+
+            pose_p_to_w = np.linalg.inv(pose_w_to_p)
+
+            r_p_to_w_as_quat = quaternion.from_rotation_matrix(pose_p_to_w[:3, :3])
+            t_p_to_w = pose_p_to_w[:3, 3]
+            poses_p_to_w.append(pose_p_to_w)
+            poses.append(t_p_to_w)
+            quaternions.append(r_p_to_w_as_quat)
+    return timestamps, poses, quaternions, poses_p_to_w
+
+
+def main(rootdir, outdir):
+    os.makedirs(outdir, exist_ok=True)
+    subdirs = ["Validation", "Training"]
+    for subdir in subdirs:
+        outsubdir = osp.join(outdir, subdir)
+        scene_dirs = sorted(
+            [
+                d
+                for d in os.listdir(osp.join(rootdir, subdir))
+                if osp.isdir(osp.join(rootdir, subdir, d))
+            ]
+        )
+
+        with Pool() as pool:
+            results = list(
+                tqdm(
+                    pool.imap(
+                        process_scene,
+                        [
+                            (rootdir, outdir, subdir, scene_subdir)
+                            for scene_subdir in scene_dirs
+                        ],
+                    ),
+                    total=len(scene_dirs),
+                )
+            )
+
+        # Filter None results and other post-processing
+        valid_scenes = [result for result in results if result is not None]
+        outlistfile = osp.join(outsubdir, "scene_list.json")
+        with open(outlistfile, "w") as f:
+            json.dump(valid_scenes, f)
+
+
+def process_scene(args):
+    rootdir, outdir, subdir, scene_subdir = args
+    # Unpack paths
+    scene_dir = osp.join(rootdir, subdir, scene_subdir)
+    outsubdir = osp.join(outdir, subdir)
+    out_scene_subdir = osp.join(outsubdir, scene_subdir)
+
+    # Validation if necessary resources exist
+    if (
+        not osp.exists(osp.join(scene_dir, "highres_depth"))
+        or not osp.exists(osp.join(scene_dir, "vga_wide"))
+        or not osp.exists(osp.join(scene_dir, "vga_wide_intrinsics"))
+        or not osp.exists(osp.join(scene_dir, "lowres_wide.traj"))
+    ):
+        return None
+
+    depth_dir = osp.join(scene_dir, "highres_depth")
+    rgb_dir = osp.join(scene_dir, "vga_wide")
+    intrinsics_dir = osp.join(scene_dir, "vga_wide_intrinsics")
+    traj_path = osp.join(scene_dir, "lowres_wide.traj")
+
+    depth_files = sorted(os.listdir(depth_dir))
+    img_files = sorted(os.listdir(rgb_dir))
+
+    out_scene_subdir = osp.join(outsubdir, scene_subdir)
+
+    # STEP 3: parse the scene and export the list of valid (K, pose, rgb, depth) and convert images
+    scene_metadata_path = osp.join(out_scene_subdir, "scene_metadata.npz")
+    if osp.isfile(scene_metadata_path):
+        print(f"Skipping {scene_subdir}")
+    else:
+        print(f"parsing {scene_subdir}")
+        # loads traj
+        timestamps, poses, quaternions, poses_cam_to_world = read_traj(traj_path)
+
+        poses = np.array(poses)
+        quaternions = np.array(quaternions, dtype=np.quaternion)
+        quaternions = quaternion.unflip_rotors(quaternions)
+        timestamps = np.array(timestamps)
+
+        all_depths = sorted(
+            [
+                (basename, basename.split(".png")[0].split("_")[1])
+                for basename in depth_files
+            ],
+            key=lambda x: float(x[1]),
+        )
+
+        selected_depths = []
+        timestamps_selected = []
+        timestamp_min = timestamps.min()
+        timestamp_max = timestamps.max()
+        for basename, frame_id in all_depths:
+            frame_id = float(frame_id)
+            if frame_id < timestamp_min or frame_id > timestamp_max:
+                continue
+            selected_depths.append((basename, frame_id))
+            timestamps_selected.append(frame_id)
+
+        sky_direction_scene, trajectories, intrinsics, images, depths = (
+            convert_scene_metadata(
+                scene_subdir,
+                intrinsics_dir,
+                timestamps,
+                quaternions,
+                poses,
+                poses_cam_to_world,
+                img_files,
+                selected_depths,
+                timestamps_selected,
+            )
+        )
+
+        if len(images) == 0:
+            print(f"Skipping {scene_subdir}")
+            return None
+
+        os.makedirs(out_scene_subdir, exist_ok=True)
+
+        os.makedirs(os.path.join(out_scene_subdir, "vga_wide"), exist_ok=True)
+        os.makedirs(os.path.join(out_scene_subdir, "highres_depth"), exist_ok=True)
+        assert isinstance(sky_direction_scene, str)
+
+        for image_path, depth_path in zip(images, depths):
+            img_out = os.path.join(
+                out_scene_subdir, "vga_wide", image_path.replace(".png", ".jpg")
+            )
+            depth_out = os.path.join(out_scene_subdir, "highres_depth", depth_path)
+            if osp.isfile(img_out) and osp.isfile(depth_out):
+                continue
+
+            vga_wide_path = osp.join(rgb_dir, image_path)
+            depth_path = osp.join(depth_dir, depth_path)
+
+            if not osp.isfile(vga_wide_path) or not osp.isfile(depth_path):
+                continue
+
+            img = Image.open(vga_wide_path)
+            depth = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED)
+
+            # rotate the image
+            if sky_direction_scene == "RIGHT":
+                try:
+                    img = img.transpose(Image.Transpose.ROTATE_90)
+                except Exception:
+                    img = img.transpose(Image.ROTATE_90)
+                depth = cv2.rotate(depth, cv2.ROTATE_90_COUNTERCLOCKWISE)
+
+            elif sky_direction_scene == "LEFT":
+                try:
+                    img = img.transpose(Image.Transpose.ROTATE_270)
+                except Exception:
+                    img = img.transpose(Image.ROTATE_270)
+                depth = cv2.rotate(depth, cv2.ROTATE_90_CLOCKWISE)
+
+            elif sky_direction_scene == "DOWN":
+                try:
+                    img = img.transpose(Image.Transpose.ROTATE_180)
+                except Exception:
+                    img = img.transpose(Image.ROTATE_180)
+                depth = cv2.rotate(depth, cv2.ROTATE_180)
+
+            W, H = img.size
+            if not osp.isfile(img_out):
+                img.save(img_out)
+
+            depth = cv2.resize(depth, (W, H), interpolation=cv2.INTER_NEAREST)
+            if not osp.isfile(
+                depth_out
+            ):  # avoid destroying the base dataset when you mess up the paths
+                cv2.imwrite(depth_out, depth)
+
+        # save at the end
+        np.savez(
+            scene_metadata_path,
+            trajectories=trajectories,
+            intrinsics=intrinsics,
+            images=images,
+        )
+
+
+def convert_scene_metadata(
+    scene_subdir,
+    intrinsics_dir,
+    timestamps,
+    quaternions,
+    poses,
+    poses_cam_to_world,
+    all_images,
+    selected_depths,
+    timestamps_selected,
+):
+    # find scene orientation
+    sky_direction_scene, rotated_to_cam = find_scene_orientation(poses_cam_to_world)
+
+    # find/compute pose for selected timestamps
+    # most images have a valid timestamp / exact pose associated
+    timestamps_selected = np.array(timestamps_selected)
+    spline = interpolate.interp1d(timestamps, poses, kind="linear", axis=0)
+    interpolated_rotations = quaternion.squad(
+        quaternions, timestamps, timestamps_selected
+    )
+    interpolated_positions = spline(timestamps_selected)
+
+    trajectories = []
+    intrinsics = []
+    images = []
+    depths = []
+    for i, (basename, frame_id) in enumerate(selected_depths):
+        intrinsic_fn = osp.join(intrinsics_dir, f"{scene_subdir}_{frame_id}.pincam")
+        search_interval = int(0.1 / 0.001)
+        for timestamp in range(-search_interval, search_interval + 1):
+            if osp.exists(intrinsic_fn):
+                break
+            intrinsic_fn = osp.join(
+                intrinsics_dir,
+                f"{scene_subdir}_{float(frame_id) + timestamp * 0.001:.3f}.pincam",
+            )
+        if not osp.exists(intrinsic_fn):
+            print(f"Skipping {intrinsic_fn}")
+            continue
+
+        image_path = "{}_{}.png".format(scene_subdir, frame_id)
+        search_interval = int(0.001 / 0.001)
+        for timestamp in range(-search_interval, search_interval + 1):
+            if image_path in all_images:
+                break
+            image_path = "{}_{}.png".format(
+                scene_subdir, float(frame_id) + timestamp * 0.001
+            )
+        if image_path not in all_images:
+            print(f"Skipping {scene_subdir} {frame_id}")
+            continue
+
+        w, h, fx, fy, hw, hh = np.loadtxt(intrinsic_fn)  # PINHOLE
+
+        pose = np.eye(4)
+        pose[:3, :3] = quaternion.as_rotation_matrix(interpolated_rotations[i])
+        pose[:3, 3] = interpolated_positions[i]
+
+        images.append(basename)
+        depths.append(basename)
+        if sky_direction_scene == "RIGHT" or sky_direction_scene == "LEFT":
+            intrinsics.append([h, w, fy, fx, hh, hw])  # swapped intrinsics
+        else:
+            intrinsics.append([w, h, fx, fy, hw, hh])
+        trajectories.append(
+            pose @ rotated_to_cam
+        )  # pose_cam_to_world @ rotated_to_cam = rotated(cam) to world
+
+    return sky_direction_scene, trajectories, intrinsics, images, depths
+
+
+def find_scene_orientation(poses_cam_to_world):
+    if len(poses_cam_to_world) > 0:
+        up_vector = sum(get_up_vectors(p) for p in poses_cam_to_world) / len(
+            poses_cam_to_world
+        )
+        right_vector = sum(get_right_vectors(p) for p in poses_cam_to_world) / len(
+            poses_cam_to_world
+        )
+        up_world = np.array([[0.0], [0.0], [1.0], [0.0]])
+    else:
+        up_vector = np.array([[0.0], [-1.0], [0.0], [0.0]])
+        right_vector = np.array([[1.0], [0.0], [0.0], [0.0]])
+        up_world = np.array([[0.0], [0.0], [1.0], [0.0]])
+
+    # value between 0, 180
+    device_up_to_world_up_angle = (
+        np.arccos(np.clip(np.dot(np.transpose(up_world), up_vector), -1.0, 1.0)).item()
+        * 180.0
+        / np.pi
+    )
+    device_right_to_world_up_angle = (
+        np.arccos(
+            np.clip(np.dot(np.transpose(up_world), right_vector), -1.0, 1.0)
+        ).item()
+        * 180.0
+        / np.pi
+    )
+
+    up_closest_to_90 = abs(device_up_to_world_up_angle - 90.0) < abs(
+        device_right_to_world_up_angle - 90.0
+    )
+    if up_closest_to_90:
+        assert abs(device_up_to_world_up_angle - 90.0) < 45.0
+        # LEFT
+        if device_right_to_world_up_angle > 90.0:
+            sky_direction_scene = "LEFT"
+            cam_to_rotated_q = quaternion.from_rotation_vector(
+                [0.0, 0.0, math.pi / 2.0]
+            )
+        else:
+            # note that in metadata.csv RIGHT does not exist, but again it's not accurate...
+            # well, turns out there are scenes oriented like this
+            # for example Training/41124801
+            sky_direction_scene = "RIGHT"
+            cam_to_rotated_q = quaternion.from_rotation_vector(
+                [0.0, 0.0, -math.pi / 2.0]
+            )
+    else:
+        # right is close to 90
+        assert abs(device_right_to_world_up_angle - 90.0) < 45.0
+        if device_up_to_world_up_angle > 90.0:
+            sky_direction_scene = "DOWN"
+            cam_to_rotated_q = quaternion.from_rotation_vector([0.0, 0.0, math.pi])
+        else:
+            sky_direction_scene = "UP"
+            cam_to_rotated_q = quaternion.quaternion(1, 0, 0, 0)
+    cam_to_rotated = np.eye(4)
+    cam_to_rotated[:3, :3] = quaternion.as_rotation_matrix(cam_to_rotated_q)
+    rotated_to_cam = np.linalg.inv(cam_to_rotated)
+    return sky_direction_scene, rotated_to_cam
+
+
+if __name__ == "__main__":
+    parser = get_parser()
+    args = parser.parse_args()
+    main(args.arkitscenes_dir, args.output_dir)

extern/CUT3R/datasets_preprocess/preprocess_bedlam.py

@@ -0,0 +1,402 @@
+#!/usr/bin/env python3
+"""
+Process Bedlam scenes by computing camera intrinsics and extrinsics
+from extracted data. The script reads per-scene CSV and image/depth files,
+computes the necessary camera parameters, and saves the resulting camera
+files (as .npz files) in an output directory.
+
+Usage:
+    python preprocess_bedlam.py --root /path/to/extracted_data \
+                             --outdir /path/to/processed_bedlam \
+                             [--num_workers 4]
+"""
+
+import os
+import cv2
+import numpy as np
+import pandas as pd
+from glob import glob
+import shutil
+import OpenEXR  # Ensure OpenEXR is installed
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from tqdm import tqdm
+import argparse
+
+# Enable OpenEXR support in OpenCV.
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+
+# Global constants
+IMG_FORMAT = ".png"
+rotate_flag = False
+SENSOR_W = 36
+SENSOR_H = 20.25
+IMG_W = 1280
+IMG_H = 720
+
+# -----------------------------------------------------------------------------
+# Helper functions for camera parameter conversion
+# -----------------------------------------------------------------------------
+
+
+def focalLength_mm2px(focalLength, dslr_sens, focalPoint):
+    focal_pixel = (focalLength / dslr_sens) * focalPoint * 2
+    return focal_pixel
+
+
+def get_cam_int(fl, sens_w, sens_h, cx, cy):
+    flx = focalLength_mm2px(fl, sens_w, cx)
+    fly = focalLength_mm2px(fl, sens_h, cy)
+    cam_mat = np.array([[flx, 0, cx], [0, fly, cy], [0, 0, 1]])
+    return cam_mat
+
+
+def unreal2cv2(points):
+    # Permute coordinates: x --> y, y --> z, z --> x
+    points = np.roll(points, 2, axis=1)
+    # Invert the y-axis
+    points = points * np.array([1.0, -1.0, 1.0])
+    return points
+
+
+def get_cam_trans(body_trans, cam_trans):
+    cam_trans = np.array(cam_trans) / 100
+    cam_trans = unreal2cv2(np.reshape(cam_trans, (1, 3)))
+    body_trans = np.array(body_trans) / 100
+    body_trans = unreal2cv2(np.reshape(body_trans, (1, 3)))
+    trans = body_trans - cam_trans
+    return trans
+
+
+def get_cam_rotmat(pitch, yaw, roll):
+    rotmat_yaw, _ = cv2.Rodrigues(np.array([[0, (yaw / 180) * np.pi, 0]], dtype=float))
+    rotmat_pitch, _ = cv2.Rodrigues(np.array([pitch / 180 * np.pi, 0, 0]).reshape(3, 1))
+    rotmat_roll, _ = cv2.Rodrigues(np.array([0, 0, roll / 180 * np.pi]).reshape(3, 1))
+    final_rotmat = rotmat_roll @ (rotmat_pitch @ rotmat_yaw)
+    return final_rotmat
+
+
+def get_global_orient(cam_pitch, cam_yaw, cam_roll):
+    pitch_rotmat, _ = cv2.Rodrigues(
+        np.array([cam_pitch / 180 * np.pi, 0, 0]).reshape(3, 1)
+    )
+    roll_rotmat, _ = cv2.Rodrigues(
+        np.array([0, 0, cam_roll / 180 * np.pi]).reshape(3, 1)
+    )
+    final_rotmat = roll_rotmat @ pitch_rotmat
+    return final_rotmat
+
+
+def convert_translation_to_opencv(x, y, z):
+    t_cv = np.array([y, -z, x])
+    return t_cv
+
+
+def rotation_matrix_unreal(yaw, pitch, roll):
+    yaw_rad = np.deg2rad(yaw)
+    pitch_rad = np.deg2rad(pitch)
+    roll_rad = np.deg2rad(roll)
+    # Yaw (left-handed)
+    R_yaw = np.array(
+        [
+            [np.cos(-yaw_rad), -np.sin(-yaw_rad), 0],
+            [np.sin(-yaw_rad), np.cos(-yaw_rad), 0],
+            [0, 0, 1],
+        ]
+    )
+    # Pitch (right-handed)
+    R_pitch = np.array(
+        [
+            [np.cos(pitch_rad), 0, np.sin(pitch_rad)],
+            [0, 1, 0],
+            [-np.sin(pitch_rad), 0, np.cos(pitch_rad)],
+        ]
+    )
+    # Roll (right-handed)
+    R_roll = np.array(
+        [
+            [1, 0, 0],
+            [0, np.cos(roll_rad), -np.sin(roll_rad)],
+            [0, np.sin(roll_rad), np.cos(roll_rad)],
+        ]
+    )
+    R_unreal = R_roll @ R_pitch @ R_yaw
+    return R_unreal
+
+
+def convert_rotation_to_opencv(R_unreal):
+    # Transformation matrix from Unreal to OpenCV coordinate system.
+    C = np.array([[0, 1, 0], [0, 0, -1], [1, 0, 0]])
+    R_cv = C @ R_unreal @ C.T
+    return R_cv
+
+
+def get_rot_unreal(yaw, pitch, roll):
+    yaw_rad = np.deg2rad(yaw)
+    pitch_rad = np.deg2rad(pitch)
+    roll_rad = np.deg2rad(roll)
+    R_yaw = np.array(
+        [
+            [np.cos(yaw_rad), -np.sin(yaw_rad), 0],
+            [np.sin(yaw_rad), np.cos(yaw_rad), 0],
+            [0, 0, 1],
+        ]
+    )
+    R_pitch = np.array(
+        [
+            [np.cos(pitch_rad), 0, -np.sin(pitch_rad)],
+            [0, 1, 0],
+            [np.sin(pitch_rad), 0, np.cos(pitch_rad)],
+        ]
+    )
+    R_roll = np.array(
+        [
+            [1, 0, 0],
+            [0, np.cos(roll_rad), np.sin(roll_rad)],
+            [0, -np.sin(roll_rad), np.cos(roll_rad)],
+        ]
+    )
+    R_unreal = R_yaw @ R_pitch @ R_roll
+    return R_unreal
+
+
+def get_extrinsics_unreal(R_unreal, t_unreal):
+    cam_trans = np.array(t_unreal)
+    ext = np.eye(4)
+    ext[:3, :3] = R_unreal
+    ext[:3, 3] = cam_trans.reshape(1, 3)
+    return ext
+
+
+def get_extrinsics_opencv(yaw, pitch, roll, x, y, z):
+    R_unreal = get_rot_unreal(yaw, pitch, roll)
+    t_unreal = np.array([x / 100.0, y / 100.0, z / 100.0])
+    T_u2wu = get_extrinsics_unreal(R_unreal, t_unreal)
+    T_opencv2unreal = np.array(
+        [[0, 0, -1, 0], [1, 0, 0, 0], [0, -1, 0, 0], [0, 0, 0, 1]], dtype=np.float32
+    )
+    T_wu2ou = np.array(
+        [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=np.float32
+    )
+    return np.linalg.inv(T_opencv2unreal @ T_u2wu @ T_wu2ou)
+
+
+# -----------------------------------------------------------------------------
+# Get camera parameters from the extracted images and CSV data.
+# -----------------------------------------------------------------------------
+
+
+def get_params(
+    image_folder,
+    fl,
+    trans_body,
+    cam_x,
+    cam_y,
+    cam_z,
+    fps,
+    cam_pitch_,
+    cam_roll_,
+    cam_yaw_,
+):
+    all_images = sorted(glob(os.path.join(image_folder, "*" + IMG_FORMAT)))
+    imgnames, cam_ext, cam_int = [], [], []
+
+    for img_ind, image_path in enumerate(all_images):
+        # Process every 5th frame.
+        if img_ind % 5 != 0:
+            continue
+        cam_ind = img_ind
+
+        cam_pitch_ind = cam_pitch_[cam_ind]
+        cam_yaw_ind = cam_yaw_[cam_ind]
+        cam_roll_ind = cam_roll_[cam_ind]
+
+        CAM_INT = get_cam_int(fl[cam_ind], SENSOR_W, SENSOR_H, IMG_W / 2.0, IMG_H / 2.0)
+
+        rot_unreal = rotation_matrix_unreal(cam_yaw_ind, cam_pitch_ind, cam_roll_ind)
+        rot_cv = convert_rotation_to_opencv(rot_unreal)
+        trans_cv = convert_translation_to_opencv(
+            cam_x[cam_ind] / 100.0, cam_y[cam_ind] / 100.0, cam_z[cam_ind] / 100.0
+        )
+        cam_ext_ = np.eye(4)
+        cam_ext_[:3, :3] = rot_cv
+        # The camera pose is computed as the inverse of the transformed translation.
+        cam_ext_[:3, 3] = -rot_cv @ trans_cv
+
+        imgnames.append(
+            os.path.join(image_path.split("/")[-2], image_path.split("/")[-1])
+        )
+        cam_ext.append(cam_ext_)
+        cam_int.append(CAM_INT)
+    return imgnames, cam_ext, cam_int
+
+
+# -----------------------------------------------------------------------------
+# Processing per sequence.
+# -----------------------------------------------------------------------------
+
+
+def process_seq(args):
+    """
+    Process a single sequence task. For each image, load the corresponding
+    depth and image files, and save the computed camera intrinsics and the inverse
+    of the extrinsic matrix (i.e. the camera pose in world coordinates) as an NPZ file.
+    """
+    (
+        scene,
+        seq_name,
+        outdir,
+        image_folder_base,
+        depth_folder_base,
+        imgnames,
+        cam_ext,
+        cam_int,
+    ) = args
+
+    out_rgb_dir = os.path.join(outdir, '_'.join([scene, seq_name]), 'rgb')
+    out_depth_dir = os.path.join(outdir, '_'.join([scene, seq_name]), 'depth')
+    out_cam_dir = os.path.join(outdir, "_".join([scene, seq_name]), "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)
+
+    assert (
+        len(imgnames) == len(cam_ext) == len(cam_int)
+    ), f"Inconsistent lengths for {scene}_{seq_name}"
+    for imgname, ext, intr in zip(imgnames, cam_ext, cam_int):
+        depthname = imgname.replace(".png", "_depth.exr")
+        imgpath = os.path.join(image_folder_base, imgname)
+        depthpath = os.path.join(depth_folder_base, depthname)
+        depth= OpenEXR.File(depthpath).parts[0].channels['Depth'].pixels
+        depth = depth.astype(np.float32)/100.0
+        
+        outimg_path = os.path.join(out_rgb_dir, os.path.basename(imgpath))
+        outdepth_path = os.path.join(out_depth_dir, os.path.basename(imgpath).replace('.png','.npy'))
+        outcam_path = os.path.join(
+            out_cam_dir, os.path.basename(imgpath).replace(".png", ".npz")
+        )
+
+        shutil.copy(imgpath, outimg_path)
+        np.save(outdepth_path, depth)
+        np.savez(outcam_path, intrinsics=intr, pose=np.linalg.inv(ext))
+    return None
+
+
+# -----------------------------------------------------------------------------
+# Main entry point.
+# -----------------------------------------------------------------------------
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Process Bedlam scenes: compute camera intrinsics and extrinsics, "
+        "and save processed camera files."
+    )
+    parser.add_argument(
+        "--root",
+        type=str,
+        required=True,
+        help="Root directory of the extracted data (scenes).",
+    )
+    parser.add_argument(
+        "--outdir", type=str, required=True, help="Output directory for processed data."
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=None,
+        help="Number of worker processes (default: os.cpu_count()//2).",
+    )
+    args = parser.parse_args()
+
+    root = args.root
+    outdir = args.outdir
+    num_workers = (
+        args.num_workers if args.num_workers is not None else (os.cpu_count() or 4) // 2
+    )
+
+    # Get scene directories from the root folder.
+    scenes = sorted(
+        [d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))]
+    )
+    # Exclude HDRI scenes.
+    hdri_scenes = [
+        "20221010_3_1000_batch01hand",
+        "20221017_3_1000_batch01hand",
+        "20221018_3-8_250_batch01hand",
+        "20221019_3_250_highbmihand",
+    ]
+    scenes = np.setdiff1d(scenes, hdri_scenes)
+
+    tasks = []
+    for scene in tqdm(scenes, desc="Collecting tasks"):
+        # Skip closeup scenes.
+        if "closeup" in scene:
+            continue
+        base_folder = os.path.join(root, scene)
+        image_folder_base = os.path.join(root, scene, "png")
+        depth_folder_base = os.path.join(root, scene, "depth")
+        csv_path = os.path.join(base_folder, "be_seq.csv")
+        if not os.path.exists(csv_path):
+            continue
+        csv_data = pd.read_csv(csv_path)
+        csv_data = csv_data.to_dict("list")
+        cam_csv_base = os.path.join(base_folder, "ground_truth", "camera")
+
+        # Look for a row in the CSV with a "sequence_name" comment.
+        for idx, comment in enumerate(csv_data.get("Comment", [])):
+            if "sequence_name" in comment:
+                seq_name = comment.split(";")[0].split("=")[-1]
+                cam_csv_path = os.path.join(cam_csv_base, seq_name + "_camera.csv")
+                if not os.path.exists(cam_csv_path):
+                    continue
+                cam_csv_data = pd.read_csv(cam_csv_path)
+                cam_csv_data = cam_csv_data.to_dict("list")
+                cam_x = cam_csv_data["x"]
+                cam_y = cam_csv_data["y"]
+                cam_z = cam_csv_data["z"]
+                cam_yaw_ = cam_csv_data["yaw"]
+                cam_pitch_ = cam_csv_data["pitch"]
+                cam_roll_ = cam_csv_data["roll"]
+                fl = cam_csv_data["focal_length"]
+                image_folder = os.path.join(image_folder_base, seq_name)
+                trans_body = None  # Not used here.
+                imgnames, cam_ext, cam_int = get_params(
+                    image_folder,
+                    fl,
+                    trans_body,
+                    cam_x,
+                    cam_y,
+                    cam_z,
+                    6,
+                    cam_pitch_=cam_pitch_,
+                    cam_roll_=cam_roll_,
+                    cam_yaw_=cam_yaw_,
+                )
+                tasks.append(
+                    (
+                        scene,
+                        seq_name,
+                        outdir,
+                        image_folder_base,
+                        depth_folder_base,
+                        imgnames,
+                        cam_ext,
+                        cam_int,
+                    )
+                )
+                # Process only the first valid sequence for this scene.
+                break
+
+    # Process each task in parallel.
+    with ProcessPoolExecutor(max_workers=num_workers) as executor:
+        futures = {executor.submit(process_seq, task): task for task in tasks}
+        for future in tqdm(
+            as_completed(futures), total=len(futures), desc="Processing sequences"
+        ):
+            error = future.result()
+            if error:
+                print(error)
+
+
+if __name__ == "__main__":
+    main()

extern/CUT3R/datasets_preprocess/preprocess_blendedmvs.py

@@ -0,0 +1,168 @@
+#!/usr/bin/env python3
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Preprocessing code for the BlendedMVS dataset
+# dataset at https://github.com/YoYo000/BlendedMVS
+# 1) Download BlendedMVS.zip
+# 2) Download BlendedMVS+.zip
+# 3) Download BlendedMVS++.zip
+# 4) Unzip everything in the same /path/to/tmp/blendedMVS/ directory
+# 5) python datasets_preprocess/preprocess_blendedMVS.py --blendedmvs_dir /path/to/tmp/blendedMVS/
+# --------------------------------------------------------
+import os
+import os.path as osp
+import re
+from tqdm import tqdm
+import numpy as np
+
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+import cv2
+
+import path_to_root  # noqa
+from datasets_preprocess.utils.parallel import parallel_threads
+from datasets_preprocess.utils import cropping  # noqa
+
+
+def get_parser():
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--blendedmvs_dir", required=True)
+    parser.add_argument("--precomputed_pairs", required=True)
+    parser.add_argument("--output_dir", default="data/blendedmvs_processed")
+    return parser
+
+
+def main(db_root, pairs_path, output_dir):
+    print(">> Listing all sequences")
+    sequences = [f for f in os.listdir(db_root) if len(f) == 24]
+    # should find 502 scenes
+    assert sequences, f"did not found any sequences at {db_root}"
+    print(f"   (found {len(sequences)} sequences)")
+
+    for i, seq in enumerate(tqdm(sequences)):
+        out_dir = osp.join(output_dir, seq)
+        os.makedirs(out_dir, exist_ok=True)
+
+        # generate the crops
+        root = osp.join(db_root, seq)
+        cam_dir = osp.join(root, "cams")
+        func_args = [
+            (root, f[:-8], out_dir)
+            for f in os.listdir(cam_dir)
+            if not f.startswith("pair")
+        ]
+        parallel_threads(load_crop_and_save, func_args, star_args=True, leave=False)
+
+    # verify that all pairs are there
+    pairs = np.load(pairs_path)
+    for seqh, seql, img1, img2, score in tqdm(pairs):
+        for view_index in [img1, img2]:
+            impath = osp.join(
+                output_dir, f"{seqh:08x}{seql:016x}", f"{view_index:08n}.jpg"
+            )
+            assert osp.isfile(impath), f"missing image at {impath=}"
+
+    print(f">> Done, saved everything in {output_dir}/")
+
+
+def load_crop_and_save(root, img, out_dir):
+    if osp.isfile(osp.join(out_dir, img + ".npz")):
+        return  # already done
+
+    # load everything
+    intrinsics_in, R_camin2world, t_camin2world = _load_pose(
+        osp.join(root, "cams", img + "_cam.txt")
+    )
+    color_image_in = cv2.cvtColor(
+        cv2.imread(osp.join(root, "blended_images", img + ".jpg"), cv2.IMREAD_COLOR),
+        cv2.COLOR_BGR2RGB,
+    )
+    depthmap_in = load_pfm_file(osp.join(root, "rendered_depth_maps", img + ".pfm"))
+
+    # do the crop
+    H, W = color_image_in.shape[:2]
+    assert H * 4 == W * 3
+    image, depthmap, intrinsics_out, R_in2out = _crop_image(
+        intrinsics_in, color_image_in, depthmap_in, (512, 384)
+    )
+
+    # write everything
+    image.save(osp.join(out_dir, img + ".jpg"), quality=80)
+    cv2.imwrite(osp.join(out_dir, img + ".exr"), depthmap)
+
+    # New camera parameters
+    R_camout2world = R_camin2world @ R_in2out.T
+    t_camout2world = t_camin2world
+    np.savez(
+        osp.join(out_dir, img + ".npz"),
+        intrinsics=intrinsics_out,
+        R_cam2world=R_camout2world,
+        t_cam2world=t_camout2world,
+    )
+
+
+def _crop_image(intrinsics_in, color_image_in, depthmap_in, resolution_out=(800, 800)):
+    image, depthmap, intrinsics_out = cropping.rescale_image_depthmap(
+        color_image_in, depthmap_in, intrinsics_in, resolution_out
+    )
+    R_in2out = np.eye(3)
+    return image, depthmap, intrinsics_out, R_in2out
+
+
+def _load_pose(path, ret_44=False):
+    f = open(path)
+    RT = np.loadtxt(f, skiprows=1, max_rows=4, dtype=np.float32)
+    assert RT.shape == (4, 4)
+    RT = np.linalg.inv(RT)  # world2cam to cam2world
+
+    K = np.loadtxt(f, skiprows=2, max_rows=3, dtype=np.float32)
+    assert K.shape == (3, 3)
+
+    if ret_44:
+        return K, RT
+    return K, RT[:3, :3], RT[:3, 3]  # , depth_uint8_to_f32
+
+
+def load_pfm_file(file_path):
+    with open(file_path, "rb") as file:
+        header = file.readline().decode("UTF-8").strip()
+
+        if header == "PF":
+            is_color = True
+        elif header == "Pf":
+            is_color = False
+        else:
+            raise ValueError("The provided file is not a valid PFM file.")
+
+        dimensions = re.match(r"^(\d+)\s(\d+)\s$", file.readline().decode("UTF-8"))
+        if dimensions:
+            img_width, img_height = map(int, dimensions.groups())
+        else:
+            raise ValueError("Invalid PFM header format.")
+
+        endian_scale = float(file.readline().decode("UTF-8").strip())
+        if endian_scale < 0:
+            dtype = "<f"  # little-endian
+        else:
+            dtype = ">f"  # big-endian
+
+        data_buffer = file.read()
+        img_data = np.frombuffer(data_buffer, dtype=dtype)
+
+        if is_color:
+            img_data = np.reshape(img_data, (img_height, img_width, 3))
+        else:
+            img_data = np.reshape(img_data, (img_height, img_width))
+
+        img_data = cv2.flip(img_data, 0)
+
+    return img_data
+
+
+if __name__ == "__main__":
+    parser = get_parser()
+    args = parser.parse_args()
+    main(args.blendedmvs_dir, args.precomputed_pairs, args.output_dir)

extern/CUT3R/datasets_preprocess/preprocess_co3d.py

@@ -0,0 +1,391 @@
+#!/usr/bin/env python3
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Script to pre-process the CO3D dataset.
+# Usage:
+# python3 datasets_preprocess/preprocess_co3d.py --co3d_dir /path/to/co3d
+# --------------------------------------------------------
+
+import argparse
+import random
+import gzip
+import json
+import os
+import os.path as osp
+
+import torch
+import PIL.Image
+import numpy as np
+import cv2
+
+from tqdm.auto import tqdm
+import matplotlib.pyplot as plt
+
+import path_to_root  # noqa
+import datasets_preprocess.utils.cropping as cropping  # noqa
+
+
+CATEGORIES = [
+    "apple",
+    "backpack",
+    "ball",
+    "banana",
+    "baseballbat",
+    "baseballglove",
+    "bench",
+    "bicycle",
+    "book",
+    "bottle",
+    "bowl",
+    "broccoli",
+    "cake",
+    "car",
+    "carrot",
+    "cellphone",
+    "chair",
+    "couch",
+    "cup",
+    "donut",
+    "frisbee",
+    "hairdryer",
+    "handbag",
+    "hotdog",
+    "hydrant",
+    "keyboard",
+    "kite",
+    "laptop",
+    "microwave",
+    "motorcycle",
+    "mouse",
+    "orange",
+    "parkingmeter",
+    "pizza",
+    "plant",
+    "remote",
+    "sandwich",
+    "skateboard",
+    "stopsign",
+    "suitcase",
+    "teddybear",
+    "toaster",
+    "toilet",
+    "toybus",
+    "toyplane",
+    "toytrain",
+    "toytruck",
+    "tv",
+    "umbrella",
+    "vase",
+    "wineglass",
+]
+CATEGORIES_IDX = {cat: i for i, cat in enumerate(CATEGORIES)}  # for seeding
+
+SINGLE_SEQUENCE_CATEGORIES = sorted(
+    set(CATEGORIES) - set(["microwave", "stopsign", "tv"])
+)
+
+
+def get_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--category", type=str, default=None)
+    parser.add_argument(
+        "--single_sequence_subset",
+        default=False,
+        action="store_true",
+        help="prepare the single_sequence_subset instead.",
+    )
+    parser.add_argument("--output_dir", type=str, default="data/co3d_processed")
+    parser.add_argument("--co3d_dir", type=str, required=True)
+    parser.add_argument("--num_sequences_per_object", type=int, default=50)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--min_quality",
+        type=float,
+        default=0.5,
+        help="Minimum viewpoint quality score.",
+    )
+
+    parser.add_argument(
+        "--img_size",
+        type=int,
+        default=512,
+        help=(
+            "lower dimension will be >= img_size * 3/4, and max dimension will be >= img_size"
+        ),
+    )
+    return parser
+
+
+def convert_ndc_to_pinhole(focal_length, principal_point, image_size):
+    focal_length = np.array(focal_length)
+    principal_point = np.array(principal_point)
+    image_size_wh = np.array([image_size[1], image_size[0]])
+    half_image_size = image_size_wh / 2
+    rescale = half_image_size.min()
+    principal_point_px = half_image_size - principal_point * rescale
+    focal_length_px = focal_length * rescale
+    fx, fy = focal_length_px[0], focal_length_px[1]
+    cx, cy = principal_point_px[0], principal_point_px[1]
+    K = np.array([[fx, 0.0, cx], [0.0, fy, cy], [0.0, 0.0, 1.0]], dtype=np.float32)
+    return K
+
+
+def opencv_from_cameras_projection(R, T, focal, p0, image_size):
+    R = torch.from_numpy(R)[None, :, :]
+    T = torch.from_numpy(T)[None, :]
+    focal = torch.from_numpy(focal)[None, :]
+    p0 = torch.from_numpy(p0)[None, :]
+    image_size = torch.from_numpy(image_size)[None, :]
+
+    R_pytorch3d = R.clone()
+    T_pytorch3d = T.clone()
+    focal_pytorch3d = focal
+    p0_pytorch3d = p0
+    T_pytorch3d[:, :2] *= -1
+    R_pytorch3d[:, :, :2] *= -1
+    tvec = T_pytorch3d
+    R = R_pytorch3d.permute(0, 2, 1)
+
+    # Retype the image_size correctly and flip to width, height.
+    image_size_wh = image_size.to(R).flip(dims=(1,))
+
+    # NDC to screen conversion.
+    scale = image_size_wh.to(R).min(dim=1, keepdim=True)[0] / 2.0
+    scale = scale.expand(-1, 2)
+    c0 = image_size_wh / 2.0
+
+    principal_point = -p0_pytorch3d * scale + c0
+    focal_length = focal_pytorch3d * scale
+
+    camera_matrix = torch.zeros_like(R)
+    camera_matrix[:, :2, 2] = principal_point
+    camera_matrix[:, 2, 2] = 1.0
+    camera_matrix[:, 0, 0] = focal_length[:, 0]
+    camera_matrix[:, 1, 1] = focal_length[:, 1]
+    return R[0], tvec[0], camera_matrix[0]
+
+
+def get_set_list(category_dir, split, is_single_sequence_subset=False):
+    listfiles = os.listdir(osp.join(category_dir, "set_lists"))
+    if is_single_sequence_subset:
+        # not all objects have manyview_dev
+        subset_list_files = [f for f in listfiles if "manyview_dev" in f]
+    else:
+        subset_list_files = [f for f in listfiles if f"fewview_train" in f]
+
+    sequences_all = []
+    for subset_list_file in subset_list_files:
+        with open(osp.join(category_dir, "set_lists", subset_list_file)) as f:
+            subset_lists_data = json.load(f)
+            sequences_all.extend(subset_lists_data[split])
+
+    return sequences_all
+
+
+def prepare_sequences(
+    category,
+    co3d_dir,
+    output_dir,
+    img_size,
+    split,
+    min_quality,
+    max_num_sequences_per_object,
+    seed,
+    is_single_sequence_subset=False,
+):
+    random.seed(seed)
+    category_dir = osp.join(co3d_dir, category)
+    category_output_dir = osp.join(output_dir, category)
+    sequences_all = get_set_list(category_dir, split, is_single_sequence_subset)
+    sequences_numbers = sorted(set(seq_name for seq_name, _, _ in sequences_all))
+
+    frame_file = osp.join(category_dir, "frame_annotations.jgz")
+    sequence_file = osp.join(category_dir, "sequence_annotations.jgz")
+
+    with gzip.open(frame_file, "r") as fin:
+        frame_data = json.loads(fin.read())
+    with gzip.open(sequence_file, "r") as fin:
+        sequence_data = json.loads(fin.read())
+
+    frame_data_processed = {}
+    for f_data in frame_data:
+        sequence_name = f_data["sequence_name"]
+        frame_data_processed.setdefault(sequence_name, {})[
+            f_data["frame_number"]
+        ] = f_data
+
+    good_quality_sequences = set()
+    for seq_data in sequence_data:
+        if seq_data["viewpoint_quality_score"] > min_quality:
+            good_quality_sequences.add(seq_data["sequence_name"])
+
+    sequences_numbers = [
+        seq_name for seq_name in sequences_numbers if seq_name in good_quality_sequences
+    ]
+    if len(sequences_numbers) < max_num_sequences_per_object:
+        selected_sequences_numbers = sequences_numbers
+    else:
+        selected_sequences_numbers = random.sample(
+            sequences_numbers, max_num_sequences_per_object
+        )
+
+    selected_sequences_numbers_dict = {
+        seq_name: [] for seq_name in selected_sequences_numbers
+    }
+    sequences_all = [
+        (seq_name, frame_number, filepath)
+        for seq_name, frame_number, filepath in sequences_all
+        if seq_name in selected_sequences_numbers_dict
+    ]
+
+    for seq_name, frame_number, filepath in tqdm(sequences_all):
+        frame_idx = int(filepath.split("/")[-1][5:-4])
+        selected_sequences_numbers_dict[seq_name].append(frame_idx)
+        mask_path = filepath.replace("images", "masks").replace(".jpg", ".png")
+        frame_data = frame_data_processed[seq_name][frame_number]
+        focal_length = frame_data["viewpoint"]["focal_length"]
+        principal_point = frame_data["viewpoint"]["principal_point"]
+        image_size = frame_data["image"]["size"]
+        K = convert_ndc_to_pinhole(focal_length, principal_point, image_size)
+        R, tvec, camera_intrinsics = opencv_from_cameras_projection(
+            np.array(frame_data["viewpoint"]["R"]),
+            np.array(frame_data["viewpoint"]["T"]),
+            np.array(focal_length),
+            np.array(principal_point),
+            np.array(image_size),
+        )
+
+        frame_data = frame_data_processed[seq_name][frame_number]
+        depth_path = os.path.join(co3d_dir, frame_data["depth"]["path"])
+        assert frame_data["depth"]["scale_adjustment"] == 1.0
+        image_path = os.path.join(co3d_dir, filepath)
+        mask_path_full = os.path.join(co3d_dir, mask_path)
+
+        input_rgb_image = PIL.Image.open(image_path).convert("RGB")
+        input_mask = plt.imread(mask_path_full)
+
+        with PIL.Image.open(depth_path) as depth_pil:
+            # the image is stored with 16-bit depth but PIL reads it as I (32 bit).
+            # we cast it to uint16, then reinterpret as float16, then cast to float32
+            input_depthmap = (
+                np.frombuffer(np.array(depth_pil, dtype=np.uint16), dtype=np.float16)
+                .astype(np.float32)
+                .reshape((depth_pil.size[1], depth_pil.size[0]))
+            )
+        depth_mask = np.stack((input_depthmap, input_mask), axis=-1)
+        H, W = input_depthmap.shape
+
+        camera_intrinsics = camera_intrinsics.numpy()
+        cx, cy = camera_intrinsics[:2, 2].round().astype(int)
+        min_margin_x = min(cx, W - cx)
+        min_margin_y = min(cy, H - cy)
+
+        # the new window will be a rectangle of size (2*min_margin_x, 2*min_margin_y) centered on (cx,cy)
+        l, t = cx - min_margin_x, cy - min_margin_y
+        r, b = cx + min_margin_x, cy + min_margin_y
+        crop_bbox = (l, t, r, b)
+        input_rgb_image, depth_mask, input_camera_intrinsics = (
+            cropping.crop_image_depthmap(
+                input_rgb_image, depth_mask, camera_intrinsics, crop_bbox
+            )
+        )
+
+        # try to set the lower dimension to img_size * 3/4 -> img_size=512 => 384
+        scale_final = ((img_size * 3 // 4) / min(H, W)) + 1e-8
+        output_resolution = np.floor(np.array([W, H]) * scale_final).astype(int)
+        if max(output_resolution) < img_size:
+            # let's put the max dimension to img_size
+            scale_final = (img_size / max(H, W)) + 1e-8
+            output_resolution = np.floor(np.array([W, H]) * scale_final).astype(int)
+
+        input_rgb_image, depth_mask, input_camera_intrinsics = (
+            cropping.rescale_image_depthmap(
+                input_rgb_image, depth_mask, input_camera_intrinsics, output_resolution
+            )
+        )
+        input_depthmap = depth_mask[:, :, 0]
+        input_mask = depth_mask[:, :, 1]
+
+        # generate and adjust camera pose
+        camera_pose = np.eye(4, dtype=np.float32)
+        camera_pose[:3, :3] = R
+        camera_pose[:3, 3] = tvec
+        camera_pose = np.linalg.inv(camera_pose)
+
+        # save crop images and depth, metadata
+        save_img_path = os.path.join(output_dir, filepath)
+        save_depth_path = os.path.join(output_dir, frame_data["depth"]["path"])
+        save_mask_path = os.path.join(output_dir, mask_path)
+        os.makedirs(os.path.split(save_img_path)[0], exist_ok=True)
+        os.makedirs(os.path.split(save_depth_path)[0], exist_ok=True)
+        os.makedirs(os.path.split(save_mask_path)[0], exist_ok=True)
+
+        input_rgb_image.save(save_img_path)
+        scaled_depth_map = (input_depthmap / np.max(input_depthmap) * 65535).astype(
+            np.uint16
+        )
+        cv2.imwrite(save_depth_path, scaled_depth_map)
+        cv2.imwrite(save_mask_path, (input_mask * 255).astype(np.uint8))
+
+        save_meta_path = save_img_path.replace("jpg", "npz")
+        np.savez(
+            save_meta_path,
+            camera_intrinsics=input_camera_intrinsics,
+            camera_pose=camera_pose,
+            maximum_depth=np.max(input_depthmap),
+        )
+
+    return selected_sequences_numbers_dict
+
+
+if __name__ == "__main__":
+    parser = get_parser()
+    args = parser.parse_args()
+    assert args.co3d_dir != args.output_dir
+    if args.category is None:
+        if args.single_sequence_subset:
+            categories = SINGLE_SEQUENCE_CATEGORIES
+        else:
+            categories = CATEGORIES
+    else:
+        categories = [args.category]
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    for split in ["train", "test"]:
+        selected_sequences_path = os.path.join(
+            args.output_dir, f"selected_seqs_{split}.json"
+        )
+        if os.path.isfile(selected_sequences_path):
+            continue
+
+        all_selected_sequences = {}
+        for category in categories:
+            category_output_dir = osp.join(args.output_dir, category)
+            os.makedirs(category_output_dir, exist_ok=True)
+            category_selected_sequences_path = os.path.join(
+                category_output_dir, f"selected_seqs_{split}.json"
+            )
+            if os.path.isfile(category_selected_sequences_path):
+                with open(category_selected_sequences_path, "r") as fid:
+                    category_selected_sequences = json.load(fid)
+            else:
+                print(f"Processing {split} - category = {category}")
+                category_selected_sequences = prepare_sequences(
+                    category=category,
+                    co3d_dir=args.co3d_dir,
+                    output_dir=args.output_dir,
+                    img_size=args.img_size,
+                    split=split,
+                    min_quality=args.min_quality,
+                    max_num_sequences_per_object=args.num_sequences_per_object,
+                    seed=args.seed + CATEGORIES_IDX[category],
+                    is_single_sequence_subset=args.single_sequence_subset,
+                )
+                with open(category_selected_sequences_path, "w") as file:
+                    json.dump(category_selected_sequences, file)
+
+            all_selected_sequences[category] = category_selected_sequences
+        with open(selected_sequences_path, "w") as file:
+            json.dump(all_selected_sequences, file)

extern/CUT3R/datasets_preprocess/preprocess_cop3d.py

@@ -0,0 +1,322 @@
+#!/usr/bin/env python3
+
+# --------------------------------------------------------
+# Script to pre-process the COP3D dataset.
+# Usage:
+#   python3 preprocess_cop3d.py --cop3d_dir /path/to/cop3d \
+#       --output_dir /path/to/processed_cop3d
+# --------------------------------------------------------
+
+import argparse
+import random
+import gzip
+import json
+import os
+import os.path as osp
+
+import torch
+import PIL.Image
+import numpy as np
+import cv2
+
+from tqdm.auto import tqdm
+import matplotlib.pyplot as plt
+
+import src.dust3r.datasets.utils.cropping as cropping
+
+# Define the object categories. (These are used for seeding.)
+CATEGORIES = ["cat", "dog"]
+CATEGORIES_IDX = {cat: i for i, cat in enumerate(CATEGORIES)}
+
+
+def get_parser():
+    """Set up the argument parser."""
+    parser = argparse.ArgumentParser(
+        description="Preprocess the CO3D dataset and output processed images, masks, and metadata."
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="",
+        help="Output directory for processed CO3D data.",
+    )
+    parser.add_argument(
+        "--cop3d_dir",
+        type=str,
+        default="",
+        help="Directory containing the raw CO3D data.",
+    )
+    parser.add_argument(
+        "--seed", type=int, default=42, help="Random seed for reproducibility."
+    )
+    parser.add_argument(
+        "--min_quality",
+        type=float,
+        default=0.5,
+        help="Minimum viewpoint quality score.",
+    )
+    parser.add_argument(
+        "--img_size",
+        type=int,
+        default=512,
+        help=(
+            "Lower dimension will be >= img_size * 3/4, and max dimension will be >= img_size"
+        ),
+    )
+    return parser
+
+
+def convert_ndc_to_pinhole(focal_length, principal_point, image_size):
+    """Convert normalized device coordinates to a pinhole camera intrinsic matrix."""
+    focal_length = np.array(focal_length)
+    principal_point = np.array(principal_point)
+    image_size_wh = np.array([image_size[1], image_size[0]])
+    half_image_size = image_size_wh / 2
+    rescale = half_image_size.min()
+    principal_point_px = half_image_size - principal_point * rescale
+    focal_length_px = focal_length * rescale
+    fx, fy = focal_length_px[0], focal_length_px[1]
+    cx, cy = principal_point_px[0], principal_point_px[1]
+    K = np.array([[fx, 0.0, cx], [0.0, fy, cy], [0.0, 0.0, 1.0]], dtype=np.float32)
+    return K
+
+
+def opencv_from_cameras_projection(R, T, focal, p0, image_size):
+    """
+    Convert camera projection parameters from CO3D (NDC) to OpenCV coordinates.
+
+    Returns:
+        R, tvec, camera_matrix: OpenCV-style rotation matrix, translation vector, and intrinsic matrix.
+    """
+    R = torch.from_numpy(R)[None, :, :]
+    T = torch.from_numpy(T)[None, :]
+    focal = torch.from_numpy(focal)[None, :]
+    p0 = torch.from_numpy(p0)[None, :]
+    image_size = torch.from_numpy(image_size)[None, :]
+
+    # Convert to PyTorch3D convention.
+    R_pytorch3d = R.clone()
+    T_pytorch3d = T.clone()
+    focal_pytorch3d = focal
+    p0_pytorch3d = p0
+    T_pytorch3d[:, :2] *= -1
+    R_pytorch3d[:, :, :2] *= -1
+    tvec = T_pytorch3d
+    R = R_pytorch3d.permute(0, 2, 1)
+
+    # Retype image_size (flip to width, height).
+    image_size_wh = image_size.to(R).flip(dims=(1,))
+
+    # Compute scale and principal point.
+    scale = image_size_wh.to(R).min(dim=1, keepdim=True)[0] / 2.0
+    scale = scale.expand(-1, 2)
+    c0 = image_size_wh / 2.0
+    principal_point = -p0_pytorch3d * scale + c0
+    focal_length = focal_pytorch3d * scale
+
+    camera_matrix = torch.zeros_like(R)
+    camera_matrix[:, :2, 2] = principal_point
+    camera_matrix[:, 2, 2] = 1.0
+    camera_matrix[:, 0, 0] = focal_length[:, 0]
+    camera_matrix[:, 1, 1] = focal_length[:, 1]
+    return R[0], tvec[0], camera_matrix[0]
+
+
+def get_set_list(category_dir, split):
+    """Obtain a list of sequences for a given category and split."""
+    listfiles = os.listdir(osp.join(category_dir, "set_lists"))
+    subset_list_files = [f for f in listfiles if "manyview" in f]
+    if len(subset_list_files) <= 0:
+        subset_list_files = [f for f in listfiles if "fewview" in f]
+
+    sequences_all = []
+    for subset_list_file in subset_list_files:
+        with open(osp.join(category_dir, "set_lists", subset_list_file)) as f:
+            subset_lists_data = json.load(f)
+            sequences_all.extend(subset_lists_data[split])
+    return sequences_all
+
+
+def prepare_sequences(
+    category, cop3d_dir, output_dir, img_size, split, min_quality, seed
+):
+    """
+    Process sequences for a given category and split.
+
+    This function loads per-frame and per-sequence annotations,
+    filters sequences based on quality, crops and rescales images,
+    and saves metadata for each frame.
+
+    Returns a dictionary mapping sequence names to lists of selected frame indices.
+    """
+    random.seed(seed)
+    category_dir = osp.join(cop3d_dir, category)
+    category_output_dir = osp.join(output_dir, category)
+    sequences_all = get_set_list(category_dir, split)
+
+    # Get unique sequence names.
+    sequences_numbers = sorted(set(seq_name for seq_name, _, _ in sequences_all))
+
+    # Load frame and sequence annotation files.
+    frame_file = osp.join(category_dir, "frame_annotations.jgz")
+    sequence_file = osp.join(category_dir, "sequence_annotations.jgz")
+
+    with gzip.open(frame_file, "r") as fin:
+        frame_data = json.loads(fin.read())
+    with gzip.open(sequence_file, "r") as fin:
+        sequence_data = json.loads(fin.read())
+
+    # Organize frame annotations per sequence.
+    frame_data_processed = {}
+    for f_data in frame_data:
+        sequence_name = f_data["sequence_name"]
+        frame_data_processed.setdefault(sequence_name, {})[
+            f_data["frame_number"]
+        ] = f_data
+
+    # Select sequences with quality above the threshold.
+    good_quality_sequences = set()
+    for seq_data in sequence_data:
+        if seq_data["viewpoint_quality_score"] > min_quality:
+            good_quality_sequences.add(seq_data["sequence_name"])
+    sequences_numbers = [
+        seq_name for seq_name in sequences_numbers if seq_name in good_quality_sequences
+    ]
+    selected_sequences_numbers = sequences_numbers
+    selected_sequences_numbers_dict = {
+        seq_name: [] for seq_name in selected_sequences_numbers
+    }
+
+    # Filter frames to only those from selected sequences.
+    sequences_all = [
+        (seq_name, frame_number, filepath)
+        for seq_name, frame_number, filepath in sequences_all
+        if seq_name in selected_sequences_numbers_dict
+    ]
+
+    # Process each frame.
+    for seq_name, frame_number, filepath in tqdm(
+        sequences_all, desc="Processing frames"
+    ):
+        frame_idx = int(filepath.split("/")[-1][5:-4])
+        selected_sequences_numbers_dict[seq_name].append(frame_idx)
+        mask_path = filepath.replace("images", "masks").replace(".jpg", ".png")
+        frame_data_entry = frame_data_processed[seq_name][frame_number]
+        focal_length = frame_data_entry["viewpoint"]["focal_length"]
+        principal_point = frame_data_entry["viewpoint"]["principal_point"]
+        image_size = frame_data_entry["image"]["size"]
+        K = convert_ndc_to_pinhole(focal_length, principal_point, image_size)
+        R, tvec, camera_intrinsics = opencv_from_cameras_projection(
+            np.array(frame_data_entry["viewpoint"]["R"]),
+            np.array(frame_data_entry["viewpoint"]["T"]),
+            np.array(focal_length),
+            np.array(principal_point),
+            np.array(image_size),
+        )
+
+        # Load input image and mask.
+        image_path = osp.join(cop3d_dir, filepath)
+        mask_path_full = osp.join(cop3d_dir, mask_path)
+        input_rgb_image = PIL.Image.open(image_path).convert("RGB")
+        input_mask = plt.imread(mask_path_full)
+        H, W = input_mask.shape
+
+        camera_intrinsics = camera_intrinsics.numpy()
+        cx, cy = camera_intrinsics[:2, 2].round().astype(int)
+        min_margin_x = min(cx, W - cx)
+        min_margin_y = min(cy, H - cy)
+        l, t = cx - min_margin_x, cy - min_margin_y
+        r, b = cx + min_margin_x, cy + min_margin_y
+        crop_bbox = (l, t, r, b)
+
+        # Crop the image, mask, and adjust intrinsics.
+        input_rgb_image, input_mask, input_camera_intrinsics = (
+            cropping.crop_image_depthmap(
+                input_rgb_image, input_mask, camera_intrinsics, crop_bbox
+            )
+        )
+        scale_final = ((img_size * 3 // 4) / min(H, W)) + 1e-8
+        output_resolution = np.floor(np.array([W, H]) * scale_final).astype(int)
+        if max(output_resolution) < img_size:
+            scale_final = (img_size / max(H, W)) + 1e-8
+            output_resolution = np.floor(np.array([W, H]) * scale_final).astype(int)
+        input_rgb_image, input_mask, input_camera_intrinsics = (
+            cropping.rescale_image_depthmap(
+                input_rgb_image, input_mask, input_camera_intrinsics, output_resolution
+            )
+        )
+
+        # Generate and adjust camera pose.
+        camera_pose = np.eye(4, dtype=np.float32)
+        camera_pose[:3, :3] = R
+        camera_pose[:3, 3] = tvec
+        camera_pose = np.linalg.inv(camera_pose)
+
+        # Save processed image and mask.
+        save_img_path = osp.join(output_dir, filepath)
+        save_mask_path = osp.join(output_dir, mask_path)
+        os.makedirs(osp.split(save_img_path)[0], exist_ok=True)
+        os.makedirs(osp.split(save_mask_path)[0], exist_ok=True)
+        input_rgb_image.save(save_img_path)
+        cv2.imwrite(save_mask_path, (input_mask * 255).astype(np.uint8))
+
+        # Save metadata (intrinsics and pose).
+        save_meta_path = save_img_path.replace("jpg", "npz")
+        np.savez(
+            save_meta_path,
+            camera_intrinsics=input_camera_intrinsics,
+            camera_pose=camera_pose,
+        )
+
+    return selected_sequences_numbers_dict
+
+
+def main():
+    parser = get_parser()
+    args = parser.parse_args()
+    assert (
+        args.cop3d_dir != args.output_dir
+    ), "Input and output directories must differ."
+    categories = CATEGORIES
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Process each split separately.
+    for split in ["train", "test"]:
+        selected_sequences_path = osp.join(
+            args.output_dir, f"selected_seqs_{split}.json"
+        )
+        if os.path.isfile(selected_sequences_path):
+            continue
+
+        all_selected_sequences = {}
+        for category in categories:
+            category_output_dir = osp.join(args.output_dir, category)
+            os.makedirs(category_output_dir, exist_ok=True)
+            category_selected_sequences_path = osp.join(
+                category_output_dir, f"selected_seqs_{split}.json"
+            )
+            if os.path.isfile(category_selected_sequences_path):
+                with open(category_selected_sequences_path, "r") as fid:
+                    category_selected_sequences = json.load(fid)
+            else:
+                print(f"Processing {split} - category = {category}")
+                category_selected_sequences = prepare_sequences(
+                    category=category,
+                    cop3d_dir=args.cop3d_dir,
+                    output_dir=args.output_dir,
+                    img_size=args.img_size,
+                    split=split,
+                    min_quality=args.min_quality,
+                    seed=args.seed + CATEGORIES_IDX[category],
+                )
+                with open(category_selected_sequences_path, "w") as file:
+                    json.dump(category_selected_sequences, file)
+
+            all_selected_sequences[category] = category_selected_sequences
+
+        with open(selected_sequences_path, "w") as file:
+            json.dump(all_selected_sequences, file)
+
+
+if __name__ == "__main__":
+    main()

extern/CUT3R/datasets_preprocess/preprocess_dl3dv.py

@@ -0,0 +1,188 @@
+import argparse
+import random
+import gzip
+import json
+import os
+import sys
+
+import os.path as osp
+
+import torch
+import PIL.Image
+from PIL import Image
+import numpy as np
+import cv2
+
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+import shutil
+from read_write_model import run
+
+import torch
+import torchvision
+
+
+def get_parser():
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dl3dv_dir", default="../DL3DV-Dense/3K/")  # TODO
+    parser.add_argument("--output_dir", default="../processed_dl3dv/3K/")  # TODO
+    return parser
+
+
+from scipy.spatial.transform import Rotation as R
+
+
+def read_array(path):
+    with open(path, "rb") as fid:
+        width, height, channels = np.genfromtxt(
+            fid, delimiter="&", max_rows=1, usecols=(0, 1, 2), dtype=int
+        )
+        fid.seek(0)
+        num_delimiter = 0
+        byte = fid.read(1)
+        while True:
+            if byte == b"&":
+                num_delimiter += 1
+                if num_delimiter >= 3:
+                    break
+            byte = fid.read(1)
+        array = np.fromfile(fid, np.float32)
+    array = array.reshape((width, height, channels), order="F")
+    return np.transpose(array, (1, 0, 2)).squeeze()
+
+
+def main(rootdir, outdir):
+    os.makedirs(outdir, exist_ok=True)
+
+    envs = [f for f in os.listdir(rootdir) if os.path.isdir(osp.join(rootdir, f))]
+    for env in tqdm(envs):
+        subseqs = [
+            f
+            for f in os.listdir(osp.join(rootdir, env))
+            if os.path.isdir(osp.join(rootdir, env, f)) and f.startswith("dense")
+        ]
+        for subseq in subseqs:
+            sparse_dir = osp.join(rootdir, env, subseq, "sparse")
+            images_dir = osp.join(rootdir, env, subseq, "images")
+            # depth_dir = osp.join(rootdir, env, subseq, "stereo", "depth_maps")
+            if (
+                (not os.path.exists(sparse_dir))
+                or (not os.path.exists(images_dir))
+                # or (not os.path.exists(depth_dir))
+            ):
+                continue
+            intrins_file = sparse_dir + "/cameras.txt"
+            poses_file = sparse_dir + "/images.txt"
+            if os.path.exists(intrins_file) and os.path.exists(poses_file):
+                continue
+            run(sparse_dir, sparse_dir)
+
+            cam_params = {}
+            with open(intrins_file, "r") as f:
+                for line in f:
+                    if line.startswith("#"):
+                        continue
+                    parts = line.strip().split()
+                    if len(parts) == 0:
+                        continue
+                    cam_id = int(parts[0])
+                    fx = float(parts[4])
+                    fy = float(parts[5])
+                    cx = float(parts[6])
+                    cy = float(parts[7])
+                    cam_params[cam_id] = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
+
+            poses = []
+            images = []
+            intrinsics = []
+
+            with open(poses_file, "r") as f:
+                for i, line in enumerate(f):
+                    if line.startswith("#"):
+                        continue
+                    parts = line.strip().split()
+                    if len(parts) == 0:
+                        continue
+                    if "." in parts[0]:
+                        continue
+
+                    img_name = parts[-1]
+                    w, x, y, z = map(float, parts[1:5])
+                    R = np.array(
+                        [
+                            [
+                                1 - 2 * y * y - 2 * z * z,
+                                2 * x * y - 2 * z * w,
+                                2 * x * z + 2 * y * w,
+                            ],
+                            [
+                                2 * x * y + 2 * z * w,
+                                1 - 2 * x * x - 2 * z * z,
+                                2 * y * z - 2 * x * w,
+                            ],
+                            [
+                                2 * x * z - 2 * y * w,
+                                2 * y * z + 2 * x * w,
+                                1 - 2 * x * x - 2 * y * y,
+                            ],
+                        ]
+                    )
+                    tx, ty, tz = map(float, parts[5:8])
+                    cam_id = int(parts[-2])
+                    pose = np.eye(4)
+                    pose[:3, :3] = R
+                    pose[:3, 3] = [tx, ty, tz]
+                    poses.append(np.linalg.inv(pose))
+                    images.append(img_name)
+                    intrinsics.append(cam_params[cam_id])
+
+            os.makedirs(osp.join(outdir, env, subseq), exist_ok=True)
+            os.makedirs(osp.join(outdir, env, subseq, "rgb"), exist_ok=True)
+            # os.makedirs(osp.join(outdir, env, subseq, "depth"), exist_ok=True)
+            os.makedirs(osp.join(outdir, env, subseq, "cam"), exist_ok=True)
+
+            for i, img_name in enumerate(tqdm(images)):
+                basename = img_name.split("/")[-1]
+                if os.path.exists(
+                    osp.join(
+                        outdir, env, subseq, "cam", basename.replace(".png", ".npz")
+                    )
+                ):
+                    print("Exist!")
+                    continue
+                img_path = os.path.join(images_dir, img_name)
+                # depth_path = os.path.join(depth_dir, img_name + ".geometric.bin")
+                if not os.path.exists(depth_path) or not os.path.exists(img_path):
+                    continue
+                try:
+                    rgb = Image.open(img_path)
+                    # depth = read_array(depth_path)
+                except:
+                    continue
+                intrinsic = intrinsics[i]
+                pose = poses[i]
+
+                # save all
+
+                rgb.save(osp.join(outdir, env, subseq, "rgb", basename))
+                # np.save(
+                #     osp.join(
+                #         outdir, env, subseq, "depth", basename.replace(".png", ".npy")
+                #     ),
+                #     depth,
+                # )
+                np.savez(
+                    osp.join(
+                        outdir, env, subseq, "cam", basename.replace(".png", ".npz")
+                    ),
+                    intrinsic=intrinsic,
+                    pose=pose,
+                )
+
+
+if __name__ == "__main__":
+    parser = get_parser()
+    args = parser.parse_args()
+    main(args.dl3dv_dir, args.output_dir)

extern/CUT3R/datasets_preprocess/preprocess_dynamic_replica.py

@@ -0,0 +1,344 @@
+#!/usr/bin/env python3
+"""
+Preprocess the Dynamic Replica dataset.
+
+This script reads frame annotations (stored in compressed JSON files),
+loads images, depth maps, optical flow, and camera parameters, and saves
+processed images, depth maps, flow files, and camera metadata (intrinsics and poses)
+to an output directory organized by split, sequence, and camera view.
+
+Usage:
+    python preprocess_dynamic_replica.py --root_dir /path/to/data_dynamic_replica \
+                                           --out_dir /path/to/processed_dynamic_replica \
+                                           [--splits train valid test] \
+                                           [--num_processes 8]
+"""
+
+import argparse
+import gzip
+import json
+import os
+import os.path as osp
+import re
+import shutil
+import time
+from collections import defaultdict
+from dataclasses import dataclass
+from multiprocessing import Pool, cpu_count
+from typing import List, Optional
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import PIL.Image
+import torch
+from PIL import Image
+from pytorch3d.implicitron.dataset.types import (
+    FrameAnnotation as ImplicitronFrameAnnotation,
+    load_dataclass,
+)
+from tqdm import tqdm
+import imageio
+
+# Enable OpenEXR support in OpenCV.
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+
+TAG_CHAR = np.array([202021.25], np.float32)
+
+
+def readFlow(fn):
+    """Read .flo file in Middlebury format."""
+    with open(fn, "rb") as f:
+        magic = np.fromfile(f, np.float32, count=1)
+        if 202021.25 != magic:
+            print("Magic number incorrect. Invalid .flo file")
+            return None
+        else:
+            w = np.fromfile(f, np.int32, count=1)
+            h = np.fromfile(f, np.int32, count=1)
+            data = np.fromfile(f, np.float32, count=2 * int(w) * int(h))
+            return np.resize(data, (int(h), int(w), 2))
+
+
+def readPFM(file):
+    with open(file, "rb") as f:
+        header = f.readline().rstrip()
+        if header == b"PF":
+            color = True
+        elif header == b"Pf":
+            color = False
+        else:
+            raise Exception("Not a PFM file.")
+
+        dim_match = re.match(rb"^(\d+)\s(\d+)\s$", f.readline())
+        if dim_match:
+            width, height = map(int, dim_match.groups())
+        else:
+            raise Exception("Malformed PFM header.")
+
+        scale = float(f.readline().rstrip())
+        endian = "<" if scale < 0 else ">"
+        if scale < 0:
+            scale = -scale
+
+        data = np.fromfile(f, endian + "f")
+        shape = (height, width, 3) if color else (height, width)
+        data = np.reshape(data, shape)
+        data = np.flipud(data)
+        return data
+
+
+def read_gen(file_name, pil=False):
+    ext = osp.splitext(file_name)[-1].lower()
+    if ext in [".png", ".jpeg", ".ppm", ".jpg"]:
+        return Image.open(file_name)
+    elif ext in [".bin", ".raw"]:
+        return np.load(file_name)
+    elif ext == ".flo":
+        return readFlow(file_name).astype(np.float32)
+    elif ext == ".pfm":
+        flow = readPFM(file_name).astype(np.float32)
+        return flow if len(flow.shape) == 2 else flow[:, :, :-1]
+    return []
+
+
+def _load_16big_png_depth(depth_png):
+    with Image.open(depth_png) as depth_pil:
+        depth = (
+            np.frombuffer(np.array(depth_pil, dtype=np.uint16), dtype=np.float16)
+            .astype(np.float32)
+            .reshape((depth_pil.size[1], depth_pil.size[0]))
+        )
+    return depth
+
+
+@dataclass
+class DynamicReplicaFrameAnnotation(ImplicitronFrameAnnotation):
+    """A dataclass used to load annotations from .json for Dynamic Replica."""
+
+    camera_name: Optional[str] = None
+    instance_id_map_path: Optional[str] = None
+    flow_forward: Optional[str] = None
+    flow_forward_mask: Optional[str] = None
+    flow_backward: Optional[str] = None
+    flow_backward_mask: Optional[str] = None
+    trajectories: Optional[str] = None
+
+
+def _get_pytorch3d_camera(entry_viewpoint, image_size, scale: float):
+    """
+    Convert the camera parameters stored in an annotation to PyTorch3D convention.
+
+    Returns:
+        R, tvec, focal, principal_point
+    """
+    assert entry_viewpoint is not None
+    principal_point = torch.tensor(entry_viewpoint.principal_point, dtype=torch.float)
+    focal_length = torch.tensor(entry_viewpoint.focal_length, dtype=torch.float)
+    half_image_size_wh_orig = (
+        torch.tensor(list(reversed(image_size)), dtype=torch.float) / 2.0
+    )
+
+    fmt = entry_viewpoint.intrinsics_format
+    if fmt.lower() == "ndc_norm_image_bounds":
+        rescale = half_image_size_wh_orig
+    elif fmt.lower() == "ndc_isotropic":
+        rescale = half_image_size_wh_orig.min()
+    else:
+        raise ValueError(f"Unknown intrinsics format: {fmt}")
+
+    principal_point_px = half_image_size_wh_orig - principal_point * rescale
+    focal_length_px = focal_length * rescale
+
+    # Prepare rotation and translation for PyTorch3D
+    R = torch.tensor(entry_viewpoint.R, dtype=torch.float)
+    T = torch.tensor(entry_viewpoint.T, dtype=torch.float)
+    R_pytorch3d = R.clone()
+    T_pytorch3d = T.clone()
+    T_pytorch3d[..., :2] *= -1
+    R_pytorch3d[..., :, :2] *= -1
+    tvec = T_pytorch3d
+
+    return R, tvec, focal_length_px, principal_point_px
+
+
+# Global configuration for splits and output.
+SPLITS = ["train", "valid", "test"]
+# (You can override the default root and out_dir via command-line arguments.)
+
+
+def process_split_data(args):
+    """
+    Process all frames for a given split.
+
+    Reads the frame annotation file for the given split, groups frames per sequence
+    and camera, and for each frame loads the image, depth map, optical flows (if available),
+    computes the camera intrinsics and pose (using _get_pytorch3d_camera), and saves the data.
+    """
+    split, root_dir, out_dir = args
+    split_dir = osp.join(root_dir, split)
+    # The frame annotations are stored in a compressed json file.
+    frame_annotations_file = osp.join(split_dir, f"frame_annotations_{split}.jgz")
+    with gzip.open(frame_annotations_file, "rt", encoding="utf8") as zipfile:
+        frame_annots_list = load_dataclass(zipfile, List[DynamicReplicaFrameAnnotation])
+
+    # Group frames by sequence and camera.
+    seq_annot = defaultdict(lambda: defaultdict(list))
+    for frame_annot in frame_annots_list:
+        seq_annot[frame_annot.sequence_name][frame_annot.camera_name].append(
+            frame_annot
+        )
+
+    # Process each sequence.
+    for seq_name in tqdm(seq_annot.keys(), desc=f"Processing split '{split}'"):
+        # For each camera (e.g., 'left', 'right'), create output directories.
+        for cam in ["left", "right"]:
+            out_img_dir = osp.join(out_dir, split, seq_name, cam, "rgb")
+            out_depth_dir = osp.join(out_dir, split, seq_name, cam, "depth")
+            out_fflow_dir = osp.join(out_dir, split, seq_name, cam, "flow_forward")
+            out_bflow_dir = osp.join(out_dir, split, seq_name, cam, "flow_backward")
+            out_cam_dir = osp.join(out_dir, split, seq_name, cam, "cam")
+            os.makedirs(out_img_dir, exist_ok=True)
+            os.makedirs(out_depth_dir, exist_ok=True)
+            os.makedirs(out_fflow_dir, exist_ok=True)
+            os.makedirs(out_bflow_dir, exist_ok=True)
+            os.makedirs(out_cam_dir, exist_ok=True)
+
+            for framedata in tqdm(
+                seq_annot[seq_name][cam], desc=f"Seq {seq_name} [{cam}]", leave=False
+            ):
+                timestamp = framedata.frame_timestamp
+                im_path = osp.join(split_dir, framedata.image.path)
+                depth_path = osp.join(split_dir, framedata.depth.path)
+                if framedata.flow_forward["path"]:
+                    flow_forward_path = osp.join(
+                        split_dir, framedata.flow_forward["path"]
+                    )
+                    flow_forward_mask_path = osp.join(
+                        split_dir, framedata.flow_forward_mask["path"]
+                    )
+                if framedata.flow_backward["path"]:
+                    flow_backward_path = osp.join(
+                        split_dir, framedata.flow_backward["path"]
+                    )
+                    flow_backward_mask_path = osp.join(
+                        split_dir, framedata.flow_backward_mask["path"]
+                    )
+
+                # Ensure required files exist.
+                assert os.path.isfile(im_path), im_path
+                assert os.path.isfile(depth_path), depth_path
+                if framedata.flow_forward["path"]:
+                    assert os.path.isfile(flow_forward_path), flow_forward_path
+                    assert os.path.isfile(
+                        flow_forward_mask_path
+                    ), flow_forward_mask_path
+                if framedata.flow_backward["path"]:
+                    assert os.path.isfile(flow_backward_path), flow_backward_path
+                    assert os.path.isfile(
+                        flow_backward_mask_path
+                    ), flow_backward_mask_path
+
+                viewpoint = framedata.viewpoint
+                # Load depth map.
+                depth = _load_16big_png_depth(depth_path)
+
+                # Process optical flow if available.
+                if framedata.flow_forward["path"]:
+                    flow_forward = cv2.imread(flow_forward_path, cv2.IMREAD_UNCHANGED)
+                    flow_forward_mask = cv2.imread(
+                        flow_forward_mask_path, cv2.IMREAD_UNCHANGED
+                    )
+                    np.savez(
+                        osp.join(out_fflow_dir, f"{timestamp}.npz"),
+                        flow=flow_forward,
+                        mask=flow_forward_mask,
+                    )
+                if framedata.flow_backward["path"]:
+                    flow_backward = cv2.imread(flow_backward_path, cv2.IMREAD_UNCHANGED)
+                    flow_backward_mask = cv2.imread(
+                        flow_backward_mask_path, cv2.IMREAD_UNCHANGED
+                    )
+                    np.savez(
+                        osp.join(out_bflow_dir, f"{timestamp}.npz"),
+                        flow=flow_backward,
+                        mask=flow_backward_mask,
+                    )
+
+                # Get camera parameters.
+                R, t, focal, pp = _get_pytorch3d_camera(
+                    viewpoint, framedata.image.size, scale=1.0
+                )
+                intrinsics = np.eye(3)
+                intrinsics[0, 0] = focal[0].item()
+                intrinsics[1, 1] = focal[1].item()
+                intrinsics[0, 2] = pp[0].item()
+                intrinsics[1, 2] = pp[1].item()
+                pose = np.eye(4)
+                # Invert the camera pose.
+                pose[:3, :3] = R.numpy().T
+                pose[:3, 3] = -R.numpy().T @ t.numpy()
+
+                # Define output file paths.
+                out_img_path = osp.join(out_img_dir, f"{timestamp}.png")
+                out_depth_path = osp.join(out_depth_dir, f"{timestamp}.npy")
+                out_cam_path = osp.join(out_cam_dir, f"{timestamp}.npz")
+
+                # Copy RGB image.
+                shutil.copy(im_path, out_img_path)
+                # Save depth.
+                np.save(out_depth_path, depth)
+                # Save camera metadata.
+                np.savez(out_cam_path, intrinsics=intrinsics, pose=pose)
+    # (Optionally, you could return some summary information.)
+    return None
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Preprocess Dynamic Replica dataset: convert raw annotations, images, "
+        "depth, and flow files to a processed format."
+    )
+    parser.add_argument(
+        "--root_dir",
+        type=str,
+        required=True,
+        help="Root directory of the Dynamic Replica data.",
+    )
+    parser.add_argument(
+        "--out_dir",
+        type=str,
+        required=True,
+        help="Output directory for processed data.",
+    )
+    parser.add_argument(
+        "--splits",
+        type=str,
+        nargs="+",
+        default=SPLITS,
+        help="List of splits to process (default: train valid test).",
+    )
+    parser.add_argument(
+        "--num_processes",
+        type=int,
+        default=cpu_count(),
+        help="Number of processes to use (default: number of CPU cores).",
+    )
+    args = parser.parse_args()
+
+    os.makedirs(args.out_dir, exist_ok=True)
+    tasks = [(split, args.root_dir, args.out_dir) for split in args.splits]
+
+    print("Processing splits:", args.splits)
+    with Pool(processes=args.num_processes) as pool:
+        list(
+            tqdm(
+                pool.imap(process_split_data, tasks),
+                total=len(tasks),
+                desc="Overall Progress",
+            )
+        )
+
+
+if __name__ == "__main__":
+    main()

extern/CUT3R/datasets_preprocess/preprocess_eden.py

@@ -0,0 +1,181 @@
+#!/usr/bin/env python3
+"""
+Preprocess the Eden dataset.
+
+This script processes the Eden dataset by copying RGB images, converting depth
+data from .mat files to .npy format, and saving camera intrinsics from .mat files
+into a structured output directory. Files are processed in parallel using
+a ProcessPoolExecutor.
+
+Usage:
+    python preprocess_eden.py --root /path/to/data_raw_videos/data_eden \
+                              --out_dir /path/to/data_raw_videos/processed_eden \
+                              [--num_workers N]
+"""
+
+import os
+import shutil
+import scipy.io
+import numpy as np
+from tqdm import tqdm
+from concurrent.futures import ProcessPoolExecutor, as_completed
+import argparse
+
+
+def process_basename(args):
+    """
+    Process a single basename: load the corresponding image, depth, and camera
+    intrinsics files, then copy/save them into the output directories.
+
+    Parameters:
+        args (tuple): Contains (seq, basename, rgb_dir, depth_dir, cam_dir,
+                      out_rgb_dir, out_depth_dir, out_cam_dir)
+    Returns:
+        None on success or an error message string on failure.
+    """
+    (
+        seq,
+        basename,
+        rgb_dir,
+        depth_dir,
+        cam_dir,
+        out_rgb_dir,
+        out_depth_dir,
+        out_cam_dir,
+    ) = args
+    # Define output 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_cam_path = os.path.join(out_cam_dir, f"{basename}.npz")
+
+    # Skip processing if the camera file has already been saved.
+    if os.path.exists(out_cam_path):
+        return None
+
+    try:
+        cam_type = "L"
+        img_file = os.path.join(rgb_dir, f"{basename}_{cam_type}.png")
+        depth_file = os.path.join(depth_dir, f"{basename}_{cam_type}.mat")
+        cam_file = os.path.join(cam_dir, f"{basename}.mat")
+
+        # Check if the required files exist.
+        if not (
+            os.path.exists(img_file)
+            and os.path.exists(depth_file)
+            and os.path.exists(cam_file)
+        ):
+            return f"Missing files for {basename} in {seq}"
+
+        # Load depth data.
+        depth_mat = scipy.io.loadmat(depth_file)
+        depth = depth_mat.get("Depth")
+        if depth is None:
+            return f"Depth data missing in {depth_file}"
+        depth = depth[..., 0]
+
+        # Load camera intrinsics.
+        cam_mat = scipy.io.loadmat(cam_file)
+        intrinsics = cam_mat.get(f"K_{cam_type}")
+        if intrinsics is None:
+            return f"Intrinsics data missing in {cam_file}"
+
+        # Copy the RGB image.
+        shutil.copyfile(img_file, out_img_path)
+        # Save the depth data.
+        np.save(out_depth_path, depth)
+        # Save the camera intrinsics.
+        np.savez(out_cam_path, intrinsics=intrinsics)
+
+    except Exception as e:
+        return f"Error processing {basename} in {seq}: {e}"
+
+    return None  # Indicate success.
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Preprocess Eden dataset: copy RGB images, process depth maps, and save camera intrinsics."
+    )
+    parser.add_argument(
+        "--root", type=str, default="", help="Root directory of the raw Eden data."
+    )
+    parser.add_argument(
+        "--out_dir",
+        type=str,
+        default="",
+        help="Output directory for processed Eden data.",
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=os.cpu_count(),
+        help="Number of worker processes to use.",
+    )
+    args = parser.parse_args()
+
+    root = args.root
+    out_dir = args.out_dir
+    # Modes typically found in the Eden dataset.
+    modes = ["clear", "cloudy", "overcast", "sunset", "twilight"]
+
+    rgb_root = os.path.join(root, "RGB")
+    depth_root = os.path.join(root, "Depth")
+    cam_root = os.path.join(root, "cam_matrix")
+
+    # Collect sequence directories by traversing the RGB root.
+    seq_dirs = []
+    for d in os.listdir(rgb_root):
+        for m in modes:
+            seq_path = os.path.join(rgb_root, d, m)
+            if os.path.isdir(seq_path):
+                # Save the relative path (e.g., "d/m").
+                seq_dirs.append(os.path.join(d, m))
+
+    all_tasks = []
+    for seq in seq_dirs:
+        rgb_dir = os.path.join(rgb_root, seq)
+        depth_dir = os.path.join(depth_root, seq)
+        cam_dir = os.path.join(cam_root, seq)
+
+        # Create output directories for this sequence.
+        # Replace any os.sep in the sequence name with an underscore.
+        seq_name = "_".join(seq.split(os.sep))
+        out_rgb_dir = os.path.join(out_dir, seq_name, "rgb")
+        out_depth_dir = os.path.join(out_dir, seq_name, "depth")
+        out_cam_dir = os.path.join(out_dir, seq_name, "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 the camera directory (assuming file extension .mat).
+        basenames = sorted([d[:-4] for d in os.listdir(cam_dir) if d.endswith(".mat")])
+
+        for basename in basenames:
+            task = (
+                seq,
+                basename,
+                rgb_dir,
+                depth_dir,
+                cam_dir,
+                out_rgb_dir,
+                out_depth_dir,
+                out_cam_dir,
+            )
+            all_tasks.append(task)
+
+    num_workers = args.num_workers
+    print(f"Processing {len(all_tasks)} tasks using {num_workers} workers...")
+    with ProcessPoolExecutor(max_workers=num_workers) as executor:
+        futures = {
+            executor.submit(process_basename, task): task[1] for task in all_tasks
+        }
+        for future in tqdm(
+            as_completed(futures), total=len(futures), desc="Processing tasks"
+        ):
+            error = future.result()
+            if error:
+                print(error)
+
+
+if __name__ == "__main__":
+    main()

extern/CUT3R/datasets_preprocess/preprocess_hoi4d.py

@@ -0,0 +1,175 @@
+#!/usr/bin/env python3
+"""
+HOI4D Preprocessing Script
+
+This script processes HOI4D data by:
+  1. Searching specific subdirectories for RGB and depth images.
+  2. Reading camera intrinsics from a .npy file (one per high-level scene).
+  3. Rescaling the RGB images and depth maps to a fixed output resolution
+     (e.g., 640x480) using the 'cropping' module.
+  4. Saving results (RGB, .npy depth, .npz camera intrinsics) in a new directory structure.
+
+Usage:
+    python preprocess_hoi4d.py \
+        --root_dir /path/to/HOI4D_release \
+        --cam_root /path/to/camera_params \
+        --out_dir /path/to/processed_hoi4d
+"""
+
+import os
+import glob
+import cv2
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+from concurrent.futures import ProcessPoolExecutor
+import argparse
+
+import src.dust3r.datasets.utils.cropping as cropping
+
+def parse_arguments():
+    """
+    Parse command-line arguments for HOI4D preprocessing.
+
+    Returns:
+        argparse.Namespace: The parsed arguments.
+    """
+    parser = argparse.ArgumentParser(
+        description="Preprocess HOI4D dataset by rescaling RGB and depth images."
+    )
+    parser.add_argument("--root_dir", required=True,
+                        help="Path to the HOI4D_release directory.")
+    parser.add_argument("--cam_root", required=True,
+                        help="Path to the directory containing camera intrinsics.")
+    parser.add_argument("--out_dir", required=True,
+                        help="Path to the directory where processed files will be saved.")
+    parser.add_argument("--max_workers", type=int, default=None,
+                        help="Number of parallel workers. Default uses half of available CPU cores.")
+    args = parser.parse_args()
+    return args
+
+def process_image(args):
+    """
+    Process a single image and depth map:
+      - Loads the image (using PIL) and depth (using OpenCV).
+      - Converts depth from mm to meters (divided by 1000).
+      - Rescales both using 'cropping.rescale_image_depthmap'.
+      - Saves the rescaled image (.png), depth (.npy), and camera intrinsics (.npz).
+
+    Args:
+        args (tuple): A tuple of:
+          (img_path, depth_path, out_img_path, out_depth_path, out_cam_path, intrinsics)
+
+    Returns:
+        None. Errors are printed to the console but do not stop the workflow.
+    """
+    img_path, depth_path, out_img_path, out_depth_path, out_cam_path, intrinsics = args
+
+    try:
+        # Load image
+        img = Image.open(img_path)
+
+        # Load depth (in mm) and convert to meters
+        depth = cv2.imread(depth_path, cv2.IMREAD_ANYDEPTH)
+        if depth is None:
+            raise ValueError(f"Could not read depth image: {depth_path}")
+        depth = depth.astype(np.float32) / 1000.0
+
+        # Rescale image and depth map
+        img_rescaled, depth_rescaled, intrinsics_rescaled = cropping.rescale_image_depthmap(
+            img, depth, intrinsics.copy(), (640, 480)
+        )
+
+        # Save processed data
+        img_rescaled.save(out_img_path)      # PNG image
+        np.save(out_depth_path, depth_rescaled)  # Depth .npy
+        np.savez(out_cam_path, intrinsics=intrinsics_rescaled)
+
+    except Exception as e:
+        print(f"Error processing {img_path}: {e}")
+
+def main():
+    args = parse_arguments()
+
+    root = args.root_dir
+    cam_root = args.cam_root
+    out_dir = args.out_dir
+    if not os.path.exists(out_dir):
+        os.makedirs(out_dir, exist_ok=True)
+
+    # Collect a list of subdirectories using a glob pattern
+    # e.g.: root/ZY2021*/H*/C*/N*/S*/s*/T*
+    scene_dirs = glob.glob(os.path.join(root, "ZY2021*", "H*", "C*", "N*", "S*", "s*", "T*"))
+
+    # Build tasks
+    tasks = []
+    for scene_dir in tqdm(scene_dirs, desc="Collecting scenes"):
+        # Build an output sub-directory name
+        # Example: "ZY202101/H1/C1/N1/S1/s1/T1" -> "ZY202101_H1_C1_N1_S1_s1_T1"
+        scene_relpath = os.path.relpath(scene_dir, root)
+        scene_name = "_".join(scene_relpath.split(os.sep))
+
+        # Load camera intrinsics from a .npy file in cam_root
+        # e.g., first token of scene_relpath might point to the relevant .npy
+        # "ZY202101" -> "cam_root/ZY202101/intrin.npy" (adjust logic as needed)
+        top_level = scene_relpath.split(os.sep)[0]
+        cam_file = os.path.join(cam_root, top_level, "intrin.npy")
+        if not os.path.isfile(cam_file):
+            print(f"Warning: Camera file not found: {cam_file}. Skipping {scene_dir}")
+            continue
+        intrinsics = np.load(cam_file)
+
+        # Directories for this sequence
+        rgb_dir = os.path.join(scene_dir, "align_rgb")
+        depth_dir = os.path.join(scene_dir, "align_depth")
+
+        # Output directories
+        out_rgb_dir = os.path.join(out_dir, scene_name, "rgb")
+        out_depth_dir = os.path.join(out_dir, scene_name, "depth")
+        out_cam_dir = os.path.join(out_dir, scene_name, "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)
+
+        # Find all image paths
+        img_paths = sorted(glob.glob(os.path.join(rgb_dir, "*.jpg")))
+
+        # Build tasks for each image
+        for img_path in img_paths:
+            basename = os.path.splitext(os.path.basename(img_path))[0]
+            depth_path = os.path.join(depth_dir, f"{basename}.png")
+
+            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")
+
+            # Skip if already processed
+            if (os.path.exists(out_img_path) and os.path.exists(out_depth_path) and
+                    os.path.exists(out_cam_path)):
+                continue
+
+            task = (
+                img_path,
+                depth_path,
+                out_img_path,
+                out_depth_path,
+                out_cam_path,
+                intrinsics
+            )
+            tasks.append(task)
+
+    # Process tasks in parallel
+    max_workers = args.max_workers
+    if max_workers is None:
+        max_workers = max(1, os.cpu_count() // 2)
+
+    with ProcessPoolExecutor(max_workers=max_workers) as executor:
+        list(tqdm(
+            executor.map(process_image, tasks),
+            total=len(tasks),
+            desc="Processing images"
+        ))
+
+
+if __name__ == "__main__":
+    main()

extern/CUT3R/datasets_preprocess/preprocess_hypersim.py

@@ -0,0 +1,268 @@
+#!/usr/bin/env python3
+"""
+Preprocess the Hypersim dataset.
+
+This script reads camera parameters from a CSV file, converts an OpenGL-style
+projection matrix into a camera intrinsic matrix, applies tone mapping, and
+saves processed RGB images, depth maps, and camera metadata into an output
+directory. Processing is done per scene and per camera view.
+
+Usage:
+    python preprocess_hypersim.py --hypersim_dir /path/to/hypersim \
+                                  --output_dir /path/to/processed_hypersim
+"""
+
+import argparse
+import os
+import shutil
+import time
+
+import cv2
+import h5py
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from PIL import Image
+from tqdm import tqdm
+
+# Ensure OpenEXR support for OpenCV.
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(
+        description="Preprocess the Hypersim dataset by converting projection "
+        "matrices, applying tone mapping, and saving processed outputs."
+    )
+    parser.add_argument(
+        "--hypersim_dir",
+        default="/path/to/hypersim",
+        help="Root directory of the Hypersim dataset.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        default="/path/to/processed_hypersim",
+        help="Output directory for processed Hypersim data.",
+    )
+    return parser
+
+
+def opengl_to_intrinsics(proj_matrix, width_pixels, height_pixels):
+    # Extract parameters from the projection matrix.
+    K00 = proj_matrix[0, 0] * width_pixels / 2.0
+    K01 = -proj_matrix[0, 1] * width_pixels / 2.0
+    K02 = (1.0 - proj_matrix[0, 2]) * width_pixels / 2.0
+    K11 = proj_matrix[1, 1] * height_pixels / 2.0
+    K12 = (1.0 + proj_matrix[1, 2]) * height_pixels / 2.0
+    return np.array([[K00, K01, K02], [0.0, K11, K12], [0.0, 0.0, 1.0]])
+
+
+def process_scene(args):
+    rootdir, outdir, scene_name = args
+    scene_outdir = os.path.join(outdir, scene_name)
+    os.makedirs(scene_outdir, exist_ok=True)
+    seq_dir = os.path.join(rootdir, scene_name)
+    seq_detail_dir = os.path.join(seq_dir, "_detail")
+    seq_images_dir = os.path.join(seq_dir, "images")
+
+    # Read global camera parameters from the CSV file.
+    all_metafile = os.path.join(rootdir, "metadata_camera_parameters.csv")
+    df_camera_parameters = pd.read_csv(all_metafile, index_col="scene_name")
+    df_ = df_camera_parameters.loc[scene_name]
+
+    width_pixels = int(df_["settings_output_img_width"])
+    height_pixels = int(df_["settings_output_img_height"])
+
+    M_proj = np.array(
+        [
+            [df_["M_proj_00"], df_["M_proj_01"], df_["M_proj_02"], df_["M_proj_03"]],
+            [df_["M_proj_10"], df_["M_proj_11"], df_["M_proj_12"], df_["M_proj_13"]],
+            [df_["M_proj_20"], df_["M_proj_21"], df_["M_proj_22"], df_["M_proj_23"]],
+            [df_["M_proj_30"], df_["M_proj_31"], df_["M_proj_32"], df_["M_proj_33"]],
+        ]
+    )
+
+    camera_intrinsics = opengl_to_intrinsics(
+        M_proj, width_pixels, height_pixels
+    ).astype(np.float32)
+    if camera_intrinsics[0, 1] != 0:
+        print(f"camera_intrinsics[0, 1] != 0: {camera_intrinsics[0, 1]}")
+        return
+
+    # Read world scale and camera IDs.
+    worldscale = (
+        pd.read_csv(
+            os.path.join(seq_detail_dir, "metadata_scene.csv"),
+            index_col="parameter_name",
+        )
+        .to_numpy()
+        .flatten()[0]
+        .astype(np.float32)
+    )
+    camera_ids = (
+        pd.read_csv(
+            os.path.join(seq_detail_dir, "metadata_cameras.csv"),
+            header=None,
+            skiprows=1,
+        )
+        .to_numpy()
+        .flatten()
+    )
+
+    # Tone mapping parameters.
+    gamma = 1.0 / 2.2  # Standard gamma correction exponent.
+    inv_gamma = 1.0 / gamma
+    percentile = 90  # Desired percentile brightness in the unmodified image.
+    brightness_nth_percentile_desired = 0.8  # Desired brightness after scaling.
+
+    for camera_id in camera_ids:
+        subscene_dir = os.path.join(scene_outdir, f"{camera_id}")
+        os.makedirs(subscene_dir, exist_ok=True)
+        camera_dir = os.path.join(seq_detail_dir, camera_id)
+        if not os.path.exists(camera_dir):
+            print(f"{camera_dir} does not exist.")
+            continue
+        color_dir = os.path.join(seq_images_dir, f"scene_{camera_id}_final_hdf5")
+        geometry_dir = os.path.join(seq_images_dir, f"scene_{camera_id}_geometry_hdf5")
+        if not (os.path.exists(color_dir) and os.path.exists(geometry_dir)):
+            print(f"{color_dir} or {geometry_dir} does not exist.")
+            continue
+
+        camera_positions_hdf5_file = os.path.join(
+            camera_dir, "camera_keyframe_positions.hdf5"
+        )
+        camera_orientations_hdf5_file = os.path.join(
+            camera_dir, "camera_keyframe_orientations.hdf5"
+        )
+
+        with h5py.File(camera_positions_hdf5_file, "r") as f:
+            camera_positions = f["dataset"][:]
+        with h5py.File(camera_orientations_hdf5_file, "r") as f:
+            camera_orientations = f["dataset"][:]
+
+        assert len(camera_positions) == len(
+            camera_orientations
+        ), f"len(camera_positions)={len(camera_positions)} != len(camera_orientations)={len(camera_orientations)}"
+
+        rgbs = sorted([f for f in os.listdir(color_dir) if f.endswith(".color.hdf5")])
+        depths = sorted(
+            [f for f in os.listdir(geometry_dir) if f.endswith(".depth_meters.hdf5")]
+        )
+        assert len(rgbs) == len(
+            depths
+        ), f"len(rgbs)={len(rgbs)} != len(depths)={len(depths)}"
+        exist_frame_ids = [int(f.split(".")[1]) for f in rgbs]
+        valid_camera_positions = camera_positions[exist_frame_ids]
+        valid_camera_orientations = camera_orientations[exist_frame_ids]
+
+        for i, (rgb, depth) in enumerate(tqdm(zip(rgbs, depths), total=len(rgbs))):
+            frame_id = int(rgb.split(".")[1])
+            assert frame_id == int(
+                depth.split(".")[1]
+            ), f"frame_id={frame_id} != {int(depth.split('.')[1])}"
+            # Tone mapping.
+            render_entity = os.path.join(
+                geometry_dir,
+                depth.replace("depth_meters.hdf5", "render_entity_id.hdf5"),
+            )
+            with h5py.File(os.path.join(color_dir, rgb), "r") as f:
+                color = f["dataset"][:]
+            with h5py.File(os.path.join(geometry_dir, depth), "r") as f:
+                distance = f["dataset"][:]
+            R_cam2world = valid_camera_orientations[i]
+            R_cam2world = R_cam2world @ np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
+            t_cam2world = valid_camera_positions[i] * worldscale
+            T_cam2world = np.eye(4)
+            T_cam2world[:3, :3] = R_cam2world
+            T_cam2world[:3, 3] = t_cam2world
+
+            if not np.isfinite(T_cam2world).all():
+                print(f"frame_id={frame_id} T_cam2world is not finite.")
+                continue
+
+            focal = (camera_intrinsics[0, 0] + camera_intrinsics[1, 1]) / 2.0
+            ImageplaneX = (
+                np.linspace(
+                    (-0.5 * width_pixels) + 0.5,
+                    (0.5 * width_pixels) - 0.5,
+                    width_pixels,
+                )
+                .reshape(1, width_pixels)
+                .repeat(height_pixels, 0)
+                .astype(np.float32)[:, :, None]
+            )
+            ImageplaneY = (
+                np.linspace(
+                    (-0.5 * height_pixels) + 0.5,
+                    (0.5 * height_pixels) - 0.5,
+                    height_pixels,
+                )
+                .reshape(height_pixels, 1)
+                .repeat(width_pixels, 1)
+                .astype(np.float32)[:, :, None]
+            )
+            ImageplaneZ = np.full([height_pixels, width_pixels, 1], focal, np.float32)
+            Imageplane = np.concatenate([ImageplaneX, ImageplaneY, ImageplaneZ], axis=2)
+
+            depth = distance / np.linalg.norm(Imageplane, axis=2) * focal
+
+            with h5py.File(render_entity, "r") as f:
+                render_entity_id = f["dataset"][:].astype(np.int32)
+            assert (render_entity_id != 0).all()
+            valid_mask = render_entity_id != -1
+
+            if np.sum(valid_mask) == 0:
+                scale = 1.0  # If there are no valid pixels, set scale to 1.0.
+            else:
+                brightness = (
+                    0.3 * color[:, :, 0] + 0.59 * color[:, :, 1] + 0.11 * color[:, :, 2]
+                )
+                brightness_valid = brightness[valid_mask]
+                eps = 0.0001  # Avoid division by zero.
+                brightness_nth_percentile_current = np.percentile(
+                    brightness_valid, percentile
+                )
+                if brightness_nth_percentile_current < eps:
+                    scale = 0.0
+                else:
+                    scale = (
+                        np.power(brightness_nth_percentile_desired, inv_gamma)
+                        / brightness_nth_percentile_current
+                    )
+
+            color = np.power(np.maximum(scale * color, 0), gamma)
+            color = np.clip(color, 0.0, 1.0)
+
+            out_rgb_path = os.path.join(subscene_dir, f"{frame_id:06d}_rgb.png")
+            Image.fromarray((color * 255).astype(np.uint8)).save(out_rgb_path)
+            out_depth_path = os.path.join(subscene_dir, f"{frame_id:06d}_depth.npy")
+            np.save(out_depth_path, depth.astype(np.float32))
+            out_cam_path = os.path.join(subscene_dir, f"{frame_id:06d}_cam.npz")
+            np.savez(
+                out_cam_path,
+                intrinsics=camera_intrinsics,
+                pose=T_cam2world.astype(np.float32),
+            )
+
+
+def main():
+    parser = get_parser()
+    args = parser.parse_args()
+
+    # Use placeholder paths to avoid personal/private information.
+    rootdir = args.hypersim_dir  # e.g., '/path/to/hypersim'
+    outdir = args.output_dir  # e.g., '/path/to/processed_hypersim'
+    os.makedirs(outdir, exist_ok=True)
+
+    import multiprocessing
+
+    scenes = sorted(
+        [f for f in os.listdir(rootdir) if os.path.isdir(os.path.join(rootdir, f))]
+    )
+    # Process each scene sequentially (or use multiprocessing if desired)
+    for scene in scenes:
+        process_scene((rootdir, outdir, scene))
+
+
+if __name__ == "__main__":
+    main()

extern/CUT3R/datasets_preprocess/preprocess_irs.py

@@ -0,0 +1,230 @@
+#!/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()

extern/CUT3R/datasets_preprocess/preprocess_mapfree.py

@@ -0,0 +1,76 @@
+import subprocess
+import os
+import argparse
+import glob
+
+
+def get_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--mapfree_dir",
+        default="mapfree/train/",
+    )
+    parser.add_argument(
+        "--colmap_dir",
+        default="mapfree/colmap",
+    )
+    parser.add_argument(
+        "--output_dir",
+        default="processed_mapfree",
+    )
+    return parser
+
+
+def run_patch_match_stereo(root_colmap_dir, root_img_dir):
+    scene_names = sorted(os.listdir(root_colmap_dir))
+    sub_dir_names = ["seq0", "seq1"]
+    for scene_name in scene_names:
+        scene_dir = os.path.join(root_colmap_dir, scene_name)
+        img_dir = os.path.join(root_img_dir, scene_name)
+        for i, sub in enumerate(sub_dir_names):
+            sub_dir = os.path.join(scene_dir, sub)
+            out_dir = os.path.join(scene_dir, f"dense{i}")
+            if not os.path.exists(sub_dir):
+                continue
+            if os.path.exists(out_dir) and os.path.exists(
+                os.path.join(out_dir, f"stereo/depth_maps/{sub}")
+            ):
+                if len(
+                    glob.glob(
+                        os.path.join(out_dir, f"stereo/depth_maps/{sub}/*geometric.bin")
+                    )
+                ) == len(glob.glob(os.path.join(img_dir, sub, "*.jpg"))):
+                    print(f"depth maps already computed, skip {sub_dir}")
+                    continue
+
+            print(sub_dir)
+            cmd = f"colmap image_undistorter \
+                    --image_path {img_dir} \
+                    --input_path {sub_dir} \
+                    --output_path {out_dir} \
+                    --output_type COLMAP;"
+
+            subprocess.call(cmd, shell=True)
+            cmd = f"rm -rf {out_dir}/images/seq{i}; rm -rf {out_dir}/sparse;"
+            cmd += f"cp -r {sub_dir} {out_dir}/sparse;"
+            cmd += f"cp -r {img_dir}/{sub} {out_dir}/images;"
+            subprocess.call(cmd, shell=True)
+
+            # we comment this because we have released the mvs results, but feel free to re-run the mvs
+
+            # cmd = f"colmap patch_match_stereo \
+            #         --workspace_path {out_dir} \
+            #         --workspace_format COLMAP \
+            #         --PatchMatchStereo.cache_size 512 \
+            #         --PatchMatchStereo.geom_consistency true"
+            # subprocess.call(cmd, shell=True)
+
+
+if __name__ == "__main__":
+    parser = get_parser()
+    args = parser.parse_args()
+    root_colmap_dir = args.colmap_dir
+    root_img_dir = args.mapfree_dir
+
+    # run patch match stereo
+    run_patch_match_stereo(root_colmap_dir, root_img_dir)

extern/CUT3R/datasets_preprocess/preprocess_mapfree2.py

@@ -0,0 +1,123 @@
+import os
+
+
+import os.path as osp
+
+from PIL import Image
+import numpy as np
+
+
+from tqdm import tqdm
+from read_write_model import run
+
+
+def get_parser():
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--mapfree_dir", default="")  # TODO
+    parser.add_argument("--output_dir", default="test_preprocess")  # TODO
+    return parser
+
+
+def main(rootdir, outdir):
+    os.makedirs(outdir, exist_ok=True)
+
+    envs = [f for f in os.listdir(rootdir) if os.path.isdir(osp.join(rootdir, f))]
+    for env in tqdm(envs):
+        subseqs = [
+            f
+            for f in os.listdir(osp.join(rootdir, env))
+            if os.path.isdir(osp.join(rootdir, env, f))
+        ]
+        for subseq in subseqs:
+            sparse_dir = osp.join(rootdir, env, subseq, "sparse")
+            images_dir = osp.join(rootdir, env, subseq, "images")
+            run(sparse_dir, sparse_dir)
+            intrins_file = sparse_dir + "/cameras.txt"
+            poses_file = sparse_dir + "/images.txt"
+
+            cam_params = {}
+            with open(intrins_file, "r") as f:
+                for line in f:
+                    if line.startswith("#"):
+                        continue
+                    parts = line.strip().split()
+                    if len(parts) == 0:
+                        continue
+                    cam_id = int(parts[0])
+                    fx = float(parts[4])
+                    fy = float(parts[5])
+                    cx = float(parts[6])
+                    cy = float(parts[7])
+                    cam_params[cam_id] = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
+
+            poses = []
+            images = []
+            intrinsics = []
+
+            with open(poses_file, "r") as f:
+                for i, line in enumerate(f):
+                    if line.startswith("#"):
+                        continue
+                    parts = line.strip().split()
+                    if len(parts) == 0:
+                        continue
+                    if "." in parts[0]:
+                        continue
+
+                    img_name = parts[-1]
+                    w, x, y, z = map(float, parts[1:5])
+                    R = np.array(
+                        [
+                            [
+                                1 - 2 * y * y - 2 * z * z,
+                                2 * x * y - 2 * z * w,
+                                2 * x * z + 2 * y * w,
+                            ],
+                            [
+                                2 * x * y + 2 * z * w,
+                                1 - 2 * x * x - 2 * z * z,
+                                2 * y * z - 2 * x * w,
+                            ],
+                            [
+                                2 * x * z - 2 * y * w,
+                                2 * y * z + 2 * x * w,
+                                1 - 2 * x * x - 2 * y * y,
+                            ],
+                        ]
+                    )
+                    tx, ty, tz = map(float, parts[5:8])
+                    cam_id = int(parts[-2])
+                    pose = np.eye(4)
+                    pose[:3, :3] = R
+                    pose[:3, 3] = [tx, ty, tz]
+                    poses.append(np.linalg.inv(pose))
+                    images.append(img_name)
+                    intrinsics.append(cam_params[cam_id])
+
+            os.makedirs(osp.join(outdir, env, subseq), exist_ok=True)
+            os.makedirs(osp.join(outdir, env, subseq, "rgb"), exist_ok=True)
+            os.makedirs(osp.join(outdir, env, subseq, "cam"), exist_ok=True)
+
+            for i, img_name in enumerate(tqdm(images)):
+                img_path = os.path.join(images_dir, img_name)
+                rgb = Image.open(img_path)
+                intrinsic = intrinsics[i]
+                pose = poses[i]
+                # save all
+                basename = img_name.split("/")[-1]
+                rgb.save(osp.join(outdir, env, subseq, "rgb", basename))
+                np.savez(
+                    osp.join(
+                        outdir, env, subseq, "cam", basename.replace(".jpg", ".npz")
+                    ),
+                    intrinsic=intrinsic,
+                    pose=pose,
+                )
+
+
+if __name__ == "__main__":
+    parser = get_parser()
+    args = parser.parse_args()
+    main(args.mapfree_dir, args.output_dir)

extern/CUT3R/datasets_preprocess/preprocess_megadepth.py

@@ -0,0 +1,229 @@
+# --------------------------------------------------------
+# Preprocessing code for the MegaDepth dataset
+# dataset at https://www.cs.cornell.edu/projects/megadepth/
+# --------------------------------------------------------
+import os
+import os.path as osp
+import collections
+from tqdm import tqdm
+import numpy as np
+
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+import cv2
+import h5py
+
+import path_to_root  # noqa
+from datasets_preprocess.utils.parallel import parallel_threads
+from datasets_preprocess.utils import cropping  # noqa
+
+
+def get_parser():
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--megadepth_dir", required=True)
+    parser.add_argument("--num_views", default=64, type=int)
+    parser.add_argument("--precomputed_sets", required=True)
+    parser.add_argument("--output_dir", default="data/dust3r_data/processed_megadepth")
+    return parser
+
+
+def main(db_root, pairs_path, output_dir, num_views):
+    os.makedirs(output_dir, exist_ok=True)
+
+    # load all pairs
+    data = np.load(pairs_path, allow_pickle=True)
+    scenes = data["scenes"]
+    images = data["images"]
+    sets = data["sets"]
+
+    # enumerate all unique images
+    todo = collections.defaultdict(set)
+    for line in sets:
+        for i in range(1, num_views + 1):
+            todo[line[0]].add(line[i])
+
+    # for each scene, load intrinsics and then parallel crops
+    for scene, im_idxs in tqdm(todo.items(), desc="Overall"):
+        scene, subscene = scenes[scene].split()
+        out_dir = osp.join(output_dir, scene, subscene)
+        os.makedirs(out_dir, exist_ok=True)
+
+        # load all camera params
+        _, pose_w2cam, intrinsics = _load_kpts_and_poses(
+            db_root, scene, subscene, intrinsics=True
+        )
+
+        in_dir = osp.join(db_root, scene, "dense" + subscene)
+        # args = [(in_dir, img, intrinsics[img], pose_w2cam[img], out_dir)
+        #         for img in [images[im_id] for im_id in im_idxs]]
+        args = [
+            (in_dir, img, intrinsics[img], pose_w2cam[img], out_dir)
+            for img in intrinsics.keys()
+            if os.path.exists(osp.join(in_dir, "imgs", img))
+        ]
+        parallel_threads(
+            resize_one_image,
+            args,
+            star_args=True,
+            front_num=0,
+            leave=False,
+            desc=f"{scene}/{subscene}",
+        )
+
+    # save pairs
+    print("Done! prepared all images in", output_dir)
+
+
+def resize_one_image(root, tag, K_pre_rectif, pose_w2cam, out_dir):
+    if osp.isfile(osp.join(out_dir, tag + ".npz")):
+        return
+
+    # load image
+    img = cv2.cvtColor(
+        cv2.imread(osp.join(root, "imgs", tag), cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB
+    )
+    H, W = img.shape[:2]
+
+    # load depth
+    with h5py.File(osp.join(root, "depths", osp.splitext(tag)[0] + ".h5"), "r") as hd5:
+        depthmap = np.asarray(hd5["depth"])
+
+    # rectify = undistort the intrinsics
+    imsize_pre, K_pre, distortion = K_pre_rectif
+    imsize_post = img.shape[1::-1]
+    K_post = cv2.getOptimalNewCameraMatrix(
+        K_pre,
+        distortion,
+        imsize_pre,
+        alpha=0,
+        newImgSize=imsize_post,
+        centerPrincipalPoint=True,
+    )[0]
+
+    # downscale
+    img_out, depthmap_out, intrinsics_out, R_in2out = _downscale_image(
+        K_post, img, depthmap, resolution_out=(800, 600)
+    )
+
+    # write everything
+    img_out.save(osp.join(out_dir, tag + ".jpg"), quality=90)
+    cv2.imwrite(osp.join(out_dir, tag + ".exr"), depthmap_out)
+
+    camout2world = np.linalg.inv(pose_w2cam)
+    camout2world[:3, :3] = camout2world[:3, :3] @ R_in2out.T
+    np.savez(
+        osp.join(out_dir, tag + ".npz"),
+        intrinsics=intrinsics_out,
+        cam2world=camout2world,
+    )
+
+
+def _downscale_image(camera_intrinsics, image, depthmap, resolution_out=(512, 384)):
+    H, W = image.shape[:2]
+    resolution_out = sorted(resolution_out)[:: +1 if W < H else -1]
+
+    image, depthmap, intrinsics_out = cropping.rescale_image_depthmap(
+        image, depthmap, camera_intrinsics, resolution_out, force=False
+    )
+    R_in2out = np.eye(3)
+
+    return image, depthmap, intrinsics_out, R_in2out
+
+
+def _load_kpts_and_poses(root, scene_id, subscene, z_only=False, intrinsics=False):
+    if intrinsics:
+        with open(
+            os.path.join(
+                root, scene_id, "sparse", "manhattan", subscene, "cameras.txt"
+            ),
+            "r",
+        ) as f:
+            raw = f.readlines()[3:]  # skip the header
+
+        camera_intrinsics = {}
+        for camera in raw:
+            camera = camera.split(" ")
+            width, height, focal, cx, cy, k0 = [float(elem) for elem in camera[2:]]
+            K = np.eye(3)
+            K[0, 0] = focal
+            K[1, 1] = focal
+            K[0, 2] = cx
+            K[1, 2] = cy
+            camera_intrinsics[int(camera[0])] = (
+                (int(width), int(height)),
+                K,
+                (k0, 0, 0, 0),
+            )
+
+    with open(
+        os.path.join(root, scene_id, "sparse", "manhattan", subscene, "images.txt"), "r"
+    ) as f:
+        raw = f.read().splitlines()[4:]  # skip the header
+
+    extract_pose = (
+        colmap_raw_pose_to_principal_axis if z_only else colmap_raw_pose_to_RT
+    )
+
+    poses = {}
+    points3D_idxs = {}
+    camera = []
+
+    for image, points in zip(raw[::2], raw[1::2]):
+        image = image.split(" ")
+        points = points.split(" ")
+
+        image_id = image[-1]
+        camera.append(int(image[-2]))
+
+        # find the principal axis
+        raw_pose = [float(elem) for elem in image[1:-2]]
+        poses[image_id] = extract_pose(raw_pose)
+
+        current_points3D_idxs = {int(i) for i in points[2::3] if i != "-1"}
+        assert -1 not in current_points3D_idxs, bb()
+        points3D_idxs[image_id] = current_points3D_idxs
+
+    if intrinsics:
+        image_intrinsics = {
+            im_id: camera_intrinsics[cam] for im_id, cam in zip(poses, camera)
+        }
+        return points3D_idxs, poses, image_intrinsics
+    else:
+        return points3D_idxs, poses
+
+
+def colmap_raw_pose_to_principal_axis(image_pose):
+    qvec = image_pose[:4]
+    qvec = qvec / np.linalg.norm(qvec)
+    w, x, y, z = qvec
+    z_axis = np.float32(
+        [2 * x * z - 2 * y * w, 2 * y * z + 2 * x * w, 1 - 2 * x * x - 2 * y * y]
+    )
+    return z_axis
+
+
+def colmap_raw_pose_to_RT(image_pose):
+    qvec = image_pose[:4]
+    qvec = qvec / np.linalg.norm(qvec)
+    w, x, y, z = qvec
+    R = np.array(
+        [
+            [1 - 2 * y * y - 2 * z * z, 2 * x * y - 2 * z * w, 2 * x * z + 2 * y * w],
+            [2 * x * y + 2 * z * w, 1 - 2 * x * x - 2 * z * z, 2 * y * z - 2 * x * w],
+            [2 * x * z - 2 * y * w, 2 * y * z + 2 * x * w, 1 - 2 * x * x - 2 * y * y],
+        ]
+    )
+    # principal_axis.append(R[2, :])
+    t = image_pose[4:7]
+    # World-to-Camera pose
+    current_pose = np.eye(4)
+    current_pose[:3, :3] = R
+    current_pose[:3, 3] = t
+    return current_pose
+
+
+if __name__ == "__main__":
+    parser = get_parser()
+    args = parser.parse_args()
+    main(args.megadepth_dir, args.precomputed_sets, args.output_dir, args.num_views)

extern/CUT3R/datasets_preprocess/preprocess_mp3d.py

@@ -0,0 +1,217 @@
+#!/usr/bin/env python3
+"""
+Preprocess the Matterport3D (MP3D) dataset.
+
+This script reads camera parameters and overlap data from a configuration file,
+processes RGB images and corresponding depth images, adjusts camera poses using a
+conversion matrix, and then saves the processed images, depth maps, and camera
+metadata into separate output directories.
+
+Usage:
+    python preprocess_mp3d.py --root_dir /path/to/data_mp3d/v1/scans \
+                              --out_dir /path/to/processed_mp3d
+"""
+
+import os
+import numpy as np
+import cv2
+import shutil
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from tqdm import tqdm
+import argparse
+
+
+def process_image(args):
+    """
+    Process a single image: reads the RGB image and depth image, normalizes the depth,
+    adjusts the camera pose using a conversion matrix, and saves the processed outputs.
+
+    Parameters:
+      args: tuple containing
+         (i, paths, K, pose, img_dir, depth_dir, out_rgb_dir, out_depth_dir, out_cam_dir, R_conv)
+         where:
+           i             - the frame index
+           paths         - tuple of (depth filename, RGB filename)
+           K             - camera intrinsics matrix (3x3 NumPy array)
+           pose          - camera pose (4x4 NumPy array)
+           img_dir       - directory containing RGB images
+           depth_dir     - directory containing depth images
+           out_rgb_dir   - output directory for processed RGB images
+           out_depth_dir - output directory for processed depth maps
+           out_cam_dir   - output directory for processed camera metadata
+           R_conv        - a 4x4 conversion matrix (NumPy array)
+    Returns:
+      None if successful, or an error string if processing fails.
+    """
+    (
+        i,
+        paths,
+        K,
+        pose,
+        img_dir,
+        depth_dir,
+        out_rgb_dir,
+        out_depth_dir,
+        out_cam_dir,
+        R_conv,
+    ) = args
+
+    depth_path, img_path = paths
+    img_path_full = os.path.join(img_dir, img_path)
+    depth_path_full = os.path.join(depth_dir, depth_path)
+
+    try:
+        # Read depth image using OpenCV (assumed to be stored with 16-bit depth)
+        depth = cv2.imread(depth_path_full, cv2.IMREAD_ANYDEPTH).astype(np.float32)
+        depth = depth / 4000.0  # Normalize depth (adjust this factor as needed)
+
+        # Adjust the camera pose with the conversion matrix
+        pose_adjusted = pose @ R_conv
+
+        # Generate output filenames using a zero-padded frame index.
+        basename = f"{i:06d}"
+        out_img_path = os.path.join(out_rgb_dir, basename + ".png")
+        out_depth_path = os.path.join(out_depth_dir, basename + ".npy")
+        out_cam_path = os.path.join(out_cam_dir, basename + ".npz")
+
+        # Copy the RGB image.
+        shutil.copyfile(img_path_full, out_img_path)
+
+        # Save the depth map.
+        np.save(out_depth_path, depth)
+
+        # Save the camera intrinsics and adjusted pose.
+        np.savez(out_cam_path, intrinsics=K, pose=pose_adjusted)
+
+    except Exception as e:
+        return f"Error processing image {img_path}: {e}"
+
+    return None
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Preprocess MP3D scans: convert and save RGB images, depth maps, and camera metadata."
+    )
+    parser.add_argument(
+        "--root_dir",
+        type=str,
+        default="/path/to/data_mp3d/v1/scans",
+        help="Root directory of the raw MP3D data.",
+    )
+    parser.add_argument(
+        "--out_dir",
+        type=str,
+        default="/path/to/processed_mp3d",
+        help="Output directory for processed MP3D data.",
+    )
+    args = parser.parse_args()
+
+    root = args.root_dir
+    out_dir = args.out_dir
+
+    # List sequence directories (each scan is stored as a separate directory).
+    seqs = sorted([d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))])
+
+    # Define a conversion matrix from MP3D to the desired coordinate system.
+    R_conv = np.array(
+        [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]], dtype=np.float32
+    )
+
+    for seq in tqdm(seqs, desc="Sequences"):
+        # The sequence directory structure assumes that images and depth files are stored
+        # under a subdirectory with the same name as the sequence.
+        seq_dir = os.path.join(root, seq, seq)
+
+        img_dir = os.path.join(seq_dir, "undistorted_color_images")
+        depth_dir = os.path.join(seq_dir, "undistorted_depth_images")
+        cam_file = os.path.join(seq_dir, "undistorted_camera_parameters", f"{seq}.conf")
+        overlap_file = os.path.join(seq_dir, "image_overlap_data", f"{seq}_iis.txt")
+
+        # Read overlap data and save it (optional).
+        overlap = []
+        with open(overlap_file, "r") as f:
+            for line in f:
+                parts = line.split()
+                overlap.append([int(parts[1]), int(parts[2]), float(parts[3])])
+        overlap = np.array(overlap)
+        os.makedirs(os.path.join(out_dir, seq), exist_ok=True)
+        np.save(os.path.join(out_dir, seq, "overlap.npy"), overlap)
+
+        # Read camera parameters from a configuration file.
+        intrinsics = []
+        camera_poses = []
+        image_files = []
+
+        with open(cam_file, "r") as file:
+            lines = file.readlines()
+        current_intrinsics = None
+        for line in lines:
+            parts = line.split()
+            if not parts:
+                continue
+            if parts[0] == "intrinsics_matrix":
+                # Extract intrinsic parameters.
+                fx, cx, fy, cy = (
+                    float(parts[1]),
+                    float(parts[3]),
+                    float(parts[5]),
+                    float(parts[6]),
+                )
+                current_intrinsics = np.array(
+                    [[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32
+                )
+            elif parts[0] == "scan":
+                # Read the image filenames and camera pose.
+                depth_image = parts[1]
+                color_image = parts[2]
+                image_files.append((depth_image, color_image))
+                matrix_values = list(map(float, parts[3:]))
+                camera_pose = np.array(matrix_values).reshape(4, 4)
+                camera_poses.append(camera_pose)
+                if current_intrinsics is not None:
+                    intrinsics.append(current_intrinsics.copy())
+
+        if not (len(image_files) == len(intrinsics) == len(camera_poses)):
+            print(f"Inconsistent data in sequence {seq}")
+            continue
+
+        # Prepare output directories.
+        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)
+
+        tasks = []
+        for i, (paths, K, pose) in enumerate(
+            zip(image_files, intrinsics, camera_poses)
+        ):
+            args_task = (
+                i,
+                paths,
+                K,
+                pose,
+                img_dir,
+                depth_dir,
+                out_rgb_dir,
+                out_depth_dir,
+                out_cam_dir,
+                R_conv,
+            )
+            tasks.append(args_task)
+
+        num_workers = os.cpu_count() // 2
+        with ProcessPoolExecutor(max_workers=num_workers) as executor:
+            futures = {executor.submit(process_image, 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()

extern/CUT3R/datasets_preprocess/preprocess_mvimgnet.py

@@ -0,0 +1,323 @@
+#!/usr/bin/env python3
+"""
+Preprocess the MVImgNet dataset.
+
+This script processes MVImgNet sequences by:
+  - Loading a sparse SFM reconstruction.
+  - Undistorting and rescaling RGB images.
+  - Converting COLMAP intrinsics between conventions.
+  - Saving the processed images and camera metadata.
+
+Usage:
+  python preprocess_mvimgnet.py --data_dir /path/to/MVImgNet_data \
+                                --pcd_dir /path/to/MVPNet \
+                                --output_dir /path/to/processed_mvimgnet
+"""
+
+import os
+import os.path as osp
+import argparse
+import numpy as np
+import open3d as o3d
+import pyrender
+import PIL.Image as Image
+import cv2
+import shutil
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+
+# Import your custom SFM processing function.
+from read_write_model import run  # Assumed to be available
+
+# Try to set up resampling filters from PIL.
+try:
+    lanczos = Image.Resampling.LANCZOS
+    bicubic = Image.Resampling.BICUBIC
+except AttributeError:
+    lanczos = Image.LANCZOS
+    bicubic = Image.BICUBIC
+
+# Conversion matrix from COLMAP (or OpenGL) to OpenCV conventions.
+OPENGL_TO_OPENCV = np.float32(
+    [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]]
+)
+
+
+# -----------------------------------------------------------------------------
+# Helper Classes and Functions
+# -----------------------------------------------------------------------------
+class ImageList:
+    """Convenience class to apply operations to a list of images."""
+
+    def __init__(self, images):
+        if not isinstance(images, (list, tuple)):
+            images = [images]
+        self.images = []
+        for image in images:
+            if not isinstance(image, Image.Image):
+                image = Image.fromarray(image)
+            self.images.append(image)
+
+    def __len__(self):
+        return len(self.images)
+
+    def to_pil(self):
+        return tuple(self.images) if len(self.images) > 1 else self.images[0]
+
+    @property
+    def size(self):
+        sizes = [im.size for im in self.images]
+        assert all(s == sizes[0] for s in sizes)
+        return sizes[0]
+
+    def resize(self, *args, **kwargs):
+        return ImageList([im.resize(*args, **kwargs) for im in self.images])
+
+    def crop(self, *args, **kwargs):
+        return ImageList([im.crop(*args, **kwargs) for im in self.images])
+
+
+def colmap_to_opencv_intrinsics(K):
+    """
+    Convert COLMAP intrinsics (with pixel centers at (0.5, 0.5)) to OpenCV convention.
+    """
+    K = K.copy()
+    K[0, 2] -= 0.5
+    K[1, 2] -= 0.5
+    return K
+
+
+def opencv_to_colmap_intrinsics(K):
+    """
+    Convert OpenCV intrinsics (with pixel centers at (0, 0)) to COLMAP convention.
+    """
+    K = K.copy()
+    K[0, 2] += 0.5
+    K[1, 2] += 0.5
+    return K
+
+
+def rescale_image_depthmap(
+    image, depthmap, camera_intrinsics, output_resolution, force=True
+):
+    """
+    Jointly rescale an image (and its depthmap) so that the output resolution is at least the desired value.
+
+    Args:
+      image: Input image (as a PIL.Image or compatible object).
+      depthmap: A corresponding depth map (or None).
+      camera_intrinsics: A 3x3 NumPy array of intrinsics.
+      output_resolution: (width, height) desired resolution.
+      force: If True, always rescale even if the image is smaller.
+
+    Returns:
+      Tuple of (rescaled image, rescaled depthmap, updated intrinsics).
+    """
+    image = ImageList(image)
+    input_resolution = np.array(image.size)  # (W, H)
+    output_resolution = np.array(output_resolution)
+    if depthmap is not None:
+        assert tuple(depthmap.shape[:2]) == image.size[::-1]
+    scale_final = max(output_resolution / image.size) + 1e-8
+    if scale_final >= 1 and not force:
+        return image.to_pil(), depthmap, camera_intrinsics
+    output_resolution = np.floor(input_resolution * scale_final).astype(int)
+    image = image.resize(
+        tuple(output_resolution), resample=lanczos if scale_final < 1 else bicubic
+    )
+    if depthmap is not None:
+        depthmap = cv2.resize(
+            depthmap, tuple(output_resolution), interpolation=cv2.INTER_NEAREST
+        )
+    camera_intrinsics = camera_matrix_of_crop(
+        camera_intrinsics, input_resolution, output_resolution, scaling=scale_final
+    )
+    return image.to_pil(), depthmap, camera_intrinsics
+
+
+def camera_matrix_of_crop(
+    input_camera_matrix,
+    input_resolution,
+    output_resolution,
+    scaling=1,
+    offset_factor=0.5,
+    offset=None,
+):
+    """
+    Update the camera intrinsics to account for a rescaling (or cropping) of the image.
+    """
+    margins = np.asarray(input_resolution) * scaling - output_resolution
+    assert np.all(margins >= 0.0)
+    if offset is None:
+        offset = offset_factor * margins
+    output_camera_matrix_colmap = opencv_to_colmap_intrinsics(input_camera_matrix)
+    output_camera_matrix_colmap[:2, :] *= scaling
+    output_camera_matrix_colmap[:2, 2] -= offset
+    output_camera_matrix = colmap_to_opencv_intrinsics(output_camera_matrix_colmap)
+    return output_camera_matrix
+
+
+def pose_from_qwxyz_txyz(elems):
+    """
+    Convert a quaternion (qw, qx, qy, qz) and translation (tx, ty, tz) to a 4x4 pose.
+    Returns the inverse of the computed pose (i.e. cam2world).
+    """
+    from scipy.spatial.transform import Rotation
+
+    qw, qx, qy, qz, tx, ty, tz = map(float, elems)
+    pose = np.eye(4)
+    pose[:3, :3] = Rotation.from_quat((qx, qy, qz, qw)).as_matrix()
+    pose[:3, 3] = (tx, ty, tz)
+    return np.linalg.inv(pose)
+
+
+def load_sfm(sfm_dir):
+    """
+    Load sparse SFM data from COLMAP output files.
+
+    Returns a tuple (img_idx, img_infos) where:
+      - img_idx: A dict mapping image filename to index.
+      - img_infos: A dict of image information (including intrinsics, file path, and camera pose).
+    """
+    with open(osp.join(sfm_dir, "cameras.txt"), "r") as f:
+        raw = f.read().splitlines()[3:]  # skip header
+    intrinsics = {}
+    for camera in raw:
+        camera = camera.split(" ")
+        intrinsics[int(camera[0])] = [camera[1]] + [float(x) for x in camera[2:]]
+    with open(osp.join(sfm_dir, "images.txt"), "r") as f:
+        raw = f.read().splitlines()
+        raw = [line for line in raw if not line.startswith("#")]
+    img_idx = {}
+    img_infos = {}
+    for image, points in zip(raw[0::2], raw[1::2]):
+        image = image.split(" ")
+        points = points.split(" ")
+        idx = image[0]
+        img_name = image[-1]
+        assert img_name not in img_idx, f"Duplicate image: {img_name}"
+        img_idx[img_name] = idx
+        current_points2D = {
+            int(i): (float(x), float(y))
+            for i, x, y in zip(points[2::3], points[0::3], points[1::3])
+            if i != "-1"
+        }
+        img_infos[idx] = dict(
+            intrinsics=intrinsics[int(image[-2])],
+            path=img_name,
+            frame_id=img_name,
+            cam_to_world=pose_from_qwxyz_txyz(image[1:-2]),
+            sparse_pts2d=current_points2D,
+        )
+    return img_idx, img_infos
+
+
+def undistort_images(intrinsics, rgb):
+    """
+    Given camera intrinsics (in COLMAP convention) and an RGB image, compute and return
+    the corresponding OpenCV intrinsics along with the (unchanged) image.
+    """
+    width = int(intrinsics[1])
+    height = int(intrinsics[2])
+    fx = intrinsics[3]
+    fy = intrinsics[4]
+    cx = intrinsics[5]
+    cy = intrinsics[6]
+    K = np.zeros([3, 3])
+    K[0, 0] = fx
+    K[0, 2] = cx
+    K[1, 1] = fy
+    K[1, 2] = cy
+    K[2, 2] = 1
+    return width, height, K, rgb
+
+
+# -----------------------------------------------------------------------------
+# Processing Functions
+# -----------------------------------------------------------------------------
+def process_sequence(category, obj, data_dir, output_dir):
+    """
+    Process a single sequence from MVImgNet.
+
+    Steps:
+      1. Load the point cloud (from the MVPNet directory) and create a mesh (using Pyrender) for visualization.
+      2. Load the SFM reconstruction from COLMAP files.
+      3. For each image in the SFM output:
+         a. Load the image.
+         b. Undistort and rescale it.
+         c. Update the camera intrinsics.
+         d. Save the processed image and camera metadata.
+    """
+
+    # Define directories.
+    seq_dir = osp.join(data_dir, "MVImgNet_by_categories", category, obj[:-4])
+    rgb_dir = osp.join(seq_dir, "images")
+    sfm_dir = osp.join(seq_dir, "sparse", "0")
+
+    output_scene_dir = osp.join(output_dir, f"{category}_{obj[:-4]}")
+    output_rgb_dir = osp.join(output_scene_dir, "rgb")
+    output_cam_dir = osp.join(output_scene_dir, "cam")
+    os.makedirs(output_rgb_dir, exist_ok=True)
+    os.makedirs(output_cam_dir, exist_ok=True)
+
+    # Run custom SFM processing.
+    run(sfm_dir, sfm_dir)
+    img_idx, img_infos = load_sfm(sfm_dir)
+
+    for imgname in img_idx:
+        idx = img_idx[imgname]
+        info = img_infos[idx]
+        rgb_path = osp.join(rgb_dir, info["path"])
+        if not osp.exists(rgb_path):
+            continue
+        rgb = np.array(Image.open(rgb_path))
+        _, _, K, rgb = undistort_images(info["intrinsics"], rgb)
+        intrinsics = colmap_to_opencv_intrinsics(K)
+        # Rescale image to a target resolution (e.g., 640x480) preserving aspect ratio.
+        image, _, intrinsics = rescale_image_depthmap(
+            rgb, None, intrinsics, (640, int(640 * 3.0 / 4))
+        )
+        intrinsics = opencv_to_colmap_intrinsics(intrinsics)
+        out_img_path = osp.join(output_rgb_dir, info["path"][:-3] + "jpg")
+        image.save(out_img_path)
+        out_cam_path = osp.join(output_cam_dir, info["path"][:-3] + "npz")
+        np.savez(out_cam_path, intrinsics=intrinsics, pose=info["cam_to_world"])
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Preprocess MVImgNet dataset: undistort, rescale images, and save camera parameters."
+    )
+    parser.add_argument(
+        "--data_dir",
+        type=str,
+        default="/path/to/MVImgNet_data",
+        help="Directory containing MVImgNet data (images and point clouds).",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="/path/to/processed_mvimgnet",
+        help="Directory where processed data will be saved.",
+    )
+    args = parser.parse_args()
+
+    data_dir = args.data_dir
+    output_dir = args.output_dir
+
+    # Get list of categories.
+    categories = sorted(
+        [
+            d
+            for d in os.listdir(osp.join(data_dir, "MVImgNet_by_categories"))
+            if osp.isdir(osp.join(data_dir, "MVImgNet_by_categories", d))
+        ]
+    )
+    for cat in categories:
+        objects = sorted(os.listdir(osp.join(data_dir, "MVImgNet_by_categories", cat)))
+        for obj in objects:
+            process_sequence(cat, obj, data_dir, output_dir)
+
+
+if __name__ == "__main__":
+    main()

extern/CUT3R/datasets_preprocess/preprocess_mvs_synth.py

@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+"""
+Preprocess the MVS Synth dataset.
+
+This script processes each sequence in a given dataset directory by:
+  - Reading the RGB image, EXR depth image, and JSON camera parameters.
+  - Computing the camera pose from the extrinsic matrix (with a conversion matrix applied).
+  - Creating a simple camera intrinsics matrix from the provided focal lengths and principal point.
+  - Copying the RGB image (as JPG), saving the depth (as a NumPy array), and saving the camera data (as a NPZ file).
+
+Usage:
+    python preprocess_mvs_synth.py --root_dir /path/to/data_mvs_synth/GTAV_720/ \
+                                   --out_dir /path/to/processed_mvs_synth \
+                                   --num_workers 32
+"""
+
+import os
+import shutil
+import json
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from tqdm import tqdm
+import numpy as np
+import cv2
+import argparse
+
+# Ensure OpenEXR support if needed
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+
+# Conversion matrix (example conversion, adjust if needed)
+R_conv = np.array(
+    [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=np.float32
+)
+
+
+def process_basename(seq, basename, root_dir, out_dir):
+    """
+    Process a single frame identified by 'basename' within a given sequence.
+
+    Reads the RGB image, depth (EXR) file, and camera parameters (JSON file),
+    computes the adjusted camera pose, builds the camera intrinsics matrix,
+    and saves the processed outputs.
+
+    Parameters:
+      seq (str): The sequence (subdirectory) name.
+      basename (str): The basename of the file (without extension).
+      root_dir (str): Root directory containing the raw data.
+      out_dir (str): Output directory where processed data will be saved.
+
+    Returns:
+      None on success, or an error string on failure.
+    """
+    try:
+        # Define input directories.
+        seq_dir = os.path.join(root_dir, seq)
+        img_dir = os.path.join(seq_dir, "images")
+        depth_dir = os.path.join(seq_dir, "depths")
+        cam_dir = os.path.join(seq_dir, "poses")
+
+        # Define input file paths.
+        img_path = os.path.join(img_dir, basename + ".png")
+        depth_path = os.path.join(depth_dir, basename + ".exr")
+        cam_path = os.path.join(cam_dir, basename + ".json")
+
+        # Define output directories.
+        out_seq_dir = os.path.join(out_dir, seq)
+        out_img_dir = os.path.join(out_seq_dir, "rgb")
+        out_depth_dir = os.path.join(out_seq_dir, "depth")
+        out_cam_dir = os.path.join(out_seq_dir, "cam")
+        os.makedirs(out_img_dir, exist_ok=True)
+        os.makedirs(out_depth_dir, exist_ok=True)
+        os.makedirs(out_cam_dir, exist_ok=True)
+
+        # Define output file paths.
+        out_img_path = os.path.join(out_img_dir, basename + ".jpg")
+        out_depth_path = os.path.join(out_depth_dir, basename + ".npy")
+        out_cam_path = os.path.join(out_cam_dir, basename + ".npz")
+
+        # Read and process camera parameters.
+        with open(cam_path, "r") as f:
+            cam_data = json.load(f)
+        c_x = cam_data["c_x"]
+        c_y = cam_data["c_y"]
+        f_x = cam_data["f_x"]
+        f_y = cam_data["f_y"]
+        extrinsic = np.array(cam_data["extrinsic"])
+        # Invert extrinsic matrix to obtain camera-to-world pose.
+        pose = np.linalg.inv(extrinsic)
+        # Apply conversion matrix.
+        pose = R_conv @ pose
+
+        # Build a simple intrinsics matrix.
+        intrinsics = np.array(
+            [[f_x, 0, c_x], [0, f_y, c_y], [0, 0, 1]], dtype=np.float32
+        )
+
+        if np.any(np.isinf(pose)) or np.any(np.isnan(pose)):
+            raise ValueError(f"Invalid pose for {basename}")
+
+        # Read depth image.
+        depth = cv2.imread(depth_path, cv2.IMREAD_ANYDEPTH).astype(np.float32)
+        depth[np.isinf(depth)] = 0.0  # Clean up any infinite values
+
+        # Save the processed data.
+        shutil.copyfile(img_path, out_img_path)
+        np.save(out_depth_path, depth)
+        np.savez(out_cam_path, intrinsics=intrinsics, pose=pose)
+
+    except Exception as e:
+        return f"Error processing {seq}/{basename}: {e}"
+
+    return None
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Preprocess MVS Synth dataset: convert images, depth, and camera data."
+    )
+    parser.add_argument(
+        "--root_dir",
+        type=str,
+        default="/path/to/data_mvs_synth/GTAV_720/",
+        help="Root directory of the raw MVS Synth data.",
+    )
+    parser.add_argument(
+        "--out_dir",
+        type=str,
+        default="/path/to/processed_mvs_synth",
+        help="Output directory for processed data.",
+    )
+    parser.add_argument(
+        "--num_workers", type=int, default=32, help="Number of parallel workers."
+    )
+    args = parser.parse_args()
+
+    root_dir = args.root_dir
+    out_dir = args.out_dir
+
+    # Get list of sequence directories.
+    seqs = sorted(
+        [d for d in os.listdir(root_dir) if os.path.isdir(os.path.join(root_dir, d))]
+    )
+
+    # Pre-create output directories for each sequence.
+    for seq in seqs:
+        out_seq_dir = os.path.join(out_dir, seq)
+        os.makedirs(os.path.join(out_seq_dir, "rgb"), exist_ok=True)
+        os.makedirs(os.path.join(out_seq_dir, "depth"), exist_ok=True)
+        os.makedirs(os.path.join(out_seq_dir, "cam"), exist_ok=True)
+
+    # Build list of processing tasks.
+    tasks = []
+    for seq in seqs:
+        seq_dir = os.path.join(root_dir, seq)
+        img_dir = os.path.join(seq_dir, "images")
+        basenames = sorted([d[:-4] for d in os.listdir(img_dir) if d.endswith(".png")])
+        for basename in basenames:
+            tasks.append((seq, basename, root_dir, out_dir))
+
+    num_workers = args.num_workers
+    print(f"Processing {len(tasks)} tasks using {num_workers} workers...")
+
+    with ProcessPoolExecutor(max_workers=num_workers) as executor:
+        futures = {executor.submit(process_basename, *task): task[1] for task in tasks}
+        for future in tqdm(
+            as_completed(futures), total=len(futures), desc="Processing"
+        ):
+            error = future.result()
+            if error:
+                print(error)
+
+
+if __name__ == "__main__":
+    main()