Initial commit

.gitattributes

@@ -0,0 +1,37 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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
+*.whl filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,159 @@
+# Local notebooks used for local debug
+notebooks_local/
+wandb/
+
+# Gradio files
+
+served_files/
+
+# hidden folders 
+.*/**
+
+# The rest is taken from https://github.com/Anttwo/SuGaR
+*.pt
+*.pth
+output*
+*.slurm
+*.pyc
+*.ply
+*.obj
+sugar_tests.ipynb
+sugar_sh_tests.ipynb
+
+# To remove
+frosting*
+extract_shell.py
+train_frosting_refined.py
+train_frosting.py
+run_frosting_viewer.py
+slurm_a100.sh
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+learnableearthparser/fast_sampler/_sampler.c

LICENSE.txt

@@ -0,0 +1,12 @@
+Copyright 2025, Dmytro Kotovenko, Olga Grebenkova, Björn Ommer
+Redistribution and use in source and binary forms, with or without modification, are permitted for non-commercial academic research and/or non-commercial personal use only provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+Any use of this software beyond the above specified conditions requires a separate license. Please contact the copyright holders to discuss license terms.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

README.md

@@ -0,0 +1,14 @@
+---
+title: EDGS
+emoji: 🎥
+colorFrom: pink
+colorTo: blue
+sdk: gradio
+sdk_version: 5.25.2
+app_file: app.py
+pinned: false
+python_version: "3.10"
+license: other
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,434 @@
+import spaces
+import torch
+import os
+import shutil
+import tempfile
+import argparse
+import gradio as gr
+import sys
+import io
+import subprocess
+from PIL import Image
+import numpy as np
+from hydra import initialize, compose
+import hydra
+from omegaconf import OmegaConf
+import time
+
+def install_submodules():
+    subprocess.check_call(['pip', 'install', './submodules/RoMa'])
+    
+
+
+STATIC_FILE_SERVING_FOLDER = "./served_files"
+MODEL_PATH = None
+os.makedirs(STATIC_FILE_SERVING_FOLDER, exist_ok=True)
+
+trainer = None
+
+class Tee(io.TextIOBase):
+
+    def __init__(self, *streams):
+        self.streams = streams
+        
+    def write(self, data):
+        for stream in self.streams:
+            stream.write(data)
+        return len(data)
+
+    def flush(self):
+        for stream in self.streams:
+            stream.flush()
+
+
+def capture_logs(func, *args, **kwargs):
+    log_capture_string = io.StringIO()
+    tee = Tee(sys.__stdout__, log_capture_string)
+
+    with contextlib.redirect_stdout(tee):
+        result = func(*args, **kwargs)
+
+    return result, log_capture_string.getvalue()
+
+
+@spaces.GPU(duration=350)
+# Training Pipeline
+def run_training_pipeline(scene_dir, 
+                          num_ref_views=16, 
+                          num_corrs_per_view=20000, 
+                          num_steps=1_000,
+                          mode_toggle="Ours (EDGS)"):
+    with initialize(config_path="./configs", version_base="1.1"):
+        cfg = compose(config_name="train")
+
+    scene_name = os.path.basename(scene_dir)
+    model_output_dir = f"./outputs/{scene_name}_trained"
+
+    cfg.wandb.mode = "disabled"
+    cfg.gs.dataset.model_path = model_output_dir
+    cfg.gs.dataset.source_path = scene_dir
+    cfg.gs.dataset.images = "images"
+
+    cfg.gs.opt.TEST_CAM_IDX_TO_LOG = 12
+    cfg.train.gs_epochs = 30000
+    
+    if mode_toggle=="Ours (EDGS)":
+        cfg.gs.opt.opacity_reset_interval = 1_000_000
+        cfg.train.reduce_opacity = True
+        cfg.train.no_densify = True
+        cfg.train.max_lr = True
+
+        cfg.init_wC.use = True
+        cfg.init_wC.matches_per_ref = num_corrs_per_view
+        cfg.init_wC.nns_per_ref = 1
+        cfg.init_wC.num_refs = num_ref_views
+        cfg.init_wC.add_SfM_init = False
+        cfg.init_wC.scaling_factor = 0.00077 * 2.
+        
+    set_seed(cfg.seed)
+    os.makedirs(cfg.gs.dataset.model_path, exist_ok=True)
+
+    global trainer
+    global MODEL_PATH
+    generator3dgs = hydra.utils.instantiate(cfg.gs, do_train_test_split=False)
+    trainer = EDGSTrainer(GS=generator3dgs, training_config=cfg.gs.opt, device=cfg.device, log_wandb=cfg.wandb.mode != 'disabled')
+
+    # Disable evaluation and saving
+    trainer.saving_iterations = []
+    trainer.evaluate_iterations = []
+
+    # Initialize
+    trainer.timer.start()
+    start_time = time.time()
+    trainer.init_with_corr(cfg.init_wC, roma_model=roma_model)
+    time_for_init = time.time()-start_time
+
+    viewpoint_cams = trainer.GS.scene.getTrainCameras()
+    path_cameras = generate_fully_smooth_cameras_with_tsp(existing_cameras=viewpoint_cams, 
+                                                          n_selected=6, 
+                                                          n_points_per_segment=30, 
+                                                          closed=False)
+    path_cameras = path_cameras + path_cameras[::-1]
+
+    path_renderings = []
+    idx = 0
+    # Visualize after init
+    for _ in range(120):
+        with torch.no_grad():
+            viewpoint_cam = path_cameras[idx]
+            idx = (idx + 1) % len(path_cameras)
+            render_pkg = trainer.GS(viewpoint_cam)
+            image = render_pkg["render"]
+            image_np = np.clip(image.detach().cpu().numpy().transpose(1, 2, 0), 0, 1)
+            image_np = (image_np * 255).astype(np.uint8)
+            path_renderings.append(put_text_on_image(img=image_np, 
+                                                     text=f"Init stage.\nTime:{time_for_init:.3f}s.   "))
+    path_renderings = path_renderings + [put_text_on_image(img=image_np, text=f"Start fitting.\nTime:{time_for_init:.3f}s.   ")]*30
+    
+    # Train and save visualizations during training.
+    start_time = time.time()
+    for _ in range(int(num_steps//10)):
+        with torch.no_grad():
+            viewpoint_cam = path_cameras[idx]
+            idx = (idx + 1) % len(path_cameras)
+            render_pkg = trainer.GS(viewpoint_cam)
+            image = render_pkg["render"]
+            image_np = np.clip(image.detach().cpu().numpy().transpose(1, 2, 0), 0, 1)
+            image_np = (image_np * 255).astype(np.uint8)
+            path_renderings.append(put_text_on_image(
+                img=image_np, 
+                text=f"Fitting stage.\nTime:{time_for_init + time.time()-start_time:.3f}s.   "))
+    
+        cfg.train.gs_epochs = 10
+        trainer.train(cfg.train)
+        print(f"Time elapsed: {(time_for_init + time.time()-start_time):.2f}s.")
+        # if (cfg.init_wC.use == False) and (time_for_init + time.time()-start_time) > 60:
+        #     break
+    final_time = time.time()
+    
+    # Add static frame. To highlight we're done
+    path_renderings += [put_text_on_image(
+        img=image_np, text=f"Done.\nTime:{time_for_init + final_time -start_time:.3f}s.   ")]*30
+    # Final rendering at the end.
+    for _ in range(len(path_cameras)):
+        with torch.no_grad():
+            viewpoint_cam = path_cameras[idx]
+            idx = (idx + 1) % len(path_cameras)
+            render_pkg = trainer.GS(viewpoint_cam)
+            image = render_pkg["render"]
+            image_np = np.clip(image.detach().cpu().numpy().transpose(1, 2, 0), 0, 1)
+            image_np = (image_np * 255).astype(np.uint8)
+            path_renderings.append(put_text_on_image(img=image_np, 
+                                                 text=f"Final result.\nTime:{time_for_init + final_time -start_time:.3f}s.   "))
+
+    trainer.save_model()
+    final_video_path = os.path.join(STATIC_FILE_SERVING_FOLDER, f"{scene_name}_final.mp4")
+    save_numpy_frames_as_mp4(frames=path_renderings, output_path=final_video_path, fps=30, center_crop=0.85)
+    MODEL_PATH = cfg.gs.dataset.model_path
+    ply_path = os.path.join(cfg.gs.dataset.model_path, f"point_cloud/iteration_{trainer.gs_step}/point_cloud.ply")
+    shutil.copy(ply_path, os.path.join(STATIC_FILE_SERVING_FOLDER, "point_cloud_final.ply"))
+
+    return final_video_path, ply_path
+
+# Gradio Interface
+def gradio_interface(input_path, num_ref_views, num_corrs, num_steps):
+    images, scene_dir = run_full_pipeline(input_path, num_ref_views, num_corrs, max_size=1024)
+    shutil.copytree(scene_dir, STATIC_FILE_SERVING_FOLDER+'/scene_colmaped',  dirs_exist_ok=True)
+    (final_video_path, ply_path), log_output = capture_logs(run_training_pipeline,
+                                                            scene_dir,
+                                                            num_ref_views,
+                                                            num_corrs,
+                                                            num_steps
+                                                            )
+    images_rgb = [img[:, :, ::-1] for img in images]
+    return images_rgb, final_video_path, scene_dir, ply_path, log_output
+
+# Dummy Render Functions
+@spaces.GPU(duration=60)
+def render_all_views(scene_dir):
+    viewpoint_cams = trainer.GS.scene.getTrainCameras()
+    path_cameras = generate_fully_smooth_cameras_with_tsp(existing_cameras=viewpoint_cams, 
+                                                          n_selected=8, 
+                                                          n_points_per_segment=60, 
+                                                          closed=False)
+    path_cameras = path_cameras + path_cameras[::-1]
+
+    path_renderings = []
+    with torch.no_grad():
+        for viewpoint_cam in path_cameras:
+            render_pkg = trainer.GS(viewpoint_cam)
+            image = render_pkg["render"]
+            image_np = np.clip(image.detach().cpu().numpy().transpose(1, 2, 0), 0, 1)
+            image_np = (image_np * 255).astype(np.uint8)
+            path_renderings.append(image_np)
+    save_numpy_frames_as_mp4(frames=path_renderings, 
+                             output_path=os.path.join(STATIC_FILE_SERVING_FOLDER, "render_all_views.mp4"), 
+                             fps=30, 
+                             center_crop=0.85)
+    
+    return os.path.join(STATIC_FILE_SERVING_FOLDER, "render_all_views.mp4")
+
+@spaces.GPU(duration=60)
+def render_circular_path(scene_dir):
+    viewpoint_cams = trainer.GS.scene.getTrainCameras()
+    path_cameras = generate_circular_camera_path(existing_cameras=viewpoint_cams, 
+                                                 N=240, 
+                                                 radius_scale=0.65,
+                                                 d=0)
+
+    path_renderings = []
+    with torch.no_grad():
+        for viewpoint_cam in path_cameras:
+            render_pkg = trainer.GS(viewpoint_cam)
+            image = render_pkg["render"]
+            image_np = np.clip(image.detach().cpu().numpy().transpose(1, 2, 0), 0, 1)
+            image_np = (image_np * 255).astype(np.uint8)
+            path_renderings.append(image_np)
+    save_numpy_frames_as_mp4(frames=path_renderings, 
+                             output_path=os.path.join(STATIC_FILE_SERVING_FOLDER, "render_circular_path.mp4"), 
+                             fps=30, 
+                             center_crop=0.85)
+    
+    return os.path.join(STATIC_FILE_SERVING_FOLDER, "render_circular_path.mp4")
+
+# Download Functions
+def download_cameras():
+    path = os.path.join(MODEL_PATH, "cameras.json")
+    return f"[📥 Download Cameras.json](file={path})"
+
+def download_model():
+    path = os.path.join(STATIC_FILE_SERVING_FOLDER, "point_cloud_final.ply")
+    return f"[📥 Download Pretrained Model (.ply)](file={path})"
+
+# Full pipeline helpers
+def run_full_pipeline(input_path, num_ref_views, num_corrs, max_size=1024):
+    tmpdirname = tempfile.mkdtemp()
+    scene_dir = os.path.join(tmpdirname, "scene")
+    os.makedirs(scene_dir, exist_ok=True)
+
+    selected_frames = process_input(input_path, num_ref_views, scene_dir, max_size)
+    run_colmap_on_scene(scene_dir)
+
+    return selected_frames, scene_dir
+
+# Preprocess Input
+def process_input(input_path, num_ref_views, output_dir, max_size=1024):
+    if isinstance(input_path, (str, os.PathLike)):
+        if os.path.isdir(input_path):
+            frames = []
+            for img_file in sorted(os.listdir(input_path)):
+                if img_file.lower().endswith(('jpg', 'jpeg', 'png')):
+                    img = Image.open(os.path.join(output_dir, img_file)).convert('RGB')
+                    img.thumbnail((1024, 1024))
+                    frames.append(np.array(img))
+        else:
+            frames = read_video_frames(video_input=input_path, max_size=max_size)
+    else:
+        frames = read_video_frames(video_input=input_path, max_size=max_size)
+
+    frames_scores = preprocess_frames(frames)
+    selected_frames_indices = select_optimal_frames(scores=frames_scores, k=min(num_ref_views, len(frames)))
+    selected_frames = [frames[frame_idx] for frame_idx in selected_frames_indices]
+
+    save_frames_to_scene_dir(frames=selected_frames, scene_dir=output_dir)
+    return selected_frames
+
+@spaces.GPU(duration=150)
+def preprocess_input(input_path, num_ref_views, max_size=1024):
+    tmpdirname = tempfile.mkdtemp()
+    scene_dir = os.path.join(tmpdirname, "scene")
+    os.makedirs(scene_dir, exist_ok=True)
+    selected_frames = process_input(input_path, num_ref_views, scene_dir, max_size)
+    run_colmap_on_scene(scene_dir)
+    return selected_frames, scene_dir
+
+def start_training(scene_dir, num_ref_views, num_corrs, num_steps):
+    return capture_logs(run_training_pipeline, scene_dir, num_ref_views, num_corrs, num_steps)
+
+# Gradio App
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column(scale=6):
+            gr.Markdown("""
+            ## <span style='font-size: 20px;'>📄 EDGS: Eliminating Densification for Efficient Convergence of 3DGS</span>
+            🔗 <a href='https://compvis.github.io/EDGS' target='_blank'>Project Page</a>
+            """, elem_id="header")
+
+    gr.Markdown("""
+    ### <span style='font-size: 22px;'>🛠️ How to Use This Demo</span>
+
+    1. Upload a **front-facing video** or **a folder of images** of a **static** scene.
+    2. Use the sliders to configure the number of reference views, correspondences, and optimization steps.
+    3. First press on preprocess Input to extract frames from video(for videos) and COLMAP frames.
+    4.Then click **🚀 Start Reconstruction** to actually launch the reconstruction pipeline.
+    5. Watch the training visualization and explore the 3D model.
+    ‼️ **If you see nothing in the 3D model viewer**, try rotating or zooming — sometimes the initial camera orientation is off.
+
+
+    ✅ Best for scenes with small camera motion.
+    ❗ For full 360° or large-scale scenes, we recommend the Colab version (see project page).
+    """, elem_id="quickstart")
+    scene_dir_state = gr.State()
+    ply_model_state = gr.State()
+    with gr.Row():
+        with gr.Column(scale=2):
+            input_file = gr.File(label="Upload Video or Images", 
+                file_types=[".mp4", ".avi", ".mov", ".png", ".jpg", ".jpeg"], 
+                file_count="multiple")
+            gr.Examples(
+                examples = [
+                    [["assets/examples/video_bakery.mp4"]],
+                    [["assets/examples/video_flowers.mp4"]],
+                    [["assets/examples/video_fruits.mp4"]],
+                    [["assets/examples/video_plant.mp4"]],
+                    [["assets/examples/video_salad.mp4"]],
+                    [["assets/examples/video_tram.mp4"]],
+                    [["assets/examples/video_tulips.mp4"]]
+                    ],
+                inputs=[input_file],
+                label="🎞️ ALternatively, try an Example Video",
+                examples_per_page=4
+            )
+            ref_slider = gr.Slider(4, 32, value=16, step=1, label="Number of Reference Views")
+            corr_slider = gr.Slider(5000, 30000, value=20000, step=1000, label="Correspondences per Reference View")
+            fit_steps_slider = gr.Slider(100, 5000, value=400, step=100, label="Number of optimization steps")
+            preprocess_button = gr.Button("📸 Preprocess Input")
+            start_button = gr.Button("🚀 Start Reconstruction", interactive=False)
+            gallery = gr.Gallery(label="Selected Reference Views", columns=4, height=300)
+
+        with gr.Column(scale=3):
+            gr.Markdown("### 🏋️ Training Visualization")
+            video_output = gr.Video(label="Training Video", autoplay=True)
+            render_all_views_button = gr.Button("🎥 Render All-Views Path")
+            render_circular_path_button = gr.Button("🎥 Render Circular Path")
+            rendered_video_output = gr.Video(label="Rendered Video", autoplay=True)
+        with gr.Column(scale=5):
+            gr.Markdown("### 🌐 Final 3D Model")
+            model3d_viewer = gr.Model3D(label="3D Model Viewer")
+
+            gr.Markdown("### 📦 Output Files")
+            with gr.Row(height=50):
+                with gr.Column():
+                    #gr.Markdown(value=f"[📥 Download .ply](file/point_cloud_final.ply)")
+                    download_cameras_button = gr.Button("📥 Download Cameras.json")
+                    download_cameras_file = gr.File(label="📄 Cameras.json")
+                with gr.Column():
+                    download_model_button = gr.Button("📥 Download Pretrained Model (.ply)")
+                    download_model_file = gr.File(label="📄 Pretrained Model (.ply)")
+
+    log_output_box = gr.Textbox(label="🖥️ Log", lines=10, interactive=False)
+
+    def on_preprocess_click(input_file, num_ref_views):
+        images, scene_dir = preprocess_input(input_file, num_ref_views)
+        return gr.update(value=[x[...,::-1] for x in images]), scene_dir, gr.update(interactive=True)
+
+    def on_start_click(scene_dir, num_ref_views, num_corrs, num_steps):
+        (video_path, ply_path), logs = start_training(scene_dir, num_ref_views, num_corrs, num_steps)
+        return video_path, ply_path, logs
+
+    preprocess_button.click(
+        fn=on_preprocess_click,
+        inputs=[input_file, ref_slider],
+        outputs=[gallery, scene_dir_state, start_button]
+    )
+
+    start_button.click(
+        fn=on_start_click,
+        inputs=[scene_dir_state, ref_slider, corr_slider, fit_steps_slider],
+        outputs=[video_output, model3d_viewer, log_output_box]
+    )
+
+    render_all_views_button.click(fn=render_all_views, inputs=[scene_dir_state], outputs=[rendered_video_output])
+    render_circular_path_button.click(fn=render_circular_path, inputs=[scene_dir_state], outputs=[rendered_video_output])
+
+    download_cameras_button.click(fn=lambda: os.path.join(MODEL_PATH, "cameras.json"), inputs=[], outputs=[download_cameras_file])
+    download_model_button.click(fn=lambda: os.path.join(STATIC_FILE_SERVING_FOLDER, "point_cloud_final.ply"), inputs=[], outputs=[download_model_file])
+
+
+    gr.Markdown("""
+    ---
+    ### <span style='font-size: 20px;'>📖 Detailed Overview</span>
+
+    If you uploaded a video, it will be automatically cut into a smaller number of frames (default: 16).
+
+    The model pipeline:
+    1. 🧠 Runs PyCOLMAP to estimate camera intrinsics & poses (~3–7 seconds for <16 images).
+    2. 🔁 Computes 2D-2D correspondences between views. More correspondences generally improve quality.
+    3. 🔧 Optimizes a 3D Gaussian Splatting model for several steps.
+
+    ### 🎥 Training Visualization
+    You will see a visualization of the entire training process in the "Training Video" pane.
+
+    ### 🌀 Rendering & 3D Model
+    - Render the scene from a circular path of novel views.
+    - Or from camera views close to the original input.
+
+    The 3D model is shown in the right viewer. You can explore it interactively:
+    - On PC: WASD keys, arrow keys, and mouse clicks
+    - On mobile: pan and pinch to zoom
+
+    🕒 Note: the 3D viewer takes a few extra seconds (~5s) to display after training ends.
+
+    ---
+    Preloaded models coming soon. (TODO)
+    """, elem_id="details")
+
+    
+
+
+if __name__ == "__main__":
+    install_submodules()
+    from source.utils_aux import set_seed
+    from source.utils_preprocess import read_video_frames, preprocess_frames, select_optimal_frames, save_frames_to_scene_dir, run_colmap_on_scene
+    from source.trainer import EDGSTrainer
+    from source.visualization import generate_circular_camera_path, save_numpy_frames_as_mp4, generate_fully_smooth_cameras_with_tsp, put_text_on_image
+    # Init RoMA model:
+    sys.path.append('../submodules/RoMa')
+    from romatch import roma_outdoor, roma_indoor
+    
+    roma_model = roma_indoor(device="cpu")
+    roma_model = roma_model.to("cuda")
+    roma_model.upsample_preds = False
+    roma_model.symmetric = False
+
+    demo.launch(share=True)

assets/examples/video_bakery.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b65c813f2ef9637579350e145fdceed333544d933278be5613c6d49468f4eab0
+size 6362238

assets/examples/video_flowers.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c81a7c28e0d59bad38d5a45f6bfa83b80990d1fb78a06f82137e8b57ec38e62b
+size 6466943

assets/examples/video_fruits.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ac3b937e155d1965a314b478e940f6fd93e90371d3bf7d62d3225d840fe8e126
+size 3356915

assets/examples/video_plant.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e67a40e62de9aacf0941d8cf33a2dd08d256fe60d0fc58f60426c251f9d8abd8
+size 13023885

assets/examples/video_salad.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:367d86071a201c124383f0a076ce644b7f86b9c2fbce6aa595c4989ebf259bfb
+size 8774427

assets/examples/video_tram.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:298c297155d3f52edcffcb6b6f9910c992b98ecfb93cfaf8fb64fb340aba1dae
+size 4697915

assets/examples/video_tulips.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c19942319b8f7e33bb03cb4a39b11797fe38572bf7157af37471b7c8573fb495
+size 7298210

assets/video_fruits_ours_full.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c5b113566d3a083b81360b549ed89f70d5e81739f83e182518f6906811311a2
+size 14839197

configs/gs/base.yaml

@@ -0,0 +1,51 @@
+_target_: source.networks.Warper3DGS
+
+verbose: True
+viewpoint_stack: !!null
+sh_degree: 3
+
+opt:
+  iterations: 30000
+  position_lr_init: 0.00016
+  position_lr_final: 1.6e-06
+  position_lr_delay_mult: 0.01
+  position_lr_max_steps: 30000
+  feature_lr: 0.0025
+  opacity_lr: 0.025
+  scaling_lr: 0.005
+  rotation_lr: 0.001
+  percent_dense: 0.01
+  lambda_dssim: 0.2
+  densification_interval: 100
+  opacity_reset_interval: 30000
+  densify_from_iter: 500
+  densify_until_iter: 15000
+  densify_grad_threshold: 0.0002
+  random_background: false
+  save_iterations: [3000, 7000, 15000, 30000]
+  batch_size: 64
+  exposure_lr_init: 0.01
+  exposure_lr_final: 0.0001
+  exposure_lr_delay_steps: 0
+  exposure_lr_delay_mult: 0.0
+
+  TRAIN_CAM_IDX_TO_LOG: 50
+  TEST_CAM_IDX_TO_LOG: 10
+
+pipe:
+  convert_SHs_python: False
+  compute_cov3D_python: False
+  debug: False
+  antialiasing: False 
+
+dataset:
+  densify_until_iter: 15000
+  source_path:  '' #path to dataset
+  model_path:  '' #path to logs
+  images: images
+  resolution: -1
+  white_background: false
+  data_device: cuda
+  eval: false
+  depths: ""
+  train_test_exp: False

configs/train.yaml

@@ -0,0 +1,38 @@
+defaults:
+  - gs: base
+  - _self_ 
+
+seed: 228
+
+wandb:
+  mode: "online" # "disabled" for no logging
+  entity: "3dcorrespondence"
+  project: "Adv3DGS"
+  group: null
+  name: null
+  tag: "debug"
+    
+train:
+  gs_epochs: 0 # number of 3dgs iterations
+  reduce_opacity: True 
+  no_densify: False # if True, the model will not be densified
+  max_lr: True 
+
+load:
+  gs: null #path to 3dgs checkpoint
+  gs_step: null #number of iterations, e.g. 7000
+
+device: "cuda:0"
+verbose: true
+
+init_wC:
+  use: True # use EDGS
+  matches_per_ref: 15_000 # number of matches per reference
+  num_refs: 180 # number of reference images
+  nns_per_ref: 3 # number of nearest neighbors per reference
+  scaling_factor: 0.001
+  proj_err_tolerance: 0.01
+  roma_model: "outdoors" # you can change this to "indoors" or "outdoors"
+  add_SfM_init : False
+
+

requirements.txt

@@ -0,0 +1,32 @@
+--extra-index-url https://download.pytorch.org/whl/cu124
+torch
+torchvision
+torchaudio 
+
+
+# Required libraries from pip
+Pillow
+huggingface_hub
+einops
+safetensors
+sympy==1.13.1
+wandb
+hydra-core
+tqdm
+torchmetrics
+lpips
+matplotlib
+rich
+plyfile
+imageio
+imageio-ffmpeg
+numpy==1.26.4  # Match conda-installed version
+opencv-python
+pycolmap
+moviepy
+plotly
+scikit-learn
+ffmpeg
+
+https://huggingface.co/spaces/magistrkoljan/test/resolve/main/wheels/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl?download=true
+https://huggingface.co/spaces/magistrkoljan/test/resolve/main/wheels/simple_knn-0.0.0-cp310-cp310-linux_x86_64.whl?download=true

source/EDGS.code-workspace

@@ -0,0 +1,11 @@
+{
+	"folders": [
+		{
+			"path": ".."
+		},
+		{
+			"path": "../../../../.."
+		}
+	],
+	"settings": {}
+}

source/corr_init.py

@@ -0,0 +1,682 @@
+import sys
+sys.path.append('../')
+sys.path.append("../submodules")
+sys.path.append('../submodules/RoMa')
+
+from matplotlib import pyplot as plt
+from PIL import Image
+import torch
+import numpy as np
+
+#from tqdm import tqdm_notebook as tqdm
+from tqdm import tqdm
+from scipy.cluster.vq import kmeans, vq
+from scipy.spatial.distance import cdist
+
+import torch.nn.functional as F
+from romatch import roma_outdoor, roma_indoor
+from utils.sh_utils import RGB2SH
+from romatch.utils import get_tuple_transform_ops
+
+
+def pairwise_distances(matrix):
+    """
+    Computes the pairwise Euclidean distances between all vectors in the input matrix.
+
+    Args:
+        matrix (torch.Tensor): Input matrix of shape [N, D], where N is the number of vectors and D is the dimensionality.
+
+    Returns:
+        torch.Tensor: Pairwise distance matrix of shape [N, N].
+    """
+    # Compute squared pairwise distances
+    squared_diff = torch.cdist(matrix, matrix, p=2)
+    return squared_diff
+
+
+def k_closest_vectors(matrix, k):
+    """
+    Finds the k-closest vectors for each vector in the input matrix based on Euclidean distance.
+
+    Args:
+        matrix (torch.Tensor): Input matrix of shape [N, D], where N is the number of vectors and D is the dimensionality.
+        k (int): Number of closest vectors to return for each vector.
+
+    Returns:
+        torch.Tensor: Indices of the k-closest vectors for each vector, excluding the vector itself.
+    """
+    # Compute pairwise distances
+    distances = pairwise_distances(matrix)
+
+    # For each vector, sort distances and get the indices of the k-closest vectors (excluding itself)
+    # Set diagonal distances to infinity to exclude the vector itself from the nearest neighbors
+    distances.fill_diagonal_(float('inf'))
+
+    # Get the indices of the k smallest distances (k-closest vectors)
+    _, indices = torch.topk(distances, k, largest=False, dim=1)
+
+    return indices
+
+
+def select_cameras_kmeans(cameras, K):
+    """
+    Selects K cameras from a set using K-means clustering.
+
+    Args:
+        cameras: NumPy array of shape (N, 16), representing N cameras with their 4x4 homogeneous matrices flattened.
+        K: Number of clusters (cameras to select).
+
+    Returns:
+        selected_indices: List of indices of the cameras closest to the cluster centers.
+    """
+    # Ensure input is a NumPy array
+    if not isinstance(cameras, np.ndarray):
+        cameras = np.asarray(cameras)
+
+    if cameras.shape[1] != 16:
+        raise ValueError("Each camera must have 16 values corresponding to a flattened 4x4 matrix.")
+
+    # Perform K-means clustering
+    cluster_centers, _ = kmeans(cameras, K)
+
+    # Assign each camera to a cluster and find distances to cluster centers
+    cluster_assignments, _ = vq(cameras, cluster_centers)
+
+    # Find the camera nearest to each cluster center
+    selected_indices = []
+    for k in range(K):
+        cluster_members = cameras[cluster_assignments == k]
+        distances = cdist([cluster_centers[k]], cluster_members)[0]
+        nearest_camera_idx = np.where(cluster_assignments == k)[0][np.argmin(distances)]
+        selected_indices.append(nearest_camera_idx)
+
+    return selected_indices
+
+
+def compute_warp_and_confidence(viewpoint_cam1, viewpoint_cam2, roma_model, device="cuda", verbose=False, output_dict={}):
+    """
+    Computes the warp and confidence between two viewpoint cameras using the roma_model.
+
+    Args:
+        viewpoint_cam1: Source viewpoint camera.
+        viewpoint_cam2: Target viewpoint camera.
+        roma_model: Pre-trained Roma model for correspondence matching.
+        device: Device to run the computation on.
+        verbose: If True, displays the images.
+
+    Returns:
+        certainty: Confidence tensor.
+        warp: Warp tensor.
+        imB: Processed image B as numpy array.
+    """
+    # Prepare images
+    imA = viewpoint_cam1.original_image.detach().cpu().numpy().transpose(1, 2, 0)
+    imB = viewpoint_cam2.original_image.detach().cpu().numpy().transpose(1, 2, 0)
+    imA = Image.fromarray(np.clip(imA * 255, 0, 255).astype(np.uint8))
+    imB = Image.fromarray(np.clip(imB * 255, 0, 255).astype(np.uint8))
+
+    if verbose:
+        fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(16, 8))
+        cax1 = ax[0].imshow(imA)
+        ax[0].set_title("Image 1")
+        cax2 = ax[1].imshow(imB)
+        ax[1].set_title("Image 2")
+        fig.colorbar(cax1, ax=ax[0])
+        fig.colorbar(cax2, ax=ax[1])
+    
+        for axis in ax:
+            axis.axis('off')
+        # Save the figure into the dictionary
+        output_dict[f'image_pair'] = fig
+   
+    # Transform images
+    ws, hs = roma_model.w_resized, roma_model.h_resized
+    test_transform = get_tuple_transform_ops(resize=(hs, ws), normalize=True)
+    im_A, im_B = test_transform((imA, imB))
+    batch = {"im_A": im_A[None].to(device), "im_B": im_B[None].to(device)}
+
+    # Forward pass through Roma model
+    corresps = roma_model.forward(batch) if not roma_model.symmetric else roma_model.forward_symmetric(batch)
+    finest_scale = 1
+    hs, ws = roma_model.upsample_res if roma_model.upsample_preds else (hs, ws)
+
+    # Process certainty and warp
+    certainty = corresps[finest_scale]["certainty"]
+    im_A_to_im_B = corresps[finest_scale]["flow"]
+    if roma_model.attenuate_cert:
+        low_res_certainty = F.interpolate(
+            corresps[16]["certainty"], size=(hs, ws), align_corners=False, mode="bilinear"
+        )
+        certainty -= 0.5 * low_res_certainty * (low_res_certainty < 0)
+
+    # Upsample predictions if needed
+    if roma_model.upsample_preds:
+        im_A_to_im_B = F.interpolate(
+            im_A_to_im_B, size=(hs, ws), align_corners=False, mode="bilinear"
+        )
+        certainty = F.interpolate(
+            certainty, size=(hs, ws), align_corners=False, mode="bilinear"
+        )
+
+    # Convert predictions to final format
+    im_A_to_im_B = im_A_to_im_B.permute(0, 2, 3, 1)
+    im_A_coords = torch.stack(torch.meshgrid(
+        torch.linspace(-1 + 1 / hs, 1 - 1 / hs, hs, device=device),
+        torch.linspace(-1 + 1 / ws, 1 - 1 / ws, ws, device=device),
+        indexing='ij'
+    ), dim=0).permute(1, 2, 0).unsqueeze(0).expand(im_A_to_im_B.size(0), -1, -1, -1)
+
+    warp = torch.cat((im_A_coords, im_A_to_im_B), dim=-1)
+    certainty = certainty.sigmoid()
+
+    return certainty[0, 0], warp[0], np.array(imB)
+
+
+def resize_batch(tensors_3d, tensors_4d, target_shape):
+    """
+    Resizes a batch of tensors with shapes [B, H, W] and [B, H, W, 4] to the target spatial dimensions.
+
+    Args:
+        tensors_3d: Tensor of shape [B, H, W].
+        tensors_4d: Tensor of shape [B, H, W, 4].
+        target_shape: Tuple (target_H, target_W) specifying the target spatial dimensions.
+
+    Returns:
+        resized_tensors_3d: Tensor of shape [B, target_H, target_W].
+        resized_tensors_4d: Tensor of shape [B, target_H, target_W, 4].
+    """
+    target_H, target_W = target_shape
+
+    # Resize [B, H, W] tensor
+    resized_tensors_3d = F.interpolate(
+        tensors_3d.unsqueeze(1), size=(target_H, target_W), mode="bilinear", align_corners=False
+    ).squeeze(1)
+
+    # Resize [B, H, W, 4] tensor
+    B, _, _, C = tensors_4d.shape
+    resized_tensors_4d = F.interpolate(
+        tensors_4d.permute(0, 3, 1, 2), size=(target_H, target_W), mode="bilinear", align_corners=False
+    ).permute(0, 2, 3, 1)
+
+    return resized_tensors_3d, resized_tensors_4d
+
+
+def aggregate_confidences_and_warps(viewpoint_stack, closest_indices, roma_model, source_idx, verbose=False, output_dict={}):
+    """
+    Aggregates confidences and warps by iterating over the nearest neighbors of the source viewpoint.
+
+    Args:
+        viewpoint_stack: Stack of viewpoint cameras.
+        closest_indices: Indices of the nearest neighbors for each viewpoint.
+        roma_model: Pre-trained Roma model.
+        source_idx: Index of the source viewpoint.
+        verbose: If True, displays intermediate results.
+
+    Returns:
+        certainties_max: Aggregated maximum confidences.
+        warps_max: Aggregated warps corresponding to maximum confidences.
+        certainties_max_idcs: Pixel-wise index of the image  from which we taken the best matching.
+        imB_compound: List of the neighboring images.
+    """
+    certainties_all, warps_all, imB_compound = [], [], []
+
+    for nn in tqdm(closest_indices[source_idx]):
+
+        viewpoint_cam1 = viewpoint_stack[source_idx]
+        viewpoint_cam2 = viewpoint_stack[nn]
+
+        certainty, warp, imB = compute_warp_and_confidence(viewpoint_cam1, viewpoint_cam2, roma_model, verbose=verbose, output_dict=output_dict)
+        certainties_all.append(certainty)
+        warps_all.append(warp)
+        imB_compound.append(imB)
+
+    certainties_all = torch.stack(certainties_all, dim=0)
+    target_shape = imB_compound[0].shape[:2]
+    if verbose: 
+        print("certainties_all.shape:", certainties_all.shape)
+        print("torch.stack(warps_all, dim=0).shape:", torch.stack(warps_all, dim=0).shape)
+        print("target_shape:", target_shape)        
+
+    certainties_all_resized, warps_all_resized = resize_batch(certainties_all,
+                                                              torch.stack(warps_all, dim=0),
+                                                              target_shape
+                                                              )
+
+    if verbose:
+        print("warps_all_resized.shape:", warps_all_resized.shape)
+        for n, cert in enumerate(certainties_all):
+            fig, ax = plt.subplots()
+            cax = ax.imshow(cert.cpu().numpy(), cmap='viridis')
+            fig.colorbar(cax, ax=ax)
+            ax.set_title("Pixel-wise Confidence")
+            output_dict[f'certainty_{n}'] = fig
+
+        for n, warp in enumerate(warps_all):
+            fig, ax = plt.subplots()
+            cax = ax.imshow(warp.cpu().numpy()[:, :, :3], cmap='viridis')
+            fig.colorbar(cax, ax=ax)
+            ax.set_title("Pixel-wise warp")
+            output_dict[f'warp_resized_{n}'] = fig
+
+        for n, cert in enumerate(certainties_all_resized):
+            fig, ax = plt.subplots()
+            cax = ax.imshow(cert.cpu().numpy(), cmap='viridis')
+            fig.colorbar(cax, ax=ax)
+            ax.set_title("Pixel-wise Confidence resized")
+            output_dict[f'certainty_resized_{n}'] = fig
+
+        for n, warp in enumerate(warps_all_resized):
+            fig, ax = plt.subplots()
+            cax = ax.imshow(warp.cpu().numpy()[:, :, :3], cmap='viridis')
+            fig.colorbar(cax, ax=ax)
+            ax.set_title("Pixel-wise warp resized")
+            output_dict[f'warp_resized_{n}'] = fig
+
+    certainties_max, certainties_max_idcs = torch.max(certainties_all_resized, dim=0)
+    H, W = certainties_max.shape
+
+    warps_max = warps_all_resized[certainties_max_idcs, torch.arange(H).unsqueeze(1), torch.arange(W)]
+
+    imA = viewpoint_cam1.original_image.detach().cpu().numpy().transpose(1, 2, 0)
+    imA = np.clip(imA * 255, 0, 255).astype(np.uint8)
+
+    return certainties_max, warps_max, certainties_max_idcs, imA, imB_compound, certainties_all_resized, warps_all_resized
+
+
+
+def extract_keypoints_and_colors(imA, imB_compound, certainties_max, certainties_max_idcs, matches, roma_model,
+                                 verbose=False, output_dict={}):
+    """
+    Extracts keypoints and corresponding colors from the source image (imA) and multiple target images (imB_compound).
+
+    Args:
+        imA: Source image as a NumPy array (H_A, W_A, C).
+        imB_compound: List of target images as NumPy arrays [(H_B, W_B, C), ...].
+        certainties_max: Tensor of pixel-wise maximum confidences.
+        certainties_max_idcs: Tensor of pixel-wise indices for the best matches.
+        matches: Matches in normalized coordinates.
+        roma_model: Roma model instance for keypoint operations.
+        verbose: if to show intermediate outputs and visualize results
+
+    Returns:
+        kptsA_np: Keypoints in imA in normalized coordinates.
+        kptsB_np: Keypoints in imB in normalized coordinates.
+        kptsA_color: Colors of keypoints in imA.
+        kptsB_color: Colors of keypoints in imB based on certainties_max_idcs.
+    """
+    H_A, W_A, _ = imA.shape
+    H, W = certainties_max.shape
+
+    # Convert matches to pixel coordinates
+    kptsA, kptsB = roma_model.to_pixel_coordinates(
+        matches, W_A, H_A, H, W  # W, H
+    )
+
+    kptsA_np = kptsA.detach().cpu().numpy()
+    kptsB_np = kptsB.detach().cpu().numpy()
+    kptsA_np = kptsA_np[:, [1, 0]]
+
+    if verbose:
+        fig, ax = plt.subplots(figsize=(12, 6))
+        cax = ax.imshow(imA)
+        ax.set_title("Reference image, imA")
+        output_dict[f'reference_image'] = fig
+
+        fig, ax = plt.subplots(figsize=(12, 6))
+        cax = ax.imshow(imB_compound[0])
+        ax.set_title("Image to compare to image, imB_compound")
+        output_dict[f'imB_compound'] = fig
+    
+        fig, ax = plt.subplots(figsize=(12, 6))
+        cax = ax.imshow(np.flipud(imA))
+        cax = ax.scatter(kptsA_np[:, 0], H_A - kptsA_np[:, 1], s=.03)
+        ax.set_title("Keypoints in imA")
+        ax.set_xlim(0, W_A)
+        ax.set_ylim(0, H_A)
+        output_dict[f'kptsA'] = fig
+
+        fig, ax = plt.subplots(figsize=(12, 6))
+        cax = ax.imshow(np.flipud(imB_compound[0]))
+        cax = ax.scatter(kptsB_np[:, 0], H_A - kptsB_np[:, 1], s=.03)
+        ax.set_title("Keypoints in imB")
+        ax.set_xlim(0, W_A)
+        ax.set_ylim(0, H_A)
+        output_dict[f'kptsB'] = fig
+
+    # Keypoints are in format (row, column) so the first value is alwain in range [0;height] and second is in range[0;width]
+
+    kptsA_np = kptsA.detach().cpu().numpy()
+    kptsB_np = kptsB.detach().cpu().numpy()
+
+    # Extract colors for keypoints in imA (vectorized)
+    # New experimental version
+    kptsA_x = np.round(kptsA_np[:, 0] / 1.).astype(int)
+    kptsA_y = np.round(kptsA_np[:, 1] / 1.).astype(int)
+    kptsA_color = imA[np.clip(kptsA_x, 0, H - 1), np.clip(kptsA_y, 0, W - 1)]
+   
+    # Create a composite image from imB_compound
+    imB_compound_np = np.stack(imB_compound, axis=0)
+    H_B, W_B, _ = imB_compound[0].shape
+
+    # Extract colors for keypoints in imB using certainties_max_idcs
+    imB_np = imB_compound_np[
+            certainties_max_idcs.detach().cpu().numpy(),
+            np.arange(H).reshape(-1, 1),
+            np.arange(W)
+        ]
+    
+    if verbose:
+        print("imB_np.shape:", imB_np.shape)
+        print("imB_np:", imB_np)
+        fig, ax = plt.subplots(figsize=(12, 6))
+        cax = ax.imshow(np.flipud(imB_np))
+        cax = ax.scatter(kptsB_np[:, 0], H_A - kptsB_np[:, 1], s=.03)
+        ax.set_title("np.flipud(imB_np[0]")
+        ax.set_xlim(0, W_A)
+        ax.set_ylim(0, H_A)
+        output_dict[f'np.flipud(imB_np[0]'] = fig
+
+
+    kptsB_x = np.round(kptsB_np[:, 0]).astype(int)
+    kptsB_y = np.round(kptsB_np[:, 1]).astype(int)
+
+    certainties_max_idcs_np = certainties_max_idcs.detach().cpu().numpy()
+    kptsB_proj_matrices_idx = certainties_max_idcs_np[np.clip(kptsA_x, 0, H - 1), np.clip(kptsA_y, 0, W - 1)]
+    kptsB_color = imB_compound_np[kptsB_proj_matrices_idx, np.clip(kptsB_y, 0, H - 1), np.clip(kptsB_x, 0, W - 1)]
+
+    # Normalize keypoints in both images
+    kptsA_np[:, 0] = kptsA_np[:, 0] / H * 2.0 - 1.0
+    kptsA_np[:, 1] = kptsA_np[:, 1] / W * 2.0 - 1.0
+    kptsB_np[:, 0] = kptsB_np[:, 0] / W_B * 2.0 - 1.0
+    kptsB_np[:, 1] = kptsB_np[:, 1] / H_B * 2.0 - 1.0
+
+    return kptsA_np[:, [1, 0]], kptsB_np, kptsB_proj_matrices_idx, kptsA_color, kptsB_color
+
+def prepare_tensor(input_array, device):
+    """
+    Converts an input array to a torch tensor, clones it, and detaches it for safe computation.
+    Args:
+        input_array (array-like): The input array to convert.
+        device (str or torch.device): The device to move the tensor to.
+    Returns:
+        torch.Tensor: A detached tensor clone of the input array on the specified device.
+    """
+    if not isinstance(input_array, torch.Tensor):
+        return torch.tensor(input_array, dtype=torch.float32).to(device).clone().detach()
+    return input_array.clone().detach().to(device).to(torch.float32)
+
+def triangulate_points(P1, P2, k1_x, k1_y, k2_x, k2_y, device="cuda"):
+    """
+    Solves for a batch of 3D points given batches of projection matrices and corresponding image points.
+
+    Parameters:
+    - P1, P2: Tensors of projection matrices of size (batch_size, 4, 4) or (4, 4)
+    - k1_x, k1_y: Tensors of shape (batch_size,)
+    - k2_x, k2_y: Tensors of shape (batch_size,)
+
+    Returns:
+    - X: A tensor containing the 3D homogeneous coordinates, shape (batch_size, 4)
+    """
+    EPS = 1e-4
+    # Ensure inputs are tensors
+
+    P1 = prepare_tensor(P1, device)
+    P2 = prepare_tensor(P2, device)
+    k1_x = prepare_tensor(k1_x, device)
+    k1_y = prepare_tensor(k1_y, device)
+    k2_x = prepare_tensor(k2_x, device)
+    k2_y =  prepare_tensor(k2_y, device)
+    batch_size = k1_x.shape[0]
+
+    # Expand P1 and P2 if they are not batched
+    if P1.ndim == 2:
+        P1 = P1.unsqueeze(0).expand(batch_size, -1, -1)
+    if P2.ndim == 2:
+        P2 = P2.unsqueeze(0).expand(batch_size, -1, -1)
+
+    # Extract columns from P1 and P2
+    P1_0 = P1[:, :, 0]  # Shape: (batch_size, 4)
+    P1_1 = P1[:, :, 1]
+    P1_2 = P1[:, :, 2]
+
+    P2_0 = P2[:, :, 0]
+    P2_1 = P2[:, :, 1]
+    P2_2 = P2[:, :, 2]
+
+    # Reshape kx and ky to (batch_size, 1)
+    k1_x = k1_x.view(-1, 1)
+    k1_y = k1_y.view(-1, 1)
+    k2_x = k2_x.view(-1, 1)
+    k2_y = k2_y.view(-1, 1)
+
+    # Construct the equations for each batch
+    # For camera 1
+    A1 = P1_0 - k1_x * P1_2  # Shape: (batch_size, 4)
+    A2 = P1_1 - k1_y * P1_2
+    # For camera 2
+    A3 = P2_0 - k2_x * P2_2
+    A4 = P2_1 - k2_y * P2_2
+
+    # Stack the equations
+    A = torch.stack([A1, A2, A3, A4], dim=1)  # Shape: (batch_size, 4, 4)
+
+    # Right-hand side (constants)
+    b = -A[:, :, 3]  # Shape: (batch_size, 4)
+    A_reduced = A[:, :, :3]  # Coefficients of x, y, z
+
+    # Solve using torch.linalg.lstsq (supports batching)
+    X_xyz = torch.linalg.lstsq(A_reduced, b.unsqueeze(2)).solution.squeeze(2)  # Shape: (batch_size, 3)
+
+    # Append 1 to get homogeneous coordinates
+    ones = torch.ones((batch_size, 1), dtype=torch.float32, device=X_xyz.device)
+    X = torch.cat([X_xyz, ones], dim=1)  # Shape: (batch_size, 4)
+
+    # Now compute the errors of projections.
+    seeked_splats_proj1 = (X.unsqueeze(1) @ P1).squeeze(1)
+    seeked_splats_proj1 = seeked_splats_proj1 / (EPS + seeked_splats_proj1[:, [3]])
+    seeked_splats_proj2 = (X.unsqueeze(1) @ P2).squeeze(1)
+    seeked_splats_proj2 = seeked_splats_proj2 / (EPS + seeked_splats_proj2[:, [3]])
+    proj1_target = torch.concat([k1_x, k1_y], dim=1)
+    proj2_target = torch.concat([k2_x, k2_y], dim=1)
+    errors_proj1 = torch.abs(seeked_splats_proj1[:, :2] - proj1_target).sum(1).detach().cpu().numpy()
+    errors_proj2 = torch.abs(seeked_splats_proj2[:, :2] - proj2_target).sum(1).detach().cpu().numpy()
+
+    return X, errors_proj1, errors_proj2
+
+
+
+def select_best_keypoints(
+        NNs_triangulated_points, NNs_errors_proj1, NNs_errors_proj2, device="cuda"):
+    """
+    From all the points fitted to  keypoints and corresponding colors from the source image (imA) and multiple target images (imB_compound).
+
+    Args:
+        NNs_triangulated_points:  torch tensor with keypoints coordinates (num_nns, num_points, dim). dim can be arbitrary,
+            usually 3 or 4(for homogeneous representation).
+        NNs_errors_proj1:  numpy array with projection error of the estimated keypoint on the reference frame (num_nns, num_points).
+        NNs_errors_proj2:  numpy array with projection error of the estimated keypoint on the neighbor frame (num_nns, num_points).
+    Returns:
+        selected_keypoints: keypoints with the best score.
+    """
+
+    NNs_errors_proj = np.maximum(NNs_errors_proj1, NNs_errors_proj2)
+
+    # Convert indices to PyTorch tensor
+    indices = torch.from_numpy(np.argmin(NNs_errors_proj, axis=0)).long().to(device)
+
+    # Create index tensor for the second dimension
+    n_indices = torch.arange(NNs_triangulated_points.shape[1]).long().to(device)
+
+    # Use advanced indexing to select elements
+    NNs_triangulated_points_selected = NNs_triangulated_points[indices, n_indices, :]  # Shape: [N, k]
+
+    return NNs_triangulated_points_selected, np.min(NNs_errors_proj, axis=0)
+
+
+
+def init_gaussians_with_corr(gaussians, scene, cfg, device, verbose = False, roma_model=None):
+    """
+    For a given input gaussians and a scene we instantiate a RoMa model(change to indoors if necessary) and process scene
+    training frames to extract correspondences. Those are used to initialize gaussians
+    Args:
+        gaussians: object gaussians of the class GaussianModel that we need to enrich with gaussians.
+        scene: object of the Scene class.
+        cfg: configuration. Use init_wC
+    Returns:
+        gaussians: inplace transforms object gaussians of the class GaussianModel.
+
+    """
+    if roma_model is None:
+        if cfg.roma_model == "indoors":
+            roma_model = roma_indoor(device=device)
+        else:
+            roma_model = roma_outdoor(device=device)
+        roma_model.upsample_preds = False
+        roma_model.symmetric = False
+    M = cfg.matches_per_ref
+    upper_thresh = roma_model.sample_thresh
+    scaling_factor = cfg.scaling_factor
+    expansion_factor = 1
+    keypoint_fit_error_tolerance = cfg.proj_err_tolerance
+    visualizations = {}
+    viewpoint_stack = scene.getTrainCameras().copy()
+    NUM_REFERENCE_FRAMES = min(cfg.num_refs, len(viewpoint_stack))
+    NUM_NNS_PER_REFERENCE = min(cfg.nns_per_ref , len(viewpoint_stack))
+    # Select cameras using K-means
+    viewpoint_cam_all = torch.stack([x.world_view_transform.flatten() for x in viewpoint_stack], axis=0)
+
+    selected_indices = select_cameras_kmeans(cameras=viewpoint_cam_all.detach().cpu().numpy(), K=NUM_REFERENCE_FRAMES)
+    selected_indices = sorted(selected_indices)
+   
+
+    # Find the k-closest vectors for each vector
+    viewpoint_cam_all = torch.stack([x.world_view_transform.flatten() for x in viewpoint_stack], axis=0)
+    closest_indices = k_closest_vectors(viewpoint_cam_all, NUM_NNS_PER_REFERENCE)
+    if verbose: print("Indices of k-closest vectors for each vector:\n", closest_indices)
+
+    closest_indices_selected = closest_indices[:, :].detach().cpu().numpy()
+
+    all_new_xyz = []
+    all_new_features_dc = []
+    all_new_features_rest = []
+    all_new_opacities = []
+    all_new_scaling = []
+    all_new_rotation = []
+
+    # Run roma_model.match once to kinda initialize the model
+    with torch.no_grad():
+        viewpoint_cam1 = viewpoint_stack[0]
+        viewpoint_cam2 = viewpoint_stack[1]
+        imA = viewpoint_cam1.original_image.detach().cpu().numpy().transpose(1, 2, 0)
+        imB = viewpoint_cam2.original_image.detach().cpu().numpy().transpose(1, 2, 0)
+        imA = Image.fromarray(np.clip(imA * 255, 0, 255).astype(np.uint8))
+        imB = Image.fromarray(np.clip(imB * 255, 0, 255).astype(np.uint8))
+        warp, certainty_warp = roma_model.match(imA, imB, device=device)
+        print("Once run full roma_model.match warp.shape:", warp.shape)
+        print("Once run full roma_model.match certainty_warp.shape:", certainty_warp.shape)
+        del warp, certainty_warp
+        torch.cuda.empty_cache()
+
+    for source_idx in tqdm(sorted(selected_indices)):
+        # 1. Compute keypoints and warping for all the neigboring views
+        with torch.no_grad():
+            # Call the aggregation function to get imA and imB_compound
+            certainties_max, warps_max, certainties_max_idcs, imA, imB_compound, certainties_all, warps_all = aggregate_confidences_and_warps(
+                viewpoint_stack=viewpoint_stack,
+                closest_indices=closest_indices_selected,
+                roma_model=roma_model,
+                source_idx=source_idx,
+                verbose=verbose, output_dict=visualizations
+            )
+
+
+        # Triangulate keypoints
+        with torch.no_grad():
+            matches = warps_max
+            certainty = certainties_max
+            certainty = certainty.clone()
+            certainty[certainty > upper_thresh] = 1
+            matches, certainty = (
+                matches.reshape(-1, 4),
+                certainty.reshape(-1),
+            )
+
+            # Select based on certainty elements with high confidence. These are basically all of
+            # kptsA_np.
+            good_samples = torch.multinomial(certainty,
+                                             num_samples=min(expansion_factor * M, len(certainty)),
+                                             replacement=False)
+
+        certainties_max, warps_max, certainties_max_idcs, imA, imB_compound, certainties_all, warps_all
+        reference_image_dict = {
+            "ref_image": imA,
+            "NNs_images": imB_compound,
+            "certainties_all": certainties_all,
+            "warps_all": warps_all,
+            "triangulated_points": [],
+            "triangulated_points_errors_proj1": [],
+            "triangulated_points_errors_proj2": []
+
+        }
+        with torch.no_grad():
+            for NN_idx in tqdm(range(len(warps_all))):
+                matches_NN = warps_all[NN_idx].reshape(-1, 4)[good_samples]
+
+                # Extract keypoints and colors
+                kptsA_np, kptsB_np, kptsB_proj_matrices_idcs, kptsA_color, kptsB_color = extract_keypoints_and_colors(
+                    imA, imB_compound, certainties_max, certainties_max_idcs, matches_NN, roma_model
+                )
+
+                proj_matrices_A = viewpoint_stack[source_idx].full_proj_transform
+                proj_matrices_B = viewpoint_stack[closest_indices_selected[source_idx, NN_idx]].full_proj_transform
+                triangulated_points, triangulated_points_errors_proj1, triangulated_points_errors_proj2 = triangulate_points(
+                    P1=torch.stack([proj_matrices_A] * M, axis=0),
+                    P2=torch.stack([proj_matrices_B] * M, axis=0),
+                    k1_x=kptsA_np[:M, 0], k1_y=kptsA_np[:M, 1],
+                    k2_x=kptsB_np[:M, 0], k2_y=kptsB_np[:M, 1])
+
+                reference_image_dict["triangulated_points"].append(triangulated_points)
+                reference_image_dict["triangulated_points_errors_proj1"].append(triangulated_points_errors_proj1)
+                reference_image_dict["triangulated_points_errors_proj2"].append(triangulated_points_errors_proj2)
+
+        with torch.no_grad():
+            NNs_triangulated_points_selected, NNs_triangulated_points_selected_proj_errors = select_best_keypoints(
+                NNs_triangulated_points=torch.stack(reference_image_dict["triangulated_points"], dim=0),
+                NNs_errors_proj1=np.stack(reference_image_dict["triangulated_points_errors_proj1"], axis=0),
+                NNs_errors_proj2=np.stack(reference_image_dict["triangulated_points_errors_proj2"], axis=0))
+
+        # 4. Save as gaussians
+        viewpoint_cam1 = viewpoint_stack[source_idx]
+        N = len(NNs_triangulated_points_selected)
+        with torch.no_grad():
+            new_xyz = NNs_triangulated_points_selected[:, :-1]
+            all_new_xyz.append(new_xyz)  # seeked_splats
+            all_new_features_dc.append(RGB2SH(torch.tensor(kptsA_color.astype(np.float32) / 255.)).unsqueeze(1))
+            all_new_features_rest.append(torch.stack([gaussians._features_rest[-1].clone().detach() * 0.] * N, dim=0))
+            # new version that sets points with large error invisible
+            # TODO: remove those points instead. However it doesn't affect the performance.
+            mask_bad_points = torch.tensor(
+                NNs_triangulated_points_selected_proj_errors > keypoint_fit_error_tolerance,
+                dtype=torch.float32).unsqueeze(1).to(device)
+            all_new_opacities.append(torch.stack([gaussians._opacity[-1].clone().detach()] * N, dim=0) * 0. - mask_bad_points * (1e1))
+
+            dist_points_to_cam1 = torch.linalg.norm(viewpoint_cam1.camera_center.clone().detach() - new_xyz,
+                                                    dim=1, ord=2)
+            #all_new_scaling.append(torch.log(((dist_points_to_cam1) / 1. * scaling_factor).unsqueeze(1).repeat(1, 3)))
+            all_new_scaling.append(gaussians.scaling_inverse_activation((dist_points_to_cam1 * scaling_factor).unsqueeze(1).repeat(1, 3)))
+            all_new_rotation.append(torch.stack([gaussians._rotation[-1].clone().detach()] * N, dim=0))
+
+    all_new_xyz = torch.cat(all_new_xyz, dim=0) 
+    all_new_features_dc = torch.cat(all_new_features_dc, dim=0)
+    new_tmp_radii = torch.zeros(all_new_xyz.shape[0])
+    prune_mask = torch.ones(all_new_xyz.shape[0], dtype=torch.bool)
+    
+    gaussians.densification_postfix(all_new_xyz[prune_mask].to(device),
+                                    all_new_features_dc[prune_mask].to(device),
+                                    torch.cat(all_new_features_rest, dim=0)[prune_mask].to(device),
+                                    torch.cat(all_new_opacities, dim=0)[prune_mask].to(device),
+                                    torch.cat(all_new_scaling, dim=0)[prune_mask].to(device),
+                                    torch.cat(all_new_rotation, dim=0)[prune_mask].to(device),
+                                    new_tmp_radii[prune_mask].to(device))
+    
+    return viewpoint_stack, closest_indices_selected, visualizations

source/corr_init_new.py

@@ -0,0 +1,904 @@
+import sys
+sys.path.append('../')
+sys.path.append("../submodules")
+sys.path.append('../submodules/RoMa')
+
+from matplotlib import pyplot as plt
+from PIL import Image
+import torch
+import numpy as np
+
+#from tqdm import tqdm_notebook as tqdm
+from tqdm import tqdm
+from scipy.cluster.vq import kmeans, vq
+from scipy.spatial.distance import cdist
+
+import torch.nn.functional as F
+from romatch import roma_outdoor, roma_indoor
+from utils.sh_utils import RGB2SH
+from romatch.utils import get_tuple_transform_ops
+
+
+def pairwise_distances(matrix):
+    """
+    Computes the pairwise Euclidean distances between all vectors in the input matrix.
+
+    Args:
+        matrix (torch.Tensor): Input matrix of shape [N, D], where N is the number of vectors and D is the dimensionality.
+
+    Returns:
+        torch.Tensor: Pairwise distance matrix of shape [N, N].
+    """
+    # Compute squared pairwise distances
+    squared_diff = torch.cdist(matrix, matrix, p=2)
+    return squared_diff
+
+
+def k_closest_vectors(matrix, k):
+    """
+    Finds the k-closest vectors for each vector in the input matrix based on Euclidean distance.
+
+    Args:
+        matrix (torch.Tensor): Input matrix of shape [N, D], where N is the number of vectors and D is the dimensionality.
+        k (int): Number of closest vectors to return for each vector.
+
+    Returns:
+        torch.Tensor: Indices of the k-closest vectors for each vector, excluding the vector itself.
+    """
+    # Compute pairwise distances
+    distances = pairwise_distances(matrix)
+
+    # For each vector, sort distances and get the indices of the k-closest vectors (excluding itself)
+    # Set diagonal distances to infinity to exclude the vector itself from the nearest neighbors
+    distances.fill_diagonal_(float('inf'))
+
+    # Get the indices of the k smallest distances (k-closest vectors)
+    _, indices = torch.topk(distances, k, largest=False, dim=1)
+
+    return indices
+
+
+def select_cameras_kmeans(cameras, K):
+    """
+    Selects K cameras from a set using K-means clustering.
+
+    Args:
+        cameras: NumPy array of shape (N, 16), representing N cameras with their 4x4 homogeneous matrices flattened.
+        K: Number of clusters (cameras to select).
+
+    Returns:
+        selected_indices: List of indices of the cameras closest to the cluster centers.
+    """
+    # Ensure input is a NumPy array
+    if not isinstance(cameras, np.ndarray):
+        cameras = np.asarray(cameras)
+
+    if cameras.shape[1] != 16:
+        raise ValueError("Each camera must have 16 values corresponding to a flattened 4x4 matrix.")
+
+    # Perform K-means clustering
+    cluster_centers, _ = kmeans(cameras, K)
+
+    # Assign each camera to a cluster and find distances to cluster centers
+    cluster_assignments, _ = vq(cameras, cluster_centers)
+
+    # Find the camera nearest to each cluster center
+    selected_indices = []
+    for k in range(K):
+        cluster_members = cameras[cluster_assignments == k]
+        distances = cdist([cluster_centers[k]], cluster_members)[0]
+        nearest_camera_idx = np.where(cluster_assignments == k)[0][np.argmin(distances)]
+        selected_indices.append(nearest_camera_idx)
+
+    return selected_indices
+
+
+def compute_warp_and_confidence(viewpoint_cam1, viewpoint_cam2, roma_model, device="cuda", verbose=False, output_dict={}):
+    """
+    Computes the warp and confidence between two viewpoint cameras using the roma_model.
+
+    Args:
+        viewpoint_cam1: Source viewpoint camera.
+        viewpoint_cam2: Target viewpoint camera.
+        roma_model: Pre-trained Roma model for correspondence matching.
+        device: Device to run the computation on.
+        verbose: If True, displays the images.
+
+    Returns:
+        certainty: Confidence tensor.
+        warp: Warp tensor.
+        imB: Processed image B as numpy array.
+    """
+    # Prepare images
+    imA = viewpoint_cam1.original_image.detach().cpu().numpy().transpose(1, 2, 0)
+    imB = viewpoint_cam2.original_image.detach().cpu().numpy().transpose(1, 2, 0)
+    imA = Image.fromarray(np.clip(imA * 255, 0, 255).astype(np.uint8))
+    imB = Image.fromarray(np.clip(imB * 255, 0, 255).astype(np.uint8))
+
+    if verbose:
+        fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(16, 8))
+        cax1 = ax[0].imshow(imA)
+        ax[0].set_title("Image 1")
+        cax2 = ax[1].imshow(imB)
+        ax[1].set_title("Image 2")
+        fig.colorbar(cax1, ax=ax[0])
+        fig.colorbar(cax2, ax=ax[1])
+    
+        for axis in ax:
+            axis.axis('off')
+        # Save the figure into the dictionary
+        output_dict[f'image_pair'] = fig
+   
+    # Transform images
+    ws, hs = roma_model.w_resized, roma_model.h_resized
+    test_transform = get_tuple_transform_ops(resize=(hs, ws), normalize=True)
+    im_A, im_B = test_transform((imA, imB))
+    batch = {"im_A": im_A[None].to(device), "im_B": im_B[None].to(device)}
+
+    # Forward pass through Roma model
+    corresps = roma_model.forward(batch) if not roma_model.symmetric else roma_model.forward_symmetric(batch)
+    finest_scale = 1
+    hs, ws = roma_model.upsample_res if roma_model.upsample_preds else (hs, ws)
+
+    # Process certainty and warp
+    certainty = corresps[finest_scale]["certainty"]
+    im_A_to_im_B = corresps[finest_scale]["flow"]
+    if roma_model.attenuate_cert:
+        low_res_certainty = F.interpolate(
+            corresps[16]["certainty"], size=(hs, ws), align_corners=False, mode="bilinear"
+        )
+        certainty -= 0.5 * low_res_certainty * (low_res_certainty < 0)
+
+    # Upsample predictions if needed
+    if roma_model.upsample_preds:
+        im_A_to_im_B = F.interpolate(
+            im_A_to_im_B, size=(hs, ws), align_corners=False, mode="bilinear"
+        )
+        certainty = F.interpolate(
+            certainty, size=(hs, ws), align_corners=False, mode="bilinear"
+        )
+
+    # Convert predictions to final format
+    im_A_to_im_B = im_A_to_im_B.permute(0, 2, 3, 1)
+    im_A_coords = torch.stack(torch.meshgrid(
+        torch.linspace(-1 + 1 / hs, 1 - 1 / hs, hs, device=device),
+        torch.linspace(-1 + 1 / ws, 1 - 1 / ws, ws, device=device),
+        indexing='ij'
+    ), dim=0).permute(1, 2, 0).unsqueeze(0).expand(im_A_to_im_B.size(0), -1, -1, -1)
+
+    warp = torch.cat((im_A_coords, im_A_to_im_B), dim=-1)
+    certainty = certainty.sigmoid()
+
+    return certainty[0, 0], warp[0], np.array(imB)
+
+
+def resize_batch(tensors_3d, tensors_4d, target_shape):
+    """
+    Resizes a batch of tensors with shapes [B, H, W] and [B, H, W, 4] to the target spatial dimensions.
+
+    Args:
+        tensors_3d: Tensor of shape [B, H, W].
+        tensors_4d: Tensor of shape [B, H, W, 4].
+        target_shape: Tuple (target_H, target_W) specifying the target spatial dimensions.
+
+    Returns:
+        resized_tensors_3d: Tensor of shape [B, target_H, target_W].
+        resized_tensors_4d: Tensor of shape [B, target_H, target_W, 4].
+    """
+    target_H, target_W = target_shape
+
+    # Resize [B, H, W] tensor
+    resized_tensors_3d = F.interpolate(
+        tensors_3d.unsqueeze(1), size=(target_H, target_W), mode="bilinear", align_corners=False
+    ).squeeze(1)
+
+    # Resize [B, H, W, 4] tensor
+    B, _, _, C = tensors_4d.shape
+    resized_tensors_4d = F.interpolate(
+        tensors_4d.permute(0, 3, 1, 2), size=(target_H, target_W), mode="bilinear", align_corners=False
+    ).permute(0, 2, 3, 1)
+
+    return resized_tensors_3d, resized_tensors_4d
+
+
+def aggregate_confidences_and_warps(viewpoint_stack, closest_indices, roma_model, source_idx, verbose=False, output_dict={}):
+    """
+    Aggregates confidences and warps by iterating over the nearest neighbors of the source viewpoint.
+
+    Args:
+        viewpoint_stack: Stack of viewpoint cameras.
+        closest_indices: Indices of the nearest neighbors for each viewpoint.
+        roma_model: Pre-trained Roma model.
+        source_idx: Index of the source viewpoint.
+        verbose: If True, displays intermediate results.
+
+    Returns:
+        certainties_max: Aggregated maximum confidences.
+        warps_max: Aggregated warps corresponding to maximum confidences.
+        certainties_max_idcs: Pixel-wise index of the image  from which we taken the best matching.
+        imB_compound: List of the neighboring images.
+    """
+    certainties_all, warps_all, imB_compound = [], [], []
+
+    for nn in tqdm(closest_indices[source_idx]):
+
+        viewpoint_cam1 = viewpoint_stack[source_idx]
+        viewpoint_cam2 = viewpoint_stack[nn]
+
+        certainty, warp, imB = compute_warp_and_confidence(viewpoint_cam1, viewpoint_cam2, roma_model, verbose=verbose, output_dict=output_dict)
+        certainties_all.append(certainty)
+        warps_all.append(warp)
+        imB_compound.append(imB)
+
+    certainties_all = torch.stack(certainties_all, dim=0)
+    target_shape = imB_compound[0].shape[:2]
+    if verbose: 
+        print("certainties_all.shape:", certainties_all.shape)
+        print("torch.stack(warps_all, dim=0).shape:", torch.stack(warps_all, dim=0).shape)
+        print("target_shape:", target_shape)        
+
+    certainties_all_resized, warps_all_resized = resize_batch(certainties_all,
+                                                              torch.stack(warps_all, dim=0),
+                                                              target_shape
+                                                              )
+
+    if verbose:
+        print("warps_all_resized.shape:", warps_all_resized.shape)
+        for n, cert in enumerate(certainties_all):
+            fig, ax = plt.subplots()
+            cax = ax.imshow(cert.cpu().numpy(), cmap='viridis')
+            fig.colorbar(cax, ax=ax)
+            ax.set_title("Pixel-wise Confidence")
+            output_dict[f'certainty_{n}'] = fig
+
+        for n, warp in enumerate(warps_all):
+            fig, ax = plt.subplots()
+            cax = ax.imshow(warp.cpu().numpy()[:, :, :3], cmap='viridis')
+            fig.colorbar(cax, ax=ax)
+            ax.set_title("Pixel-wise warp")
+            output_dict[f'warp_resized_{n}'] = fig
+
+        for n, cert in enumerate(certainties_all_resized):
+            fig, ax = plt.subplots()
+            cax = ax.imshow(cert.cpu().numpy(), cmap='viridis')
+            fig.colorbar(cax, ax=ax)
+            ax.set_title("Pixel-wise Confidence resized")
+            output_dict[f'certainty_resized_{n}'] = fig
+
+        for n, warp in enumerate(warps_all_resized):
+            fig, ax = plt.subplots()
+            cax = ax.imshow(warp.cpu().numpy()[:, :, :3], cmap='viridis')
+            fig.colorbar(cax, ax=ax)
+            ax.set_title("Pixel-wise warp resized")
+            output_dict[f'warp_resized_{n}'] = fig
+
+    certainties_max, certainties_max_idcs = torch.max(certainties_all_resized, dim=0)
+    H, W = certainties_max.shape
+
+    warps_max = warps_all_resized[certainties_max_idcs, torch.arange(H).unsqueeze(1), torch.arange(W)]
+
+
+    return certainties_max, warps_max, certainties_max_idcs, imB_compound, certainties_all_resized, warps_all_resized
+
+
+
+def extract_keypoints_and_colors(imA, imB_compound, certainties_max, certainties_max_idcs, matches, roma_model,
+                                 verbose=False, output_dict={}):
+    """
+    Extracts keypoints and corresponding colors from the source image (imA) and multiple target images (imB_compound).
+
+    Args:
+        imA: Source image as a NumPy array (H_A, W_A, C).
+        imB_compound: List of target images as NumPy arrays [(H_B, W_B, C), ...].
+        certainties_max: Tensor of pixel-wise maximum confidences.
+        certainties_max_idcs: Tensor of pixel-wise indices for the best matches.
+        matches: Matches in normalized coordinates.
+        roma_model: Roma model instance for keypoint operations.
+        verbose: if to show intermediate outputs and visualize results
+
+    Returns:
+        kptsA_np: Keypoints in imA in normalized coordinates.
+        kptsB_np: Keypoints in imB in normalized coordinates.
+        kptsA_color: Colors of keypoints in imA.
+        kptsB_color: Colors of keypoints in imB based on certainties_max_idcs.
+    """
+    H_A, W_A, _ = imA.shape
+    H, W = certainties_max.shape
+
+    # Convert matches to pixel coordinates
+    kptsA, kptsB = roma_model.to_pixel_coordinates(
+        matches, W_A, H_A, H, W  # W, H
+    )
+
+    kptsA_np = kptsA.detach().cpu().numpy()
+    kptsB_np = kptsB.detach().cpu().numpy()
+    kptsA_np = kptsA_np[:, [1, 0]]
+
+    if verbose:
+        fig, ax = plt.subplots(figsize=(12, 6))
+        cax = ax.imshow(imA)
+        ax.set_title("Reference image, imA")
+        output_dict[f'reference_image'] = fig
+
+        fig, ax = plt.subplots(figsize=(12, 6))
+        cax = ax.imshow(imB_compound[0])
+        ax.set_title("Image to compare to image, imB_compound")
+        output_dict[f'imB_compound'] = fig
+    
+        fig, ax = plt.subplots(figsize=(12, 6))
+        cax = ax.imshow(np.flipud(imA))
+        cax = ax.scatter(kptsA_np[:, 0], H_A - kptsA_np[:, 1], s=.03)
+        ax.set_title("Keypoints in imA")
+        ax.set_xlim(0, W_A)
+        ax.set_ylim(0, H_A)
+        output_dict[f'kptsA'] = fig
+
+        fig, ax = plt.subplots(figsize=(12, 6))
+        cax = ax.imshow(np.flipud(imB_compound[0]))
+        cax = ax.scatter(kptsB_np[:, 0], H_A - kptsB_np[:, 1], s=.03)
+        ax.set_title("Keypoints in imB")
+        ax.set_xlim(0, W_A)
+        ax.set_ylim(0, H_A)
+        output_dict[f'kptsB'] = fig
+
+    # Keypoints are in format (row, column) so the first value is alwain in range [0;height] and second is in range[0;width]
+
+    kptsA_np = kptsA.detach().cpu().numpy()
+    kptsB_np = kptsB.detach().cpu().numpy()
+
+    # Extract colors for keypoints in imA (vectorized)
+    # New experimental version
+    kptsA_x = np.round(kptsA_np[:, 0] / 1.).astype(int)
+    kptsA_y = np.round(kptsA_np[:, 1] / 1.).astype(int)
+    kptsA_color = imA[np.clip(kptsA_x, 0, H - 1), np.clip(kptsA_y, 0, W - 1)]
+   
+    # Create a composite image from imB_compound
+    imB_compound_np = np.stack(imB_compound, axis=0)
+    H_B, W_B, _ = imB_compound[0].shape
+
+    # Extract colors for keypoints in imB using certainties_max_idcs
+    imB_np = imB_compound_np[
+            certainties_max_idcs.detach().cpu().numpy(),
+            np.arange(H).reshape(-1, 1),
+            np.arange(W)
+        ]
+    
+    if verbose:
+        print("imB_np.shape:", imB_np.shape)
+        print("imB_np:", imB_np)
+        fig, ax = plt.subplots(figsize=(12, 6))
+        cax = ax.imshow(np.flipud(imB_np))
+        cax = ax.scatter(kptsB_np[:, 0], H_A - kptsB_np[:, 1], s=.03)
+        ax.set_title("np.flipud(imB_np[0]")
+        ax.set_xlim(0, W_A)
+        ax.set_ylim(0, H_A)
+        output_dict[f'np.flipud(imB_np[0]'] = fig
+
+
+    kptsB_x = np.round(kptsB_np[:, 0]).astype(int)
+    kptsB_y = np.round(kptsB_np[:, 1]).astype(int)
+
+    certainties_max_idcs_np = certainties_max_idcs.detach().cpu().numpy()
+    kptsB_proj_matrices_idx = certainties_max_idcs_np[np.clip(kptsA_x, 0, H - 1), np.clip(kptsA_y, 0, W - 1)]
+    kptsB_color = imB_compound_np[kptsB_proj_matrices_idx, np.clip(kptsB_y, 0, H - 1), np.clip(kptsB_x, 0, W - 1)]
+
+    # Normalize keypoints in both images
+    kptsA_np[:, 0] = kptsA_np[:, 0] / H * 2.0 - 1.0
+    kptsA_np[:, 1] = kptsA_np[:, 1] / W * 2.0 - 1.0
+    kptsB_np[:, 0] = kptsB_np[:, 0] / W_B * 2.0 - 1.0
+    kptsB_np[:, 1] = kptsB_np[:, 1] / H_B * 2.0 - 1.0
+
+    return kptsA_np[:, [1, 0]], kptsB_np, kptsB_proj_matrices_idx, kptsA_color, kptsB_color
+
+def prepare_tensor(input_array, device):
+    """
+    Converts an input array to a torch tensor, clones it, and detaches it for safe computation.
+    Args:
+        input_array (array-like): The input array to convert.
+        device (str or torch.device): The device to move the tensor to.
+    Returns:
+        torch.Tensor: A detached tensor clone of the input array on the specified device.
+    """
+    if not isinstance(input_array, torch.Tensor):
+        return torch.tensor(input_array, dtype=torch.float32).to(device).clone().detach()
+    return input_array.clone().detach().to(device).to(torch.float32)
+
+def triangulate_points(P1, P2, k1_x, k1_y, k2_x, k2_y, device="cuda"):
+    """
+    Solves for a batch of 3D points given batches of projection matrices and corresponding image points.
+
+    Parameters:
+    - P1, P2: Tensors of projection matrices of size (batch_size, 4, 4) or (4, 4)
+    - k1_x, k1_y: Tensors of shape (batch_size,)
+    - k2_x, k2_y: Tensors of shape (batch_size,)
+
+    Returns:
+    - X: A tensor containing the 3D homogeneous coordinates, shape (batch_size, 4)
+    """
+    EPS = 1e-4
+    # Ensure inputs are tensors
+
+    P1 = prepare_tensor(P1, device)
+    P2 = prepare_tensor(P2, device)
+    k1_x = prepare_tensor(k1_x, device)
+    k1_y = prepare_tensor(k1_y, device)
+    k2_x = prepare_tensor(k2_x, device)
+    k2_y =  prepare_tensor(k2_y, device)
+    batch_size = k1_x.shape[0]
+
+    # Expand P1 and P2 if they are not batched
+    if P1.ndim == 2:
+        P1 = P1.unsqueeze(0).expand(batch_size, -1, -1)
+    if P2.ndim == 2:
+        P2 = P2.unsqueeze(0).expand(batch_size, -1, -1)
+
+    # Extract columns from P1 and P2
+    P1_0 = P1[:, :, 0]  # Shape: (batch_size, 4)
+    P1_1 = P1[:, :, 1]
+    P1_2 = P1[:, :, 2]
+
+    P2_0 = P2[:, :, 0]
+    P2_1 = P2[:, :, 1]
+    P2_2 = P2[:, :, 2]
+
+    # Reshape kx and ky to (batch_size, 1)
+    k1_x = k1_x.view(-1, 1)
+    k1_y = k1_y.view(-1, 1)
+    k2_x = k2_x.view(-1, 1)
+    k2_y = k2_y.view(-1, 1)
+
+    # Construct the equations for each batch
+    # For camera 1
+    A1 = P1_0 - k1_x * P1_2  # Shape: (batch_size, 4)
+    A2 = P1_1 - k1_y * P1_2
+    # For camera 2
+    A3 = P2_0 - k2_x * P2_2
+    A4 = P2_1 - k2_y * P2_2
+
+    # Stack the equations
+    A = torch.stack([A1, A2, A3, A4], dim=1)  # Shape: (batch_size, 4, 4)
+
+    # Right-hand side (constants)
+    b = -A[:, :, 3]  # Shape: (batch_size, 4)
+    A_reduced = A[:, :, :3]  # Coefficients of x, y, z
+
+    # Solve using torch.linalg.lstsq (supports batching)
+    X_xyz = torch.linalg.lstsq(A_reduced, b.unsqueeze(2)).solution.squeeze(2)  # Shape: (batch_size, 3)
+
+    # Append 1 to get homogeneous coordinates
+    ones = torch.ones((batch_size, 1), dtype=torch.float32, device=X_xyz.device)
+    X = torch.cat([X_xyz, ones], dim=1)  # Shape: (batch_size, 4)
+
+    # Now compute the errors of projections.
+    seeked_splats_proj1 = (X.unsqueeze(1) @ P1).squeeze(1)
+    seeked_splats_proj1 = seeked_splats_proj1 / (EPS + seeked_splats_proj1[:, [3]])
+    seeked_splats_proj2 = (X.unsqueeze(1) @ P2).squeeze(1)
+    seeked_splats_proj2 = seeked_splats_proj2 / (EPS + seeked_splats_proj2[:, [3]])
+    proj1_target = torch.concat([k1_x, k1_y], dim=1)
+    proj2_target = torch.concat([k2_x, k2_y], dim=1)
+    errors_proj1 = torch.abs(seeked_splats_proj1[:, :2] - proj1_target).sum(1).detach().cpu().numpy()
+    errors_proj2 = torch.abs(seeked_splats_proj2[:, :2] - proj2_target).sum(1).detach().cpu().numpy()
+
+    return X, errors_proj1, errors_proj2
+
+
+
+def select_best_keypoints(
+        NNs_triangulated_points, NNs_errors_proj1, NNs_errors_proj2, device="cuda"):
+    """
+    From all the points fitted to  keypoints and corresponding colors from the source image (imA) and multiple target images (imB_compound).
+
+    Args:
+        NNs_triangulated_points:  torch tensor with keypoints coordinates (num_nns, num_points, dim). dim can be arbitrary,
+            usually 3 or 4(for homogeneous representation).
+        NNs_errors_proj1:  numpy array with projection error of the estimated keypoint on the reference frame (num_nns, num_points).
+        NNs_errors_proj2:  numpy array with projection error of the estimated keypoint on the neighbor frame (num_nns, num_points).
+    Returns:
+        selected_keypoints: keypoints with the best score.
+    """
+
+    NNs_errors_proj = np.maximum(NNs_errors_proj1, NNs_errors_proj2)
+
+    # Convert indices to PyTorch tensor
+    indices = torch.from_numpy(np.argmin(NNs_errors_proj, axis=0)).long().to(device)
+
+    # Create index tensor for the second dimension
+    n_indices = torch.arange(NNs_triangulated_points.shape[1]).long().to(device)
+
+    # Use advanced indexing to select elements
+    NNs_triangulated_points_selected = NNs_triangulated_points[indices, n_indices, :]  # Shape: [N, k]
+
+    return NNs_triangulated_points_selected, np.min(NNs_errors_proj, axis=0)
+
+
+
+import time
+from collections import defaultdict
+from tqdm import tqdm
+
+# def init_gaussians_with_corr_profiled(gaussians, scene, cfg, device, verbose=False, roma_model=None):
+#     timings = defaultdict(list)  # To accumulate timings
+
+#     if roma_model is None:
+#         if cfg.roma_model == "indoors":
+#             roma_model = roma_indoor(device=device)
+#         else:
+#             roma_model = roma_outdoor(device=device)
+#         roma_model.upsample_preds = False
+#         roma_model.symmetric = False
+
+#     M = cfg.matches_per_ref
+#     upper_thresh = roma_model.sample_thresh
+#     scaling_factor = cfg.scaling_factor
+#     expansion_factor = 1
+#     keypoint_fit_error_tolerance = cfg.proj_err_tolerance
+#     visualizations = {}
+#     viewpoint_stack = scene.getTrainCameras().copy()
+#     NUM_REFERENCE_FRAMES = min(cfg.num_refs, len(viewpoint_stack))
+#     NUM_NNS_PER_REFERENCE = min(cfg.nns_per_ref, len(viewpoint_stack))
+
+#     viewpoint_cam_all = torch.stack([x.world_view_transform.flatten() for x in viewpoint_stack], axis=0)
+
+#     selected_indices = select_cameras_kmeans(cameras=viewpoint_cam_all.detach().cpu().numpy(), K=NUM_REFERENCE_FRAMES)
+#     selected_indices = sorted(selected_indices)
+
+#     viewpoint_cam_all = torch.stack([x.world_view_transform.flatten() for x in viewpoint_stack], axis=0)
+#     closest_indices = k_closest_vectors(viewpoint_cam_all, NUM_NNS_PER_REFERENCE)
+#     closest_indices_selected = closest_indices[:, :].detach().cpu().numpy()
+
+#     all_new_xyz = []
+#     all_new_features_dc = []
+#     all_new_features_rest = []
+#     all_new_opacities = []
+#     all_new_scaling = []
+#     all_new_rotation = []
+
+#     # Dummy first pass to initialize model
+#     with torch.no_grad():
+#         viewpoint_cam1 = viewpoint_stack[0]
+#         viewpoint_cam2 = viewpoint_stack[1]
+#         imA = viewpoint_cam1.original_image.detach().cpu().numpy().transpose(1, 2, 0)
+#         imB = viewpoint_cam2.original_image.detach().cpu().numpy().transpose(1, 2, 0)
+#         imA = Image.fromarray(np.clip(imA * 255, 0, 255).astype(np.uint8))
+#         imB = Image.fromarray(np.clip(imB * 255, 0, 255).astype(np.uint8))
+#         warp, certainty_warp = roma_model.match(imA, imB, device=device)
+#         del warp, certainty_warp
+#         torch.cuda.empty_cache()
+
+#     # Main Loop over source_idx
+#     for source_idx in tqdm(sorted(selected_indices), desc="Profiling source frames"):
+
+#         # =================== Step 1: Aggregate Confidences and Warps ===================
+#         start = time.time()
+#         viewpoint_cam1 = viewpoint_stack[source_idx]
+#         viewpoint_cam2 = viewpoint_stack[closest_indices_selected[source_idx,0]]
+#         imA = viewpoint_cam1.original_image.detach().cpu().numpy().transpose(1, 2, 0)
+#         imB = viewpoint_cam2.original_image.detach().cpu().numpy().transpose(1, 2, 0)
+#         imA = Image.fromarray(np.clip(imA * 255, 0, 255).astype(np.uint8))
+#         imB = Image.fromarray(np.clip(imB * 255, 0, 255).astype(np.uint8))
+#         warp, certainty_warp = roma_model.match(imA, imB, device=device)
+
+#         certainties_max, warps_max, certainties_max_idcs, imB_compound, certainties_all, warps_all = aggregate_confidences_and_warps(
+#             viewpoint_stack=viewpoint_stack,
+#             closest_indices=closest_indices_selected,
+#             roma_model=roma_model,
+#             source_idx=source_idx,
+#             verbose=verbose, 
+#             output_dict=visualizations
+#         )
+
+#         certainties_max = certainty_warp
+#         with torch.no_grad():
+#             warps_all = warps.unsqueeze(0)
+        
+#         timings['aggregation_warp_certainty'].append(time.time() - start)
+
+#         # =================== Step 2: Good Samples Selection ===================
+#         start = time.time()
+#         certainty = certainties_max.reshape(-1).clone()
+#         certainty[certainty > upper_thresh] = 1
+#         good_samples = torch.multinomial(certainty, num_samples=min(expansion_factor * M, len(certainty)), replacement=False)
+#         timings['good_samples_selection'].append(time.time() - start)
+
+#         # =================== Step 3: Triangulate Keypoints for Each NN ===================
+#         reference_image_dict = {
+#             "triangulated_points": [],
+#             "triangulated_points_errors_proj1": [],
+#             "triangulated_points_errors_proj2": []
+#         }
+
+#         start = time.time()
+#         for NN_idx in range(len(warps_all)):
+#             matches_NN = warps_all[NN_idx].reshape(-1, 4)[good_samples]
+
+#             # Extract keypoints and colors
+#             kptsA_np, kptsB_np, kptsB_proj_matrices_idcs, kptsA_color, kptsB_color = extract_keypoints_and_colors(
+#                 imA, imB_compound, certainties_max, certainties_max_idcs, matches_NN, roma_model
+#             )
+
+#             proj_matrices_A = viewpoint_stack[source_idx].full_proj_transform
+#             proj_matrices_B = viewpoint_stack[closest_indices_selected[source_idx, NN_idx]].full_proj_transform
+#             triangulated_points, triangulated_points_errors_proj1, triangulated_points_errors_proj2 = triangulate_points(
+#                 P1=torch.stack([proj_matrices_A] * M, axis=0),
+#                 P2=torch.stack([proj_matrices_B] * M, axis=0),
+#                 k1_x=kptsA_np[:M, 0], k1_y=kptsA_np[:M, 1],
+#                 k2_x=kptsB_np[:M, 0], k2_y=kptsB_np[:M, 1])
+
+#             reference_image_dict["triangulated_points"].append(triangulated_points)
+#             reference_image_dict["triangulated_points_errors_proj1"].append(triangulated_points_errors_proj1)
+#             reference_image_dict["triangulated_points_errors_proj2"].append(triangulated_points_errors_proj2)
+#         timings['triangulation_per_NN'].append(time.time() - start)
+
+#         # =================== Step 4: Select Best Triangulated Points ===================
+#         start = time.time()
+#         NNs_triangulated_points_selected, NNs_triangulated_points_selected_proj_errors = select_best_keypoints(
+#             NNs_triangulated_points=torch.stack(reference_image_dict["triangulated_points"], dim=0),
+#             NNs_errors_proj1=np.stack(reference_image_dict["triangulated_points_errors_proj1"], axis=0),
+#             NNs_errors_proj2=np.stack(reference_image_dict["triangulated_points_errors_proj2"], axis=0))
+#         timings['select_best_keypoints'].append(time.time() - start)
+
+#         # =================== Step 5: Create New Gaussians ===================
+#         start = time.time()
+#         viewpoint_cam1 = viewpoint_stack[source_idx]
+#         N = len(NNs_triangulated_points_selected)
+#         new_xyz = NNs_triangulated_points_selected[:, :-1]
+#         all_new_xyz.append(new_xyz)
+#         all_new_features_dc.append(RGB2SH(torch.tensor(kptsA_color.astype(np.float32) / 255.)).unsqueeze(1))
+#         all_new_features_rest.append(torch.stack([gaussians._features_rest[-1].clone().detach() * 0.] * N, dim=0))
+
+#         mask_bad_points = torch.tensor(
+#             NNs_triangulated_points_selected_proj_errors > keypoint_fit_error_tolerance,
+#             dtype=torch.float32).unsqueeze(1).to(device)
+
+#         all_new_opacities.append(torch.stack([gaussians._opacity[-1].clone().detach()] * N, dim=0) * 0. - mask_bad_points * (1e1))
+
+#         dist_points_to_cam1 = torch.linalg.norm(viewpoint_cam1.camera_center.clone().detach() - new_xyz, dim=1, ord=2)
+#         all_new_scaling.append(gaussians.scaling_inverse_activation((dist_points_to_cam1 * scaling_factor).unsqueeze(1).repeat(1, 3)))
+#         all_new_rotation.append(torch.stack([gaussians._rotation[-1].clone().detach()] * N, dim=0))
+#         timings['save_gaussians'].append(time.time() - start)
+
+#     # =================== Final Densification Postfix ===================
+#     start = time.time()
+#     all_new_xyz = torch.cat(all_new_xyz, dim=0) 
+#     all_new_features_dc = torch.cat(all_new_features_dc, dim=0)
+#     new_tmp_radii = torch.zeros(all_new_xyz.shape[0])
+#     prune_mask = torch.ones(all_new_xyz.shape[0], dtype=torch.bool)
+
+#     gaussians.densification_postfix(
+#         all_new_xyz[prune_mask].to(device),
+#         all_new_features_dc[prune_mask].to(device),
+#         torch.cat(all_new_features_rest, dim=0)[prune_mask].to(device),
+#         torch.cat(all_new_opacities, dim=0)[prune_mask].to(device),
+#         torch.cat(all_new_scaling, dim=0)[prune_mask].to(device),
+#         torch.cat(all_new_rotation, dim=0)[prune_mask].to(device),
+#         new_tmp_radii[prune_mask].to(device)
+#     )
+#     timings['final_densification_postfix'].append(time.time() - start)
+
+#     # =================== Print Profiling Results ===================
+#     print("\n=== Profiling Summary (average per frame) ===")
+#     for key, times in timings.items():
+#         print(f"{key:35s}: {sum(times) / len(times):.4f} sec (total {sum(times):.2f} sec)")
+
+#     return viewpoint_stack, closest_indices_selected, visualizations
+
+
+
+def extract_keypoints_and_colors_single(imA, imB, matches, roma_model, verbose=False, output_dict={}):
+    """
+    Extracts keypoints and corresponding colors from a source image (imA) and a single target image (imB).
+
+    Args:
+        imA: Source image as a NumPy array (H_A, W_A, C).
+        imB: Target image as a NumPy array (H_B, W_B, C).
+        matches: Matches in normalized coordinates (torch.Tensor).
+        roma_model: Roma model instance for keypoint operations.
+        verbose: If True, outputs intermediate visualizations.
+    Returns:
+        kptsA_np: Keypoints in imA (normalized).
+        kptsB_np: Keypoints in imB (normalized).
+        kptsA_color: Colors of keypoints in imA.
+        kptsB_color: Colors of keypoints in imB.
+    """
+    H_A, W_A, _ = imA.shape
+    H_B, W_B, _ = imB.shape
+
+    # Convert matches to pixel coordinates
+    # Matches format: (B, 4) = (x1_norm, y1_norm, x2_norm, y2_norm)
+    kptsA = matches[:, :2]  # [N, 2]
+    kptsB = matches[:, 2:]  # [N, 2]
+
+    # Scale normalized coordinates [-1,1] to pixel coordinates
+    kptsA_pix = torch.zeros_like(kptsA)
+    kptsB_pix = torch.zeros_like(kptsB)
+
+    # Important! [Normalized to pixel space]
+    kptsA_pix[:, 0] = (kptsA[:, 0] + 1) * (W_A - 1) / 2
+    kptsA_pix[:, 1] = (kptsA[:, 1] + 1) * (H_A - 1) / 2
+
+    kptsB_pix[:, 0] = (kptsB[:, 0] + 1) * (W_B - 1) / 2
+    kptsB_pix[:, 1] = (kptsB[:, 1] + 1) * (H_B - 1) / 2
+
+    kptsA_np = kptsA_pix.detach().cpu().numpy()
+    kptsB_np = kptsB_pix.detach().cpu().numpy()
+
+    # Extract colors
+    kptsA_x = np.round(kptsA_np[:, 0]).astype(int)
+    kptsA_y = np.round(kptsA_np[:, 1]).astype(int)
+    kptsB_x = np.round(kptsB_np[:, 0]).astype(int)
+    kptsB_y = np.round(kptsB_np[:, 1]).astype(int)
+
+    kptsA_color = imA[np.clip(kptsA_y, 0, H_A-1), np.clip(kptsA_x, 0, W_A-1)]
+    kptsB_color = imB[np.clip(kptsB_y, 0, H_B-1), np.clip(kptsB_x, 0, W_B-1)]
+
+    # Normalize keypoints into [-1, 1] for downstream triangulation
+    kptsA_np_norm = np.zeros_like(kptsA_np)
+    kptsB_np_norm = np.zeros_like(kptsB_np)
+
+    kptsA_np_norm[:, 0] = kptsA_np[:, 0] / (W_A - 1) * 2.0 - 1.0
+    kptsA_np_norm[:, 1] = kptsA_np[:, 1] / (H_A - 1) * 2.0 - 1.0
+
+    kptsB_np_norm[:, 0] = kptsB_np[:, 0] / (W_B - 1) * 2.0 - 1.0
+    kptsB_np_norm[:, 1] = kptsB_np[:, 1] / (H_B - 1) * 2.0 - 1.0
+
+    return kptsA_np_norm, kptsB_np_norm, kptsA_color, kptsB_color
+
+
+
+def init_gaussians_with_corr_profiled(gaussians, scene, cfg, device, verbose=False, roma_model=None):
+    timings = defaultdict(list)
+
+    if roma_model is None:
+        if cfg.roma_model == "indoors":
+            roma_model = roma_indoor(device=device)
+        else:
+            roma_model = roma_outdoor(device=device)
+        roma_model.upsample_preds = False
+        roma_model.symmetric = False
+
+    M = cfg.matches_per_ref
+    upper_thresh = roma_model.sample_thresh
+    scaling_factor = cfg.scaling_factor
+    expansion_factor = 1
+    keypoint_fit_error_tolerance = cfg.proj_err_tolerance
+    visualizations = {}
+    viewpoint_stack = scene.getTrainCameras().copy()
+    NUM_REFERENCE_FRAMES = min(cfg.num_refs, len(viewpoint_stack))
+    NUM_NNS_PER_REFERENCE = 1  # Only ONE neighbor now!
+
+    viewpoint_cam_all = torch.stack([x.world_view_transform.flatten() for x in viewpoint_stack], axis=0)
+
+    selected_indices = select_cameras_kmeans(cameras=viewpoint_cam_all.detach().cpu().numpy(), K=NUM_REFERENCE_FRAMES)
+    selected_indices = sorted(selected_indices)
+
+    viewpoint_cam_all = torch.stack([x.world_view_transform.flatten() for x in viewpoint_stack], axis=0)
+    closest_indices = k_closest_vectors(viewpoint_cam_all, NUM_NNS_PER_REFERENCE)
+    closest_indices_selected = closest_indices[:, :].detach().cpu().numpy()
+
+    all_new_xyz = []
+    all_new_features_dc = []
+    all_new_features_rest = []
+    all_new_opacities = []
+    all_new_scaling = []
+    all_new_rotation = []
+
+    # Dummy first pass to initialize model
+    with torch.no_grad():
+        viewpoint_cam1 = viewpoint_stack[0]
+        viewpoint_cam2 = viewpoint_stack[1]
+        imA = viewpoint_cam1.original_image.detach().cpu().numpy().transpose(1, 2, 0)
+        imB = viewpoint_cam2.original_image.detach().cpu().numpy().transpose(1, 2, 0)
+        imA = Image.fromarray(np.clip(imA * 255, 0, 255).astype(np.uint8))
+        imB = Image.fromarray(np.clip(imB * 255, 0, 255).astype(np.uint8))
+        warp, certainty_warp = roma_model.match(imA, imB, device=device)
+        del warp, certainty_warp
+        torch.cuda.empty_cache()
+
+    # Main Loop over source_idx
+    for source_idx in tqdm(sorted(selected_indices), desc="Profiling source frames"):
+
+        # =================== Step 1: Compute Warp and Certainty ===================
+        start = time.time()
+        viewpoint_cam1 = viewpoint_stack[source_idx]
+        NNs=closest_indices_selected.shape[1]
+        viewpoint_cam2 = viewpoint_stack[closest_indices_selected[source_idx, np.random.randint(NNs)]]
+        imA = viewpoint_cam1.original_image.detach().cpu().numpy().transpose(1, 2, 0)
+        imB = viewpoint_cam2.original_image.detach().cpu().numpy().transpose(1, 2, 0)
+        imA = Image.fromarray(np.clip(imA * 255, 0, 255).astype(np.uint8))
+        imB = Image.fromarray(np.clip(imB * 255, 0, 255).astype(np.uint8))
+        warp, certainty_warp = roma_model.match(imA, imB, device=device)
+
+        certainties_max = certainty_warp  # New manual sampling
+        timings['aggregation_warp_certainty'].append(time.time() - start)
+
+        # =================== Step 2: Good Samples Selection ===================
+        start = time.time()
+        certainty = certainties_max.reshape(-1).clone()
+        certainty[certainty > upper_thresh] = 1
+        good_samples = torch.multinomial(certainty, num_samples=min(expansion_factor * M, len(certainty)), replacement=False)
+        timings['good_samples_selection'].append(time.time() - start)
+
+        # =================== Step 3: Triangulate Keypoints ===================
+        reference_image_dict = {
+            "triangulated_points": [],
+            "triangulated_points_errors_proj1": [],
+            "triangulated_points_errors_proj2": []
+        }
+
+        start = time.time()
+        matches_NN = warp.reshape(-1, 4)[good_samples]
+
+        # Convert matches to pixel coordinates
+        kptsA_np, kptsB_np, kptsA_color, kptsB_color = extract_keypoints_and_colors_single(
+            np.array(imA).astype(np.uint8), 
+            np.array(imB).astype(np.uint8), 
+            matches_NN, 
+            roma_model
+        )
+
+        proj_matrices_A = viewpoint_stack[source_idx].full_proj_transform
+        proj_matrices_B = viewpoint_stack[closest_indices_selected[source_idx, 0]].full_proj_transform
+
+        triangulated_points, triangulated_points_errors_proj1, triangulated_points_errors_proj2 = triangulate_points(
+            P1=torch.stack([proj_matrices_A] * M, axis=0),
+            P2=torch.stack([proj_matrices_B] * M, axis=0),
+            k1_x=kptsA_np[:M, 0], k1_y=kptsA_np[:M, 1],
+            k2_x=kptsB_np[:M, 0], k2_y=kptsB_np[:M, 1])
+
+        reference_image_dict["triangulated_points"].append(triangulated_points)
+        reference_image_dict["triangulated_points_errors_proj1"].append(triangulated_points_errors_proj1)
+        reference_image_dict["triangulated_points_errors_proj2"].append(triangulated_points_errors_proj2)
+        timings['triangulation_per_NN'].append(time.time() - start)
+
+        # =================== Step 4: Select Best Triangulated Points ===================
+        start = time.time()
+        NNs_triangulated_points_selected, NNs_triangulated_points_selected_proj_errors = select_best_keypoints(
+            NNs_triangulated_points=torch.stack(reference_image_dict["triangulated_points"], dim=0),
+            NNs_errors_proj1=np.stack(reference_image_dict["triangulated_points_errors_proj1"], axis=0),
+            NNs_errors_proj2=np.stack(reference_image_dict["triangulated_points_errors_proj2"], axis=0))
+        timings['select_best_keypoints'].append(time.time() - start)
+
+        # =================== Step 5: Create New Gaussians ===================
+        start = time.time()
+        viewpoint_cam1 = viewpoint_stack[source_idx]
+        N = len(NNs_triangulated_points_selected)
+        new_xyz = NNs_triangulated_points_selected[:, :-1]
+        all_new_xyz.append(new_xyz)
+        all_new_features_dc.append(RGB2SH(torch.tensor(kptsA_color.astype(np.float32) / 255.)).unsqueeze(1))
+        all_new_features_rest.append(torch.stack([gaussians._features_rest[-1].clone().detach() * 0.] * N, dim=0))
+
+        mask_bad_points = torch.tensor(
+            NNs_triangulated_points_selected_proj_errors > keypoint_fit_error_tolerance,
+            dtype=torch.float32).unsqueeze(1).to(device)
+
+        all_new_opacities.append(torch.stack([gaussians._opacity[-1].clone().detach()] * N, dim=0) * 0. - mask_bad_points * (1e1))
+
+        dist_points_to_cam1 = torch.linalg.norm(viewpoint_cam1.camera_center.clone().detach() - new_xyz, dim=1, ord=2)
+        all_new_scaling.append(gaussians.scaling_inverse_activation((dist_points_to_cam1 * scaling_factor).unsqueeze(1).repeat(1, 3)))
+        all_new_rotation.append(torch.stack([gaussians._rotation[-1].clone().detach()] * N, dim=0))
+        timings['save_gaussians'].append(time.time() - start)
+
+    # =================== Final Densification Postfix ===================
+    start = time.time()
+    all_new_xyz = torch.cat(all_new_xyz, dim=0) 
+    all_new_features_dc = torch.cat(all_new_features_dc, dim=0)
+    new_tmp_radii = torch.zeros(all_new_xyz.shape[0])
+    prune_mask = torch.ones(all_new_xyz.shape[0], dtype=torch.bool)
+
+    gaussians.densification_postfix(
+        all_new_xyz[prune_mask].to(device),
+        all_new_features_dc[prune_mask].to(device),
+        torch.cat(all_new_features_rest, dim=0)[prune_mask].to(device),
+        torch.cat(all_new_opacities, dim=0)[prune_mask].to(device),
+        torch.cat(all_new_scaling, dim=0)[prune_mask].to(device),
+        torch.cat(all_new_rotation, dim=0)[prune_mask].to(device),
+        new_tmp_radii[prune_mask].to(device)
+    )
+    timings['final_densification_postfix'].append(time.time() - start)
+
+    # =================== Print Profiling Results ===================
+    print("\n=== Profiling Summary (average per frame) ===")
+    for key, times in timings.items():
+        print(f"{key:35s}: {sum(times) / len(times):.4f} sec (total {sum(times):.2f} sec)")
+
+    return viewpoint_stack, closest_indices_selected, visualizations

source/data_utils.py

@@ -0,0 +1,28 @@
+def scene_cameras_train_test_split(scene, verbose=False):
+    """
+    Iterate over resolutions in the scene. For each resolution check if this resolution has test_cameras
+    if it doesn't then extract every 8th camera from the train and put it to the test set. This follows the
+    evaluation protocol suggested by Kerbl et al. in the seminal work on 3DGS. All changes to the input
+    object scene are inplace changes.
+    :param scene: Scene Class object from the gaussian-splatting.scene module
+    :param verbose: Print initial and final stage of the function
+    :return:  None
+
+    """
+    if verbose: print("Preparing train and test sets split...")
+    for resolution in scene.train_cameras.keys():
+        if len(scene.test_cameras[resolution]) == 0:
+            if verbose:
+                print(f"Found no test_cameras for resolution {resolution}. Move every 8th camera out ouf total "+\
+                      f"{len(scene.train_cameras[resolution])} train cameras to the test set now")
+            N = len(scene.train_cameras[resolution])
+            scene.test_cameras[resolution] = [scene.train_cameras[resolution][idx] for idx in range(0, N) 
+                                              if idx % 8 == 0]
+            scene.train_cameras[resolution] = [scene.train_cameras[resolution][idx] for idx in range(0, N)
+                                               if idx % 8 != 0]
+            if verbose:
+                print(f"Done. Now train and test sets contain each {len(scene.train_cameras[resolution])} and " + \
+                      f"{len(scene.test_cameras[resolution])} cameras respectively.")
+
+
+    return

source/losses.py

@@ -0,0 +1,100 @@
+# Code is copied from the gaussian-splatting/utils/loss_utils.py
+
+import torch
+import torch.nn.functional as F
+from torch.autograd import Variable
+from math import exp
+
+def l1_loss(network_output, gt, mean=True):
+    return torch.abs((network_output - gt)).mean() if mean else torch.abs((network_output - gt))
+
+def l2_loss(network_output, gt):
+    return ((network_output - gt) ** 2).mean()
+
+def gaussian(window_size, sigma):
+    gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)])
+    return gauss / gauss.sum()
+
+def create_window(window_size, channel):
+    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
+    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
+    return window
+
+def ssim(img1, img2, window_size=11, size_average=True, mask = None):
+    channel = img1.size(-3)
+    window = create_window(window_size, channel)
+
+    if img1.is_cuda:
+        window = window.cuda(img1.get_device())
+    window = window.type_as(img1)
+
+    return _ssim(img1, img2, window, window_size, channel, size_average, mask)
+
+def _ssim(img1, img2, window, window_size, channel, size_average=True, mask = None):
+    mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel)
+    mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel)
+
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1 * mu2
+
+    sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq
+    sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq
+    sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2
+
+    C1 = 0.01 ** 2
+    C2 = 0.03 ** 2
+
+    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
+
+    if mask is not None:
+        ssim_map = ssim_map * mask
+
+    if size_average:
+        return ssim_map.mean()
+    else:
+        return ssim_map.mean(1).mean(1).mean(1)
+
+
+def mse(img1, img2):
+    return (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True)
+
+def psnr(img1, img2):
+    """
+    Computes the Peak Signal-to-Noise Ratio (PSNR) between two single images. NOT BATCHED!
+    Args:
+        img1 (torch.Tensor): The first image tensor, with pixel values scaled between 0 and 1.
+                             Shape should be (channels, height, width).
+        img2 (torch.Tensor): The second image tensor with the same shape as img1, used for comparison.
+
+    Returns:
+        torch.Tensor: A scalar tensor containing the PSNR value in decibels (dB).
+    """
+    mse = (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True)
+    return 20 * torch.log10(1.0 / torch.sqrt(mse))
+
+
+def tv_loss(image):
+    """
+    Computes the total variation (TV) loss for an image of shape [3, H, W].
+
+    Args:
+        image (torch.Tensor): Input image of shape [3, H, W]
+
+    Returns:
+        torch.Tensor: Scalar value representing the total variation loss.
+    """
+    # Ensure the image has the correct dimensions
+    assert image.ndim == 3 and image.shape[0] == 3, "Input must be of shape [3, H, W]"
+
+    # Compute the difference between adjacent pixels in the x-direction (width)
+    diff_x = torch.abs(image[:, :, 1:] - image[:, :, :-1])
+
+    # Compute the difference between adjacent pixels in the y-direction (height)
+    diff_y = torch.abs(image[:, 1:, :] - image[:, :-1, :])
+
+    # Sum the total variation in both directions
+    tv_loss_value = torch.mean(diff_x) + torch.mean(diff_y)
+
+    return tv_loss_value

source/networks.py

@@ -0,0 +1,52 @@
+import torch
+
+import sys
+sys.path.append('./submodules/gaussian-splatting/')
+
+from random import randint
+from scene import Scene, GaussianModel
+from gaussian_renderer import render
+from source.data_utils import scene_cameras_train_test_split
+
+class Warper3DGS(torch.nn.Module):
+    def __init__(self, sh_degree,  opt, pipe, dataset, viewpoint_stack, verbose,
+                 do_train_test_split=True):
+        super(Warper3DGS, self).__init__()
+        """
+        Init Warper using all the objects necessary for rendering gaussian splats.
+        Here we merely link class objects to the objects instantiated outsided the class.
+        """
+        print("ready!!!7")
+        self.gaussians = GaussianModel(sh_degree)
+        print("ready!!!8")
+        self.gaussians.tmp_radii = torch.zeros((self.gaussians.get_xyz.shape[0]), device="cuda")
+        self.render = render
+        self.gs_config_opt = opt
+        bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]
+        self.bg = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
+        self.pipe = pipe
+        print("ready!!!")
+        self.scene = Scene(dataset, self.gaussians, shuffle=False)
+        print("ready2")
+        if do_train_test_split:
+            scene_cameras_train_test_split(self.scene, verbose=verbose)
+
+        self.gaussians.training_setup(opt)
+        self.viewpoint_stack = viewpoint_stack
+        if not self.viewpoint_stack:
+            self.viewpoint_stack = self.scene.getTrainCameras().copy()
+
+    def forward(self, viewpoint_cam=None):
+        """
+        For a provided camera viewpoint_cam we render gaussians from this viewpoint.
+        If no camera provided then we use the self.viewpoint_stack (list of cameras).
+        If the latter is empty we reinitialize it using the self.scene object.
+        """
+        if not viewpoint_cam:
+            if not self.viewpoint_stack:
+                self.viewpoint_stack = self.scene.getTrainCameras().copy()
+            viewpoint_cam = self.viewpoint_stack[randint(0, len(self.viewpoint_stack) - 1)]
+
+        render_pkg = self.render(viewpoint_cam, self.gaussians, self.pipe, self.bg)
+        return render_pkg
+

source/timer.py

@@ -0,0 +1,24 @@
+import time
+class Timer:
+    def __init__(self):
+        self.start_time = None
+        self.elapsed = 0
+        self.paused = False
+
+    def start(self):
+        if self.start_time is None:
+            self.start_time = time.time()
+        elif self.paused:
+            self.start_time = time.time() - self.elapsed
+            self.paused = False
+
+    def pause(self):
+        if not self.paused:
+            self.elapsed = time.time() - self.start_time
+            self.paused = True
+
+    def get_elapsed_time(self):
+        if self.paused:
+            return self.elapsed
+        else:
+            return time.time() - self.start_time

source/trainer.py

@@ -0,0 +1,262 @@
+import torch
+from random import randint
+from tqdm.rich import trange
+from tqdm import tqdm as tqdm
+from source.networks import Warper3DGS
+import wandb
+import sys
+
+sys.path.append('./submodules/gaussian-splatting/')
+import lpips
+from source.losses import ssim, l1_loss, psnr
+from rich.console import Console
+from rich.theme import Theme
+
+custom_theme = Theme({
+    "info": "dim cyan",
+    "warning": "magenta",
+    "danger": "bold red"
+})
+
+#from source.corr_init import init_gaussians_with_corr
+from source.corr_init_new import init_gaussians_with_corr_profiled as init_gaussians_with_corr
+from source.utils_aux import log_samples
+
+from source.timer import Timer
+
+class EDGSTrainer:
+    def __init__(self,
+                 GS: Warper3DGS,
+                 training_config,
+                 dataset_white_background=False,
+                 device=torch.device('cuda'),
+                 log_wandb=True,
+                 ):
+        self.GS = GS
+        self.scene = GS.scene
+        self.viewpoint_stack = GS.viewpoint_stack
+        self.gaussians = GS.gaussians
+
+        self.training_config = training_config
+        self.GS_optimizer = GS.gaussians.optimizer
+        self.dataset_white_background = dataset_white_background
+
+        self.training_step = 1
+        self.gs_step = 0
+        self.CONSOLE = Console(width=120, theme=custom_theme)
+        self.saving_iterations = training_config.save_iterations
+        self.evaluate_iterations = None
+        self.batch_size = training_config.batch_size
+        self.ema_loss_for_log = 0.0
+
+        # Logs in the format {step:{"loss1":loss1_value, "loss2":loss2_value}}
+        self.logs_losses = {}
+        self.lpips = lpips.LPIPS(net='vgg').to(device)
+        self.device = device
+        self.timer = Timer()
+        self.log_wandb = log_wandb
+
+    def load_checkpoints(self, load_cfg):
+        # Load 3DGS checkpoint
+        if load_cfg.gs:
+            self.gs.gaussians.restore(
+                torch.load(f"{load_cfg.gs}/chkpnt{load_cfg.gs_step}.pth")[0],
+                self.training_config)
+            self.GS_optimizer = self.GS.gaussians.optimizer
+            self.CONSOLE.print(f"3DGS loaded from checkpoint for iteration {load_cfg.gs_step}",
+                               style="info")
+            self.training_step += load_cfg.gs_step
+            self.gs_step += load_cfg.gs_step
+
+    def train(self, train_cfg):
+        # 3DGS training
+        self.CONSOLE.print("Train 3DGS for {} iterations".format(train_cfg.gs_epochs), style="info")    
+        with trange(self.training_step, self.training_step + train_cfg.gs_epochs, desc="[green]Train gaussians") as progress_bar:
+            for self.training_step in progress_bar:
+                radii = self.train_step_gs(max_lr=train_cfg.max_lr, no_densify=train_cfg.no_densify)
+                with torch.no_grad():
+                    if train_cfg.no_densify:
+                        self.prune(radii)
+                    else:
+                        self.densify_and_prune(radii)
+                    if train_cfg.reduce_opacity:
+                        # Slightly reduce opacity every few steps:
+                        if self.gs_step < self.training_config.densify_until_iter and self.gs_step % 10 == 0:
+                            opacities_new = torch.log(torch.exp(self.GS.gaussians._opacity.data) * 0.99)
+                            self.GS.gaussians._opacity.data = opacities_new
+                    self.timer.pause()
+                    # Progress bar
+                    if self.training_step % 10 == 0:
+                        progress_bar.set_postfix({"[red]Loss": f"{self.ema_loss_for_log:.{7}f}"}, refresh=True)
+                    # Log and save
+                    if self.training_step in self.saving_iterations:
+                        self.save_model()
+                    if self.evaluate_iterations is not None:
+                        if self.training_step in self.evaluate_iterations:
+                            self.evaluate()
+                    else:
+                        if (self.training_step <= 3000 and self.training_step % 500 == 0) or \
+                            (self.training_step > 3000 and self.training_step % 1000 == 228) :
+                            self.evaluate()
+
+                    self.timer.start()
+
+
+    def evaluate(self):
+        torch.cuda.empty_cache()
+        log_gen_images, log_real_images = [], []
+        validation_configs = ({'name': 'test', 'cameras': self.scene.getTestCameras(), 'cam_idx': self.training_config.TEST_CAM_IDX_TO_LOG},
+                              {'name': 'train',
+                               'cameras': [self.scene.getTrainCameras()[idx % len(self.scene.getTrainCameras())] for idx in
+                                           range(0, 150, 5)], 'cam_idx': 10})
+        if self.log_wandb:
+            wandb.log({f"Number of Gaussians": len(self.GS.gaussians._xyz)}, step=self.training_step)
+        for config in validation_configs:
+            if config['cameras'] and len(config['cameras']) > 0:
+                l1_test = 0.0
+                psnr_test = 0.0
+                ssim_test = 0.0
+                lpips_splat_test = 0.0
+                for idx, viewpoint in enumerate(config['cameras']):
+                    image = torch.clamp(self.GS(viewpoint)["render"], 0.0, 1.0)
+                    gt_image = torch.clamp(viewpoint.original_image.to(self.device), 0.0, 1.0)
+                    l1_test += l1_loss(image, gt_image).double()
+                    psnr_test += psnr(image.unsqueeze(0), gt_image.unsqueeze(0)).double()
+                    ssim_test += ssim(image, gt_image).double()
+                    lpips_splat_test += self.lpips(image, gt_image).detach().double()
+                    if idx in [config['cam_idx']]:
+                        log_gen_images.append(image)
+                        log_real_images.append(gt_image)
+                psnr_test /= len(config['cameras'])
+                l1_test /= len(config['cameras'])
+                ssim_test /= len(config['cameras'])
+                lpips_splat_test /= len(config['cameras'])
+                if self.log_wandb:
+                    wandb.log({f"{config['name']}/L1": l1_test.item(), f"{config['name']}/PSNR": psnr_test.item(), \
+                            f"{config['name']}/SSIM": ssim_test.item(), f"{config['name']}/LPIPS_splat": lpips_splat_test.item()}, step = self.training_step)
+                self.CONSOLE.print("\n[ITER {}], #{} gaussians, Evaluating {}: L1={:.6f},  PSNR={:.6f}, SSIM={:.6f}, LPIPS_splat={:.6f} ".format(
+                    self.training_step, len(self.GS.gaussians._xyz), config['name'], l1_test.item(), psnr_test.item(), ssim_test.item(), lpips_splat_test.item()), style="info")
+        if self.log_wandb:
+            with torch.no_grad():
+                log_samples(torch.stack((log_real_images[0],log_gen_images[0])) , [], self.training_step, caption="Real and Generated Samples")
+                wandb.log({"time": self.timer.get_elapsed_time()}, step=self.training_step)
+        torch.cuda.empty_cache()
+
+    def train_step_gs(self, max_lr = False, no_densify = False):
+        self.gs_step += 1
+        if max_lr:
+            self.GS.gaussians.update_learning_rate(max(self.gs_step, 8_000))
+        else:
+            self.GS.gaussians.update_learning_rate(self.gs_step)
+        # Every 1000 its we increase the levels of SH up to a maximum degree
+        if self.gs_step % 1000 == 0:
+            self.GS.gaussians.oneupSHdegree()
+
+        # Pick a random Camera
+        if not self.viewpoint_stack:
+            self.viewpoint_stack = self.scene.getTrainCameras().copy()
+        viewpoint_cam = self.viewpoint_stack.pop(randint(0, len(self.viewpoint_stack) - 1))
+      
+        render_pkg = self.GS(viewpoint_cam=viewpoint_cam)
+        image = render_pkg["render"]
+        # Loss
+        gt_image = viewpoint_cam.original_image.to(self.device)
+        L1_loss = l1_loss(image, gt_image)
+
+        ssim_loss = (1.0 - ssim(image, gt_image))
+        loss = (1.0 - self.training_config.lambda_dssim) * L1_loss + \
+               self.training_config.lambda_dssim * ssim_loss
+        self.timer.pause() 
+        self.logs_losses[self.training_step] = {"loss": loss.item(),
+                                                "L1_loss": L1_loss.item(),
+                                                "ssim_loss": ssim_loss.item()}
+        
+        if self.log_wandb:
+            for k, v in self.logs_losses[self.training_step].items():
+                wandb.log({f"train/{k}": v}, step=self.training_step)
+        self.ema_loss_for_log = 0.4 * self.logs_losses[self.training_step]["loss"] + 0.6 * self.ema_loss_for_log
+        self.timer.start()
+        self.GS_optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        with torch.no_grad():
+            if self.gs_step < self.training_config.densify_until_iter and not no_densify:
+                self.GS.gaussians.max_radii2D[render_pkg["visibility_filter"]] = torch.max(
+                    self.GS.gaussians.max_radii2D[render_pkg["visibility_filter"]],
+                    render_pkg["radii"][render_pkg["visibility_filter"]])
+                self.GS.gaussians.add_densification_stats(render_pkg["viewspace_points"],
+                                                                     render_pkg["visibility_filter"])
+
+        # Optimizer step
+        self.GS_optimizer.step()
+        self.GS_optimizer.zero_grad(set_to_none=True)
+        return render_pkg["radii"]
+
+    def densify_and_prune(self, radii = None):
+        # Densification or pruning
+        if self.gs_step < self.training_config.densify_until_iter:
+            if (self.gs_step > self.training_config.densify_from_iter) and \
+                    (self.gs_step % self.training_config.densification_interval == 0):
+                size_threshold = 20 if self.gs_step > self.training_config.opacity_reset_interval else None
+                self.GS.gaussians.densify_and_prune(self.training_config.densify_grad_threshold,
+                                                               0.005,
+                                                               self.GS.scene.cameras_extent,
+                                                               size_threshold, radii)
+            if self.gs_step % self.training_config.opacity_reset_interval == 0 or (
+                    self.dataset_white_background and self.gs_step == self.training_config.densify_from_iter):
+                self.GS.gaussians.reset_opacity()             
+
+          
+
+    def save_model(self):
+        print("\n[ITER {}] Saving Gaussians".format(self.gs_step))
+        self.scene.save(self.gs_step)
+        print("\n[ITER {}] Saving Checkpoint".format(self.gs_step))
+        torch.save((self.GS.gaussians.capture(), self.gs_step),
+                self.scene.model_path + "/chkpnt" + str(self.gs_step) + ".pth")
+
+
+    def init_with_corr(self, cfg, verbose=False, roma_model=None): 
+        """
+        Initializes image with matchings. Also removes SfM init points.
+        Args:
+            cfg: configuration part named init_wC. Check train.yaml
+            verbose: whether you want to print intermediate results. Useful for debug.
+            roma_model: optionally you can pass here preinit RoMA model to avoid reinit 
+                it every time.  
+        """
+        if not cfg.use:
+            return None
+        N_splats_at_init = len(self.GS.gaussians._xyz)
+        print("N_splats_at_init:", N_splats_at_init)
+        camera_set, selected_indices, visualization_dict = init_gaussians_with_corr(
+            self.GS.gaussians, 
+            self.scene, 
+            cfg, 
+            self.device,                                                                                    
+            verbose=verbose,
+            roma_model=roma_model)
+
+        # Remove SfM points and leave only matchings inits
+        if not cfg.add_SfM_init:
+            with torch.no_grad():
+                N_splats_after_init = len(self.GS.gaussians._xyz)
+                print("N_splats_after_init:", N_splats_after_init)
+                self.gaussians.tmp_radii = torch.zeros(self.gaussians._xyz.shape[0]).to(self.device)
+                mask = torch.concat([torch.ones(N_splats_at_init, dtype=torch.bool),
+                                    torch.zeros(N_splats_after_init-N_splats_at_init, dtype=torch.bool)],
+                                axis=0)
+                self.GS.gaussians.prune_points(mask)
+        with torch.no_grad():
+            gaussians =  self.gaussians
+            gaussians._scaling =  gaussians.scaling_inverse_activation(gaussians.scaling_activation(gaussians._scaling)*0.5)
+        return visualization_dict
+    
+
+    def prune(self, radii, min_opacity=0.005):
+        self.GS.gaussians.tmp_radii = radii
+        if self.gs_step < self.training_config.densify_until_iter:
+            prune_mask = (self.GS.gaussians.get_opacity < min_opacity).squeeze()
+            self.GS.gaussians.prune_points(prune_mask)
+            torch.cuda.empty_cache()
+        self.GS.gaussians.tmp_radii = None
+

source/utils_aux.py

@@ -0,0 +1,92 @@
+# Perlin noise code taken from https://gist.github.com/adefossez/0646dbe9ed4005480a2407c62aac8869
+from types import SimpleNamespace
+import random
+import numpy as np
+import torch
+import torchvision
+import wandb
+import random
+import torchvision.transforms as T
+import torchvision.transforms.functional as F
+import torch
+from PIL import Image
+
+def parse_dict_to_namespace(dict_nested):
+    """Turns nested dictionary into nested namespaces"""
+    if type(dict_nested) != dict and type(dict_nested) != list: return dict_nested
+    x = SimpleNamespace()
+    for key, val in dict_nested.items():
+        if type(val) == dict:
+            setattr(x, key, parse_dict_to_namespace(val))
+        elif type(val) == list:
+            setattr(x, key, [parse_dict_to_namespace(v) for v in val])
+        else:
+            setattr(x, key, val)
+    return x
+
+def set_seed(seed=42, cuda=True):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if cuda:
+        torch.cuda.manual_seed_all(seed)
+
+
+
+def log_samples(samples, scores, iteration, caption="Real Samples"):
+    # Create a grid of images
+    grid = torchvision.utils.make_grid(samples)
+
+    # Log the images and scores to wandb
+    wandb.log({
+        f"{caption}_images": [wandb.Image(grid, caption=f"{caption}: {scores}")],
+    }, step = iteration)
+
+
+
+def pairwise_distances(matrix):
+    """
+    Computes the pairwise Euclidean distances between all vectors in the input matrix.
+
+    Args:
+        matrix (torch.Tensor): Input matrix of shape [N, D], where N is the number of vectors and D is the dimensionality.
+
+    Returns:
+        torch.Tensor: Pairwise distance matrix of shape [N, N].
+    """
+    # Compute squared pairwise distances
+    squared_diff = torch.cdist(matrix, matrix, p=2)
+    return squared_diff
+
+def k_closest_vectors(matrix, k):
+    """
+    Finds the k-closest vectors for each vector in the input matrix based on Euclidean distance.
+
+    Args:
+        matrix (torch.Tensor): Input matrix of shape [N, D], where N is the number of vectors and D is the dimensionality.
+        k (int): Number of closest vectors to return for each vector.
+
+    Returns:
+        torch.Tensor: Indices of the k-closest vectors for each vector, excluding the vector itself.
+    """
+    # Compute pairwise distances
+    distances = pairwise_distances(matrix)
+
+    # For each vector, sort distances and get the indices of the k-closest vectors (excluding itself)
+    # Set diagonal distances to infinity to exclude the vector itself from the nearest neighbors
+    distances.fill_diagonal_(float('inf'))
+    
+    # Get the indices of the k smallest distances (k-closest vectors)
+    _, indices = torch.topk(distances, k, largest=False, dim=1)
+
+    return indices
+
+def process_image(image_tensor):
+    image_np = image_tensor.detach().cpu().numpy().transpose(1, 2, 0)
+    return Image.fromarray(np.clip(image_np * 255, 0, 255).astype(np.uint8))
+
+
+def normalize_keypoints(kpts_np, width, height):
+    kpts_np[:, 0] = kpts_np[:, 0] / width * 2. - 1.
+    kpts_np[:, 1] = kpts_np[:, 1] / height * 2. - 1.
+    return kpts_np

source/utils_preprocess.py

@@ -0,0 +1,334 @@
+# This file contains function for video or image collection preprocessing.
+# For video we do the preprocessing and select k sharpest frames.
+# Afterwards scene is constructed 
+import cv2
+import numpy as np
+from tqdm import tqdm
+import pycolmap
+import os
+import time
+import tempfile
+from moviepy import VideoFileClip
+from matplotlib import pyplot as plt
+from PIL import Image
+import cv2
+from tqdm import tqdm
+
+WORKDIR = "../outputs/"
+
+
+def get_rotation_moviepy(video_path):
+    clip = VideoFileClip(video_path)
+    rotation = 0
+
+    try:
+        displaymatrix = clip.reader.infos['inputs'][0]['streams'][2]['metadata'].get('displaymatrix', '')
+        if 'rotation of' in displaymatrix:
+            angle = float(displaymatrix.strip().split('rotation of')[-1].split('degrees')[0])
+            rotation = int(angle) % 360
+            
+    except Exception as e:
+        print(f"No displaymatrix rotation found: {e}")
+
+    clip.reader.close()
+    #if clip.audio:
+    #    clip.audio.reader.close_proc()
+
+    return rotation
+
+def resize_max_side(frame, max_size):
+    h, w = frame.shape[:2]
+    scale = max_size / max(h, w)
+    if scale < 1:
+        frame = cv2.resize(frame, (int(w * scale), int(h * scale)))
+    return frame
+
+def read_video_frames(video_input, k=1, max_size=1024):
+    """
+    Extracts every k-th frame from a video or list of images, resizes to max size, and returns frames as list.
+
+    Parameters:
+        video_input (str, file-like, or list): Path to video file, file-like object, or list of image files.
+        k (int): Interval for frame extraction (every k-th frame).
+        max_size (int): Maximum size for width or height after resizing.
+
+    Returns:
+        frames (list): List of resized frames (numpy arrays).
+    """
+    # Handle list of image files (not single video in a list)
+    if isinstance(video_input, list):
+        # If it's a single video in a list, treat it as video
+        if len(video_input) == 1 and video_input[0].name.endswith(('.mp4', '.avi', '.mov')):
+            video_input = video_input[0]  # unwrap single video file
+        else:
+            # Treat as list of images
+            frames = []
+            for img_file in video_input:
+                img = Image.open(img_file.name).convert("RGB")
+                img.thumbnail((max_size, max_size))
+                frames.append(np.array(img)[...,::-1])
+            return frames
+
+    # Handle file-like or path
+    if hasattr(video_input, 'name'):
+        video_path = video_input.name
+    elif isinstance(video_input, (str, os.PathLike)):
+        video_path = str(video_input)
+    else:
+        raise ValueError("Unsupported video input type. Must be a filepath, file-like object, or list of images.")
+
+    
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        raise ValueError(f"Error: Could not open video {video_path}.")
+
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    frame_count = 0
+    frames = []
+
+    with tqdm(total=total_frames // k, desc="Processing Video", unit="frame") as pbar:
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            if frame_count % k == 0:
+                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                h, w = frame.shape[:2]
+                scale = max(h, w) / max_size
+                if scale > 1:
+                    frame = cv2.resize(frame, (int(w / scale), int(h / scale)))
+                frames.append(frame[...,[2,1,0]])
+                pbar.update(1)
+            frame_count += 1
+
+    cap.release()
+    return frames
+
+def resize_max_side(frame, max_size):
+    """
+    Resizes the frame so that its largest side equals max_size, maintaining aspect ratio.
+    """
+    height, width = frame.shape[:2]
+    max_dim = max(height, width)
+    
+    if max_dim <= max_size:
+        return frame  # No need to resize
+
+    scale = max_size / max_dim
+    new_width = int(width * scale)
+    new_height = int(height * scale)
+
+    resized_frame = cv2.resize(frame, (new_width, new_height), interpolation=cv2.INTER_AREA)
+    return resized_frame
+
+
+
+def variance_of_laplacian(image):
+	# compute the Laplacian of the image and then return the focus
+	# measure, which is simply the variance of the Laplacian
+	return cv2.Laplacian(image, cv2.CV_64F).var()
+    
+def process_all_frames(IMG_FOLDER = '/scratch/datasets/hq_data/night2_all_frames',
+                       to_visualize=False,
+                       save_images=True):
+    dict_scores = {}
+    for idx, img_name in tqdm(enumerate(sorted([x for x in os.listdir(IMG_FOLDER) if '.png' in x]))):
+        
+        img = cv2.imread(os.path.join(IMG_FOLDER, img_name))#[250:, 100:]
+        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        fm = variance_of_laplacian(gray) + \
+                variance_of_laplacian(cv2.resize(gray, (0,0), fx=0.75, fy=0.75)) + \
+                variance_of_laplacian(cv2.resize(gray, (0,0), fx=0.5, fy=0.5)) + \
+                variance_of_laplacian(cv2.resize(gray, (0,0), fx=0.25, fy=0.25))
+        if to_visualize:
+            plt.figure()
+            plt.title(f"Laplacian score: {fm:.2f}")
+            plt.imshow(img[..., [2,1,0]])
+            plt.show()
+        dict_scores[idx] = {"idx" : idx, 
+                            "img_name" : img_name,
+                            "score" : fm}
+        if save_images:
+            dict_scores[idx]["img"] = img
+        
+    return dict_scores
+
+def select_optimal_frames(scores, k):
+    """
+    Selects a minimal subset of frames while ensuring no gaps exceed k.
+
+    Args:
+        scores (list of float): List of scores where index represents frame number.
+        k (int): Maximum allowed gap between selected frames.
+
+    Returns:
+        list of int: Indices of selected frames.
+    """
+    n = len(scores)
+    selected = [0, n-1]
+    i = 0  # Start at the first frame
+
+    while i < n:
+        # Find the best frame to select within the next k frames
+        best_idx = max(range(i, min(i + k + 1, n)), key=lambda x: scores[x], default=None)
+
+        if best_idx is None:
+            break  # No more frames left
+
+        selected.append(best_idx)
+        i = best_idx + k + 1  # Move forward, ensuring gaps stay within k
+
+    return sorted(selected)
+
+
+def variance_of_laplacian(image):
+    """
+    Compute the variance of Laplacian as a focus measure.
+    """
+    return cv2.Laplacian(image, cv2.CV_64F).var()
+
+def preprocess_frames(frames, verbose=False):
+    """
+    Compute sharpness scores for a list of frames using multi-scale Laplacian variance.
+
+    Args:
+        frames (list of np.ndarray): List of frames (BGR images).
+        verbose (bool): If True, print scores.
+
+    Returns:
+        list of float: Sharpness scores for each frame.
+    """
+    scores = []
+
+    for idx, frame in enumerate(tqdm(frames, desc="Scoring frames")):
+        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+
+        fm = (
+            variance_of_laplacian(gray) +
+            variance_of_laplacian(cv2.resize(gray, (0, 0), fx=0.75, fy=0.75)) +
+            variance_of_laplacian(cv2.resize(gray, (0, 0), fx=0.5, fy=0.5)) +
+            variance_of_laplacian(cv2.resize(gray, (0, 0), fx=0.25, fy=0.25))
+        )
+        
+        if verbose:
+            print(f"Frame {idx}: Sharpness Score = {fm:.2f}")
+
+        scores.append(fm)
+
+    return scores
+
+def select_optimal_frames(scores, k):
+    """
+    Selects k frames by splitting into k segments and picking the sharpest frame from each.
+
+    Args:
+        scores (list of float): List of sharpness scores.
+        k (int): Number of frames to select.
+
+    Returns:
+        list of int: Indices of selected frames.  
+    """
+    n = len(scores)
+    selected_indices = []
+    segment_size = n // k
+
+    for i in range(k):
+        start = i * segment_size
+        end = (i + 1) * segment_size if i < k - 1 else n  # Last chunk may be larger
+        segment_scores = scores[start:end]
+        
+        if len(segment_scores) == 0:
+            continue  # Safety check if some segment is empty
+        
+        best_in_segment = start + np.argmax(segment_scores)
+        selected_indices.append(best_in_segment)
+
+    return sorted(selected_indices)
+
+def save_frames_to_scene_dir(frames, scene_dir):
+    """
+    Saves a list of frames into the target scene directory under 'images/' subfolder.
+
+    Args:
+        frames (list of np.ndarray): List of frames (BGR images) to save.
+        scene_dir (str): Target path where 'images/' subfolder will be created.
+    """
+    images_dir = os.path.join(scene_dir, "images")
+    os.makedirs(images_dir, exist_ok=True)
+
+    for idx, frame in enumerate(frames):
+        filename = os.path.join(images_dir, f"{idx:08d}.png")  # 00000000.png, 00000001.png, etc.
+        cv2.imwrite(filename, frame)
+
+    print(f"Saved {len(frames)} frames to {images_dir}")
+
+
+def run_colmap_on_scene(scene_dir):
+    """
+    Runs feature extraction, matching, and mapping on all images inside scene_dir/images using pycolmap.
+
+    Args:
+        scene_dir (str): Path to scene directory containing 'images' folder.
+    
+    TODO: if the function hasn't managed to match all the frames either increase image size,
+    increase number of features or just remove those frames from the folder scene_dir/images
+    """
+    start_time = time.time()
+    print(f"Running COLMAP pipeline on all images inside {scene_dir}")
+
+    # Setup paths
+    database_path = os.path.join(scene_dir, "database.db")
+    sparse_path = os.path.join(scene_dir, "sparse")
+    image_dir = os.path.join(scene_dir, "images")
+    
+    # Make sure output directories exist
+    os.makedirs(sparse_path, exist_ok=True)
+
+    # Step 1: Feature Extraction
+    pycolmap.extract_features(
+        database_path,
+        image_dir,
+        sift_options={
+            "max_num_features": 512 * 2,
+            "max_image_size": 512 * 1,
+        }
+    )
+    print(f"Finished feature extraction in {(time.time() - start_time):.2f}s.")
+
+    # Step 2: Feature Matching
+    pycolmap.match_exhaustive(database_path)
+    print(f"Finished feature matching in {(time.time() - start_time):.2f}s.")
+
+    # Step 3: Mapping
+    pipeline_options = pycolmap.IncrementalPipelineOptions()
+    pipeline_options.min_num_matches = 15
+    pipeline_options.multiple_models = True
+    pipeline_options.max_num_models = 50
+    pipeline_options.max_model_overlap = 20
+    pipeline_options.min_model_size = 10
+    pipeline_options.extract_colors = True
+    pipeline_options.num_threads = 8
+    pipeline_options.mapper.init_min_num_inliers = 30
+    pipeline_options.mapper.init_max_error = 8.0
+    pipeline_options.mapper.init_min_tri_angle = 5.0
+
+    reconstruction = pycolmap.incremental_mapping(
+        database_path=database_path,
+        image_path=image_dir,
+        output_path=sparse_path,
+        options=pipeline_options,
+    )
+    print(f"Finished incremental mapping in {(time.time() - start_time):.2f}s.")
+
+    # Step 4: Post-process Cameras to SIMPLE_PINHOLE
+    recon_path = os.path.join(sparse_path, "0")
+    reconstruction = pycolmap.Reconstruction(recon_path)
+
+    for cam in reconstruction.cameras.values():
+        cam.model = 'SIMPLE_PINHOLE'
+        cam.params = cam.params[:3]  # Keep only [f, cx, cy]
+
+    reconstruction.write(recon_path)
+
+    print(f"Total pipeline time: {(time.time() - start_time):.2f}s.")
+

source/vggt_to_colmap.py

@@ -0,0 +1,598 @@
+# Reuse code taken from the implementation of atakan-topaloglu:
+#  https://github.com/atakan-topaloglu/vggt/blob/main/vggt_to_colmap.py
+
+import os
+import argparse
+import numpy as np
+import torch
+import glob
+import struct
+from scipy.spatial.transform import Rotation
+import sys
+from PIL import Image
+import cv2
+import requests
+import tempfile
+
+sys.path.append("submodules/vggt/")
+
+from vggt.models.vggt import VGGT
+from vggt.utils.load_fn import load_and_preprocess_images
+from vggt.utils.pose_enc import pose_encoding_to_extri_intri
+from vggt.utils.geometry import unproject_depth_map_to_point_map
+
+def load_model(device=None):
+    """Load and initialize the VGGT model."""
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {device}")
+    
+    model = VGGT.from_pretrained("facebook/VGGT-1B")
+
+    # model = VGGT()
+    # _URL = "https://huggingface.co/facebook/VGGT-1B/resolve/main/model.pt"
+    # model.load_state_dict(torch.hub.load_state_dict_from_url(_URL))
+    
+    model.eval()
+    model = model.to(device)
+    return model, device
+
+def process_images(image_dir, model, device):
+    """Process images with VGGT and return predictions."""
+    image_names = glob.glob(os.path.join(image_dir, "*"))
+    image_names = sorted([f for f in image_names if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
+    print(f"Found {len(image_names)} images")
+    
+    if len(image_names) == 0:
+        raise ValueError(f"No images found in {image_dir}")
+
+    original_images = []
+    for img_path in image_names:
+        img = Image.open(img_path).convert('RGB')
+        original_images.append(np.array(img))
+    
+    images = load_and_preprocess_images(image_names).to(device)
+    print(f"Preprocessed images shape: {images.shape}")
+    
+    print("Running inference...")
+    dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
+    
+    with torch.no_grad():
+        with torch.cuda.amp.autocast(dtype=dtype):
+            predictions = model(images)
+
+    print("Converting pose encoding to camera parameters...")
+    extrinsic, intrinsic = pose_encoding_to_extri_intri(predictions["pose_enc"], images.shape[-2:])
+    predictions["extrinsic"] = extrinsic
+    predictions["intrinsic"] = intrinsic
+
+    for key in predictions.keys():
+        if isinstance(predictions[key], torch.Tensor):
+            predictions[key] = predictions[key].cpu().numpy().squeeze(0)  # remove batch dimension
+    
+    print("Computing 3D points from depth maps...")
+    depth_map = predictions["depth"]  # (S, H, W, 1)
+    world_points = unproject_depth_map_to_point_map(depth_map, predictions["extrinsic"], predictions["intrinsic"])
+    predictions["world_points_from_depth"] = world_points
+    
+    predictions["original_images"] = original_images
+    
+    S, H, W = world_points.shape[:3]
+    normalized_images = np.zeros((S, H, W, 3), dtype=np.float32)
+    
+    for i, img in enumerate(original_images):
+        resized_img = cv2.resize(img, (W, H))
+        normalized_images[i] = resized_img / 255.0
+    
+    predictions["images"] = normalized_images
+    
+    return predictions, image_names
+
+def extrinsic_to_colmap_format(extrinsics):
+    """Convert extrinsic matrices to COLMAP format (quaternion + translation)."""
+    num_cameras = extrinsics.shape[0]
+    quaternions = []
+    translations = []
+    
+    for i in range(num_cameras):
+        # VGGT's extrinsic is camera-to-world (R|t) format
+        R = extrinsics[i, :3, :3]
+        t = extrinsics[i, :3, 3]
+        
+        # Convert rotation matrix to quaternion
+        # COLMAP quaternion format is [qw, qx, qy, qz]
+        rot = Rotation.from_matrix(R)
+        quat = rot.as_quat()  # scipy returns [x, y, z, w]
+        quat = np.array([quat[3], quat[0], quat[1], quat[2]])  # Convert to [w, x, y, z]
+        
+        quaternions.append(quat)
+        translations.append(t)
+    
+    return np.array(quaternions), np.array(translations)
+
+def download_file_from_url(url, filename):
+    """Downloads a file from a URL, handling redirects."""
+    try:
+        response = requests.get(url, allow_redirects=False)
+        response.raise_for_status() 
+
+        if response.status_code == 302:  
+            redirect_url = response.headers["Location"]
+            response = requests.get(redirect_url, stream=True)
+            response.raise_for_status()
+        else:
+            response = requests.get(url, stream=True)
+            response.raise_for_status()
+
+        with open(filename, "wb") as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+        print(f"Downloaded {filename} successfully.")
+        return True
+
+    except requests.exceptions.RequestException as e:
+        print(f"Error downloading file: {e}")
+        return False
+
+def segment_sky(image_path, onnx_session, mask_filename=None):
+    """
+    Segments sky from an image using an ONNX model.
+    """
+    image = cv2.imread(image_path)
+
+    result_map = run_skyseg(onnx_session, [320, 320], image)
+    result_map_original = cv2.resize(result_map, (image.shape[1], image.shape[0]))
+
+    # Fix: Invert the mask so that 255 = non-sky, 0 = sky
+    # The model outputs low values for sky, high values for non-sky
+    output_mask = np.zeros_like(result_map_original)
+    output_mask[result_map_original < 32] = 255  # Use threshold of 32
+
+    if mask_filename is not None:
+        os.makedirs(os.path.dirname(mask_filename), exist_ok=True)
+        cv2.imwrite(mask_filename, output_mask)
+    
+    return output_mask
+
+def run_skyseg(onnx_session, input_size, image):
+    """
+    Runs sky segmentation inference using ONNX model.
+    """
+    import copy
+    
+    temp_image = copy.deepcopy(image)
+    resize_image = cv2.resize(temp_image, dsize=(input_size[0], input_size[1]))
+    x = cv2.cvtColor(resize_image, cv2.COLOR_BGR2RGB)
+    x = np.array(x, dtype=np.float32)
+    mean = [0.485, 0.456, 0.406]
+    std = [0.229, 0.224, 0.225]
+    x = (x / 255 - mean) / std
+    x = x.transpose(2, 0, 1)
+    x = x.reshape(-1, 3, input_size[0], input_size[1]).astype("float32")
+
+    input_name = onnx_session.get_inputs()[0].name
+    output_name = onnx_session.get_outputs()[0].name
+    onnx_result = onnx_session.run([output_name], {input_name: x})
+
+    onnx_result = np.array(onnx_result).squeeze()
+    min_value = np.min(onnx_result)
+    max_value = np.max(onnx_result)
+    onnx_result = (onnx_result - min_value) / (max_value - min_value)
+    onnx_result *= 255
+    onnx_result = onnx_result.astype("uint8")
+
+    return onnx_result
+
+def filter_and_prepare_points(predictions, conf_threshold, mask_sky=False, mask_black_bg=False, 
+                             mask_white_bg=False, stride=1, prediction_mode="Depthmap and Camera Branch"):
+    """
+    Filter points based on confidence and prepare for COLMAP format.
+    Implementation matches the conventions in the original VGGT code.
+    """
+    
+    if "Pointmap" in prediction_mode:
+        print("Using Pointmap Branch")
+        if "world_points" in predictions:
+            pred_world_points = predictions["world_points"]
+            pred_world_points_conf = predictions.get("world_points_conf", np.ones_like(pred_world_points[..., 0]))
+        else:
+            print("Warning: world_points not found in predictions, falling back to depth-based points")
+            pred_world_points = predictions["world_points_from_depth"]
+            pred_world_points_conf = predictions.get("depth_conf", np.ones_like(pred_world_points[..., 0]))
+    else:
+        print("Using Depthmap and Camera Branch")
+        pred_world_points = predictions["world_points_from_depth"]
+        pred_world_points_conf = predictions.get("depth_conf", np.ones_like(pred_world_points[..., 0]))
+
+    colors_rgb = predictions["images"] 
+    
+    S, H, W = pred_world_points.shape[:3]
+    if colors_rgb.shape[:3] != (S, H, W):
+        print(f"Reshaping colors_rgb from {colors_rgb.shape} to match {(S, H, W, 3)}")
+        reshaped_colors = np.zeros((S, H, W, 3), dtype=np.float32)
+        for i in range(S):
+            if i < len(colors_rgb):
+                reshaped_colors[i] = cv2.resize(colors_rgb[i], (W, H))
+        colors_rgb = reshaped_colors
+    
+    colors_rgb = (colors_rgb * 255).astype(np.uint8)
+    
+    if mask_sky:
+        print("Applying sky segmentation mask")
+        try:
+            import onnxruntime
+         
+            with tempfile.TemporaryDirectory() as temp_dir:
+                print(f"Created temporary directory for sky segmentation: {temp_dir}")
+                temp_images_dir = os.path.join(temp_dir, "images")
+                sky_masks_dir = os.path.join(temp_dir, "sky_masks")
+                os.makedirs(temp_images_dir, exist_ok=True)
+                os.makedirs(sky_masks_dir, exist_ok=True)
+                
+                image_list = []
+                for i, img in enumerate(colors_rgb):
+                    img_path = os.path.join(temp_images_dir, f"image_{i:04d}.png")
+                    image_list.append(img_path)
+                    cv2.imwrite(img_path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
+                
+           
+                skyseg_path = os.path.join(temp_dir, "skyseg.onnx")
+                if not os.path.exists("skyseg.onnx"): 
+                    print("Downloading skyseg.onnx...")
+                    download_success = download_file_from_url(
+                        "https://huggingface.co/JianyuanWang/skyseg/resolve/main/skyseg.onnx", 
+                        skyseg_path
+                    )
+                    if not download_success:
+                        print("Failed to download skyseg model, skipping sky filtering")
+                        mask_sky = False
+                else:
+            
+                    import shutil
+                    shutil.copy("skyseg.onnx", skyseg_path)
+                
+                if mask_sky:  
+                    skyseg_session = onnxruntime.InferenceSession(skyseg_path)
+                    sky_mask_list = []
+                    
+                    for img_path in image_list:
+                        mask_path = os.path.join(sky_masks_dir, os.path.basename(img_path))
+                        sky_mask = segment_sky(img_path, skyseg_session, mask_path)
+           
+                        if sky_mask.shape[0] != H or sky_mask.shape[1] != W:
+                            sky_mask = cv2.resize(sky_mask, (W, H))
+                        
+                        sky_mask_list.append(sky_mask)
+                    
+                    sky_mask_array = np.array(sky_mask_list)
+                    
+                    sky_mask_binary = (sky_mask_array > 0.1).astype(np.float32)
+                    pred_world_points_conf = pred_world_points_conf * sky_mask_binary
+                    print(f"Applied sky mask, shape: {sky_mask_binary.shape}")
+                
+        except (ImportError, Exception) as e:
+            print(f"Error in sky segmentation: {e}")
+            mask_sky = False
+    
+    vertices_3d = pred_world_points.reshape(-1, 3)
+    conf = pred_world_points_conf.reshape(-1)
+    colors_rgb_flat = colors_rgb.reshape(-1, 3)
+
+    
+
+    if len(conf) != len(colors_rgb_flat):
+        print(f"WARNING: Shape mismatch between confidence ({len(conf)}) and colors ({len(colors_rgb_flat)})")
+        min_size = min(len(conf), len(colors_rgb_flat))
+        conf = conf[:min_size]
+        vertices_3d = vertices_3d[:min_size]
+        colors_rgb_flat = colors_rgb_flat[:min_size]
+    
+    if conf_threshold == 0.0:
+        conf_thres_value = 0.0
+    else:
+        conf_thres_value = np.percentile(conf, conf_threshold)
+    
+    print(f"Using confidence threshold: {conf_threshold}% (value: {conf_thres_value:.4f})")
+    conf_mask = (conf >= conf_thres_value) & (conf > 1e-5)
+    
+    if mask_black_bg:
+        print("Filtering black background")
+        black_bg_mask = colors_rgb_flat.sum(axis=1) >= 16
+        conf_mask = conf_mask & black_bg_mask
+    
+    if mask_white_bg:
+        print("Filtering white background")
+        white_bg_mask = ~((colors_rgb_flat[:, 0] > 240) & (colors_rgb_flat[:, 1] > 240) & (colors_rgb_flat[:, 2] > 240))
+        conf_mask = conf_mask & white_bg_mask
+    
+    filtered_vertices = vertices_3d[conf_mask]
+    filtered_colors = colors_rgb_flat[conf_mask]
+    
+    if len(filtered_vertices) == 0:
+        print("Warning: No points remaining after filtering. Using default point.")
+        filtered_vertices = np.array([[0, 0, 0]])
+        filtered_colors = np.array([[200, 200, 200]])
+    
+    print(f"Filtered to {len(filtered_vertices)} points")
+    
+    points3D = []
+    point_indices = {}
+    image_points2D = [[] for _ in range(len(pred_world_points))]
+    
+    print(f"Preparing points for COLMAP format with stride {stride}...")
+    
+    total_points = 0
+    for img_idx in range(S):
+        for y in range(0, H, stride):
+            for x in range(0, W, stride):
+                flat_idx = img_idx * H * W + y * W + x
+                
+                if flat_idx >= len(conf):
+                    continue
+                
+                if conf[flat_idx] < conf_thres_value or conf[flat_idx] <= 1e-5:
+                    continue
+                
+                if mask_black_bg and colors_rgb_flat[flat_idx].sum() < 16:
+                    continue
+                
+                if mask_white_bg and all(colors_rgb_flat[flat_idx] > 240):
+                    continue
+                
+                point3D = vertices_3d[flat_idx]
+                rgb = colors_rgb_flat[flat_idx]
+                
+                if not np.all(np.isfinite(point3D)):
+                    continue
+                
+                point_hash = hash_point(point3D, scale=100)
+                
+                if point_hash not in point_indices:
+                    point_idx = len(points3D)
+                    point_indices[point_hash] = point_idx
+                    
+                    point_entry = {
+                        "id": point_idx,
+                        "xyz": point3D,
+                        "rgb": rgb,
+                        "error": 1.0,
+                        "track": [(img_idx, len(image_points2D[img_idx]))]
+                    }
+                    points3D.append(point_entry)
+                    total_points += 1
+                else:
+                    point_idx = point_indices[point_hash]
+                    points3D[point_idx]["track"].append((img_idx, len(image_points2D[img_idx])))
+                
+                image_points2D[img_idx].append((x, y, point_indices[point_hash]))
+    
+    print(f"Prepared {len(points3D)} 3D points with {sum(len(pts) for pts in image_points2D)} observations for COLMAP")
+    return points3D, image_points2D
+
+def hash_point(point, scale=100):
+    """Create a hash for a 3D point by quantizing coordinates."""
+    quantized = tuple(np.round(point * scale).astype(int))
+    return hash(quantized)
+
+def write_colmap_cameras_txt(file_path, intrinsics, image_width, image_height):
+    """Write camera intrinsics to COLMAP cameras.txt format."""
+    with open(file_path, 'w') as f:
+        f.write("# Camera list with one line of data per camera:\n")
+        f.write("#   CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n")
+        f.write(f"# Number of cameras: {len(intrinsics)}\n")
+        
+        for i, intrinsic in enumerate(intrinsics):
+            camera_id = i + 1  # COLMAP uses 1-indexed camera IDs
+            model = "PINHOLE" 
+            
+            fx = intrinsic[0, 0]
+            fy = intrinsic[1, 1]
+            cx = intrinsic[0, 2]
+            cy = intrinsic[1, 2]
+            
+            f.write(f"{camera_id} {model} {image_width} {image_height} {fx} {fy} {cx} {cy}\n")
+
+def write_colmap_images_txt(file_path, quaternions, translations, image_points2D, image_names):
+    """Write camera poses and keypoints to COLMAP images.txt format."""
+    with open(file_path, 'w') as f:
+        f.write("# Image list with two lines of data per image:\n")
+        f.write("#   IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n")
+        f.write("#   POINTS2D[] as (X, Y, POINT3D_ID)\n")
+        
+        num_points = sum(len(points) for points in image_points2D)
+        avg_points = num_points / len(image_points2D) if image_points2D else 0
+        f.write(f"# Number of images: {len(quaternions)}, mean observations per image: {avg_points:.1f}\n")
+        
+        for i in range(len(quaternions)):
+            image_id = i + 1 
+            camera_id = i + 1  
+          
+            qw, qx, qy, qz = quaternions[i]
+            tx, ty, tz = translations[i]
+            
+            f.write(f"{image_id} {qw} {qx} {qy} {qz} {tx} {ty} {tz} {camera_id} {os.path.basename(image_names[i])}\n")
+            
+            points_line = " ".join([f"{x} {y} {point3d_id+1}" for x, y, point3d_id in image_points2D[i]])
+            f.write(f"{points_line}\n")
+
+def write_colmap_points3D_txt(file_path, points3D):
+    """Write 3D points and tracks to COLMAP points3D.txt format."""
+    with open(file_path, 'w') as f:
+        f.write("# 3D point list with one line of data per point:\n")
+        f.write("#   POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n")
+        
+        avg_track_length = sum(len(point["track"]) for point in points3D) / len(points3D) if points3D else 0
+        f.write(f"# Number of points: {len(points3D)}, mean track length: {avg_track_length:.4f}\n")
+        
+        for point in points3D:
+            point_id = point["id"] + 1  
+            x, y, z = point["xyz"]
+            r, g, b = point["rgb"]
+            error = point["error"]
+            
+            track = " ".join([f"{img_id+1} {point2d_idx}" for img_id, point2d_idx in point["track"]])
+            
+            f.write(f"{point_id} {x} {y} {z} {int(r)} {int(g)} {int(b)} {error} {track}\n")
+
+def write_colmap_cameras_bin(file_path, intrinsics, image_width, image_height):
+    """Write camera intrinsics to COLMAP cameras.bin format."""
+    with open(file_path, 'wb') as fid:
+        # Write number of cameras (uint64)
+        fid.write(struct.pack('<Q', len(intrinsics)))
+        
+        for i, intrinsic in enumerate(intrinsics):
+            camera_id = i + 1
+            model_id = 1 
+            
+            fx = float(intrinsic[0, 0])
+            fy = float(intrinsic[1, 1])
+            cx = float(intrinsic[0, 2])
+            cy = float(intrinsic[1, 2])
+            
+            # Camera ID (uint32)
+            fid.write(struct.pack('<I', camera_id))
+            # Model ID (uint32)
+            fid.write(struct.pack('<I', model_id))
+            # Width (uint64)
+            fid.write(struct.pack('<Q', image_width))
+            # Height (uint64)
+            fid.write(struct.pack('<Q', image_height))
+            
+            # Parameters (double)
+            fid.write(struct.pack('<dddd', fx, fy, cx, cy))
+
+def write_colmap_images_bin(file_path, quaternions, translations, image_points2D, image_names):
+    """Write camera poses and keypoints to COLMAP images.bin format."""
+    with open(file_path, 'wb') as fid:
+        # Write number of images (uint64)
+        fid.write(struct.pack('<Q', len(quaternions)))
+        
+        for i in range(len(quaternions)):
+            image_id = i + 1
+            camera_id = i + 1
+            
+            qw, qx, qy, qz = quaternions[i].astype(float)
+            tx, ty, tz = translations[i].astype(float)
+            
+            image_name = os.path.basename(image_names[i]).encode()
+            points = image_points2D[i]
+            
+            # Image ID (uint32)
+            fid.write(struct.pack('<I', image_id))
+            # Quaternion (double): qw, qx, qy, qz
+            fid.write(struct.pack('<dddd', qw, qx, qy, qz))
+            # Translation (double): tx, ty, tz
+            fid.write(struct.pack('<ddd', tx, ty, tz))
+            # Camera ID (uint32)
+            fid.write(struct.pack('<I', camera_id))
+            # Image name
+            fid.write(struct.pack('<I', len(image_name)))
+            fid.write(image_name)
+            
+            # Write number of 2D points (uint64)
+            fid.write(struct.pack('<Q', len(points)))
+            
+            # Write 2D points: x, y, point3D_id
+            for x, y, point3d_id in points:
+                fid.write(struct.pack('<dd', float(x), float(y)))
+                fid.write(struct.pack('<Q', point3d_id + 1))
+
+def write_colmap_points3D_bin(file_path, points3D):
+    """Write 3D points and tracks to COLMAP points3D.bin format."""
+    with open(file_path, 'wb') as fid:
+        # Write number of points (uint64)
+        fid.write(struct.pack('<Q', len(points3D)))
+        
+        for point in points3D:
+            point_id = point["id"] + 1
+            x, y, z = point["xyz"].astype(float)
+            r, g, b = point["rgb"].astype(np.uint8)
+            error = float(point["error"])
+            track = point["track"]
+            
+            # Point ID (uint64)
+            fid.write(struct.pack('<Q', point_id))
+            # Position (double): x, y, z
+            fid.write(struct.pack('<ddd', x, y, z))
+            # Color (uint8): r, g, b
+            fid.write(struct.pack('<BBB', int(r), int(g), int(b)))
+            # Error (double)
+            fid.write(struct.pack('<d', error))
+            
+            # Track: list of (image_id, point2D_idx)
+            fid.write(struct.pack('<Q', len(track)))
+            for img_id, point2d_idx in track:
+                fid.write(struct.pack('<II', img_id + 1, point2d_idx))
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert images to COLMAP format using VGGT")
+    parser.add_argument("--image_dir", type=str, required=True, 
+                        help="Directory containing input images")
+    parser.add_argument("--output_dir", type=str, default="colmap_output", 
+                        help="Directory to save COLMAP files")
+    parser.add_argument("--conf_threshold", type=float, default=50.0, 
+                        help="Confidence threshold (0-100%) for including points")
+    parser.add_argument("--mask_sky", action="store_true",
+                        help="Filter out points likely to be sky")
+    parser.add_argument("--mask_black_bg", action="store_true",
+                        help="Filter out points with very dark/black color")
+    parser.add_argument("--mask_white_bg", action="store_true",
+                        help="Filter out points with very bright/white color")
+    parser.add_argument("--binary", action="store_true", 
+                        help="Output binary COLMAP files instead of text")
+    parser.add_argument("--stride", type=int, default=1, 
+                        help="Stride for point sampling (higher = fewer points)")
+    parser.add_argument("--prediction_mode", type=str, default="Depthmap and Camera Branch",
+                        choices=["Depthmap and Camera Branch", "Pointmap Branch"],
+                        help="Which prediction branch to use")
+    
+    args = parser.parse_args()
+    
+    os.makedirs(args.output_dir, exist_ok=True)
+    
+    model, device = load_model()
+    
+    predictions, image_names = process_images(args.image_dir, model, device)
+    
+    print("Converting camera parameters to COLMAP format...")
+    quaternions, translations = extrinsic_to_colmap_format(predictions["extrinsic"])
+    
+    print(f"Filtering points with confidence threshold {args.conf_threshold}% and stride {args.stride}...")
+    points3D, image_points2D = filter_and_prepare_points(
+        predictions, 
+        args.conf_threshold, 
+        mask_sky=args.mask_sky, 
+        mask_black_bg=args.mask_black_bg,
+        mask_white_bg=args.mask_white_bg,
+        stride=args.stride,
+        prediction_mode=args.prediction_mode
+    )
+    
+    height, width = predictions["depth"].shape[1:3]
+
+    print(f"Writing {'binary' if args.binary else 'text'} COLMAP files to {args.output_dir}...")
+    if args.binary:
+        write_colmap_cameras_bin(
+            os.path.join(args.output_dir, "cameras.bin"), 
+            predictions["intrinsic"], width, height)
+        write_colmap_images_bin(
+            os.path.join(args.output_dir, "images.bin"), 
+            quaternions, translations, image_points2D, image_names)
+        write_colmap_points3D_bin(
+            os.path.join(args.output_dir, "points3D.bin"), 
+            points3D)
+    else:
+        write_colmap_cameras_txt(
+            os.path.join(args.output_dir, "cameras.txt"), 
+            predictions["intrinsic"], width, height)
+        write_colmap_images_txt(
+            os.path.join(args.output_dir, "images.txt"), 
+            quaternions, translations, image_points2D, image_names)
+        write_colmap_points3D_txt(
+            os.path.join(args.output_dir, "points3D.txt"), 
+            points3D)
+    
+    print(f"COLMAP files successfully written to {args.output_dir}")
+
+if __name__ == "__main__":
+    main()

source/visualization.py

@@ -0,0 +1,1072 @@
+from matplotlib import pyplot as plt
+import numpy as np
+import torch
+
+import numpy as np
+from typing import List
+import sys
+sys.path.append('./submodules/gaussian-splatting/')
+from scene.cameras import Camera
+from PIL import Image
+import imageio
+from scipy.interpolate import splprep, splev
+
+import cv2
+import numpy as np
+import plotly.graph_objects as go
+import numpy as np
+from scipy.spatial.transform import Rotation as R, Slerp
+from scipy.spatial import distance_matrix
+from sklearn.decomposition import PCA
+from scipy.interpolate import splprep, splev
+from typing import List
+from sklearn.mixture import GaussianMixture
+
+def render_gaussians_rgb(generator3DGS, viewpoint_cam, visualize=False):
+    """
+    Simply render gaussians from the generator3DGS from the viewpoint_cam.
+    Args:
+        generator3DGS : instance of the Generator3DGS class from the networks.py file
+        viewpoint_cam : camera instance
+        visualize : boolean flag. If True, will call pyplot function and render image inplace
+    Returns:
+        uint8 numpy array with shape (H, W, 3) representing the image
+    """
+    with torch.no_grad():
+        render_pkg = generator3DGS(viewpoint_cam)
+        image = render_pkg["render"]
+        image_np = image.clone().detach().cpu().numpy().transpose(1, 2, 0)
+
+        # Clip values to be in the range [0, 1]
+        image_np = np.clip(image_np * 255, 0, 255).astype(np.uint8)
+        if visualize:
+            plt.figure(figsize=(12, 8))
+            plt.imshow(image_np)
+            plt.show()
+
+        return image_np
+
+def render_gaussians_D_scores(generator3DGS, viewpoint_cam, mask=None, mask_channel=0, visualize=False):
+    """
+        Simply render D_scores of gaussians from the generator3DGS from the viewpoint_cam.
+        Args:
+            generator3DGS : instance of the Generator3DGS class from the networks.py file
+            viewpoint_cam : camera instance
+            visualize : boolean flag. If True, will call pyplot function and render image inplace
+            mask : optional mask to highlight specific gaussians. Must be of shape (N) where N is the numnber
+                of gaussians in generator3DGS.gaussians. Must be a torch tensor of floats, please scale according
+                to how much color you want to have. Recommended mask value is 10.
+            mask_channel: to which color channel should we add mask
+        Returns:
+            uint8 numpy array with shape (H, W, 3) representing the generator3DGS.gaussians.D_scores rendered as colors
+        """
+    with torch.no_grad():
+        # Visualize D_scores
+        generator3DGS.gaussians._features_dc = generator3DGS.gaussians._features_dc * 1e-4 + \
+                                               torch.stack([generator3DGS.gaussians.D_scores] * 3, axis=-1)
+        generator3DGS.gaussians._features_rest = generator3DGS.gaussians._features_rest * 1e-4
+        if mask is not None:
+            generator3DGS.gaussians._features_dc[..., mask_channel] += mask.unsqueeze(-1)
+        render_pkg = generator3DGS(viewpoint_cam)
+        image = render_pkg["render"]
+        image_np = image.clone().detach().cpu().numpy().transpose(1, 2, 0)
+
+        # Clip values to be in the range [0, 1]
+        image_np = np.clip(image_np * 255, 0, 255).astype(np.uint8)
+        if visualize:
+            plt.figure(figsize=(12, 8))
+            plt.imshow(image_np)
+            plt.show()
+
+        if mask is not None:
+            generator3DGS.gaussians._features_dc[..., mask_channel] -= mask.unsqueeze(-1)
+
+        generator3DGS.gaussians._features_dc = (generator3DGS.gaussians._features_dc - \
+                                                     torch.stack([generator3DGS.gaussians.D_scores] * 3, axis=-1)) * 1e4
+        generator3DGS.gaussians._features_rest = generator3DGS.gaussians._features_rest * 1e4
+
+        return image_np
+    
+
+
+def normalize(v):
+    """
+    Normalize a vector to unit length.
+
+    Parameters:
+        v (np.ndarray): Input vector.
+
+    Returns:
+        np.ndarray: Unit vector in the same direction as `v`.
+    """
+    return v / np.linalg.norm(v)
+
+def look_at_rotation(camera_position: np.ndarray, target: np.ndarray, world_up=np.array([0, 1, 0])):
+    """
+    Compute a rotation matrix for a camera looking at a target point.
+
+    Parameters:
+        camera_position (np.ndarray): The 3D position of the camera.
+        target (np.ndarray): The point the camera should look at.
+        world_up (np.ndarray): A vector that defines the global 'up' direction.
+
+    Returns:
+        np.ndarray: A 3x3 rotation matrix (camera-to-world) with columns [right, up, forward].
+    """
+    z_axis = normalize(target - camera_position)         # Forward direction
+    x_axis = normalize(np.cross(world_up, z_axis))       # Right direction
+    y_axis = np.cross(z_axis, x_axis)                    # Recomputed up
+    return np.stack([x_axis, y_axis, z_axis], axis=1)
+
+    
+def generate_circular_camera_path(existing_cameras: List[Camera], N: int = 12, radius_scale: float = 1.0, d: float = 2.0) -> List[Camera]:
+    """
+    Generate a circular path of cameras around an existing camera group, 
+    with each new camera oriented to look at the average viewing direction.
+
+    Parameters:
+        existing_cameras (List[Camera]): List of existing camera objects to estimate average orientation and layout.
+        N (int): Number of new cameras to generate along the circular path.
+        radius_scale (float): Scale factor to adjust the radius of the circle.
+        d (float): Distance ahead of each camera used to estimate its look-at point.
+
+    Returns:
+        List[Camera]: A list of newly generated Camera objects forming a circular path and oriented toward a shared view center.
+    """
+    # Step 1: Compute average camera position
+    center = np.mean([cam.T for cam in existing_cameras], axis=0)
+
+    # Estimate where each camera is looking
+    # d denotes how far ahead each camera sees — you can scale this
+    look_targets = [cam.T + cam.R[:, 2] * d for cam in existing_cameras]
+    center_of_view = np.mean(look_targets, axis=0)
+
+    # Step 2: Define circular plane basis using fixed up vector
+    avg_forward = normalize(np.mean([cam.R[:, 2] for cam in existing_cameras], axis=0))
+    up_guess = np.array([0, 1, 0])
+    right = normalize(np.cross(avg_forward, up_guess))
+    up = normalize(np.cross(right, avg_forward))
+
+    # Step 3: Estimate radius
+    avg_radius = np.mean([np.linalg.norm(cam.T - center) for cam in existing_cameras]) * radius_scale
+
+    # Step 4: Create cameras on a circular path
+    angles = np.linspace(0, 2 * np.pi, N, endpoint=False)
+    reference_cam = existing_cameras[0]
+    new_cameras = []
+
+    
+    for i, a in enumerate(angles):
+        position = center + avg_radius * (np.cos(a) * right + np.sin(a) * up)
+
+        if d < 1e-5 or radius_scale < 1e-5:
+            # Use same orientation as the first camera
+            R = reference_cam.R.copy()
+        else:
+            # Change orientation
+            R = look_at_rotation(position, center_of_view)
+        new_cameras.append(Camera(
+            R=R, 
+            T=position,                                   # New position
+            FoVx=reference_cam.FoVx,
+            FoVy=reference_cam.FoVy,
+            resolution=(reference_cam.image_width, reference_cam.image_height),
+            colmap_id=-1,
+            depth_params=None,
+            image=Image.fromarray(np.zeros((reference_cam.image_height, reference_cam.image_width, 3), dtype=np.uint8)),
+            invdepthmap=None,
+            image_name=f"circular_a={a:.3f}",
+            uid=i
+        ))
+
+    return new_cameras
+
+
+def save_numpy_frames_as_gif(frames, output_path="animation.gif", duration=100):
+    """
+    Save a list of RGB NumPy frames as a looping GIF animation.
+
+    Parameters:
+        frames (List[np.ndarray]): List of RGB images as uint8 NumPy arrays (shape HxWx3).
+        output_path (str): Path to save the output GIF.
+        duration (int): Duration per frame in milliseconds.
+
+    Returns:
+        None
+    """
+    pil_frames = [Image.fromarray(f) for f in frames]
+    pil_frames[0].save(
+        output_path,
+        save_all=True,
+        append_images=pil_frames[1:],
+        duration=duration,  # duration per frame in ms
+        loop=0
+    )
+    print(f"GIF saved to: {output_path}")
+
+def center_crop_frame(frame: np.ndarray, crop_fraction: float) -> np.ndarray:
+    """
+    Crop the central region of the frame by the given fraction.
+
+    Parameters:
+        frame (np.ndarray): Input RGB image (H, W, 3).
+        crop_fraction (float): Fraction of the original size to retain (e.g., 0.8 keeps 80%).
+
+    Returns:
+        np.ndarray: Cropped RGB image.
+    """
+    if crop_fraction >= 1.0:
+        return frame
+
+    h, w, _ = frame.shape
+    new_h, new_w = int(h * crop_fraction), int(w * crop_fraction)
+    start_y = (h - new_h) // 2
+    start_x = (w - new_w) // 2
+    return frame[start_y:start_y + new_h, start_x:start_x + new_w, :]
+
+
+
+def generate_smooth_closed_camera_path(existing_cameras: List[Camera], N: int = 120, d: float = 2.0, s=.25) -> List[Camera]:
+    """
+    Generate a smooth, closed path interpolating the positions of existing cameras.
+
+    Parameters:
+        existing_cameras (List[Camera]): List of existing cameras.
+        N (int): Number of points (cameras) to sample along the smooth path.
+        d (float): Distance ahead for estimating the center of view.
+
+    Returns:
+        List[Camera]: A list of smoothly moving Camera objects along a closed loop.
+    """
+    # Step 1: Extract camera positions
+    positions = np.array([cam.T for cam in existing_cameras])
+    
+    # Step 2: Estimate center of view
+    look_targets = [cam.T + cam.R[:, 2] * d for cam in existing_cameras]
+    center_of_view = np.mean(look_targets, axis=0)
+
+    # Step 3: Fit a smooth closed spline through the positions
+    positions = np.vstack([positions, positions[0]])  # close the loop
+    tck, u = splprep(positions.T, s=s, per=True)  # periodic=True for closed loop
+
+    # Step 4: Sample points along the spline
+    u_fine = np.linspace(0, 1, N)
+    smooth_path = np.stack(splev(u_fine, tck), axis=-1)
+
+    # Step 5: Generate cameras along the smooth path
+    reference_cam = existing_cameras[0]
+    new_cameras = []
+
+    for i, pos in enumerate(smooth_path):
+        R = look_at_rotation(pos, center_of_view)
+        new_cameras.append(Camera(
+            R=R,
+            T=pos,
+            FoVx=reference_cam.FoVx,
+            FoVy=reference_cam.FoVy,
+            resolution=(reference_cam.image_width, reference_cam.image_height),
+            colmap_id=-1,
+            depth_params=None,
+            image=Image.fromarray(np.zeros((reference_cam.image_height, reference_cam.image_width, 3), dtype=np.uint8)),
+            invdepthmap=None,
+            image_name=f"smooth_path_i={i}",
+            uid=i
+        ))
+
+    return new_cameras
+
+
+def save_numpy_frames_as_mp4(frames, output_path="animation.mp4", fps=10, center_crop: float = 1.0):
+    """
+    Save a list of RGB NumPy frames as an MP4 video with optional center cropping.
+
+    Parameters:
+        frames (List[np.ndarray]): List of RGB images as uint8 NumPy arrays (shape HxWx3).
+        output_path (str): Path to save the output MP4.
+        fps (int): Frames per second for playback speed.
+        center_crop (float): Fraction (0 < center_crop <= 1.0) of central region to retain. 
+                             Use 1.0 for no cropping; 0.8 to crop to 80% center region.
+
+    Returns:
+        None
+    """
+    with imageio.get_writer(output_path, fps=fps, codec='libx264', quality=8) as writer:
+        for frame in frames:
+            cropped = center_crop_frame(frame, center_crop)
+            writer.append_data(cropped)
+    print(f"MP4 saved to: {output_path}")
+
+
+    
+def put_text_on_image(img: np.ndarray, text: str) -> np.ndarray:
+    """
+    Draws multiline white text on a copy of the input image, positioned near the bottom
+    and around 80% of the image width. Handles '\n' characters to split text into multiple lines.
+
+    Args:
+        img (np.ndarray): Input image as a (H, W, 3) uint8 numpy array.
+        text (str): Text string to draw on the image. Newlines '\n' are treated as line breaks.
+
+    Returns:
+        np.ndarray: The output image with the text drawn on it.
+    
+    Notes:
+        - The function automatically adjusts line spacing and prevents text from going outside the image.
+        - Text is drawn in white with small font size (0.5) for minimal visual impact.
+    """
+    img = img.copy()
+    height, width, _ = img.shape
+    
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    font_scale = 1.
+    color = (255, 255, 255)
+    thickness = 2
+    line_spacing = 5  # extra pixels between lines
+    
+    lines = text.split('\n')
+    
+    # Precompute the maximum text width to adjust starting x
+    max_text_width = max(cv2.getTextSize(line, font, font_scale, thickness)[0][0] for line in lines)
+    
+    x = int(0.8 * width)
+    x = min(x, width - max_text_width - 30)  # margin on right
+    #x = int(0.03 * width)
+    
+    # Start near the bottom, but move up depending on number of lines
+    total_text_height = len(lines) * (cv2.getTextSize('A', font, font_scale, thickness)[0][1] + line_spacing)
+    y_start = int(height*0.9) - total_text_height  # 30 pixels from bottom
+
+    for i, line in enumerate(lines):
+        y = y_start + i * (cv2.getTextSize(line, font, font_scale, thickness)[0][1] + line_spacing)
+        cv2.putText(img, line, (x, y), font, font_scale, color, thickness, cv2.LINE_AA)
+    
+    return img
+
+
+
+
+def catmull_rom_spline(P0, P1, P2, P3, n_points=20):
+    """
+    Compute Catmull-Rom spline segment between P1 and P2.
+    """
+    t = np.linspace(0, 1, n_points)[:, None]
+
+    M = 0.5 * np.array([
+        [-1,  3, -3, 1],
+        [ 2, -5,  4, -1],
+        [-1,  0,  1, 0],
+        [ 0,  2,  0, 0]
+    ])
+
+    G = np.stack([P0, P1, P2, P3], axis=0)
+    T = np.concatenate([t**3, t**2, t, np.ones_like(t)], axis=1)
+
+    return T @ M @ G
+
+def sort_cameras_pca(existing_cameras: List[Camera]):
+    """
+    Sort cameras along the main PCA axis.
+    """
+    positions = np.array([cam.T for cam in existing_cameras])
+    pca = PCA(n_components=1)
+    scores = pca.fit_transform(positions)
+    sorted_indices = np.argsort(scores[:, 0])
+    return sorted_indices
+
+def generate_fully_smooth_cameras(existing_cameras: List[Camera], 
+                                  n_selected: int = 30, 
+                                  n_points_per_segment: int = 20, 
+                                  d: float = 2.0,
+                                  closed: bool = False) -> List[Camera]:
+    """
+    Generate a fully smooth camera path using PCA ordering, global Catmull-Rom spline for positions, and global SLERP for orientations.
+
+    Args:
+        existing_cameras (List[Camera]): List of input cameras.
+        n_selected (int): Number of cameras to select after sorting.
+        n_points_per_segment (int): Number of interpolated points per spline segment.
+        d (float): Distance ahead for estimating center of view.
+        closed (bool): Whether to close the path.
+
+    Returns:
+        List[Camera]: List of smoothly moving Camera objects.
+    """
+    # 1. Sort cameras along PCA axis
+    sorted_indices = sort_cameras_pca(existing_cameras)
+    sorted_cameras = [existing_cameras[i] for i in sorted_indices]
+    positions = np.array([cam.T for cam in sorted_cameras])
+
+    # 2. Subsample uniformly
+    idx = np.linspace(0, len(positions) - 1, n_selected).astype(int)
+    sampled_positions = positions[idx]
+    sampled_cameras = [sorted_cameras[i] for i in idx]
+
+    # 3. Prepare for Catmull-Rom
+    if closed:
+        sampled_positions = np.vstack([sampled_positions[-1], sampled_positions, sampled_positions[0], sampled_positions[1]])
+    else:
+        sampled_positions = np.vstack([sampled_positions[0], sampled_positions, sampled_positions[-1], sampled_positions[-1]])
+
+    # 4. Generate smooth path positions
+    path_positions = []
+    for i in range(1, len(sampled_positions) - 2):
+        segment = catmull_rom_spline(sampled_positions[i-1], sampled_positions[i], sampled_positions[i+1], sampled_positions[i+2], n_points_per_segment)
+        path_positions.append(segment)
+    path_positions = np.concatenate(path_positions, axis=0)
+
+    # 5. Global SLERP for rotations
+    rotations = R.from_matrix([cam.R for cam in sampled_cameras])
+    key_times = np.linspace(0, 1, len(rotations))
+    slerp = Slerp(key_times, rotations)
+
+    query_times = np.linspace(0, 1, len(path_positions))
+    interpolated_rotations = slerp(query_times)
+
+    # 6. Generate Camera objects
+    reference_cam = existing_cameras[0]
+    smooth_cameras = []
+
+    for i, pos in enumerate(path_positions):
+        R_interp = interpolated_rotations[i].as_matrix()
+
+        smooth_cameras.append(Camera(
+            R=R_interp,
+            T=pos,
+            FoVx=reference_cam.FoVx,
+            FoVy=reference_cam.FoVy,
+            resolution=(reference_cam.image_width, reference_cam.image_height),
+            colmap_id=-1,
+            depth_params=None,
+            image=Image.fromarray(np.zeros((reference_cam.image_height, reference_cam.image_width, 3), dtype=np.uint8)),
+            invdepthmap=None,
+            image_name=f"fully_smooth_path_i={i}",
+            uid=i
+        ))
+
+    return smooth_cameras
+
+
+def plot_cameras_and_smooth_path_with_orientation(existing_cameras: List[Camera], smooth_cameras: List[Camera], scale: float = 0.1):
+    """
+    Plot input cameras and smooth path cameras with their orientations in 3D.
+
+    Args:
+        existing_cameras (List[Camera]): List of original input cameras.
+        smooth_cameras (List[Camera]): List of smooth path cameras.
+        scale (float): Length of orientation arrows.
+
+    Returns:
+        None
+    """
+    # Input cameras
+    input_positions = np.array([cam.T for cam in existing_cameras])
+
+    # Smooth cameras
+    smooth_positions = np.array([cam.T for cam in smooth_cameras])
+
+    fig = go.Figure()
+
+    # Plot input camera positions
+    fig.add_trace(go.Scatter3d(
+        x=input_positions[:, 0], y=input_positions[:, 1], z=input_positions[:, 2],
+        mode='markers',
+        marker=dict(size=4, color='blue'),
+        name='Input Cameras'
+    ))
+
+    # Plot smooth path positions
+    fig.add_trace(go.Scatter3d(
+        x=smooth_positions[:, 0], y=smooth_positions[:, 1], z=smooth_positions[:, 2],
+        mode='lines+markers',
+        line=dict(color='red', width=3),
+        marker=dict(size=2, color='red'),
+        name='Smooth Path Cameras'
+    ))
+
+    # Plot input camera orientations
+    for cam in existing_cameras:
+        origin = cam.T
+        forward = cam.R[:, 2]  # Forward direction
+
+        fig.add_trace(go.Cone(
+            x=[origin[0]], y=[origin[1]], z=[origin[2]],
+            u=[forward[0]], v=[forward[1]], w=[forward[2]],
+            colorscale=[[0, 'blue'], [1, 'blue']],
+            sizemode="absolute",
+            sizeref=scale,
+            anchor="tail",
+            showscale=False,
+            name='Input Camera Direction'
+        ))
+
+    # Plot smooth camera orientations
+    for cam in smooth_cameras:
+        origin = cam.T
+        forward = cam.R[:, 2]  # Forward direction
+
+        fig.add_trace(go.Cone(
+            x=[origin[0]], y=[origin[1]], z=[origin[2]],
+            u=[forward[0]], v=[forward[1]], w=[forward[2]],
+            colorscale=[[0, 'red'], [1, 'red']],
+            sizemode="absolute",
+            sizeref=scale,
+            anchor="tail",
+            showscale=False,
+            name='Smooth Camera Direction'
+        ))
+
+    fig.update_layout(
+        scene=dict(
+            xaxis_title='X',
+            yaxis_title='Y',
+            zaxis_title='Z',
+            aspectmode='data'
+        ),
+        title="Input Cameras and Smooth Path with Orientations",
+        margin=dict(l=0, r=0, b=0, t=30)
+    )
+
+    fig.show()
+
+
+def solve_tsp_nearest_neighbor(points: np.ndarray):
+    """
+    Solve TSP approximately using nearest neighbor heuristic.
+
+    Args:
+        points (np.ndarray): (N, 3) array of points.
+
+    Returns:
+        List[int]: Optimal visiting order of points.
+    """
+    N = points.shape[0]
+    dist = distance_matrix(points, points)
+    visited = [0]
+    unvisited = set(range(1, N))
+
+    while unvisited:
+        last = visited[-1]
+        next_city = min(unvisited, key=lambda city: dist[last, city])
+        visited.append(next_city)
+        unvisited.remove(next_city)
+
+    return visited
+
+def solve_tsp_2opt(points: np.ndarray, n_iter: int = 1000) -> np.ndarray:
+    """
+    Solve TSP approximately using Nearest Neighbor + 2-Opt.
+
+    Args:
+        points (np.ndarray): Array of shape (N, D) with points.
+        n_iter (int): Number of 2-opt iterations.
+
+    Returns:
+        np.ndarray: Ordered list of indices.
+    """
+    n_points = points.shape[0]
+
+    # === 1. Start with Nearest Neighbor
+    unvisited = list(range(n_points))
+    current = unvisited.pop(0)
+    path = [current]
+
+    while unvisited:
+        dists = np.linalg.norm(points[unvisited] - points[current], axis=1)
+        next_idx = unvisited[np.argmin(dists)]
+        unvisited.remove(next_idx)
+        path.append(next_idx)
+        current = next_idx
+
+    # === 2. Apply 2-Opt improvements
+    def path_length(path):
+        return np.sum(np.linalg.norm(points[path[i]] - points[path[i+1]], axis=0) for i in range(len(path)-1))
+
+    best_length = path_length(path)
+    improved = True
+
+    for _ in range(n_iter):
+        if not improved:
+            break
+        improved = False
+        for i in range(1, n_points - 2):
+            for j in range(i + 1, n_points):
+                if j - i == 1: continue
+                new_path = path[:i] + path[i:j][::-1] + path[j:]
+                new_length = path_length(new_path)
+                if new_length < best_length:
+                    path = new_path
+                    best_length = new_length
+                    improved = True
+                    break
+            if improved:
+                break
+
+    return np.array(path)
+
+def generate_fully_smooth_cameras_with_tsp(existing_cameras: List[Camera], 
+                                           n_selected: int = 30, 
+                                           n_points_per_segment: int = 20, 
+                                           d: float = 2.0,
+                                           closed: bool = False) -> List[Camera]:
+    """
+    Generate a fully smooth camera path using TSP ordering, global Catmull-Rom spline for positions, and global SLERP for orientations.
+
+    Args:
+        existing_cameras (List[Camera]): List of input cameras.
+        n_selected (int): Number of cameras to select after ordering.
+        n_points_per_segment (int): Number of interpolated points per spline segment.
+        d (float): Distance ahead for estimating center of view.
+        closed (bool): Whether to close the path.
+
+    Returns:
+        List[Camera]: List of smoothly moving Camera objects.
+    """
+    positions = np.array([cam.T for cam in existing_cameras])
+
+    # 1. Solve approximate TSP
+    order = solve_tsp_nearest_neighbor(positions)
+    ordered_cameras = [existing_cameras[i] for i in order]
+    ordered_positions = positions[order]
+
+    # 2. Subsample uniformly
+    idx = np.linspace(0, len(ordered_positions) - 1, n_selected).astype(int)
+    sampled_positions = ordered_positions[idx]
+    sampled_cameras = [ordered_cameras[i] for i in idx]
+
+    # 3. Prepare for Catmull-Rom
+    if closed:
+        sampled_positions = np.vstack([sampled_positions[-1], sampled_positions, sampled_positions[0], sampled_positions[1]])
+    else:
+        sampled_positions = np.vstack([sampled_positions[0], sampled_positions, sampled_positions[-1], sampled_positions[-1]])
+
+    # 4. Generate smooth path positions
+    path_positions = []
+    for i in range(1, len(sampled_positions) - 2):
+        segment = catmull_rom_spline(sampled_positions[i-1], sampled_positions[i], sampled_positions[i+1], sampled_positions[i+2], n_points_per_segment)
+        path_positions.append(segment)
+    path_positions = np.concatenate(path_positions, axis=0)
+
+    # 5. Global SLERP for rotations
+    rotations = R.from_matrix([cam.R for cam in sampled_cameras])
+    key_times = np.linspace(0, 1, len(rotations))
+    slerp = Slerp(key_times, rotations)
+
+    query_times = np.linspace(0, 1, len(path_positions))
+    interpolated_rotations = slerp(query_times)
+
+    # 6. Generate Camera objects
+    reference_cam = existing_cameras[0]
+    smooth_cameras = []
+
+    for i, pos in enumerate(path_positions):
+        R_interp = interpolated_rotations[i].as_matrix()
+
+        smooth_cameras.append(Camera(
+            R=R_interp,
+            T=pos,
+            FoVx=reference_cam.FoVx,
+            FoVy=reference_cam.FoVy,
+            resolution=(reference_cam.image_width, reference_cam.image_height),
+            colmap_id=-1,
+            depth_params=None,
+            image=Image.fromarray(np.zeros((reference_cam.image_height, reference_cam.image_width, 3), dtype=np.uint8)),
+            invdepthmap=None,
+            image_name=f"fully_smooth_path_i={i}",
+            uid=i
+        ))
+
+    return smooth_cameras
+
+from typing import List
+import numpy as np
+from sklearn.mixture import GaussianMixture
+from scipy.spatial.transform import Rotation as R, Slerp
+from PIL import Image
+
+def generate_clustered_smooth_cameras_with_tsp(existing_cameras: List[Camera], 
+                                                n_selected: int = 30, 
+                                                n_points_per_segment: int = 20, 
+                                                d: float = 2.0,
+                                                n_clusters: int = 5,
+                                                closed: bool = False) -> List[Camera]:
+    """
+    Generate a fully smooth camera path using clustering + TSP between nearest cluster centers + TSP inside clusters.
+    Positions are normalized before clustering and denormalized before generating final cameras.
+
+    Args:
+        existing_cameras (List[Camera]): List of input cameras.
+        n_selected (int): Number of cameras to select after ordering.
+        n_points_per_segment (int): Number of interpolated points per spline segment.
+        d (float): Distance ahead for estimating center of view.
+        n_clusters (int): Number of GMM clusters.
+        closed (bool): Whether to close the path.
+
+    Returns:
+        List[Camera]: Smooth path of Camera objects.
+    """
+    # Extract positions and rotations
+    positions = np.array([cam.T for cam in existing_cameras])
+    rotations = np.array([R.from_matrix(cam.R).as_quat() for cam in existing_cameras])
+
+    # === Normalize positions
+    mean_pos = np.mean(positions, axis=0)
+    scale_pos = np.std(positions, axis=0)
+    scale_pos[scale_pos == 0] = 1.0  # avoid division by zero
+
+    positions_normalized = (positions - mean_pos) / scale_pos
+
+    # === Features for clustering (only positions, not rotations)
+    features = positions_normalized
+
+    # === 1. GMM clustering
+    gmm = GaussianMixture(n_components=n_clusters, covariance_type='full', random_state=42)
+    cluster_labels = gmm.fit_predict(features)
+
+    clusters = {}
+    cluster_centers = []
+
+    for cluster_id in range(n_clusters):
+        cluster_indices = np.where(cluster_labels == cluster_id)[0]
+        if len(cluster_indices) == 0:
+            continue
+        clusters[cluster_id] = cluster_indices
+        cluster_center = np.mean(features[cluster_indices], axis=0)
+        cluster_centers.append(cluster_center)
+
+    cluster_centers = np.stack(cluster_centers)
+
+    # === 2. Remap cluster centers to nearest existing cameras
+    if False:
+        mapped_centers = []
+        for center in cluster_centers:
+            dists = np.linalg.norm(features - center, axis=1)
+            nearest_idx = np.argmin(dists)
+            mapped_centers.append(features[nearest_idx])
+        mapped_centers = np.stack(mapped_centers)
+        cluster_centers = mapped_centers
+    # === 3. Solve TSP between mapped cluster centers
+    cluster_order = solve_tsp_2opt(cluster_centers)
+
+    # === 4. For each cluster, solve TSP inside cluster
+    final_indices = []
+    for cluster_id in cluster_order:
+        cluster_indices = clusters[cluster_id]
+        cluster_positions = features[cluster_indices]
+
+        if len(cluster_positions) == 1:
+            final_indices.append(cluster_indices[0])
+            continue
+
+        local_order = solve_tsp_nearest_neighbor(cluster_positions)
+        ordered_cluster_indices = cluster_indices[local_order]
+        final_indices.extend(ordered_cluster_indices)
+
+    ordered_cameras = [existing_cameras[i] for i in final_indices]
+    ordered_positions = positions_normalized[final_indices]
+
+    # === 5. Subsample uniformly
+    idx = np.linspace(0, len(ordered_positions) - 1, n_selected).astype(int)
+    sampled_positions = ordered_positions[idx]
+    sampled_cameras = [ordered_cameras[i] for i in idx]
+
+    # === 6. Prepare for Catmull-Rom spline
+    if closed:
+        sampled_positions = np.vstack([sampled_positions[-1], sampled_positions, sampled_positions[0], sampled_positions[1]])
+    else:
+        sampled_positions = np.vstack([sampled_positions[0], sampled_positions, sampled_positions[-1], sampled_positions[-1]])
+
+    # === 7. Smooth path positions
+    path_positions = []
+    for i in range(1, len(sampled_positions) - 2):
+        segment = catmull_rom_spline(sampled_positions[i-1], sampled_positions[i], sampled_positions[i+1], sampled_positions[i+2], n_points_per_segment)
+        path_positions.append(segment)
+    path_positions = np.concatenate(path_positions, axis=0)
+
+    # === 8. Denormalize
+    path_positions = path_positions * scale_pos + mean_pos
+
+    # === 9. SLERP for rotations
+    rotations = R.from_matrix([cam.R for cam in sampled_cameras])
+    key_times = np.linspace(0, 1, len(rotations))
+    slerp = Slerp(key_times, rotations)
+
+    query_times = np.linspace(0, 1, len(path_positions))
+    interpolated_rotations = slerp(query_times)
+
+    # === 10. Generate Camera objects
+    reference_cam = existing_cameras[0]
+    smooth_cameras = []
+
+    for i, pos in enumerate(path_positions):
+        R_interp = interpolated_rotations[i].as_matrix()
+
+        smooth_cameras.append(Camera(
+            R=R_interp,
+            T=pos,
+            FoVx=reference_cam.FoVx,
+            FoVy=reference_cam.FoVy,
+            resolution=(reference_cam.image_width, reference_cam.image_height),
+            colmap_id=-1,
+            depth_params=None,
+            image=Image.fromarray(np.zeros((reference_cam.image_height, reference_cam.image_width, 3), dtype=np.uint8)),
+            invdepthmap=None,
+            image_name=f"clustered_smooth_path_i={i}",
+            uid=i
+        ))
+
+    return smooth_cameras
+
+
+# def generate_clustered_path(existing_cameras: List[Camera], 
+#                              n_points_per_segment: int = 20, 
+#                              d: float = 2.0,
+#                              n_clusters: int = 5,
+#                              closed: bool = False) -> List[Camera]:
+#     """
+#     Generate a smooth camera path using GMM clustering and TSP on cluster centers.
+
+#     Args:
+#         existing_cameras (List[Camera]): List of input cameras.
+#         n_points_per_segment (int): Number of interpolated points per spline segment.
+#         d (float): Distance ahead for estimating center of view.
+#         n_clusters (int): Number of GMM clusters (zones).
+#         closed (bool): Whether to close the path.
+
+#     Returns:
+#         List[Camera]: Smooth path of Camera objects.
+#     """
+#     # Extract positions and rotations
+#     positions = np.array([cam.T for cam in existing_cameras])
+
+#     # === Normalize positions
+#     mean_pos = np.mean(positions, axis=0)
+#     scale_pos = np.std(positions, axis=0)
+#     scale_pos[scale_pos == 0] = 1.0
+
+#     positions_normalized = (positions - mean_pos) / scale_pos
+
+#     # === 1. GMM clustering (only positions)
+#     gmm = GaussianMixture(n_components=n_clusters, covariance_type='full', random_state=42)
+#     cluster_labels = gmm.fit_predict(positions_normalized)
+
+#     cluster_centers = []
+#     for cluster_id in range(n_clusters):
+#         cluster_indices = np.where(cluster_labels == cluster_id)[0]
+#         if len(cluster_indices) == 0:
+#             continue
+#         cluster_center = np.mean(positions_normalized[cluster_indices], axis=0)
+#         cluster_centers.append(cluster_center)
+
+#     cluster_centers = np.stack(cluster_centers)
+
+#     # === 2. Solve TSP between cluster centers
+#     cluster_order = solve_tsp_2opt(cluster_centers)
+
+#     # === 3. Reorder cluster centers
+#     ordered_centers = cluster_centers[cluster_order]
+
+#     # === 4. Prepare Catmull-Rom spline
+#     if closed:
+#         ordered_centers = np.vstack([ordered_centers[-1], ordered_centers, ordered_centers[0], ordered_centers[1]])
+#     else:
+#         ordered_centers = np.vstack([ordered_centers[0], ordered_centers, ordered_centers[-1], ordered_centers[-1]])
+
+#     # === 5. Generate smooth path positions
+#     path_positions = []
+#     for i in range(1, len(ordered_centers) - 2):
+#         segment = catmull_rom_spline(ordered_centers[i-1], ordered_centers[i], ordered_centers[i+1], ordered_centers[i+2], n_points_per_segment)
+#         path_positions.append(segment)
+#     path_positions = np.concatenate(path_positions, axis=0)
+
+#     # === 6. Denormalize back
+#     path_positions = path_positions * scale_pos + mean_pos
+
+#     # === 7. Generate dummy rotations (constant forward facing)
+#     reference_cam = existing_cameras[0]
+#     default_rotation = R.from_matrix(reference_cam.R)
+
+#     # For simplicity, fixed rotation for all
+#     smooth_cameras = []
+
+#     for i, pos in enumerate(path_positions):
+#         R_interp = default_rotation.as_matrix()
+
+#         smooth_cameras.append(Camera(
+#             R=R_interp,
+#             T=pos,
+#             FoVx=reference_cam.FoVx,
+#             FoVy=reference_cam.FoVy,
+#             resolution=(reference_cam.image_width, reference_cam.image_height),
+#             colmap_id=-1,
+#             depth_params=None,
+#             image=Image.fromarray(np.zeros((reference_cam.image_height, reference_cam.image_width, 3), dtype=np.uint8)),
+#             invdepthmap=None,
+#             image_name=f"cluster_path_i={i}",
+#             uid=i
+#         ))
+
+#     return smooth_cameras
+
+from typing import List
+import numpy as np
+from sklearn.cluster import KMeans
+from scipy.spatial.transform import Rotation as R, Slerp
+from PIL import Image
+
+def generate_clustered_path(existing_cameras: List[Camera], 
+                             n_points_per_segment: int = 20, 
+                             d: float = 2.0,
+                             n_clusters: int = 5,
+                             closed: bool = False) -> List[Camera]:
+    """
+    Generate a smooth camera path using K-Means clustering and TSP on cluster centers.
+
+    Args:
+        existing_cameras (List[Camera]): List of input cameras.
+        n_points_per_segment (int): Number of interpolated points per spline segment.
+        d (float): Distance ahead for estimating center of view.
+        n_clusters (int): Number of KMeans clusters (zones).
+        closed (bool): Whether to close the path.
+
+    Returns:
+        List[Camera]: Smooth path of Camera objects.
+    """
+    # Extract positions
+    positions = np.array([cam.T for cam in existing_cameras])
+
+    # === Normalize positions
+    mean_pos = np.mean(positions, axis=0)
+    scale_pos = np.std(positions, axis=0)
+    scale_pos[scale_pos == 0] = 1.0
+
+    positions_normalized = (positions - mean_pos) / scale_pos
+
+    # === 1. K-Means clustering (only positions)
+    kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init='auto')
+    cluster_labels = kmeans.fit_predict(positions_normalized)
+
+    cluster_centers = []
+    for cluster_id in range(n_clusters):
+        cluster_indices = np.where(cluster_labels == cluster_id)[0]
+        if len(cluster_indices) == 0:
+            continue
+        cluster_center = np.mean(positions_normalized[cluster_indices], axis=0)
+        cluster_centers.append(cluster_center)
+
+    cluster_centers = np.stack(cluster_centers)
+
+    # === 2. Solve TSP between cluster centers
+    cluster_order = solve_tsp_2opt(cluster_centers)
+
+    # === 3. Reorder cluster centers
+    ordered_centers = cluster_centers[cluster_order]
+
+    # === 4. Prepare Catmull-Rom spline
+    if closed:
+        ordered_centers = np.vstack([ordered_centers[-1], ordered_centers, ordered_centers[0], ordered_centers[1]])
+    else:
+        ordered_centers = np.vstack([ordered_centers[0], ordered_centers, ordered_centers[-1], ordered_centers[-1]])
+
+    # === 5. Generate smooth path positions
+    path_positions = []
+    for i in range(1, len(ordered_centers) - 2):
+        segment = catmull_rom_spline(ordered_centers[i-1], ordered_centers[i], ordered_centers[i+1], ordered_centers[i+2], n_points_per_segment)
+        path_positions.append(segment)
+    path_positions = np.concatenate(path_positions, axis=0)
+
+    # === 6. Denormalize back
+    path_positions = path_positions * scale_pos + mean_pos
+
+    # === 7. Generate dummy rotations (constant forward facing)
+    reference_cam = existing_cameras[0]
+    default_rotation = R.from_matrix(reference_cam.R)
+
+    # For simplicity, fixed rotation for all
+    smooth_cameras = []
+
+    for i, pos in enumerate(path_positions):
+        R_interp = default_rotation.as_matrix()
+
+        smooth_cameras.append(Camera(
+            R=R_interp,
+            T=pos,
+            FoVx=reference_cam.FoVx,
+            FoVy=reference_cam.FoVy,
+            resolution=(reference_cam.image_width, reference_cam.image_height),
+            colmap_id=-1,
+            depth_params=None,
+            image=Image.fromarray(np.zeros((reference_cam.image_height, reference_cam.image_width, 3), dtype=np.uint8)),
+            invdepthmap=None,
+            image_name=f"cluster_path_i={i}",
+            uid=i
+        ))
+
+    return smooth_cameras
+
+
+
+
+def visualize_image_with_points(image, points):
+    """
+    Visualize an image with points overlaid on top. This is useful for correspondences visualizations
+
+    Parameters:
+    - image: PIL Image object
+    - points: Numpy array of shape [N, 2] containing (x, y) coordinates of points
+
+    Returns:
+    - None (displays the visualization)
+    """
+
+    # Convert PIL image to numpy array
+    img_array = np.array(image)
+
+    # Create a figure and axis
+    fig, ax = plt.subplots(figsize=(7,7))
+
+    # Display the image
+    ax.imshow(img_array)
+
+    # Scatter plot the points on top of the image
+    ax.scatter(points[:, 0], points[:, 1], color='red', marker='o', s=1)
+
+    # Show the plot
+    plt.show()
+
+
+def visualize_correspondences(image1, points1, image2, points2):
+    """
+    Visualize two images concatenated horizontally with key points and correspondences.
+
+    Parameters:
+    - image1: PIL Image object (left image)
+    - points1: Numpy array of shape [N, 2] containing (x, y) coordinates of key points for image1
+    - image2: PIL Image object (right image)
+    - points2: Numpy array of shape [N, 2] containing (x, y) coordinates of key points for image2
+
+    Returns:
+    - None (displays the visualization)
+    """
+
+    # Concatenate images horizontally
+    concatenated_image = np.concatenate((np.array(image1), np.array(image2)), axis=1)
+
+    # Create a figure and axis
+    fig, ax = plt.subplots(figsize=(10,10))
+
+    # Display the concatenated image
+    ax.imshow(concatenated_image)
+
+    # Plot key points on the left image
+    ax.scatter(points1[:, 0], points1[:, 1], color='red', marker='o', s=10)
+
+    # Plot key points on the right image
+    ax.scatter(points2[:, 0] + image1.width, points2[:, 1], color='blue', marker='o', s=10)
+
+    # Draw lines connecting corresponding key points
+    for i in range(len(points1)):
+        ax.plot([points1[i, 0], points2[i, 0] + image1.width], [points1[i, 1], points2[i, 1]])#, color='green')
+
+    # Show the plot
+    plt.show()
+

submodules/RoMa/.gitignore

@@ -0,0 +1,11 @@
+*.egg-info*
+*.vscode*
+*__pycache__*
+vis*
+workspace*
+.venv
+.DS_Store
+jobs/*
+*ignore_me*
+*.pth
+wandb*

submodules/RoMa/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Johan Edstedt
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

submodules/RoMa/README.md

@@ -0,0 +1,123 @@
+# 
+<p align="center">
+  <h1 align="center"> <ins>RoMa</ins> 🏛️:<br> Robust Dense Feature Matching <br> ⭐CVPR 2024⭐</h1>
+  <p align="center">
+    <a href="https://scholar.google.com/citations?user=Ul-vMR0AAAAJ">Johan Edstedt</a>
+    ·
+    <a href="https://scholar.google.com/citations?user=HS2WuHkAAAAJ">Qiyu Sun</a>
+    ·
+    <a href="https://scholar.google.com/citations?user=FUE3Wd0AAAAJ">Georg Bökman</a>
+    ·
+    <a href="https://scholar.google.com/citations?user=6WRQpCQAAAAJ">Mårten Wadenbäck</a>
+    ·
+    <a href="https://scholar.google.com/citations?user=lkWfR08AAAAJ">Michael Felsberg</a>
+  </p>
+  <h2 align="center"><p>
+    <a href="https://arxiv.org/abs/2305.15404" align="center">Paper</a> | 
+    <a href="https://parskatt.github.io/RoMa" align="center">Project Page</a>
+  </p></h2>
+  <div align="center"></div>
+</p>
+<br/>
+<p align="center">
+    <img src="https://github.com/Parskatt/RoMa/assets/22053118/15d8fea7-aa6d-479f-8a93-350d950d006b" alt="example" width=80%>
+    <br>
+    <em>RoMa is the robust dense feature matcher capable of estimating pixel-dense warps and reliable certainties for almost any image pair.</em>
+</p>
+
+## Setup/Install
+In your python environment (tested on Linux python 3.10), run:
+```bash
+pip install -e .
+```
+## Demo / How to Use
+We provide two demos in the [demos folder](demo).
+Here's the gist of it:
+```python
+from romatch import roma_outdoor
+roma_model = roma_outdoor(device=device)
+# Match
+warp, certainty = roma_model.match(imA_path, imB_path, device=device)
+# Sample matches for estimation
+matches, certainty = roma_model.sample(warp, certainty)
+# Convert to pixel coordinates (RoMa produces matches in [-1,1]x[-1,1])
+kptsA, kptsB = roma_model.to_pixel_coordinates(matches, H_A, W_A, H_B, W_B)
+# Find a fundamental matrix (or anything else of interest)
+F, mask = cv2.findFundamentalMat(
+    kptsA.cpu().numpy(), kptsB.cpu().numpy(), ransacReprojThreshold=0.2, method=cv2.USAC_MAGSAC, confidence=0.999999, maxIters=10000
+)
+```
+
+**New**: You can also match arbitrary keypoints with RoMa. See [match_keypoints](romatch/models/matcher.py) in RegressionMatcher.
+
+## Settings
+
+### Resolution
+By default RoMa uses an initial resolution of (560,560) which is then upsampled to (864,864). 
+You can change this at construction (see roma_outdoor kwargs).
+You can also change this later, by changing the roma_model.w_resized, roma_model.h_resized, and roma_model.upsample_res.
+
+### Sampling
+roma_model.sample_thresh controls the thresholding used when sampling matches for estimation. In certain cases a lower or higher threshold may improve results.
+
+
+## Reproducing Results
+The experiments in the paper are provided in the [experiments folder](experiments).
+
+### Training
+1. First follow the instructions provided here: https://github.com/Parskatt/DKM for downloading and preprocessing datasets.
+2. Run the relevant experiment, e.g.,
+```bash
+torchrun --nproc_per_node=4 --nnodes=1 --rdzv_backend=c10d experiments/roma_outdoor.py
+```
+### Testing
+```bash
+python experiments/roma_outdoor.py --only_test --benchmark mega-1500
+```
+## License
+All our code except DINOv2 is MIT license.
+DINOv2 has an Apache 2 license [DINOv2](https://github.com/facebookresearch/dinov2/blob/main/LICENSE).
+
+## Acknowledgement
+Our codebase builds on the code in [DKM](https://github.com/Parskatt/DKM).
+
+## Tiny RoMa
+If you find that RoMa is too heavy, you might want to try Tiny RoMa which is built on top of XFeat.
+```python
+from romatch import tiny_roma_v1_outdoor
+tiny_roma_model = tiny_roma_v1_outdoor(device=device)
+```
+Mega1500:
+|  | AUC@5 | AUC@10 | AUC@20 |
+|----------|----------|----------|----------|
+| XFeat    | 46.4    | 58.9    | 69.2    |
+| XFeat*    |  51.9   | 67.2    | 78.9    |
+| Tiny RoMa v1    | 56.4 | 69.5 | 79.5     |
+| RoMa    |  -   | -    | -    |
+
+Mega-8-Scenes (See DKM):
+|  | AUC@5 | AUC@10 | AUC@20 |
+|----------|----------|----------|----------|
+| XFeat    | -    | -    | -    |
+| XFeat*    |  50.1   | 64.4    | 75.2    |
+| Tiny RoMa v1    | 57.7 | 70.5 | 79.6     |
+| RoMa    |  -   | -    | -    |
+
+IMC22 :'):
+|  | mAA@10 |
+|----------|----------|
+| XFeat    | 42.1    |
+| XFeat*    |  -   |
+| Tiny RoMa v1    | 42.2 |
+| RoMa    |  -   |
+
+## BibTeX
+If you find our models useful, please consider citing our paper!
+```
+@article{edstedt2024roma,
+title={{RoMa: Robust Dense Feature Matching}},
+author={Edstedt, Johan and Sun, Qiyu and Bökman, Georg and Wadenbäck, Mårten and Felsberg, Michael},
+journal={IEEE Conference on Computer Vision and Pattern Recognition},
+year={2024}
+}
+```

submodules/RoMa/data/.gitignore

@@ -0,0 +1,2 @@
+*
+!.gitignore

submodules/RoMa/demo/demo_3D_effect.py

@@ -0,0 +1,47 @@
+from PIL import Image
+import torch
+import torch.nn.functional as F
+import numpy as np
+from romatch.utils.utils import tensor_to_pil
+
+from romatch import roma_outdoor
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+if torch.backends.mps.is_available():
+    device = torch.device('mps')
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+    parser = ArgumentParser()
+    parser.add_argument("--im_A_path", default="assets/toronto_A.jpg", type=str)
+    parser.add_argument("--im_B_path", default="assets/toronto_B.jpg", type=str)
+    parser.add_argument("--save_path", default="demo/gif/roma_warp_toronto", type=str)
+
+    args, _ = parser.parse_known_args()
+    im1_path = args.im_A_path
+    im2_path = args.im_B_path
+    save_path = args.save_path
+
+    # Create model
+    roma_model = roma_outdoor(device=device, coarse_res=560, upsample_res=(864, 1152))
+    roma_model.symmetric = False
+
+    H, W = roma_model.get_output_resolution()
+
+    im1 = Image.open(im1_path).resize((W, H))
+    im2 = Image.open(im2_path).resize((W, H))
+
+    # Match
+    warp, certainty = roma_model.match(im1_path, im2_path, device=device)
+    # Sampling not needed, but can be done with model.sample(warp, certainty)
+    x1 = (torch.tensor(np.array(im1)) / 255).to(device).permute(2, 0, 1)
+    x2 = (torch.tensor(np.array(im2)) / 255).to(device).permute(2, 0, 1)
+
+    coords_A, coords_B = warp[...,:2], warp[...,2:]
+    for i, x in enumerate(np.linspace(0,2*np.pi,200)):
+        t = (1 + np.cos(x))/2
+        interp_warp = (1-t)*coords_A + t*coords_B
+        im2_transfer_rgb = F.grid_sample(
+        x2[None], interp_warp[None], mode="bilinear", align_corners=False
+        )[0]
+        tensor_to_pil(im2_transfer_rgb, unnormalize=False).save(f"{save_path}_{i:03d}.jpg")

submodules/RoMa/demo/demo_fundamental.py

@@ -0,0 +1,34 @@
+from PIL import Image
+import torch
+import cv2
+from romatch import roma_outdoor
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+if torch.backends.mps.is_available():
+    device = torch.device('mps')
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+    parser = ArgumentParser()
+    parser.add_argument("--im_A_path", default="assets/sacre_coeur_A.jpg", type=str)
+    parser.add_argument("--im_B_path", default="assets/sacre_coeur_B.jpg", type=str)
+
+    args, _ = parser.parse_known_args()
+    im1_path = args.im_A_path
+    im2_path = args.im_B_path
+
+    # Create model
+    roma_model = roma_outdoor(device=device)
+
+
+    W_A, H_A = Image.open(im1_path).size
+    W_B, H_B = Image.open(im2_path).size
+
+    # Match
+    warp, certainty = roma_model.match(im1_path, im2_path, device=device)
+    # Sample matches for estimation
+    matches, certainty = roma_model.sample(warp, certainty)
+    kpts1, kpts2 = roma_model.to_pixel_coordinates(matches, H_A, W_A, H_B, W_B)    
+    F, mask = cv2.findFundamentalMat(
+        kpts1.cpu().numpy(), kpts2.cpu().numpy(), ransacReprojThreshold=0.2, method=cv2.USAC_MAGSAC, confidence=0.999999, maxIters=10000
+    )

submodules/RoMa/demo/demo_match.py

@@ -0,0 +1,50 @@
+import os
+os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'
+import torch
+from PIL import Image
+import torch.nn.functional as F
+import numpy as np
+from romatch.utils.utils import tensor_to_pil
+
+from romatch import roma_outdoor
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+if torch.backends.mps.is_available():
+    device = torch.device('mps')
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+    parser = ArgumentParser()
+    parser.add_argument("--im_A_path", default="assets/toronto_A.jpg", type=str)
+    parser.add_argument("--im_B_path", default="assets/toronto_B.jpg", type=str)
+    parser.add_argument("--save_path", default="demo/roma_warp_toronto.jpg", type=str)
+
+    args, _ = parser.parse_known_args()
+    im1_path = args.im_A_path
+    im2_path = args.im_B_path
+    save_path = args.save_path
+
+    # Create model
+    roma_model = roma_outdoor(device=device, coarse_res=560, upsample_res=(864, 1152))
+
+    H, W = roma_model.get_output_resolution()
+
+    im1 = Image.open(im1_path).resize((W, H))
+    im2 = Image.open(im2_path).resize((W, H))
+
+    # Match
+    warp, certainty = roma_model.match(im1_path, im2_path, device=device)
+    # Sampling not needed, but can be done with model.sample(warp, certainty)
+    x1 = (torch.tensor(np.array(im1)) / 255).to(device).permute(2, 0, 1)
+    x2 = (torch.tensor(np.array(im2)) / 255).to(device).permute(2, 0, 1)
+
+    im2_transfer_rgb = F.grid_sample(
+    x2[None], warp[:,:W, 2:][None], mode="bilinear", align_corners=False
+    )[0]
+    im1_transfer_rgb = F.grid_sample(
+    x1[None], warp[:, W:, :2][None], mode="bilinear", align_corners=False
+    )[0]
+    warp_im = torch.cat((im2_transfer_rgb,im1_transfer_rgb),dim=2)
+    white_im = torch.ones((H,2*W),device=device)
+    vis_im = certainty * warp_im + (1 - certainty) * white_im
+    tensor_to_pil(vis_im, unnormalize=False).save(save_path)

submodules/RoMa/demo/demo_match_opencv_sift.py

@@ -0,0 +1,43 @@
+from PIL import Image
+import numpy as np
+
+import numpy as np
+import cv2 as cv
+import matplotlib.pyplot as plt
+
+
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+    parser = ArgumentParser()
+    parser.add_argument("--im_A_path", default="assets/toronto_A.jpg", type=str)
+    parser.add_argument("--im_B_path", default="assets/toronto_B.jpg", type=str)
+    parser.add_argument("--save_path", default="demo/roma_warp_toronto.jpg", type=str)
+
+    args, _ = parser.parse_known_args()
+    im1_path = args.im_A_path
+    im2_path = args.im_B_path
+    save_path = args.save_path
+
+    img1 = cv.imread(im1_path,cv.IMREAD_GRAYSCALE)          # queryImage
+    img2 = cv.imread(im2_path,cv.IMREAD_GRAYSCALE) # trainImage
+    # Initiate SIFT detector
+    sift = cv.SIFT_create()
+    # find the keypoints and descriptors with SIFT
+    kp1, des1 = sift.detectAndCompute(img1,None)
+    kp2, des2 = sift.detectAndCompute(img2,None)
+    # BFMatcher with default params
+    bf = cv.BFMatcher()
+    matches = bf.knnMatch(des1,des2,k=2)
+    # Apply ratio test
+    good = []
+    for m,n in matches:
+        if m.distance < 0.75*n.distance:
+            good.append([m])
+    # cv.drawMatchesKnn expects list of lists as matches.
+    draw_params = dict(matchColor = (255,0,0), # draw matches in red color
+                   singlePointColor = None,
+                   flags = 2)
+
+    img3 = cv.drawMatchesKnn(img1,kp1,img2,kp2,good,None,**draw_params)
+    Image.fromarray(img3).save("demo/sift_matches.png")

submodules/RoMa/demo/demo_match_tiny.py

@@ -0,0 +1,77 @@
+import os
+os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'
+import torch
+from PIL import Image
+import torch.nn.functional as F
+import numpy as np
+from romatch.utils.utils import tensor_to_pil
+
+from romatch import tiny_roma_v1_outdoor
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+if torch.backends.mps.is_available():
+    device = torch.device('mps')
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+    parser = ArgumentParser()
+    parser.add_argument("--im_A_path", default="assets/sacre_coeur_A.jpg", type=str)
+    parser.add_argument("--im_B_path", default="assets/sacre_coeur_B.jpg", type=str)
+    parser.add_argument("--save_A_path", default="demo/tiny_roma_warp_A.jpg", type=str)
+    parser.add_argument("--save_B_path", default="demo/tiny_roma_warp_B.jpg", type=str)
+
+    args, _ = parser.parse_known_args()
+    im1_path = args.im_A_path
+    im2_path = args.im_B_path
+
+    # Create model
+    roma_model = tiny_roma_v1_outdoor(device=device)
+
+    # Match
+    warp, certainty1 = roma_model.match(im1_path, im2_path)
+    
+    h1, w1 = warp.shape[:2]
+    
+    # maybe im1.size != im2.size
+    im1 = Image.open(im1_path).resize((w1, h1))
+    im2 = Image.open(im2_path)
+    x1 = (torch.tensor(np.array(im1)) / 255).to(device).permute(2, 0, 1)
+    x2 = (torch.tensor(np.array(im2)) / 255).to(device).permute(2, 0, 1)
+    
+    h2, w2 = x2.shape[1:]
+    g1_p2x = w2 / 2 * (warp[..., 2] + 1)
+    g1_p2y = h2 / 2 * (warp[..., 3] + 1)
+    g2_p1x = torch.zeros((h2, w2), dtype=torch.float32).to(device) - 2
+    g2_p1y = torch.zeros((h2, w2), dtype=torch.float32).to(device) - 2
+
+    x, y = torch.meshgrid(
+        torch.arange(w1, device=device),
+        torch.arange(h1, device=device),
+        indexing="xy",
+    )
+    g2x = torch.round(g1_p2x[y, x]).long()
+    g2y = torch.round(g1_p2y[y, x]).long()
+    idx_x = torch.bitwise_and(0 <= g2x, g2x < w2)
+    idx_y = torch.bitwise_and(0 <= g2y, g2y < h2)
+    idx = torch.bitwise_and(idx_x, idx_y)
+    g2_p1x[g2y[idx], g2x[idx]] = x[idx].float() * 2 / w1 - 1
+    g2_p1y[g2y[idx], g2x[idx]] = y[idx].float() * 2 / h1 - 1
+
+    certainty2 = F.grid_sample(
+        certainty1[None][None],
+        torch.stack([g2_p1x, g2_p1y], dim=2)[None],
+        mode="bilinear",
+        align_corners=False,
+    )[0]
+    
+    white_im1 = torch.ones((h1, w1), device = device)
+    white_im2 = torch.ones((h2, w2), device = device)
+    
+    certainty1 = F.avg_pool2d(certainty1[None], kernel_size=5, stride=1, padding=2)[0]
+    certainty2 = F.avg_pool2d(certainty2[None], kernel_size=5, stride=1, padding=2)[0]
+    
+    vis_im1 = certainty1 * x1 + (1 - certainty1) * white_im1
+    vis_im2 = certainty2 * x2 + (1 - certainty2) * white_im2
+    
+    tensor_to_pil(vis_im1, unnormalize=False).save(args.save_A_path)
+    tensor_to_pil(vis_im2, unnormalize=False).save(args.save_B_path)

submodules/RoMa/demo/gif/.gitignore

@@ -0,0 +1,2 @@
+*
+!.gitignore

submodules/RoMa/experiments/eval_roma_outdoor.py

@@ -0,0 +1,57 @@
+import json
+
+from romatch.benchmarks import MegadepthDenseBenchmark
+from romatch.benchmarks import MegaDepthPoseEstimationBenchmark, HpatchesHomogBenchmark
+from romatch.benchmarks import Mega1500PoseLibBenchmark
+
+def test_mega_8_scenes(model, name):
+    mega_8_scenes_benchmark = MegaDepthPoseEstimationBenchmark("data/megadepth",
+                                                scene_names=['mega_8_scenes_0019_0.1_0.3.npz',
+                                                    'mega_8_scenes_0025_0.1_0.3.npz',
+                                                    'mega_8_scenes_0021_0.1_0.3.npz',
+                                                    'mega_8_scenes_0008_0.1_0.3.npz',
+                                                    'mega_8_scenes_0032_0.1_0.3.npz',
+                                                    'mega_8_scenes_1589_0.1_0.3.npz',
+                                                    'mega_8_scenes_0063_0.1_0.3.npz',
+                                                    'mega_8_scenes_0024_0.1_0.3.npz',
+                                                    'mega_8_scenes_0019_0.3_0.5.npz',
+                                                    'mega_8_scenes_0025_0.3_0.5.npz',
+                                                    'mega_8_scenes_0021_0.3_0.5.npz',
+                                                    'mega_8_scenes_0008_0.3_0.5.npz',
+                                                    'mega_8_scenes_0032_0.3_0.5.npz',
+                                                    'mega_8_scenes_1589_0.3_0.5.npz',
+                                                    'mega_8_scenes_0063_0.3_0.5.npz',
+                                                    'mega_8_scenes_0024_0.3_0.5.npz'])
+    mega_8_scenes_results = mega_8_scenes_benchmark.benchmark(model, model_name=name)
+    print(mega_8_scenes_results)
+    json.dump(mega_8_scenes_results, open(f"results/mega_8_scenes_{name}.json", "w"))
+
+def test_mega1500(model, name):
+    mega1500_benchmark = MegaDepthPoseEstimationBenchmark("data/megadepth")
+    mega1500_results = mega1500_benchmark.benchmark(model, model_name=name)
+    json.dump(mega1500_results, open(f"results/mega1500_{name}.json", "w"))
+
+def test_mega1500_poselib(model, name):
+    mega1500_benchmark = Mega1500PoseLibBenchmark("data/megadepth")
+    mega1500_results = mega1500_benchmark.benchmark(model, model_name=name)
+    json.dump(mega1500_results, open(f"results/mega1500_{name}.json", "w"))
+
+def test_mega_dense(model, name):
+    megadense_benchmark = MegadepthDenseBenchmark("data/megadepth", num_samples = 1000)
+    megadense_results = megadense_benchmark.benchmark(model)
+    json.dump(megadense_results, open(f"results/mega_dense_{name}.json", "w"))
+    
+def test_hpatches(model, name):
+    hpatches_benchmark = HpatchesHomogBenchmark("data/hpatches")
+    hpatches_results = hpatches_benchmark.benchmark(model)
+    json.dump(hpatches_results, open(f"results/hpatches_{name}.json", "w"))
+
+
+if __name__ == "__main__":
+    from romatch import roma_outdoor
+    device = "cuda"
+    model = roma_outdoor(device = device, coarse_res = 672, upsample_res = 1344)
+    experiment_name = "roma_latest"
+    test_mega1500(model, experiment_name)
+    #test_mega1500_poselib(model, experiment_name)
+    

submodules/RoMa/experiments/eval_tiny_roma_v1_outdoor.py

@@ -0,0 +1,84 @@
+import torch
+import os
+from pathlib import Path
+import json
+from romatch.benchmarks import ScanNetBenchmark
+from romatch.benchmarks import Mega1500PoseLibBenchmark, ScanNetPoselibBenchmark
+from romatch.benchmarks import MegaDepthPoseEstimationBenchmark
+
+def test_mega_8_scenes(model, name):
+    mega_8_scenes_benchmark = MegaDepthPoseEstimationBenchmark("data/megadepth",
+                                                scene_names=['mega_8_scenes_0019_0.1_0.3.npz',
+                                                    'mega_8_scenes_0025_0.1_0.3.npz',
+                                                    'mega_8_scenes_0021_0.1_0.3.npz',
+                                                    'mega_8_scenes_0008_0.1_0.3.npz',
+                                                    'mega_8_scenes_0032_0.1_0.3.npz',
+                                                    'mega_8_scenes_1589_0.1_0.3.npz',
+                                                    'mega_8_scenes_0063_0.1_0.3.npz',
+                                                    'mega_8_scenes_0024_0.1_0.3.npz',
+                                                    'mega_8_scenes_0019_0.3_0.5.npz',
+                                                    'mega_8_scenes_0025_0.3_0.5.npz',
+                                                    'mega_8_scenes_0021_0.3_0.5.npz',
+                                                    'mega_8_scenes_0008_0.3_0.5.npz',
+                                                    'mega_8_scenes_0032_0.3_0.5.npz',
+                                                    'mega_8_scenes_1589_0.3_0.5.npz',
+                                                    'mega_8_scenes_0063_0.3_0.5.npz',
+                                                    'mega_8_scenes_0024_0.3_0.5.npz'])
+    mega_8_scenes_results = mega_8_scenes_benchmark.benchmark(model, model_name=name)
+    print(mega_8_scenes_results)
+    json.dump(mega_8_scenes_results, open(f"results/mega_8_scenes_{name}.json", "w"))
+
+def test_mega1500(model, name):
+    mega1500_benchmark = MegaDepthPoseEstimationBenchmark("data/megadepth")
+    mega1500_results = mega1500_benchmark.benchmark(model, model_name=name)
+    json.dump(mega1500_results, open(f"results/mega1500_{name}.json", "w"))
+
+def test_mega1500_poselib(model, name):
+    #model.exact_softmax = True
+    mega1500_benchmark = Mega1500PoseLibBenchmark("data/megadepth", num_ransac_iter = 1, test_every = 1)
+    mega1500_results = mega1500_benchmark.benchmark(model, model_name=name)
+    json.dump(mega1500_results, open(f"results/mega1500_poselib_{name}.json", "w"))
+
+def test_mega_8_scenes_poselib(model, name):
+    mega1500_benchmark = Mega1500PoseLibBenchmark("data/megadepth", num_ransac_iter = 1, test_every = 1,
+                                                  scene_names=['mega_8_scenes_0019_0.1_0.3.npz',
+                                                    'mega_8_scenes_0025_0.1_0.3.npz',
+                                                    'mega_8_scenes_0021_0.1_0.3.npz',
+                                                    'mega_8_scenes_0008_0.1_0.3.npz',
+                                                    'mega_8_scenes_0032_0.1_0.3.npz',
+                                                    'mega_8_scenes_1589_0.1_0.3.npz',
+                                                    'mega_8_scenes_0063_0.1_0.3.npz',
+                                                    'mega_8_scenes_0024_0.1_0.3.npz',
+                                                    'mega_8_scenes_0019_0.3_0.5.npz',
+                                                    'mega_8_scenes_0025_0.3_0.5.npz',
+                                                    'mega_8_scenes_0021_0.3_0.5.npz',
+                                                    'mega_8_scenes_0008_0.3_0.5.npz',
+                                                    'mega_8_scenes_0032_0.3_0.5.npz',
+                                                    'mega_8_scenes_1589_0.3_0.5.npz',
+                                                    'mega_8_scenes_0063_0.3_0.5.npz',
+                                                    'mega_8_scenes_0024_0.3_0.5.npz'])
+    mega1500_results = mega1500_benchmark.benchmark(model, model_name=name)
+    json.dump(mega1500_results, open(f"results/mega_8_scenes_poselib_{name}.json", "w"))
+
+def test_scannet_poselib(model, name):
+    scannet_benchmark = ScanNetPoselibBenchmark("data/scannet")
+    scannet_results = scannet_benchmark.benchmark(model)
+    json.dump(scannet_results, open(f"results/scannet_{name}.json", "w"))
+
+def test_scannet(model, name):
+    scannet_benchmark = ScanNetBenchmark("data/scannet")
+    scannet_results = scannet_benchmark.benchmark(model)
+    json.dump(scannet_results, open(f"results/scannet_{name}.json", "w"))
+
+if __name__ == "__main__":
+    os.environ["TORCH_CUDNN_V8_API_ENABLED"] = "1" # For BF16 computations
+    os.environ["OMP_NUM_THREADS"] = "16"
+    torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
+    from romatch import tiny_roma_v1_outdoor
+
+    experiment_name = Path(__file__).stem
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = tiny_roma_v1_outdoor(device)
+    #test_mega1500_poselib(model, experiment_name)
+    test_mega_8_scenes_poselib(model, experiment_name)
+ 

submodules/RoMa/experiments/roma_indoor.py

@@ -0,0 +1,320 @@
+import os
+import torch
+from argparse import ArgumentParser
+
+from torch import nn
+from torch.utils.data import ConcatDataset
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+import json
+import wandb
+from tqdm import tqdm
+
+from romatch.benchmarks import MegadepthDenseBenchmark
+from romatch.datasets.megadepth import MegadepthBuilder
+from romatch.datasets.scannet import ScanNetBuilder
+from romatch.losses.robust_loss import RobustLosses
+from romatch.benchmarks import MegadepthDenseBenchmark, ScanNetBenchmark
+from romatch.train.train import train_k_steps
+from romatch.models.matcher import *
+from romatch.models.transformer import Block, TransformerDecoder, MemEffAttention
+from romatch.models.encoders import *
+from romatch.checkpointing import CheckPoint
+
+resolutions = {"low":(448, 448), "medium":(14*8*5, 14*8*5), "high":(14*8*6, 14*8*6)}
+
+def get_model(pretrained_backbone=True, resolution = "medium", **kwargs):
+    gp_dim = 512
+    feat_dim = 512
+    decoder_dim = gp_dim + feat_dim
+    cls_to_coord_res = 64
+    coordinate_decoder = TransformerDecoder(
+        nn.Sequential(*[Block(decoder_dim, 8, attn_class=MemEffAttention) for _ in range(5)]), 
+        decoder_dim, 
+        cls_to_coord_res**2 + 1,
+        is_classifier=True,
+        amp = True,
+        pos_enc = False,)
+    dw = True
+    hidden_blocks = 8
+    kernel_size = 5
+    displacement_emb = "linear"
+    disable_local_corr_grad = True
+    
+    conv_refiner = nn.ModuleDict(
+        {
+            "16": ConvRefiner(
+                2 * 512+128+(2*7+1)**2,
+                2 * 512+128+(2*7+1)**2,
+                2 + 1,
+                kernel_size=kernel_size,
+                dw=dw,
+                hidden_blocks=hidden_blocks,
+                displacement_emb=displacement_emb,
+                displacement_emb_dim=128,
+                local_corr_radius = 7,
+                corr_in_other = True,
+                amp = True,
+                disable_local_corr_grad = disable_local_corr_grad,
+                bn_momentum = 0.01,
+            ),
+            "8": ConvRefiner(
+                2 * 512+64+(2*3+1)**2,
+                2 * 512+64+(2*3+1)**2,
+                2 + 1,
+                kernel_size=kernel_size,
+                dw=dw,
+                hidden_blocks=hidden_blocks,
+                displacement_emb=displacement_emb,
+                displacement_emb_dim=64,
+                local_corr_radius = 3,
+                corr_in_other = True,
+                amp = True,
+                disable_local_corr_grad = disable_local_corr_grad,
+                bn_momentum = 0.01,
+            ),
+            "4": ConvRefiner(
+                2 * 256+32+(2*2+1)**2,
+                2 * 256+32+(2*2+1)**2,
+                2 + 1,
+                kernel_size=kernel_size,
+                dw=dw,
+                hidden_blocks=hidden_blocks,
+                displacement_emb=displacement_emb,
+                displacement_emb_dim=32,
+                local_corr_radius = 2,
+                corr_in_other = True,
+                amp = True,
+                disable_local_corr_grad = disable_local_corr_grad,
+                bn_momentum = 0.01,
+            ),
+            "2": ConvRefiner(
+                2 * 64+16,
+                128+16,
+                2 + 1,
+                kernel_size=kernel_size,
+                dw=dw,
+                hidden_blocks=hidden_blocks,
+                displacement_emb=displacement_emb,
+                displacement_emb_dim=16,
+                amp = True,
+                disable_local_corr_grad = disable_local_corr_grad,
+                bn_momentum = 0.01,
+            ),
+            "1": ConvRefiner(
+                2 * 9 + 6,
+                24,
+                2 + 1,
+                kernel_size=kernel_size,
+                dw=dw,
+                hidden_blocks = hidden_blocks,
+                displacement_emb = displacement_emb,
+                displacement_emb_dim = 6,
+                amp = True,
+                disable_local_corr_grad = disable_local_corr_grad,
+                bn_momentum = 0.01,
+            ),
+        }
+    )
+    kernel_temperature = 0.2
+    learn_temperature = False
+    no_cov = True
+    kernel = CosKernel
+    only_attention = False
+    basis = "fourier"
+    gp16 = GP(
+        kernel,
+        T=kernel_temperature,
+        learn_temperature=learn_temperature,
+        only_attention=only_attention,
+        gp_dim=gp_dim,
+        basis=basis,
+        no_cov=no_cov,
+    )
+    gps = nn.ModuleDict({"16": gp16})
+    proj16 = nn.Sequential(nn.Conv2d(1024, 512, 1, 1), nn.BatchNorm2d(512))
+    proj8 = nn.Sequential(nn.Conv2d(512, 512, 1, 1), nn.BatchNorm2d(512))
+    proj4 = nn.Sequential(nn.Conv2d(256, 256, 1, 1), nn.BatchNorm2d(256))
+    proj2 = nn.Sequential(nn.Conv2d(128, 64, 1, 1), nn.BatchNorm2d(64))
+    proj1 = nn.Sequential(nn.Conv2d(64, 9, 1, 1), nn.BatchNorm2d(9))
+    proj = nn.ModuleDict({
+        "16": proj16,
+        "8": proj8,
+        "4": proj4,
+        "2": proj2,
+        "1": proj1,
+        })
+    displacement_dropout_p = 0.0
+    gm_warp_dropout_p = 0.0
+    decoder = Decoder(coordinate_decoder, 
+                      gps, 
+                      proj, 
+                      conv_refiner, 
+                      detach=True, 
+                      scales=["16", "8", "4", "2", "1"], 
+                      displacement_dropout_p = displacement_dropout_p,
+                      gm_warp_dropout_p = gm_warp_dropout_p)
+    h,w = resolutions[resolution]
+    encoder = CNNandDinov2(
+        cnn_kwargs = dict(
+            pretrained=pretrained_backbone,
+            amp = True),
+        amp = True,
+        use_vgg = True,
+    )
+    matcher = RegressionMatcher(encoder, decoder, h=h, w=w, alpha=1, beta=0,**kwargs)
+    return matcher
+
+def train(args):
+    dist.init_process_group('nccl')
+    #torch._dynamo.config.verbose=True
+    gpus = int(os.environ['WORLD_SIZE'])
+    # create model and move it to GPU with id rank
+    rank = dist.get_rank()
+    print(f"Start running DDP on rank {rank}")
+    device_id = rank % torch.cuda.device_count()
+    romatch.LOCAL_RANK = device_id
+    torch.cuda.set_device(device_id)
+    
+    resolution = args.train_resolution
+    wandb_log = not args.dont_log_wandb
+    experiment_name = os.path.splitext(os.path.basename(__file__))[0]
+    wandb_mode = "online" if wandb_log and rank == 0 and False else "disabled"
+    wandb.init(project="romatch", entity=args.wandb_entity, name=experiment_name, reinit=False, mode = wandb_mode)
+    checkpoint_dir = "workspace/checkpoints/"
+    h,w = resolutions[resolution]
+    model = get_model(pretrained_backbone=True, resolution=resolution, attenuate_cert = False).to(device_id)
+    # Num steps
+    global_step = 0
+    batch_size = args.gpu_batch_size
+    step_size = gpus*batch_size
+    romatch.STEP_SIZE = step_size
+    
+    N = (32 * 250000)  # 250k steps of batch size 32
+    # checkpoint every
+    k = 25000 // romatch.STEP_SIZE
+
+    # Data
+    mega = MegadepthBuilder(data_root="data/megadepth", loftr_ignore=True, imc21_ignore = True)
+    use_horizontal_flip_aug = True
+    rot_prob = 0
+    depth_interpolation_mode = "bilinear"
+    megadepth_train1 = mega.build_scenes(
+        split="train_loftr", min_overlap=0.01, shake_t=32, use_horizontal_flip_aug = use_horizontal_flip_aug, rot_prob = rot_prob,
+        ht=h,wt=w,
+    )
+    megadepth_train2 = mega.build_scenes(
+        split="train_loftr", min_overlap=0.35, shake_t=32, use_horizontal_flip_aug = use_horizontal_flip_aug, rot_prob = rot_prob,
+        ht=h,wt=w,
+    )
+    megadepth_train = ConcatDataset(megadepth_train1 + megadepth_train2)
+    mega_ws = mega.weight_scenes(megadepth_train, alpha=0.75)
+    
+    scannet = ScanNetBuilder(data_root="data/scannet")
+    scannet_train = scannet.build_scenes(split="train", ht=h, wt=w, use_horizontal_flip_aug = use_horizontal_flip_aug)
+    scannet_train = ConcatDataset(scannet_train)
+    scannet_ws = scannet.weight_scenes(scannet_train, alpha=0.75)
+
+    # Loss and optimizer
+    depth_loss_scannet = RobustLosses(
+        ce_weight=0.0, 
+        local_dist={1:4, 2:4, 4:8, 8:8},
+        local_largest_scale=8,
+        depth_interpolation_mode=depth_interpolation_mode,
+        alpha = 0.5,
+        c = 1e-4,)
+    # Loss and optimizer
+    depth_loss_mega = RobustLosses(
+        ce_weight=0.01, 
+        local_dist={1:4, 2:4, 4:8, 8:8},
+        local_largest_scale=8,
+        depth_interpolation_mode=depth_interpolation_mode,
+        alpha = 0.5,
+        c = 1e-4,)
+    parameters = [
+        {"params": model.encoder.parameters(), "lr": romatch.STEP_SIZE * 5e-6 / 8},
+        {"params": model.decoder.parameters(), "lr": romatch.STEP_SIZE * 1e-4 / 8},
+    ]
+    optimizer = torch.optim.AdamW(parameters, weight_decay=0.01)
+    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+        optimizer, milestones=[(9*N/romatch.STEP_SIZE)//10])
+    megadense_benchmark = MegadepthDenseBenchmark("data/megadepth", num_samples = 1000, h=h,w=w)
+    checkpointer = CheckPoint(checkpoint_dir, experiment_name)
+    model, optimizer, lr_scheduler, global_step = checkpointer.load(model, optimizer, lr_scheduler, global_step)
+    romatch.GLOBAL_STEP = global_step
+    ddp_model = DDP(model, device_ids=[device_id], find_unused_parameters = False, gradient_as_bucket_view=True)
+    grad_scaler = torch.cuda.amp.GradScaler(growth_interval=1_000_000)
+    grad_clip_norm = 0.01
+    for n in range(romatch.GLOBAL_STEP, N, k * romatch.STEP_SIZE):
+        mega_sampler = torch.utils.data.WeightedRandomSampler(
+            mega_ws, num_samples = batch_size * k, replacement=False
+        )
+        mega_dataloader = iter(
+            torch.utils.data.DataLoader(
+                megadepth_train,
+                batch_size = batch_size,
+                sampler = mega_sampler,
+                num_workers = 8,
+            )
+        )
+        scannet_ws_sampler = torch.utils.data.WeightedRandomSampler(
+            scannet_ws, num_samples=batch_size * k, replacement=False
+        )
+        scannet_dataloader = iter(
+            torch.utils.data.DataLoader(
+                scannet_train,
+                batch_size=batch_size,
+                sampler=scannet_ws_sampler,
+                num_workers=gpus * 8,
+            )
+        )
+        for n_k in tqdm(range(n, n + 2 * k, 2),disable = romatch.RANK > 0):
+            train_k_steps(
+                n_k, 1, mega_dataloader, ddp_model, depth_loss_mega, optimizer, lr_scheduler, grad_scaler, grad_clip_norm = grad_clip_norm, progress_bar=False
+            )
+            train_k_steps(
+                n_k + 1, 1, scannet_dataloader, ddp_model, depth_loss_scannet, optimizer, lr_scheduler, grad_scaler, grad_clip_norm = grad_clip_norm, progress_bar=False
+            )
+        checkpointer.save(model, optimizer, lr_scheduler, romatch.GLOBAL_STEP)
+        wandb.log(megadense_benchmark.benchmark(model), step = romatch.GLOBAL_STEP)
+
+def test_scannet(model, name, resolution, sample_mode):
+    scannet_benchmark = ScanNetBenchmark("data/scannet")
+    scannet_results = scannet_benchmark.benchmark(model)
+    json.dump(scannet_results, open(f"results/scannet_{name}.json", "w"))
+
+if __name__ == "__main__":
+    import warnings
+    warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
+    warnings.filterwarnings('ignore')#, category=UserWarning)#, message='WARNING batched routines are designed for small sizes.')
+    os.environ["TORCH_CUDNN_V8_API_ENABLED"] = "1" # For BF16 computations
+    os.environ["OMP_NUM_THREADS"] = "16"
+    
+    import romatch
+    parser = ArgumentParser()
+    parser.add_argument("--test", action='store_true')
+    parser.add_argument("--debug_mode", action='store_true')
+    parser.add_argument("--dont_log_wandb", action='store_true')
+    parser.add_argument("--train_resolution", default='medium')
+    parser.add_argument("--gpu_batch_size", default=4, type=int)
+    parser.add_argument("--wandb_entity", required = False)
+
+    args, _ = parser.parse_known_args()
+    romatch.DEBUG_MODE = args.debug_mode
+    if not args.test:
+        train(args)
+    experiment_name = os.path.splitext(os.path.basename(__file__))[0]
+    checkpoint_dir = "workspace/"
+    checkpoint_name = checkpoint_dir + experiment_name + ".pth"
+    test_resolution = "medium"
+    sample_mode = "threshold_balanced"
+    symmetric = True
+    upsample_preds = False
+    attenuate_cert = True
+
+    model = get_model(pretrained_backbone=False, resolution = test_resolution, sample_mode = sample_mode, upsample_preds = upsample_preds, symmetric=symmetric, name=experiment_name, attenuate_cert = attenuate_cert)
+    model = model.cuda()
+    states = torch.load(checkpoint_name)
+    model.load_state_dict(states["model"])
+    test_scannet(model, experiment_name, resolution = test_resolution, sample_mode = sample_mode)

submodules/RoMa/experiments/train_roma_outdoor.py

@@ -0,0 +1,307 @@
+import os
+import torch
+from argparse import ArgumentParser
+
+from torch import nn
+from torch.utils.data import ConcatDataset
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+import json
+import wandb
+
+from romatch.benchmarks import MegadepthDenseBenchmark
+from romatch.datasets.megadepth import MegadepthBuilder
+from romatch.losses.robust_loss import RobustLosses
+from romatch.benchmarks import MegaDepthPoseEstimationBenchmark, MegadepthDenseBenchmark, HpatchesHomogBenchmark
+
+from romatch.train.train import train_k_steps
+from romatch.models.matcher import *
+from romatch.models.transformer import Block, TransformerDecoder, MemEffAttention
+from romatch.models.encoders import *
+from romatch.checkpointing import CheckPoint
+
+resolutions = {"low":(448, 448), "medium":(14*8*5, 14*8*5), "high":(14*8*6, 14*8*6)}
+
+def get_model(pretrained_backbone=True, resolution = "medium", **kwargs):
+    import warnings
+    warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
+    gp_dim = 512
+    feat_dim = 512
+    decoder_dim = gp_dim + feat_dim
+    cls_to_coord_res = 64
+    coordinate_decoder = TransformerDecoder(
+        nn.Sequential(*[Block(decoder_dim, 8, attn_class=MemEffAttention) for _ in range(5)]), 
+        decoder_dim, 
+        cls_to_coord_res**2 + 1,
+        is_classifier=True,
+        amp = True,
+        pos_enc = False,)
+    dw = True
+    hidden_blocks = 8
+    kernel_size = 5
+    displacement_emb = "linear"
+    disable_local_corr_grad = True
+    
+    conv_refiner = nn.ModuleDict(
+        {
+            "16": ConvRefiner(
+                2 * 512+128+(2*7+1)**2,
+                2 * 512+128+(2*7+1)**2,
+                2 + 1,
+                kernel_size=kernel_size,
+                dw=dw,
+                hidden_blocks=hidden_blocks,
+                displacement_emb=displacement_emb,
+                displacement_emb_dim=128,
+                local_corr_radius = 7,
+                corr_in_other = True,
+                amp = True,
+                disable_local_corr_grad = disable_local_corr_grad,
+                bn_momentum = 0.01,
+            ),
+            "8": ConvRefiner(
+                2 * 512+64+(2*3+1)**2,
+                2 * 512+64+(2*3+1)**2,
+                2 + 1,
+                kernel_size=kernel_size,
+                dw=dw,
+                hidden_blocks=hidden_blocks,
+                displacement_emb=displacement_emb,
+                displacement_emb_dim=64,
+                local_corr_radius = 3,
+                corr_in_other = True,
+                amp = True,
+                disable_local_corr_grad = disable_local_corr_grad,
+                bn_momentum = 0.01,
+            ),
+            "4": ConvRefiner(
+                2 * 256+32+(2*2+1)**2,
+                2 * 256+32+(2*2+1)**2,
+                2 + 1,
+                kernel_size=kernel_size,
+                dw=dw,
+                hidden_blocks=hidden_blocks,
+                displacement_emb=displacement_emb,
+                displacement_emb_dim=32,
+                local_corr_radius = 2,
+                corr_in_other = True,
+                amp = True,
+                disable_local_corr_grad = disable_local_corr_grad,
+                bn_momentum = 0.01,
+            ),
+            "2": ConvRefiner(
+                2 * 64+16,
+                128+16,
+                2 + 1,
+                kernel_size=kernel_size,
+                dw=dw,
+                hidden_blocks=hidden_blocks,
+                displacement_emb=displacement_emb,
+                displacement_emb_dim=16,
+                amp = True,
+                disable_local_corr_grad = disable_local_corr_grad,
+                bn_momentum = 0.01,
+            ),
+            "1": ConvRefiner(
+                2 * 9 + 6,
+                24,
+                2 + 1,
+                kernel_size=kernel_size,
+                dw=dw,
+                hidden_blocks = hidden_blocks,
+                displacement_emb = displacement_emb,
+                displacement_emb_dim = 6,
+                amp = True,
+                disable_local_corr_grad = disable_local_corr_grad,
+                bn_momentum = 0.01,
+            ),
+        }
+    )
+    kernel_temperature = 0.2
+    learn_temperature = False
+    no_cov = True
+    kernel = CosKernel
+    only_attention = False
+    basis = "fourier"
+    gp16 = GP(
+        kernel,
+        T=kernel_temperature,
+        learn_temperature=learn_temperature,
+        only_attention=only_attention,
+        gp_dim=gp_dim,
+        basis=basis,
+        no_cov=no_cov,
+    )
+    gps = nn.ModuleDict({"16": gp16})
+    proj16 = nn.Sequential(nn.Conv2d(1024, 512, 1, 1), nn.BatchNorm2d(512))
+    proj8 = nn.Sequential(nn.Conv2d(512, 512, 1, 1), nn.BatchNorm2d(512))
+    proj4 = nn.Sequential(nn.Conv2d(256, 256, 1, 1), nn.BatchNorm2d(256))
+    proj2 = nn.Sequential(nn.Conv2d(128, 64, 1, 1), nn.BatchNorm2d(64))
+    proj1 = nn.Sequential(nn.Conv2d(64, 9, 1, 1), nn.BatchNorm2d(9))
+    proj = nn.ModuleDict({
+        "16": proj16,
+        "8": proj8,
+        "4": proj4,
+        "2": proj2,
+        "1": proj1,
+        })
+    displacement_dropout_p = 0.0
+    gm_warp_dropout_p = 0.0
+    decoder = Decoder(coordinate_decoder, 
+                      gps, 
+                      proj, 
+                      conv_refiner, 
+                      detach=True, 
+                      scales=["16", "8", "4", "2", "1"], 
+                      displacement_dropout_p = displacement_dropout_p,
+                      gm_warp_dropout_p = gm_warp_dropout_p)
+    h,w = resolutions[resolution]
+    encoder = CNNandDinov2(
+        cnn_kwargs = dict(
+            pretrained=pretrained_backbone,
+            amp = True),
+        amp = True,
+        use_vgg = True,
+    )
+    matcher = RegressionMatcher(encoder, decoder, h=h, w=w,**kwargs)
+    return matcher
+
+def train(args):
+    dist.init_process_group('nccl')
+    #torch._dynamo.config.verbose=True
+    gpus = int(os.environ['WORLD_SIZE'])
+    # create model and move it to GPU with id rank
+    rank = dist.get_rank()
+    print(f"Start running DDP on rank {rank}")
+    device_id = rank % torch.cuda.device_count()
+    romatch.LOCAL_RANK = device_id
+    torch.cuda.set_device(device_id)
+    
+    resolution = args.train_resolution
+    wandb_log = not args.dont_log_wandb
+    experiment_name = os.path.splitext(os.path.basename(__file__))[0]
+    wandb_mode = "online" if wandb_log and rank == 0 else "disabled"
+    wandb.init(project="romatch", entity=args.wandb_entity, name=experiment_name, reinit=False, mode = wandb_mode)
+    checkpoint_dir = "workspace/checkpoints/"
+    h,w = resolutions[resolution]
+    model = get_model(pretrained_backbone=True, resolution=resolution, attenuate_cert = False).to(device_id)
+    # Num steps
+    global_step = 0
+    batch_size = args.gpu_batch_size
+    step_size = gpus*batch_size
+    romatch.STEP_SIZE = step_size
+    
+    N = (32 * 250000)  # 250k steps of batch size 32
+    # checkpoint every
+    k = 25000 // romatch.STEP_SIZE
+
+    # Data
+    mega = MegadepthBuilder(data_root="data/megadepth", loftr_ignore=True, imc21_ignore = True)
+    use_horizontal_flip_aug = True
+    rot_prob = 0
+    depth_interpolation_mode = "bilinear"
+    megadepth_train1 = mega.build_scenes(
+        split="train_loftr", min_overlap=0.01, shake_t=32, use_horizontal_flip_aug = use_horizontal_flip_aug, rot_prob = rot_prob,
+        ht=h,wt=w,
+    )
+    megadepth_train2 = mega.build_scenes(
+        split="train_loftr", min_overlap=0.35, shake_t=32, use_horizontal_flip_aug = use_horizontal_flip_aug, rot_prob = rot_prob,
+        ht=h,wt=w,
+    )
+    megadepth_train = ConcatDataset(megadepth_train1 + megadepth_train2)
+    mega_ws = mega.weight_scenes(megadepth_train, alpha=0.75)
+    # Loss and optimizer
+    depth_loss = RobustLosses(
+        ce_weight=0.01, 
+        local_dist={1:4, 2:4, 4:8, 8:8},
+        local_largest_scale=8,
+        depth_interpolation_mode=depth_interpolation_mode,
+        alpha = 0.5,
+        c = 1e-4,)
+    parameters = [
+        {"params": model.encoder.parameters(), "lr": romatch.STEP_SIZE * 5e-6 / 8},
+        {"params": model.decoder.parameters(), "lr": romatch.STEP_SIZE * 1e-4 / 8},
+    ]
+    optimizer = torch.optim.AdamW(parameters, weight_decay=0.01)
+    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+        optimizer, milestones=[(9*N/romatch.STEP_SIZE)//10])
+    megadense_benchmark = MegadepthDenseBenchmark("data/megadepth", num_samples = 1000, h=h,w=w)
+    checkpointer = CheckPoint(checkpoint_dir, experiment_name)
+    model, optimizer, lr_scheduler, global_step = checkpointer.load(model, optimizer, lr_scheduler, global_step)
+    romatch.GLOBAL_STEP = global_step
+    ddp_model = DDP(model, device_ids=[device_id], find_unused_parameters = False, gradient_as_bucket_view=True)
+    grad_scaler = torch.cuda.amp.GradScaler(growth_interval=1_000_000)
+    grad_clip_norm = 0.01
+    for n in range(romatch.GLOBAL_STEP, N, k * romatch.STEP_SIZE):
+        mega_sampler = torch.utils.data.WeightedRandomSampler(
+            mega_ws, num_samples = batch_size * k, replacement=False
+        )
+        mega_dataloader = iter(
+            torch.utils.data.DataLoader(
+                megadepth_train,
+                batch_size = batch_size,
+                sampler = mega_sampler,
+                num_workers = 8,
+            )
+        )
+        train_k_steps(
+            n, k, mega_dataloader, ddp_model, depth_loss, optimizer, lr_scheduler, grad_scaler, grad_clip_norm = grad_clip_norm,
+        )
+        checkpointer.save(model, optimizer, lr_scheduler, romatch.GLOBAL_STEP)
+        wandb.log(megadense_benchmark.benchmark(model), step = romatch.GLOBAL_STEP)
+
+def test_mega_8_scenes(model, name):
+    mega_8_scenes_benchmark = MegaDepthPoseEstimationBenchmark("data/megadepth",
+                                                scene_names=['mega_8_scenes_0019_0.1_0.3.npz',
+                                                    'mega_8_scenes_0025_0.1_0.3.npz',
+                                                    'mega_8_scenes_0021_0.1_0.3.npz',
+                                                    'mega_8_scenes_0008_0.1_0.3.npz',
+                                                    'mega_8_scenes_0032_0.1_0.3.npz',
+                                                    'mega_8_scenes_1589_0.1_0.3.npz',
+                                                    'mega_8_scenes_0063_0.1_0.3.npz',
+                                                    'mega_8_scenes_0024_0.1_0.3.npz',
+                                                    'mega_8_scenes_0019_0.3_0.5.npz',
+                                                    'mega_8_scenes_0025_0.3_0.5.npz',
+                                                    'mega_8_scenes_0021_0.3_0.5.npz',
+                                                    'mega_8_scenes_0008_0.3_0.5.npz',
+                                                    'mega_8_scenes_0032_0.3_0.5.npz',
+                                                    'mega_8_scenes_1589_0.3_0.5.npz',
+                                                    'mega_8_scenes_0063_0.3_0.5.npz',
+                                                    'mega_8_scenes_0024_0.3_0.5.npz'])
+    mega_8_scenes_results = mega_8_scenes_benchmark.benchmark(model, model_name=name)
+    print(mega_8_scenes_results)
+    json.dump(mega_8_scenes_results, open(f"results/mega_8_scenes_{name}.json", "w"))
+
+def test_mega1500(model, name):
+    mega1500_benchmark = MegaDepthPoseEstimationBenchmark("data/megadepth")
+    mega1500_results = mega1500_benchmark.benchmark(model, model_name=name)
+    json.dump(mega1500_results, open(f"results/mega1500_{name}.json", "w"))
+
+def test_mega_dense(model, name):
+    megadense_benchmark = MegadepthDenseBenchmark("data/megadepth", num_samples = 1000)
+    megadense_results = megadense_benchmark.benchmark(model)
+    json.dump(megadense_results, open(f"results/mega_dense_{name}.json", "w"))
+    
+def test_hpatches(model, name):
+    hpatches_benchmark = HpatchesHomogBenchmark("data/hpatches")
+    hpatches_results = hpatches_benchmark.benchmark(model)
+    json.dump(hpatches_results, open(f"results/hpatches_{name}.json", "w"))
+
+
+if __name__ == "__main__":
+    os.environ["TORCH_CUDNN_V8_API_ENABLED"] = "1" # For BF16 computations
+    os.environ["OMP_NUM_THREADS"] = "16"
+    torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
+    import romatch
+    parser = ArgumentParser()
+    parser.add_argument("--only_test", action='store_true')
+    parser.add_argument("--debug_mode", action='store_true')
+    parser.add_argument("--dont_log_wandb", action='store_true')
+    parser.add_argument("--train_resolution", default='medium')
+    parser.add_argument("--gpu_batch_size", default=8, type=int)
+    parser.add_argument("--wandb_entity", required = False)
+
+    args, _ = parser.parse_known_args()
+    romatch.DEBUG_MODE = args.debug_mode
+    if not args.only_test:
+        train(args)

submodules/RoMa/experiments/train_tiny_roma_v1_outdoor.py

@@ -0,0 +1,498 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import os
+import torch
+from argparse import ArgumentParser
+from pathlib import Path
+import math
+import numpy as np
+
+from torch import nn
+from torch.utils.data import ConcatDataset
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+import json
+import wandb
+from PIL import Image
+from torchvision.transforms import ToTensor
+
+from romatch.benchmarks import MegadepthDenseBenchmark, ScanNetBenchmark
+from romatch.benchmarks import Mega1500PoseLibBenchmark, ScanNetPoselibBenchmark
+from romatch.datasets.megadepth import MegadepthBuilder
+from romatch.losses.robust_loss_tiny_roma import RobustLosses
+from romatch.benchmarks import MegaDepthPoseEstimationBenchmark, MegadepthDenseBenchmark, HpatchesHomogBenchmark
+from romatch.train.train import train_k_steps
+from romatch.checkpointing import CheckPoint
+
+resolutions = {"low":(448, 448), "medium":(14*8*5, 14*8*5), "high":(14*8*6, 14*8*6), "xfeat": (600,800), "big": (768, 1024)}
+
+def kde(x, std = 0.1):
+    # use a gaussian kernel to estimate density
+    x = x.half() # Do it in half precision TODO: remove hardcoding
+    scores = (-torch.cdist(x,x)**2/(2*std**2)).exp()
+    density = scores.sum(dim=-1)
+    return density
+
+class BasicLayer(nn.Module):
+    """
+        Basic Convolutional Layer: Conv2d -> BatchNorm -> ReLU
+    """
+    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, bias=False, relu = True):
+        super().__init__()
+        self.layer = nn.Sequential(
+                                        nn.Conv2d( in_channels, out_channels, kernel_size, padding = padding, stride=stride, dilation=dilation, bias = bias),
+                                        nn.BatchNorm2d(out_channels, affine=False),
+                                        nn.ReLU(inplace = True) if relu else nn.Identity()
+                                    )
+
+    def forward(self, x):
+        return self.layer(x)
+
+class XFeatModel(nn.Module):
+    """
+        Implementation of architecture described in 
+        "XFeat: Accelerated Features for Lightweight Image Matching, CVPR 2024."
+    """
+
+    def __init__(self, xfeat = None, 
+                 freeze_xfeat = True, 
+                 sample_mode = "threshold_balanced", 
+                 symmetric = False, 
+                 exact_softmax = False):
+        super().__init__()
+        if xfeat is None:
+            xfeat = torch.hub.load('verlab/accelerated_features', 'XFeat', pretrained = True, top_k = 4096).net
+            del xfeat.heatmap_head, xfeat.keypoint_head, xfeat.fine_matcher
+        if freeze_xfeat:
+            xfeat.train(False)
+            self.xfeat = [xfeat]# hide params from ddp
+        else:
+            self.xfeat = nn.ModuleList([xfeat])
+        self.freeze_xfeat = freeze_xfeat
+        match_dim = 256
+        self.coarse_matcher = nn.Sequential(
+            BasicLayer(64+64+2, match_dim,),
+            BasicLayer(match_dim, match_dim,), 
+            BasicLayer(match_dim, match_dim,), 
+            BasicLayer(match_dim, match_dim,), 
+            nn.Conv2d(match_dim, 3, kernel_size=1, bias=True, padding=0))
+        fine_match_dim = 64
+        self.fine_matcher = nn.Sequential(
+            BasicLayer(24+24+2, fine_match_dim,),
+            BasicLayer(fine_match_dim, fine_match_dim,), 
+            BasicLayer(fine_match_dim, fine_match_dim,), 
+            BasicLayer(fine_match_dim, fine_match_dim,), 
+            nn.Conv2d(fine_match_dim, 3, kernel_size=1, bias=True, padding=0),)
+        self.sample_mode = sample_mode
+        self.sample_thresh = 0.2
+        self.symmetric = symmetric
+        self.exact_softmax = exact_softmax
+    
+    @property
+    def device(self):
+        return self.fine_matcher[-1].weight.device
+    
+    def preprocess_tensor(self, x):
+        """ Guarantee that image is divisible by 32 to avoid aliasing artifacts. """
+        H, W = x.shape[-2:]
+        _H, _W = (H//32) * 32, (W//32) * 32
+        rh, rw = H/_H, W/_W
+
+        x = F.interpolate(x, (_H, _W), mode='bilinear', align_corners=False)
+        return x, rh, rw        
+    
+    def forward_single(self, x):
+        with torch.inference_mode(self.freeze_xfeat or not self.training):
+            xfeat = self.xfeat[0]
+            with torch.no_grad():
+                x = x.mean(dim=1, keepdim = True)
+                x = xfeat.norm(x)
+
+            #main backbone
+            x1 = xfeat.block1(x)
+            x2 = xfeat.block2(x1 + xfeat.skip1(x))
+            x3 = xfeat.block3(x2)
+            x4 = xfeat.block4(x3)
+            x5 = xfeat.block5(x4)
+            x4 = F.interpolate(x4, (x3.shape[-2], x3.shape[-1]), mode='bilinear')
+            x5 = F.interpolate(x5, (x3.shape[-2], x3.shape[-1]), mode='bilinear')
+            feats = xfeat.block_fusion( x3 + x4 + x5 )
+        if self.freeze_xfeat:
+            return x2.clone(), feats.clone()
+        return x2, feats
+
+    def to_pixel_coordinates(self, coords, H_A, W_A, H_B = None, W_B = None):
+        if coords.shape[-1] == 2:
+            return self._to_pixel_coordinates(coords, H_A, W_A) 
+        
+        if isinstance(coords, (list, tuple)):
+            kpts_A, kpts_B = coords[0], coords[1]
+        else:
+            kpts_A, kpts_B = coords[...,:2], coords[...,2:]
+        return self._to_pixel_coordinates(kpts_A, H_A, W_A), self._to_pixel_coordinates(kpts_B, H_B, W_B)
+
+    def _to_pixel_coordinates(self, coords, H, W):
+        kpts = torch.stack((W/2 * (coords[...,0]+1), H/2 * (coords[...,1]+1)),axis=-1)
+        return kpts
+    
+    def pos_embed(self, corr_volume: torch.Tensor):
+        B, H1, W1, H0, W0 = corr_volume.shape 
+        grid = torch.stack(
+                torch.meshgrid(
+                    torch.linspace(-1+1/W1,1-1/W1, W1), 
+                    torch.linspace(-1+1/H1,1-1/H1, H1), 
+                    indexing = "xy"), 
+                dim = -1).float().to(corr_volume).reshape(H1*W1, 2)
+        down = 4
+        if not self.training and not self.exact_softmax:
+            grid_lr = torch.stack(
+                torch.meshgrid(
+                    torch.linspace(-1+down/W1,1-down/W1, W1//down), 
+                    torch.linspace(-1+down/H1,1-down/H1, H1//down), 
+                    indexing = "xy"), 
+                dim = -1).float().to(corr_volume).reshape(H1*W1 //down**2, 2)
+            cv = corr_volume
+            best_match = cv.reshape(B,H1*W1,H0,W0).amax(dim=1) # B, HW, H, W
+            P_lowres = torch.cat((cv[:,::down,::down].reshape(B,H1*W1 // down**2,H0,W0), best_match[:,None]),dim=1).softmax(dim=1)
+            pos_embeddings = torch.einsum('bchw,cd->bdhw', P_lowres[:,:-1], grid_lr)
+            pos_embeddings += P_lowres[:,-1] * grid[best_match].permute(0,3,1,2)
+        else:
+            P = corr_volume.reshape(B,H1*W1,H0,W0).softmax(dim=1) # B, HW, H, W
+            pos_embeddings = torch.einsum('bchw,cd->bdhw', P, grid)
+        return pos_embeddings
+    
+    def visualize_warp(self, warp, certainty, im_A = None, im_B = None, 
+                       im_A_path = None, im_B_path = None, symmetric = True, save_path = None, unnormalize = False):
+        device = warp.device
+        H,W2,_ = warp.shape
+        W = W2//2 if symmetric else W2
+        if im_A is None:
+            from PIL import Image
+            im_A, im_B = Image.open(im_A_path).convert("RGB"), Image.open(im_B_path).convert("RGB")
+        if not isinstance(im_A, torch.Tensor):
+            im_A = im_A.resize((W,H))
+            im_B = im_B.resize((W,H))    
+            x_B = (torch.tensor(np.array(im_B)) / 255).to(device).permute(2, 0, 1)
+            if symmetric:
+                x_A = (torch.tensor(np.array(im_A)) / 255).to(device).permute(2, 0, 1)
+        else:
+            if symmetric:
+                x_A = im_A
+            x_B = im_B
+        im_A_transfer_rgb = F.grid_sample(
+        x_B[None], warp[:,:W, 2:][None], mode="bilinear", align_corners=False
+        )[0]
+        if symmetric:
+            im_B_transfer_rgb = F.grid_sample(
+            x_A[None], warp[:, W:, :2][None], mode="bilinear", align_corners=False
+            )[0]
+            warp_im = torch.cat((im_A_transfer_rgb,im_B_transfer_rgb),dim=2)
+            white_im = torch.ones((H,2*W),device=device)
+        else:
+            warp_im = im_A_transfer_rgb
+            white_im = torch.ones((H, W), device = device)
+        vis_im = certainty * warp_im + (1 - certainty) * white_im
+        if save_path is not None:
+            from romatch.utils import tensor_to_pil
+            tensor_to_pil(vis_im, unnormalize=unnormalize).save(save_path)
+        return vis_im
+     
+    def corr_volume(self, feat0, feat1):
+        """
+            input:
+                feat0 -> torch.Tensor(B, C, H, W)
+                feat1 -> torch.Tensor(B, C, H, W)
+            return:
+                corr_volume -> torch.Tensor(B, H, W, H, W)
+        """
+        B, C, H0, W0 = feat0.shape
+        B, C, H1, W1 = feat1.shape
+        feat0 = feat0.view(B, C, H0*W0)
+        feat1 = feat1.view(B, C, H1*W1)
+        corr_volume = torch.einsum('bci,bcj->bji', feat0, feat1).reshape(B, H1, W1, H0 , W0)/math.sqrt(C) #16*16*16
+        return corr_volume
+    
+    @torch.inference_mode()
+    def match_from_path(self, im0_path, im1_path):
+        device = self.device
+        im0 = ToTensor()(Image.open(im0_path))[None].to(device)
+        im1 = ToTensor()(Image.open(im1_path))[None].to(device)
+        return self.match(im0, im1, batched = False)
+    
+    @torch.inference_mode()
+    def match(self, im0, im1, *args, batched = True):
+        # stupid
+        if isinstance(im0, (str, Path)):
+            return self.match_from_path(im0, im1)
+        elif isinstance(im0, Image.Image):
+            batched = False
+            device = self.device
+            im0 = ToTensor()(im0)[None].to(device)
+            im1 = ToTensor()(im1)[None].to(device)
+ 
+        B,C,H0,W0 = im0.shape
+        B,C,H1,W1 = im1.shape
+        self.train(False)
+        corresps = self.forward({"im_A":im0, "im_B":im1})
+        #return 1,1
+        flow = F.interpolate(
+            corresps[4]["flow"], 
+            size = (H0, W0), 
+            mode = "bilinear", align_corners = False).permute(0,2,3,1).reshape(B,H0,W0,2)
+        grid = torch.stack(
+            torch.meshgrid(
+                torch.linspace(-1+1/W0,1-1/W0, W0), 
+                torch.linspace(-1+1/H0,1-1/H0, H0), 
+                indexing = "xy"), 
+            dim = -1).float().to(flow.device).expand(B, H0, W0, 2)
+        
+        certainty = F.interpolate(corresps[4]["certainty"], size = (H0,W0), mode = "bilinear", align_corners = False)
+        warp, cert = torch.cat((grid, flow), dim = -1), certainty[:,0].sigmoid()
+        if batched:
+            return warp, cert
+        else:
+            return warp[0], cert[0]
+
+    def sample(
+        self,
+        matches,
+        certainty,
+        num=10000,
+    ):
+        if "threshold" in self.sample_mode:
+            upper_thresh = self.sample_thresh
+            certainty = certainty.clone()
+            certainty[certainty > upper_thresh] = 1
+        matches, certainty = (
+            matches.reshape(-1, 4),
+            certainty.reshape(-1),
+        )
+        expansion_factor = 4 if "balanced" in self.sample_mode else 1
+        good_samples = torch.multinomial(certainty, 
+                          num_samples = min(expansion_factor*num, len(certainty)), 
+                          replacement=False)
+        good_matches, good_certainty = matches[good_samples], certainty[good_samples]
+        if "balanced" not in self.sample_mode:
+            return good_matches, good_certainty
+        density = kde(good_matches, std=0.1)
+        p = 1 / (density+1)
+        p[density < 10] = 1e-7 # Basically should have at least 10 perfect neighbours, or around 100 ok ones
+        balanced_samples = torch.multinomial(p, 
+                          num_samples = min(num,len(good_certainty)), 
+                          replacement=False)
+        return good_matches[balanced_samples], good_certainty[balanced_samples]
+            
+    def forward(self, batch):
+        """
+            input:
+                x -> torch.Tensor(B, C, H, W) grayscale or rgb images
+            return:
+
+        """
+        im0 = batch["im_A"]
+        im1 = batch["im_B"]
+        corresps = {}
+        im0, rh0, rw0 = self.preprocess_tensor(im0)
+        im1, rh1, rw1 = self.preprocess_tensor(im1)
+        B, C, H0, W0 = im0.shape
+        B, C, H1, W1 = im1.shape
+        to_normalized = torch.tensor((2/W1, 2/H1, 1)).to(im0.device)[None,:,None,None]
+ 
+        if im0.shape[-2:] == im1.shape[-2:]:
+            x = torch.cat([im0, im1], dim=0)
+            x = self.forward_single(x)
+            feats_x0_c, feats_x1_c = x[1].chunk(2)
+            feats_x0_f, feats_x1_f = x[0].chunk(2)
+        else:
+            feats_x0_f, feats_x0_c = self.forward_single(im0)
+            feats_x1_f, feats_x1_c = self.forward_single(im1)
+        corr_volume = self.corr_volume(feats_x0_c, feats_x1_c)
+        coarse_warp = self.pos_embed(corr_volume)
+        coarse_matches = torch.cat((coarse_warp, torch.zeros_like(coarse_warp[:,-1:])), dim=1)
+        feats_x1_c_warped = F.grid_sample(feats_x1_c, coarse_matches.permute(0, 2, 3, 1)[...,:2], mode = 'bilinear', align_corners = False)
+        coarse_matches_delta = self.coarse_matcher(torch.cat((feats_x0_c, feats_x1_c_warped, coarse_warp), dim=1))
+        coarse_matches = coarse_matches + coarse_matches_delta * to_normalized
+        corresps[8] = {"flow": coarse_matches[:,:2], "certainty": coarse_matches[:,2:]}
+        coarse_matches_up = F.interpolate(coarse_matches, size = feats_x0_f.shape[-2:], mode = "bilinear", align_corners = False)        
+        coarse_matches_up_detach = coarse_matches_up.detach()#note the detach
+        feats_x1_f_warped = F.grid_sample(feats_x1_f, coarse_matches_up_detach.permute(0, 2, 3, 1)[...,:2], mode = 'bilinear', align_corners = False)
+        fine_matches_delta = self.fine_matcher(torch.cat((feats_x0_f, feats_x1_f_warped, coarse_matches_up_detach[:,:2]), dim=1))
+        fine_matches = coarse_matches_up_detach+fine_matches_delta * to_normalized
+        corresps[4] = {"flow": fine_matches[:,:2], "certainty": fine_matches[:,2:]}
+        return corresps
+    
+
+
+
+
+def train(args):
+    rank = 0
+    gpus = 1
+    device_id = rank % torch.cuda.device_count()
+    romatch.LOCAL_RANK = 0
+    torch.cuda.set_device(device_id)
+        
+    resolution = "big"
+    wandb_log = not args.dont_log_wandb
+    experiment_name = Path(__file__).stem
+    wandb_mode = "online" if wandb_log and rank == 0 else "disabled"
+    wandb.init(project="romatch", entity=args.wandb_entity, name=experiment_name, reinit=False, mode = wandb_mode)
+    checkpoint_dir = "workspace/checkpoints/"
+    h,w = resolutions[resolution]
+    model = XFeatModel(freeze_xfeat = False).to(device_id)
+    # Num steps
+    global_step = 0
+    batch_size = args.gpu_batch_size
+    step_size = gpus*batch_size
+    romatch.STEP_SIZE = step_size
+    
+    N = 2_000_000  # 2M pairs
+    # checkpoint every
+    k = 25000 // romatch.STEP_SIZE
+
+    # Data
+    mega = MegadepthBuilder(data_root="data/megadepth", loftr_ignore=True, imc21_ignore = True)
+    use_horizontal_flip_aug = True
+    normalize = False # don't imgnet normalize
+    rot_prob = 0
+    depth_interpolation_mode = "bilinear"
+    megadepth_train1 = mega.build_scenes(
+        split="train_loftr", min_overlap=0.01, shake_t=32, use_horizontal_flip_aug = use_horizontal_flip_aug, rot_prob = rot_prob,
+        ht=h,wt=w, normalize = normalize
+    )
+    megadepth_train2 = mega.build_scenes(
+        split="train_loftr", min_overlap=0.35, shake_t=32, use_horizontal_flip_aug = use_horizontal_flip_aug, rot_prob = rot_prob,
+        ht=h,wt=w, normalize = normalize
+    )
+    megadepth_train = ConcatDataset(megadepth_train1 + megadepth_train2)
+    mega_ws = mega.weight_scenes(megadepth_train, alpha=0.75)
+    # Loss and optimizer
+    depth_loss = RobustLosses(
+        ce_weight=0.01, 
+        local_dist={4:4},
+        depth_interpolation_mode=depth_interpolation_mode,
+        alpha = {4:0.15, 8:0.15},
+        c = 1e-4,
+        epe_mask_prob_th = 0.001,
+        )
+    parameters = [
+        {"params": model.parameters(), "lr": romatch.STEP_SIZE * 1e-4 / 8},
+    ]
+    optimizer = torch.optim.AdamW(parameters, weight_decay=0.01)
+    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+        optimizer, milestones=[(9*N/romatch.STEP_SIZE)//10])
+    #megadense_benchmark = MegadepthDenseBenchmark("data/megadepth", num_samples = 1000, h=h,w=w)
+    mega1500_benchmark = Mega1500PoseLibBenchmark("data/megadepth", num_ransac_iter = 1, test_every = 30)
+
+    checkpointer = CheckPoint(checkpoint_dir, experiment_name)
+    model, optimizer, lr_scheduler, global_step = checkpointer.load(model, optimizer, lr_scheduler, global_step)
+    romatch.GLOBAL_STEP = global_step
+    grad_scaler = torch.cuda.amp.GradScaler(growth_interval=1_000_000)
+    grad_clip_norm = 0.01
+    #megadense_benchmark.benchmark(model)
+    for n in range(romatch.GLOBAL_STEP, N, k * romatch.STEP_SIZE):
+        mega_sampler = torch.utils.data.WeightedRandomSampler(
+            mega_ws, num_samples = batch_size * k, replacement=False
+        )
+        mega_dataloader = iter(
+            torch.utils.data.DataLoader(
+                megadepth_train,
+                batch_size = batch_size,
+                sampler = mega_sampler,
+                num_workers = 8,
+            )
+        )
+        train_k_steps(
+            n, k, mega_dataloader, model, depth_loss, optimizer, lr_scheduler, grad_scaler, grad_clip_norm = grad_clip_norm,
+        )
+        checkpointer.save(model, optimizer, lr_scheduler, romatch.GLOBAL_STEP)
+        wandb.log(mega1500_benchmark.benchmark(model, model_name=experiment_name), step = romatch.GLOBAL_STEP)
+
+def test_mega_8_scenes(model, name):
+    mega_8_scenes_benchmark = MegaDepthPoseEstimationBenchmark("data/megadepth",
+                                                scene_names=['mega_8_scenes_0019_0.1_0.3.npz',
+                                                    'mega_8_scenes_0025_0.1_0.3.npz',
+                                                    'mega_8_scenes_0021_0.1_0.3.npz',
+                                                    'mega_8_scenes_0008_0.1_0.3.npz',
+                                                    'mega_8_scenes_0032_0.1_0.3.npz',
+                                                    'mega_8_scenes_1589_0.1_0.3.npz',
+                                                    'mega_8_scenes_0063_0.1_0.3.npz',
+                                                    'mega_8_scenes_0024_0.1_0.3.npz',
+                                                    'mega_8_scenes_0019_0.3_0.5.npz',
+                                                    'mega_8_scenes_0025_0.3_0.5.npz',
+                                                    'mega_8_scenes_0021_0.3_0.5.npz',
+                                                    'mega_8_scenes_0008_0.3_0.5.npz',
+                                                    'mega_8_scenes_0032_0.3_0.5.npz',
+                                                    'mega_8_scenes_1589_0.3_0.5.npz',
+                                                    'mega_8_scenes_0063_0.3_0.5.npz',
+                                                    'mega_8_scenes_0024_0.3_0.5.npz'])
+    mega_8_scenes_results = mega_8_scenes_benchmark.benchmark(model, model_name=name)
+    print(mega_8_scenes_results)
+    json.dump(mega_8_scenes_results, open(f"results/mega_8_scenes_{name}.json", "w"))
+
+def test_mega1500(model, name):
+    mega1500_benchmark = MegaDepthPoseEstimationBenchmark("data/megadepth")
+    mega1500_results = mega1500_benchmark.benchmark(model, model_name=name)
+    json.dump(mega1500_results, open(f"results/mega1500_{name}.json", "w"))
+
+def test_mega1500_poselib(model, name):
+    mega1500_benchmark = Mega1500PoseLibBenchmark("data/megadepth", num_ransac_iter = 1, test_every = 1)
+    mega1500_results = mega1500_benchmark.benchmark(model, model_name=name)
+    json.dump(mega1500_results, open(f"results/mega1500_poselib_{name}.json", "w"))
+
+def test_mega_8_scenes_poselib(model, name):
+    mega1500_benchmark = Mega1500PoseLibBenchmark("data/megadepth", num_ransac_iter = 1, test_every = 1,
+                                                  scene_names=['mega_8_scenes_0019_0.1_0.3.npz',
+                                                    'mega_8_scenes_0025_0.1_0.3.npz',
+                                                    'mega_8_scenes_0021_0.1_0.3.npz',
+                                                    'mega_8_scenes_0008_0.1_0.3.npz',
+                                                    'mega_8_scenes_0032_0.1_0.3.npz',
+                                                    'mega_8_scenes_1589_0.1_0.3.npz',
+                                                    'mega_8_scenes_0063_0.1_0.3.npz',
+                                                    'mega_8_scenes_0024_0.1_0.3.npz',
+                                                    'mega_8_scenes_0019_0.3_0.5.npz',
+                                                    'mega_8_scenes_0025_0.3_0.5.npz',
+                                                    'mega_8_scenes_0021_0.3_0.5.npz',
+                                                    'mega_8_scenes_0008_0.3_0.5.npz',
+                                                    'mega_8_scenes_0032_0.3_0.5.npz',
+                                                    'mega_8_scenes_1589_0.3_0.5.npz',
+                                                    'mega_8_scenes_0063_0.3_0.5.npz',
+                                                    'mega_8_scenes_0024_0.3_0.5.npz'])
+    mega1500_results = mega1500_benchmark.benchmark(model, model_name=name)
+    json.dump(mega1500_results, open(f"results/mega_8_scenes_poselib_{name}.json", "w"))
+
+def test_scannet_poselib(model, name):
+    scannet_benchmark = ScanNetPoselibBenchmark("data/scannet")
+    scannet_results = scannet_benchmark.benchmark(model)
+    json.dump(scannet_results, open(f"results/scannet_{name}.json", "w"))
+
+def test_scannet(model, name):
+    scannet_benchmark = ScanNetBenchmark("data/scannet")
+    scannet_results = scannet_benchmark.benchmark(model)
+    json.dump(scannet_results, open(f"results/scannet_{name}.json", "w"))
+
+if __name__ == "__main__":
+    os.environ["TORCH_CUDNN_V8_API_ENABLED"] = "1" # For BF16 computations
+    os.environ["OMP_NUM_THREADS"] = "16"
+    torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
+    import romatch
+    parser = ArgumentParser()
+    parser.add_argument("--only_test", action='store_true')
+    parser.add_argument("--debug_mode", action='store_true')
+    parser.add_argument("--dont_log_wandb", action='store_true')
+    parser.add_argument("--train_resolution", default='medium')
+    parser.add_argument("--gpu_batch_size", default=8, type=int)
+    parser.add_argument("--wandb_entity", required = False)
+
+    args, _ = parser.parse_known_args()
+    romatch.DEBUG_MODE = args.debug_mode
+    if not args.only_test:
+        train(args)
+
+    experiment_name = "tiny_roma_v1_outdoor"#Path(__file__).stem
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = XFeatModel(freeze_xfeat=False, exact_softmax=False).to(device)
+    model.load_state_dict(torch.load(f"{experiment_name}.pth"))
+    test_mega1500_poselib(model, experiment_name)
+    

submodules/RoMa/requirements.txt

@@ -0,0 +1,14 @@
+torch
+einops
+torchvision
+opencv-python
+kornia
+albumentations
+loguru
+tqdm
+matplotlib
+h5py
+wandb
+timm
+poselib
+#xformers # Optional, used for memefficient attention

submodules/RoMa/romatch/__init__.py

@@ -0,0 +1,8 @@
+import os
+from .models import roma_outdoor, tiny_roma_v1_outdoor, roma_indoor
+
+DEBUG_MODE = False
+RANK = int(os.environ.get('RANK', default = 0))
+GLOBAL_STEP = 0
+STEP_SIZE = 1
+LOCAL_RANK = -1

submodules/RoMa/romatch/benchmarks/__init__.py

@@ -0,0 +1,6 @@
+from .hpatches_sequences_homog_benchmark import HpatchesHomogBenchmark
+from .scannet_benchmark import ScanNetBenchmark
+from .megadepth_pose_estimation_benchmark import MegaDepthPoseEstimationBenchmark
+from .megadepth_dense_benchmark import MegadepthDenseBenchmark
+from .megadepth_pose_estimation_benchmark_poselib import Mega1500PoseLibBenchmark
+#from .scannet_benchmark_poselib import ScanNetPoselibBenchmark

submodules/RoMa/romatch/benchmarks/hpatches_sequences_homog_benchmark.py

@@ -0,0 +1,113 @@
+from PIL import Image
+import numpy as np
+
+import os
+
+from tqdm import tqdm
+from romatch.utils import pose_auc
+import cv2
+
+
+class HpatchesHomogBenchmark:
+    """Hpatches grid goes from [0,n-1] instead of [0.5,n-0.5]"""
+
+    def __init__(self, dataset_path) -> None:
+        seqs_dir = "hpatches-sequences-release"
+        self.seqs_path = os.path.join(dataset_path, seqs_dir)
+        self.seq_names = sorted(os.listdir(self.seqs_path))
+        # Ignore seqs is same as LoFTR.
+        self.ignore_seqs = set(
+            [
+                "i_contruction",
+                "i_crownnight",
+                "i_dc",
+                "i_pencils",
+                "i_whitebuilding",
+                "v_artisans",
+                "v_astronautis",
+                "v_talent",
+            ]
+        )
+
+    def convert_coordinates(self, im_A_coords, im_A_to_im_B, wq, hq, wsup, hsup):
+        offset = 0.5  # Hpatches assumes that the center of the top-left pixel is at [0,0] (I think)
+        im_A_coords = (
+            np.stack(
+                (
+                    wq * (im_A_coords[..., 0] + 1) / 2,
+                    hq * (im_A_coords[..., 1] + 1) / 2,
+                ),
+                axis=-1,
+            )
+            - offset
+        )
+        im_A_to_im_B = (
+            np.stack(
+                (
+                    wsup * (im_A_to_im_B[..., 0] + 1) / 2,
+                    hsup * (im_A_to_im_B[..., 1] + 1) / 2,
+                ),
+                axis=-1,
+            )
+            - offset
+        )
+        return im_A_coords, im_A_to_im_B
+
+    def benchmark(self, model, model_name = None):
+        n_matches = []
+        homog_dists = []
+        for seq_idx, seq_name in tqdm(
+            enumerate(self.seq_names), total=len(self.seq_names)
+        ):
+            im_A_path = os.path.join(self.seqs_path, seq_name, "1.ppm")
+            im_A = Image.open(im_A_path)
+            w1, h1 = im_A.size
+            for im_idx in range(2, 7):
+                im_B_path = os.path.join(self.seqs_path, seq_name, f"{im_idx}.ppm")
+                im_B = Image.open(im_B_path)
+                w2, h2 = im_B.size
+                H = np.loadtxt(
+                    os.path.join(self.seqs_path, seq_name, "H_1_" + str(im_idx))
+                )
+                dense_matches, dense_certainty = model.match(
+                    im_A_path, im_B_path
+                )
+                good_matches, _ = model.sample(dense_matches, dense_certainty, 5000)
+                pos_a, pos_b = self.convert_coordinates(
+                    good_matches[:, :2], good_matches[:, 2:], w1, h1, w2, h2
+                )
+                try:
+                    H_pred, inliers = cv2.findHomography(
+                        pos_a,
+                        pos_b,
+                        method = cv2.RANSAC,
+                        confidence = 0.99999,
+                        ransacReprojThreshold = 3 * min(w2, h2) / 480,
+                    )
+                except:
+                    H_pred = None
+                if H_pred is None:
+                    H_pred = np.zeros((3, 3))
+                    H_pred[2, 2] = 1.0
+                corners = np.array(
+                    [[0, 0, 1], [0, h1 - 1, 1], [w1 - 1, 0, 1], [w1 - 1, h1 - 1, 1]]
+                )
+                real_warped_corners = np.dot(corners, np.transpose(H))
+                real_warped_corners = (
+                    real_warped_corners[:, :2] / real_warped_corners[:, 2:]
+                )
+                warped_corners = np.dot(corners, np.transpose(H_pred))
+                warped_corners = warped_corners[:, :2] / warped_corners[:, 2:]
+                mean_dist = np.mean(
+                    np.linalg.norm(real_warped_corners - warped_corners, axis=1)
+                ) / (min(w2, h2) / 480.0)
+                homog_dists.append(mean_dist)
+
+        n_matches = np.array(n_matches)
+        thresholds = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+        auc = pose_auc(np.array(homog_dists), thresholds)
+        return {
+            "hpatches_homog_auc_3": auc[2],
+            "hpatches_homog_auc_5": auc[4],
+            "hpatches_homog_auc_10": auc[9],
+        }

submodules/RoMa/romatch/benchmarks/megadepth_dense_benchmark.py

@@ -0,0 +1,106 @@
+import torch
+import numpy as np
+import tqdm
+from romatch.datasets import MegadepthBuilder
+from romatch.utils import warp_kpts
+from torch.utils.data import ConcatDataset
+import romatch
+
+class MegadepthDenseBenchmark:
+    def __init__(self, data_root="data/megadepth", h = 384, w = 512, num_samples = 2000) -> None:
+        mega = MegadepthBuilder(data_root=data_root)
+        self.dataset = ConcatDataset(
+            mega.build_scenes(split="test_loftr", ht=h, wt=w)
+        )  # fixed resolution of 384,512
+        self.num_samples = num_samples
+
+    def geometric_dist(self, depth1, depth2, T_1to2, K1, K2, dense_matches):
+        b, h1, w1, d = dense_matches.shape
+        with torch.no_grad():
+            x1 = dense_matches[..., :2].reshape(b, h1 * w1, 2)
+            mask, x2 = warp_kpts(
+                x1.double(),
+                depth1.double(),
+                depth2.double(),
+                T_1to2.double(),
+                K1.double(),
+                K2.double(),
+            )
+            x2 = torch.stack(
+                (w1 * (x2[..., 0] + 1) / 2, h1 * (x2[..., 1] + 1) / 2), dim=-1
+            )
+            prob = mask.float().reshape(b, h1, w1)
+        x2_hat = dense_matches[..., 2:]
+        x2_hat = torch.stack(
+            (w1 * (x2_hat[..., 0] + 1) / 2, h1 * (x2_hat[..., 1] + 1) / 2), dim=-1
+        )
+        gd = (x2_hat - x2.reshape(b, h1, w1, 2)).norm(dim=-1)
+        gd = gd[prob == 1]
+        pck_1 = (gd < 1.0).float().mean()
+        pck_3 = (gd < 3.0).float().mean()
+        pck_5 = (gd < 5.0).float().mean()
+        return gd, pck_1, pck_3, pck_5, prob
+
+    def benchmark(self, model, batch_size=8):
+        model.train(False)
+        with torch.no_grad():
+            gd_tot = 0.0
+            pck_1_tot = 0.0
+            pck_3_tot = 0.0
+            pck_5_tot = 0.0
+            sampler = torch.utils.data.WeightedRandomSampler(
+                torch.ones(len(self.dataset)), replacement=False, num_samples=self.num_samples
+            )
+            B = batch_size
+            dataloader = torch.utils.data.DataLoader(
+                self.dataset, batch_size=B, num_workers=batch_size, sampler=sampler
+            )
+            for idx, data in tqdm.tqdm(enumerate(dataloader), disable = romatch.RANK > 0):
+                im_A, im_B, depth1, depth2, T_1to2, K1, K2 = (
+                    data["im_A"].cuda(),
+                    data["im_B"].cuda(),
+                    data["im_A_depth"].cuda(),
+                    data["im_B_depth"].cuda(),
+                    data["T_1to2"].cuda(),
+                    data["K1"].cuda(),
+                    data["K2"].cuda(),
+                )
+                matches, certainty = model.match(im_A, im_B, batched=True)
+                gd, pck_1, pck_3, pck_5, prob = self.geometric_dist(
+                    depth1, depth2, T_1to2, K1, K2, matches
+                )
+                if romatch.DEBUG_MODE:
+                    from romatch.utils.utils import tensor_to_pil
+                    import torch.nn.functional as F
+                    path = "vis"
+                    H, W = model.get_output_resolution()
+                    white_im = torch.ones((B,1,H,W),device="cuda")
+                    im_B_transfer_rgb = F.grid_sample(
+                        im_B.cuda(), matches[:,:,:W, 2:], mode="bilinear", align_corners=False
+                    )
+                    warp_im = im_B_transfer_rgb
+                    c_b = certainty[:,None]#(certainty*0.9 + 0.1*torch.ones_like(certainty))[:,None]
+                    vis_im = c_b * warp_im + (1 - c_b) * white_im
+                    for b in range(B):
+                        import os
+                        os.makedirs(f"{path}/{model.name}/{idx}_{b}_{H}_{W}",exist_ok=True)
+                        tensor_to_pil(vis_im[b], unnormalize=True).save(
+                            f"{path}/{model.name}/{idx}_{b}_{H}_{W}/warp.jpg")
+                        tensor_to_pil(im_A[b].cuda(), unnormalize=True).save(
+                            f"{path}/{model.name}/{idx}_{b}_{H}_{W}/im_A.jpg")
+                        tensor_to_pil(im_B[b].cuda(), unnormalize=True).save(
+                            f"{path}/{model.name}/{idx}_{b}_{H}_{W}/im_B.jpg")
+
+
+                gd_tot, pck_1_tot, pck_3_tot, pck_5_tot = (
+                    gd_tot + gd.mean(),
+                    pck_1_tot + pck_1,
+                    pck_3_tot + pck_3,
+                    pck_5_tot + pck_5,
+                )
+        return {
+            "epe": gd_tot.item() / len(dataloader),
+            "mega_pck_1": pck_1_tot.item() / len(dataloader),
+            "mega_pck_3": pck_3_tot.item() / len(dataloader),
+            "mega_pck_5": pck_5_tot.item() / len(dataloader),
+        }

submodules/RoMa/romatch/benchmarks/megadepth_pose_estimation_benchmark.py

@@ -0,0 +1,118 @@
+import numpy as np
+import torch
+from romatch.utils import *
+from PIL import Image
+from tqdm import tqdm
+import torch.nn.functional as F
+import romatch
+import kornia.geometry.epipolar as kepi
+
+class MegaDepthPoseEstimationBenchmark:
+    def __init__(self, data_root="data/megadepth", scene_names = None) -> None:
+        if scene_names is None:
+            self.scene_names = [
+                "0015_0.1_0.3.npz",
+                "0015_0.3_0.5.npz",
+                "0022_0.1_0.3.npz",
+                "0022_0.3_0.5.npz",
+                "0022_0.5_0.7.npz",
+            ]
+        else:
+            self.scene_names = scene_names
+        self.scenes = [
+            np.load(f"{data_root}/{scene}", allow_pickle=True)
+            for scene in self.scene_names
+        ]
+        self.data_root = data_root
+
+    def benchmark(self, model, model_name = None):
+        with torch.no_grad():
+            data_root = self.data_root
+            tot_e_t, tot_e_R, tot_e_pose = [], [], []
+            thresholds = [5, 10, 20]
+            for scene_ind in range(len(self.scenes)):
+                import os
+                scene_name = os.path.splitext(self.scene_names[scene_ind])[0]
+                scene = self.scenes[scene_ind]
+                pairs = scene["pair_infos"]
+                intrinsics = scene["intrinsics"]
+                poses = scene["poses"]
+                im_paths = scene["image_paths"]
+                pair_inds = range(len(pairs))
+                for pairind in tqdm(pair_inds):
+                    idx1, idx2 = pairs[pairind][0]
+                    K1 = intrinsics[idx1].copy()
+                    T1 = poses[idx1].copy()
+                    R1, t1 = T1[:3, :3], T1[:3, 3]
+                    K2 = intrinsics[idx2].copy()
+                    T2 = poses[idx2].copy()
+                    R2, t2 = T2[:3, :3], T2[:3, 3]
+                    R, t = compute_relative_pose(R1, t1, R2, t2)
+                    T1_to_2 = np.concatenate((R,t[:,None]), axis=-1)
+                    im_A_path = f"{data_root}/{im_paths[idx1]}"
+                    im_B_path = f"{data_root}/{im_paths[idx2]}"
+                    dense_matches, dense_certainty = model.match(
+                        im_A_path, im_B_path, K1.copy(), K2.copy(), T1_to_2.copy()
+                    )
+                    sparse_matches,_ = model.sample(
+                        dense_matches, dense_certainty, 5_000
+                    )
+                    
+                    im_A = Image.open(im_A_path)
+                    w1, h1 = im_A.size
+                    im_B = Image.open(im_B_path)
+                    w2, h2 = im_B.size
+                    if True: # Note: we keep this true as it was used in DKM/RoMa papers. There is very little difference compared to setting to False. 
+                        scale1 = 1200 / max(w1, h1)
+                        scale2 = 1200 / max(w2, h2)
+                        w1, h1 = scale1 * w1, scale1 * h1
+                        w2, h2 = scale2 * w2, scale2 * h2
+                        K1, K2 = K1.copy(), K2.copy()
+                        K1[:2] = K1[:2] * scale1
+                        K2[:2] = K2[:2] * scale2
+
+                    kpts1, kpts2 = model.to_pixel_coordinates(sparse_matches, h1, w1, h2, w2)
+                    kpts1, kpts2 = kpts1.cpu().numpy(), kpts2.cpu().numpy()
+                    for _ in range(5):
+                        shuffling = np.random.permutation(np.arange(len(kpts1)))
+                        kpts1 = kpts1[shuffling]
+                        kpts2 = kpts2[shuffling]
+                        try:
+                            threshold = 0.5 
+                            norm_threshold = threshold / (np.mean(np.abs(K1[:2, :2])) + np.mean(np.abs(K2[:2, :2])))
+                            R_est, t_est, mask = estimate_pose(
+                                kpts1,
+                                kpts2,
+                                K1,
+                                K2,
+                                norm_threshold,
+                                conf=0.99999,
+                            )
+                            T1_to_2_est = np.concatenate((R_est, t_est), axis=-1)  #
+                            e_t, e_R = compute_pose_error(T1_to_2_est, R, t)
+                            e_pose = max(e_t, e_R)
+                        except Exception as e:
+                            print(repr(e))
+                            e_t, e_R = 90, 90
+                            e_pose = max(e_t, e_R)
+                        tot_e_t.append(e_t)
+                        tot_e_R.append(e_R)
+                        tot_e_pose.append(e_pose)
+            tot_e_pose = np.array(tot_e_pose)
+            auc = pose_auc(tot_e_pose, thresholds)
+            acc_5 = (tot_e_pose < 5).mean()
+            acc_10 = (tot_e_pose < 10).mean()
+            acc_15 = (tot_e_pose < 15).mean()
+            acc_20 = (tot_e_pose < 20).mean()
+            map_5 = acc_5
+            map_10 = np.mean([acc_5, acc_10])
+            map_20 = np.mean([acc_5, acc_10, acc_15, acc_20])
+            print(f"{model_name} auc: {auc}")
+            return {
+                "auc_5": auc[0],
+                "auc_10": auc[1],
+                "auc_20": auc[2],
+                "map_5": map_5,
+                "map_10": map_10,
+                "map_20": map_20,
+            }