CodecSplat demo

.gitattributes

@@ -83,3 +83,45 @@ examples/scene_b/21.jpg filter=lfs diff=lfs merge=lfs -text
 examples/scene_b/22.jpg filter=lfs diff=lfs merge=lfs -text
 examples/scene_b/23.jpg filter=lfs diff=lfs merge=lfs -text
 examples/scene_b/24.jpg filter=lfs diff=lfs merge=lfs -text
+examples/counter/counter_00.jpg filter=lfs diff=lfs merge=lfs -text
+examples/counter/counter_01.jpg filter=lfs diff=lfs merge=lfs -text
+examples/counter/counter_02.jpg filter=lfs diff=lfs merge=lfs -text
+examples/counter/counter_03.jpg filter=lfs diff=lfs merge=lfs -text
+examples/counter/counter_04.jpg filter=lfs diff=lfs merge=lfs -text
+examples/counter/counter_05.jpg filter=lfs diff=lfs merge=lfs -text
+examples/counter/counter_06.jpg filter=lfs diff=lfs merge=lfs -text
+examples/counter/counter_07.jpg filter=lfs diff=lfs merge=lfs -text
+examples/counter/counter_08.jpg filter=lfs diff=lfs merge=lfs -text
+examples/counter/counter_09.jpg filter=lfs diff=lfs merge=lfs -text
+examples/counter/counter_10.jpg filter=lfs diff=lfs merge=lfs -text
+examples/counter/counter_11.jpg filter=lfs diff=lfs merge=lfs -text
+examples/counter/counter_12.jpg filter=lfs diff=lfs merge=lfs -text
+examples/counter/counter_13.jpg filter=lfs diff=lfs merge=lfs -text
+examples/garden/garden_00.jpg filter=lfs diff=lfs merge=lfs -text
+examples/garden/garden_01.jpg filter=lfs diff=lfs merge=lfs -text
+examples/garden/garden_02.jpg filter=lfs diff=lfs merge=lfs -text
+examples/garden/garden_03.jpg filter=lfs diff=lfs merge=lfs -text
+examples/garden/garden_04.jpg filter=lfs diff=lfs merge=lfs -text
+examples/garden/garden_05.jpg filter=lfs diff=lfs merge=lfs -text
+examples/garden/garden_06.jpg filter=lfs diff=lfs merge=lfs -text
+examples/garden/garden_07.jpg filter=lfs diff=lfs merge=lfs -text
+examples/garden/garden_08.jpg filter=lfs diff=lfs merge=lfs -text
+examples/garden/garden_09.jpg filter=lfs diff=lfs merge=lfs -text
+examples/garden/garden_10.jpg filter=lfs diff=lfs merge=lfs -text
+examples/garden/garden_11.jpg filter=lfs diff=lfs merge=lfs -text
+examples/garden/garden_12.jpg filter=lfs diff=lfs merge=lfs -text
+examples/garden/garden_13.jpg filter=lfs diff=lfs merge=lfs -text
+examples/kitchen/kitchen_00.jpg filter=lfs diff=lfs merge=lfs -text
+examples/kitchen/kitchen_01.jpg filter=lfs diff=lfs merge=lfs -text
+examples/kitchen/kitchen_02.jpg filter=lfs diff=lfs merge=lfs -text
+examples/kitchen/kitchen_03.jpg filter=lfs diff=lfs merge=lfs -text
+examples/kitchen/kitchen_04.jpg filter=lfs diff=lfs merge=lfs -text
+examples/kitchen/kitchen_05.jpg filter=lfs diff=lfs merge=lfs -text
+examples/kitchen/kitchen_06.jpg filter=lfs diff=lfs merge=lfs -text
+examples/kitchen/kitchen_07.jpg filter=lfs diff=lfs merge=lfs -text
+examples/kitchen/kitchen_08.jpg filter=lfs diff=lfs merge=lfs -text
+examples/kitchen/kitchen_09.jpg filter=lfs diff=lfs merge=lfs -text
+examples/kitchen/kitchen_10.jpg filter=lfs diff=lfs merge=lfs -text
+examples/kitchen/kitchen_11.jpg filter=lfs diff=lfs merge=lfs -text
+examples/kitchen/kitchen_12.jpg filter=lfs diff=lfs merge=lfs -text
+examples/kitchen/kitchen_13.jpg filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Pengpeng Yu
+
+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.

README.md

@@ -6,7 +6,7 @@ colorTo: yellow
 sdk: gradio
 sdk_version: 6.30.0
 app_file: app.py
-short_description: Compress photos into a tiny 3D Gaussian splat bitstream
+short_description: Compress a 3D Gaussian scene into a few KB
 python_version: "3.12"
 startup_duration_timeout: 1h
 models:
@@ -15,40 +15,35 @@ models:
 
 # CodecSplat: Ultra-Compact Latent Coding for Feed-Forward 3D Gaussian Splatting
 
-Demo of [**CodecSplat**](https://huggingface.co/papers/2605.25563) (Yu et al.). CodecSplat
-entropy-codes the 2D features that a feed-forward 3DGS model (DepthSplat-style) uses for depth
-and Gaussian prediction. A scene becomes a bitstream of a few tens of KB, which is then decoded
-into renderable 3D Gaussians.
-
-How the Space works:
-
-1. **COLMAP** (`pycolmap`, CPU) estimates intrinsics and camera poses for the uploaded photos.
-   The reconstruction is normalised and undistorted to PINHOLE.
-2. 8 context views are picked by farthest-point sampling over the cameras. They are resized and
-   center-cropped to 256×448 with the authors' crop shim.
-3. CodecSplat (`codecsplat-base-dl3dv-256x448-view8`, the authors' DL3DV evaluation preset)
-   runs the **complete bitstream path** from `ModelWrapper.test_step`: encode context →
-   compress to a real bitstream at the chosen λ → decompress → predict Gaussians, using bf16
-   autocast like `trainer.precision=bf16-mixed`.
-4. The authors' `splatting_cuda` decoder (Inria `diff-gaussian-rasterization`, prebuilt sm_120
-   wheel) renders held-out views (PSNR is reported
-   against the real photos) and a fly-through video. The decoded Gaussians are shown in an
-   interactive `gr.Model3D` splat viewer, and the `.bin` bitstream is available to download.
-
-Notes:
-* The authors' evaluation scripts render with gsplat (antialiased). This Space uses the
-  repo's other supported decoder, `model/decoder=splatting_cuda`, because a prebuilt Blackwell
-  wheel exists for it. Renders can differ very slightly from the paper's numbers.
-* COLMAP runs in a separate process with 16 threads. On the ~192-CPU ZeroGPU hosts COLMAP's
-  default thread count makes the mapper fail to register images.
-* The viewer gets a compact copy of the Gaussians: DC colour only, near-transparent splats
-  pruned, re-centred so the scene is framed.
-
-Code: https://github.com/pengpeng-yu/CodecSplat (MIT) · Weights: https://huggingface.co/pengpeng-yu/CodecSplat (MIT)
-
-## Example scenes
-
-The example frames are from the DL3DV demo scene that the DepthSplat authors redistribute in
-[`haofeixu/depthsplat`](https://huggingface.co/datasets/haofeixu/depthsplat) (MIT-labelled).
-The imagery itself comes from [DL3DV-10K](https://dl3dv-10k.github.io/DL3DV-10K/), credit to
-the DL3DV authors.
+Interactive demo of [CodecSplat](https://huggingface.co/papers/2605.25563)
+([code](https://github.com/pengpeng-yu/CodecSplat), MIT;
+[weights](https://huggingface.co/pengpeng-yu/CodecSplat)).
+
+Upload overlapping photos of a static scene. The Space:
+
+1. runs **COLMAP** (`pycolmap`, CPU) to estimate intrinsics and camera poses,
+2. picks 8 context views by farthest-point sampling and resizes them to 256x448
+   (the `codecsplat-base-dl3dv-256x448-view8` checkpoint),
+3. runs the authors' **complete bitstream path** (`EncoderCodecSplat.codec_forward`):
+   it encodes the latent features into a real entropy-coded scene bitstream at the
+   chosen rate point λ (16-1024), then decodes that bitstream back to 3D Gaussians,
+4. renders held-out views (with PSNR/SSIM against the real photos) and a
+   fly-through with `gsplat`, and shows the Gaussians in an interactive viewer.
+
+You can download both the `.bin` bitstream and the decoded `.ply`.
+
+## Implementation notes
+
+* Model construction uses the exact Hydra overrides of the authors'
+  `scripts/test_dl3dv_view8_256x448.sh` (bf16 autocast, gsplat antialiased rasterizer).
+* COLMAP stage and scene loader are adapted from the
+  [ReSplat Space](https://huggingface.co/spaces/hugging-apps/resplat-recurrent-gs) / the ReSplat
+  reference `scripts/infer_colmap.py` (same DepthSplat camera conventions). The reconstruction is
+  normalised to a fixed extent so that the DL3DV near/far planes (0.5 / 200) are meaningful.
+* `gsplat`'s CUDA kernels are JIT-compiled once at startup (no prebuilt sm_120 wheel).
+
+## Example assets
+
+`examples/kitchen` and `examples/garden` are 14 frames each from the
+[Mip-NeRF 360](https://jonbarron.info/mipnerf360/) dataset (Barron et al.), downscaled to 960px,
+as redistributed by [`nvs-bench/mipnerf360`](https://huggingface.co/datasets/nvs-bench/mipnerf360) (MIT).

app.py

@@ -1,27 +1,36 @@
 """CodecSplat: Ultra-Compact Latent Coding for Feed-Forward 3D Gaussian Splatting.
 
-Upload overlapping photos of a static scene. The Space runs COLMAP (pycolmap, CPU) to get
-camera poses, then CodecSplat (authors' DL3DV 8-view 256x448 checkpoint):
-  context views -> 2D latent features -> learned entropy coder -> real bitstream
-  -> decoder -> 3D Gaussians -> Gaussian-splatting rasterizer.
-You get an interactive splat viewer, a fly-through video, held-out novel views, the
-size of the compressed scene bitstream (and the bitstream itself).
+Upload overlapping photos of a static scene. The Space recovers the cameras with
+COLMAP (pycolmap, CPU), then CodecSplat (DL3DV 8-view 256x448 checkpoint):
+
+  1. encodes the 8 context views into an intermediate 2D latent,
+  2. entropy-codes that latent (+ cameras) into a real scene bitstream at the
+     chosen rate point lambda,
+  3. decodes the bitstream back into 3D Gaussians,
+  4. renders held-out views / a fly-through with gsplat.
+
+The inference path follows the authors' evaluation (`scripts/test_dl3dv_view8_256x448.sh`
+-> `ModelWrapper.test_step` -> `EncoderCodecSplat.codec_forward`), with the exact same
+encoder/decoder overrides.
 """
 
 from __future__ import annotations
 
+import ctypes
+import glob
 import os
 import sys
 import tempfile
 import time
+import types
 from pathlib import Path
 
 APP_DIR = Path(__file__).parent.resolve()
 sys.path.insert(0, str(APP_DIR))
 
-# pycolmap's bundled OpenBLAS segfaults ("Bad memory unallocation") when many matcher
-# threads each spin up a multi-threaded BLAS; keep BLAS single-threaded.
-os.environ["OPENBLAS_NUM_THREADS"] = "1"
+os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
+os.environ.setdefault("MAX_JOBS", str(min(32, (os.cpu_count() or 8))))
+os.environ.setdefault("WANDB_MODE", "disabled")
 
 import spaces  # noqa: E402  (must precede any CUDA-touching import)
 import torch  # noqa: E402
@@ -29,224 +38,425 @@ import torch  # noqa: E402
 torch.set_float32_matmul_precision("high")
 
 
-def _nvidia_fds(stage: str) -> None:
-    """Diagnostics: open /dev/nvidia* handles in the (GPU-less) main process mean the
-    CUDA driver was initialised there, which breaks ZeroGPU's forked workers."""
-    try:
-        fds = []
-        for fd in os.listdir("/proc/self/fd"):
-            try:
-                t = os.readlink(f"/proc/self/fd/{fd}")
-            except OSError:
-                continue
-            if "nvidia" in t:
-                fds.append(t)
-        print(f"[diag] {stage}: nvidia fds={sorted(set(fds))}", flush=True)
-    except Exception as exc:
-        print(f"[diag] {stage}: {exc}", flush=True)
+# ---------------------------------------------------------------------------
+# 1. gsplat CUDA backend (JIT-compiled once at startup against the ZeroGPU
+#    image's CUDA toolkit -- gsplat has no prebuilt sm_120 wheel).
+# ---------------------------------------------------------------------------
+def _pick_cuda_home() -> str | None:
+    torch_cuda = (torch.version.cuda or "").split(".")[0]
+    candidates: list[str] = []
+    for pattern in ("/cuda-image/usr/local/cuda-*", "/usr/local/cuda-*"):
+        candidates += sorted(glob.glob(pattern), reverse=True)
+    candidates += ["/cuda-image/usr/local/cuda", "/usr/local/cuda"]
+    usable = [c for c in candidates if os.path.isfile(os.path.join(c, "bin", "nvcc"))]
+    if not usable:
+        return None
+    for c in usable:
+        if torch_cuda and f"cuda-{torch_cuda}" in c:
+            return c
+    return usable[0]
+
+
+CUDA_HOME = _pick_cuda_home()
+print(f"[setup] torch {torch.__version__} (cuda {torch.version.cuda}) | CUDA_HOME={CUDA_HOME}", flush=True)
+if CUDA_HOME:
+    os.environ["CUDA_HOME"] = CUDA_HOME
+    os.environ["CUDA_PATH"] = CUDA_HOME
+    os.environ["PATH"] = os.path.join(CUDA_HOME, "bin") + os.pathsep + os.environ.get("PATH", "")
+    _libdir = os.path.join(CUDA_HOME, "lib64")
+    os.environ["LD_LIBRARY_PATH"] = _libdir + os.pathsep + os.environ.get("LD_LIBRARY_PATH", "")
+    for _so in sorted(glob.glob(os.path.join(_libdir, "libcudart.so*")), reverse=True):
+        try:
+            ctypes.CDLL(_so, mode=ctypes.RTLD_GLOBAL)
+            break
+        except OSError:
+            continue
+
+import torch.utils.cpp_extension as _cppext  # noqa: E402
+
+_cppext._check_cuda_version = lambda *a, **k: None
+
+print("[setup] compiling / loading gsplat CUDA kernels...", flush=True)
+_t0 = time.time()
+from gsplat.cuda._backend import _C as _GSPLAT_C  # noqa: E402
+
+if _GSPLAT_C is None:
+    raise RuntimeError(f"gsplat CUDA backend unavailable (CUDA_HOME={CUDA_HOME}).")
+print(f"[setup] gsplat backend ready in {time.time() - _t0:.1f}s", flush=True)
 
+# ---------------------------------------------------------------------------
+# 2. Compatibility shims for the vendored reference code
+# ---------------------------------------------------------------------------
+import matplotlib.cm  # noqa: E402
+import matplotlib.pyplot as _plt  # noqa: E402
 
-_nvidia_fds("after torch import")
+if not hasattr(matplotlib.cm, "get_cmap"):
+    matplotlib.cm.get_cmap = _plt.get_cmap
 
+if "skvideo" not in sys.modules:  # training-only import in some helpers
+    _sk = types.ModuleType("skvideo")
+    _sk.__path__ = []
+    _skio = types.ModuleType("skvideo.io")
+    _sk.io = _skio
+    sys.modules["skvideo"] = _sk
+    sys.modules["skvideo.io"] = _skio
 
 # ---------------------------------------------------------------------------
-# 2. Model
+# 3. Model
 # ---------------------------------------------------------------------------
+import gradio as gr  # noqa: E402
 import imageio.v2 as imageio  # noqa: E402
 import numpy as np  # noqa: E402
-import gradio as gr  # noqa: E402
 from huggingface_hub import hf_hub_download  # noqa: E402
 from PIL import Image  # noqa: E402
+from scipy.spatial.transform import Rotation, Slerp  # noqa: E402
 
-import codecsplat_pipeline as cp  # noqa: E402
-
-_nvidia_fds("after pipeline import")
-
+from colmap_io import load_colmap_scene, select_context_views, select_target_views  # noqa: E402
 from colmap_sfm import IMAGE_EXT, run_sfm  # noqa: E402
 
 MODEL_REPO = "pengpeng-yu/CodecSplat"
 CHECKPOINT_FILE = "codecsplat-base-dl3dv-256x448-view8-075062d0.pth"
 NUM_CONTEXT = 8
-NUM_TARGET = 8
-MAX_INPUT_IMAGES = 60
+IMAGE_SHAPE = (256, 448)
+NEAR, FAR = 0.5, 200.0  # config/experiment/dl3dv.yaml
+
+# Exactly the overrides of scripts/test_dl3dv_view8_256x448.sh
+OVERRIDES = [
+    "+experiment=dl3dv",
+    "mode=test",
+    f"dataset.image_shape=[{IMAGE_SHAPE[0]},{IMAGE_SHAPE[1]}]",
+    "model.encoder.num_scales=2",
+    "model.encoder.upsample_factor=4",
+    "model.encoder.lowest_feature_resolution=8",
+    "model.encoder.monodepth_vit_type=vitb",
+    "model.encoder.enable_feature_codec=true",
+    "model.encoder.enable_depth_guided_mv_refiner=true",
+    "model.encoder.enable_post_vae_mv_refiner=true",
+    "model.encoder.supervise_intermediate_depth=false",
+    "model/decoder=splatting_gsplat",
+    "model.decoder.use_covariances=true",
+    "model.decoder.rasterize_mode=antialiased",
+    "model.decoder.eps2d=0.1",
+    "model.decoder.packed=true",
+    "model.decoder.scale_invariant=false",
+    "trainer.precision=bf16-mixed",
+    "wandb.mode=disabled",
+    "output_dir=/tmp/codecsplat_out",
+]
+
+
+def build_model(checkpoint: str):
+    from hydra import compose, initialize_config_dir
+    from hydra.core.global_hydra import GlobalHydra
+
+    GlobalHydra.instance().clear()
+    with initialize_config_dir(config_dir=str(APP_DIR / "config"), version_base=None):
+        cfg_dict = compose(config_name="main", overrides=OVERRIDES)
+
+    from src.config import load_typed_root_config
+    from src.global_cfg import set_cfg
+    from src.model.decoder import get_decoder
+    from src.model.encoder import get_encoder
+
+    set_cfg(cfg_dict)
+    cfg = load_typed_root_config(cfg_dict)
+    encoder = get_encoder(cfg.model.encoder)
+    decoder = get_decoder(cfg.model.decoder, cfg.dataset)
+
+    state = torch.load(checkpoint, map_location="cpu", weights_only=False)
+    if "state_dict" in state:
+        state = state["state_dict"]
+    enc_sd = {k[len("encoder."):]: v for k, v in state.items() if k.startswith("encoder.")}
+    dec_sd = {k[len("decoder."):]: v for k, v in state.items() if k.startswith("decoder.")}
+    missing, unexpected = encoder.load_state_dict(enc_sd, strict=False)
+    print(f"[setup] encoder: {len(enc_sd)} tensors, missing={len(missing)} unexpected={len(unexpected)}", flush=True)
+    if missing:
+        print(f"[setup]   missing (first 10): {missing[:10]}", flush=True)
+    if unexpected:
+        print(f"[setup]   unexpected (first 10): {unexpected[:10]}", flush=True)
+    if dec_sd:
+        decoder.load_state_dict(dec_sd, strict=False)
+    encoder.eval()
+    decoder.eval()
+    return encoder, decoder, encoder.get_data_shim()
+
 
 print("[setup] downloading checkpoint...", flush=True)
-ckpt_path = hf_hub_download(MODEL_REPO, CHECKPOINT_FILE)
+CHECKPOINT_PATH = hf_hub_download(MODEL_REPO, CHECKPOINT_FILE)
 print("[setup] building CodecSplat...", flush=True)
-ENCODER = cp.build_encoder(ckpt_path).to("cuda")
-DECODER = cp.build_decoder().to("cuda")
+ENCODER, DECODER, DATA_SHIM = build_model(CHECKPOINT_PATH)
+ENCODER.to("cuda")
+DECODER.to("cuda")
 print("[setup] model ready", flush=True)
-_nvidia_fds("after model build")
 
+from src.dataset.shims.crop_shim import rescale_and_crop  # noqa: E402
+from src.evaluation.metrics import compute_psnr, compute_ssim  # noqa: E402
+from src.model.ply_export import export_ply  # noqa: E402
+
+MAX_INPUT_IMAGES = 40
+VIEWER_MAX_GAUSSIANS = 600_000
+VIEWER_MIN_OPACITY = 0.02
 
-def _u8(t: torch.Tensor) -> np.ndarray:
+
+# ---------------------------------------------------------------------------
+# 4. Helpers
+# ---------------------------------------------------------------------------
+def _to_uint8(t) -> np.ndarray:
     return (t.clamp(0, 1).permute(1, 2, 0).float().cpu().numpy() * 255).round().astype(np.uint8)
 
 
+def _load_views(paths, intrinsics, shape):
+    """Load images + normalized intrinsics and rescale/center-crop them to `shape`
+    with the same crop shim the DL3DV loader uses."""
+    imgs = []
+    for p in paths:
+        im = np.asarray(Image.open(p).convert("RGB"), dtype=np.float32) / 255.0
+        imgs.append(torch.from_numpy(im).permute(2, 0, 1))
+    h0 = min(i.shape[1] for i in imgs)
+    w0 = min(i.shape[2] for i in imgs)
+    imgs = torch.stack([i[:, :h0, :w0] for i in imgs])
+    K = torch.tensor(np.asarray(intrinsics), dtype=torch.float32)
+    return rescale_and_crop(imgs, K, shape, depths=None)
+
+
+def _interpolate_path(c2w: torch.Tensor, frames_per_segment: int = 12) -> torch.Tensor:
+    """Smooth fly-through visiting the context cameras in order (slerp + lerp)."""
+    c2w = c2w.cpu().numpy()
+    out = []
+    for a, b in zip(c2w[:-1], c2w[1:]):
+        slerp = Slerp([0, 1], Rotation.from_matrix(np.stack([a[:3, :3], b[:3, :3]])))
+        for t in np.linspace(0, 1, frames_per_segment, endpoint=False):
+            m = np.eye(4, dtype=np.float32)
+            m[:3, :3] = slerp([t]).as_matrix()[0]
+            m[:3, 3] = (1 - t) * a[:3, 3] + t * b[:3, 3]
+            out.append(m)
+    out.append(c2w[-1])
+    return torch.from_numpy(np.stack(out)).float()
+
+
+def export_viewer_ply(gaussians, extrinsics, path: Path) -> str:
+    """Compact, well-framed copy of the Gaussians for the in-browser splat viewer
+    (DC colour only, near-transparent splats pruned, re-centred and re-scaled)."""
+    means = gaussians.means[0]
+    scales = gaussians.scales[0]
+    rotations = gaussians.rotations[0]
+    harmonics = gaussians.harmonics[0]
+    opacities = gaussians.opacities[0]
+
+    keep = opacities > VIEWER_MIN_OPACITY
+    if int(keep.sum()) < 1000:
+        keep = torch.ones_like(opacities, dtype=torch.bool)
+    idx = torch.nonzero(keep, as_tuple=False).squeeze(1)
+    if idx.numel() > VIEWER_MAX_GAUSSIANS:
+        idx = idx[torch.randperm(idx.numel(), device=idx.device)[:VIEWER_MAX_GAUSSIANS]]
+    means, scales = means[idx].float(), scales[idx].float()
+    rotations, harmonics, opacities = rotations[idx], harmonics[idx], opacities[idx]
+
+    rot = extrinsics[:3, :3].float().inverse().to(means.device)
+    world = (rot @ means.T).T
+    center = world.median(dim=0).values
+    radius = torch.quantile((world - center).norm(dim=-1), 0.7)
+    s = float(3.0 / max(float(radius), 1e-6))
+    means = means * s - rot.T @ (center * s)
+    scales = scales * s
+    export_ply(extrinsics, means, scales, rotations, harmonics[..., :1], opacities, path,
+               align_to_view=True)
+    return str(path)
+
+
 # ---------------------------------------------------------------------------
-# 3. GPU part: encode -> bitstream -> decode -> render
+# 5. GPU stage: encode -> bitstream -> decode -> render
 # ---------------------------------------------------------------------------
-@spaces.GPU(duration=60)
-def codecsplat_gpu(image_paths: list[str], c2w: np.ndarray, K: np.ndarray, lmb: float = 1024.0):
-    ctx_idx, tgt_idx = cp.select_views(c2w, NUM_CONTEXT, NUM_TARGET)
-    batch = cp.to_device(cp.build_example(image_paths, c2w, K, ctx_idx, tgt_idx), "cuda")
+@spaces.GPU(duration=40)
+def codecsplat_run(scene_dir: str, lmb: float = 1024.0):
+    """Run CodecSplat on a COLMAP scene directory (undistorted images + sparse model)."""
+    out = Path(scene_dir) / "out"
+    out.mkdir(parents=True, exist_ok=True)
+
+    scene = load_colmap_scene(scene_dir, "sparse", "images")
+    n = len(scene["image_paths"])
+    if n < 2:
+        raise gr.Error(f"Only {n} usable image(s) after COLMAP -- need at least 2.")
+
+    w0, h0 = Image.open(scene["image_paths"][0]).size
+    shape = IMAGE_SHAPE if w0 >= h0 else IMAGE_SHAPE[::-1]
+
+    num_context = min(NUM_CONTEXT, n)
+    ctx_idx = select_context_views(scene["c2w"], num_context, "fps")
+    tgt_idx = select_target_views(n, ctx_idx, "remaining", 8)
+    if len(tgt_idx) == 0:
+        tgt_idx = ctx_idx.copy()
+
+    ctx_img, ctx_K = _load_views([scene["image_paths"][i] for i in ctx_idx], scene["intrinsics"][ctx_idx], shape)
+    tgt_img, tgt_K = _load_views([scene["image_paths"][i] for i in tgt_idx], scene["intrinsics"][tgt_idx], shape)
+
+    c2w = torch.tensor(scene["c2w"], dtype=torch.float32)
+    pivot = torch.inverse(c2w[ctx_idx[len(ctx_idx) // 2]])
+    c2w = pivot[None] @ c2w  # express all cameras in the middle context view's frame
+    ctx_c2w, tgt_c2w = c2w[ctx_idx], c2w[tgt_idx]
+
+    def views(img, K, E):
+        v = len(E)
+        return {
+            "image": img[None].cuda(),
+            "intrinsics": K[None].cuda(),
+            "extrinsics": E[None].cuda(),
+            "near": torch.full((1, v), NEAR, device="cuda"),
+            "far": torch.full((1, v), FAR, device="cuda"),
+            "index": torch.arange(v, device="cuda")[None],
+        }
+
+    batch = DATA_SHIM({
+        "context": views(ctx_img, ctx_K, ctx_c2w),
+        "target": views(tgt_img, tgt_K, tgt_c2w),
+        "scene": ["user_scene"],
+    })
     h, w = batch["context"]["image"].shape[-2:]
 
+    ENCODER.feature_codec.cfg.lmb = float(lmb)
+    t0 = time.time()
+    with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
+        enc = ENCODER(batch["context"], 0, return_bitstream=True)
+        gaussians = enc["gaussians"]
+
+        def render(E, K, chunk=8):
+            frames = []
+            for i in range(0, E.shape[1], chunk):
+                o = DECODER.forward(
+                    gaussians, E[:, i:i + chunk], K[:, i:i + chunk],
+                    torch.full_like(E[:, i:i + chunk, 0, 0], NEAR),
+                    torch.full_like(E[:, i:i + chunk, 0, 0], FAR),
+                    (h, w), depth_mode=None,
+                )
+                frames.append(o.color[0].float())
+            return torch.cat(frames)
+
+        rendered = render(batch["target"]["extrinsics"], batch["target"]["intrinsics"])
+        path_E = _interpolate_path(batch["context"]["extrinsics"][0])[None].cuda()
+        path_K = batch["context"]["intrinsics"][:, :1].expand(-1, path_E.shape[1], -1, -1)
+        video = render(path_E, path_K)
     torch.cuda.synchronize()
-    t0 = time.perf_counter()
-    out = cp.encode_decode(ENCODER, batch, lmb)
-    torch.cuda.synchronize()
-    t_codec = time.perf_counter() - t0
-    g = out["gaussians"]
-
-    # held-out novel views (+ PSNR against the real photos)
-    renders, score = None, float("nan")
-    if len(tgt_idx) > 0:
-        renders = cp.render(DECODER, g, batch["target"]["extrinsics"][0],
-                            batch["target"]["intrinsics"][0], (h, w))
-        score = cp.psnr(batch["target"]["image"][0].float(), renders)
-
-    # smooth fly-through through the context cameras
-    traj_c2w, traj_K = cp.interpolate_trajectory(
-        batch["context"]["extrinsics"][0], batch["context"]["intrinsics"][0]
-    )
-    video = cp.render(DECODER, g, traj_c2w.cuda(), traj_K.cuda(), (h, w))
-    t_total = time.perf_counter() - t0
-
-    gauss = {
-        "means": g.means[0].float().cpu(),
-        "scales": g.scales[0].float().cpu(),
-        "rotations": g.rotations[0].float().cpu(),
-        "harmonics": g.harmonics[0].float().cpu(),
-        "opacities": g.opacities[0].float().cpu(),
-    }
-    return {
-        "bitstream": out["bitstream"],
-        "num_bytes": float(out["codec_num_bytes"]),
-        "camera_bytes": float(out["codec_camera_num_bytes"]),
-        "feature_bytes": float(out["codec_feature_num_bytes"]),
-        "gaussians": gauss,
-        "context_c2w": batch["context"]["extrinsics"][0].cpu(),
-        "context_images": [_u8(x) for x in batch["context"]["image"][0]],
-        "target_images": [_u8(x) for x in batch["target"]["image"][0]],
-        "renders": [] if renders is None else [_u8(x) for x in renders],
-        "video": [_u8(x) for x in video],
-        "psnr": score,
-        "t_codec": t_codec,
-        "t_total": t_total,
-        "hw": (int(h), int(w)),
-        "ctx_idx": ctx_idx.tolist(),
-        "tgt_idx": tgt_idx.tolist(),
+    t_run = time.time() - t0
+
+    gt = batch["target"]["image"][0].float()
+    psnr = float(compute_psnr(gt, rendered.clamp(0, 1)).mean())
+    ssim = float(compute_ssim(gt, rendered.clamp(0, 1)).mean())
+
+    bitstream = enc["bitstream"]
+    bin_path = out / f"codecsplat_scene_lmb{int(lmb)}.bin"
+    bin_path.write_bytes(bitstream)
+
+    video_path = out / "flythrough.mp4"
+    imageio.mimwrite(video_path, [_to_uint8(f) for f in video], fps=24, quality=8, macro_block_size=1)
+
+    gallery = []
+    for i, (r, g) in enumerate(zip(rendered, gt)):
+        p = out / f"heldout_{i:02d}.png"
+        Image.fromarray(np.concatenate([_to_uint8(r), _to_uint8(g)], axis=1)).save(p)
+        gallery.append((str(p), f"held-out view {int(tgt_idx[i])}: CodecSplat (left) | real photo (right)"))
+
+    ref_E = batch["context"]["extrinsics"][0, len(ctx_idx) // 2]
+    ply_path = out / "codecsplat_gaussians.ply"
+    export_ply(ref_E, gaussians.means[0], gaussians.scales[0], gaussians.rotations[0],
+               gaussians.harmonics[0], gaussians.opacities[0], ply_path, align_to_view=True)
+    viewer_ply = export_viewer_ply(gaussians, ref_E, out / "codecsplat_viewer.ply")
+
+    stats = {
+        "num_bytes": len(bitstream),
+        "camera_bytes": float(enc["codec_camera_num_bytes"]),
+        "feature_bytes": float(enc["codec_feature_num_bytes"]),
+        "num_gaussians": int(gaussians.means.shape[1]),
+        "ply_bytes": os.path.getsize(ply_path),
+        "num_context": len(ctx_idx),
+        "num_target": len(tgt_idx),
+        "resolution": f"{h}x{w}",
+        "psnr": psnr,
+        "ssim": ssim,
+        "t_run": t_run,
+        "t_encode": float(enc["codec_encode_context_seconds"] + enc["codec_compress_seconds"]),
+        "t_decode": float(enc["codec_decompress_seconds"]),
     }
+    print(f"[run] lmb={lmb} bytes={len(bitstream)} psnr={psnr:.2f} t={t_run:.1f}s", flush=True)
+    return viewer_ply, str(video_path), gallery, str(bin_path), str(ply_path), stats
 
 
-# ---------------------------------------------------------------------------
-# 4. Full pipeline (CPU COLMAP -> GPU CodecSplat -> CPU export)
-# ---------------------------------------------------------------------------
-def _scene_arrays(frames: list[dict]):
-    if len(frames) < 2:
-        raise gr.Error("COLMAP registered fewer than 2 usable images.")
-    paths, c2ws, Ks = [], [], []
-    for f in frames:
-        w2c = np.eye(4, dtype=np.float64)
-        w2c[:3, :4] = np.asarray(f["w2c"], dtype=np.float64)
-        paths.append(f["path"])
-        c2ws.append(np.linalg.inv(w2c).astype(np.float32))
-        Ks.append(np.asarray(f["K"], dtype=np.float32))
-    return paths, np.stack(c2ws), np.stack(Ks)
-
-
-def reconstruct(images, lmb: float = 1024.0, matcher: str = "exhaustive",
-                progress=gr.Progress(track_tqdm=False)):
-    """Reconstruct a compressed 3D Gaussian scene from multi-view photos with CodecSplat.
+def reconstruct(
+    images,
+    lmb: float = 1024.0,
+    matcher: str = "exhaustive",
+    progress=gr.Progress(track_tqdm=False),
+):
+    """COLMAP + CodecSplat: compress a scene into a tiny bitstream and decode it to 3D Gaussians.
 
     Args:
-        images: 3-60 overlapping photos of one static scene (camera poses are estimated with COLMAP).
-        lmb: rate-distortion trade-off lambda (16 = smallest bitstream, 1024 = best quality).
-        matcher: COLMAP matching mode, "exhaustive" (unordered photos) or "sequential" (video frames).
+        images: Overlapping photos of one static scene (8-30 works well).
+        lmb: Rate point lambda in [16, 1024]. Higher = more bytes, better quality.
+        matcher: COLMAP matching -- "exhaustive" (unordered photos) or "sequential" (video frames).
 
     Returns:
-        Interactive splat (.ply), fly-through video, held-out novel views (render vs. photo),
-        selected input views, the scene bitstream (.bin) and a summary.
+        Interactive 3D Gaussian splat, a rendered fly-through video, held-out view comparisons,
+        the compressed scene bitstream, the decoded Gaussians as .ply, and a summary.
     """
     if not images:
-        raise gr.Error("Please upload at least 3 overlapping photos of the same scene.")
+        raise gr.Error("Please upload at least 3 overlapping images of the same scene.")
     if isinstance(images, (str, Path)):
         images = [images]
     paths = [str(im if isinstance(im, (str, Path)) else getattr(im, "name", im)) for im in images]
     paths = sorted([p for p in paths if p.lower().endswith(IMAGE_EXT)], key=os.path.basename)
     if len(paths) < 3:
         raise gr.Error(f"Got {len(paths)} usable image(s); COLMAP needs at least 3 overlapping photos.")
-    note = ""
     if len(paths) > MAX_INPUT_IMAGES:
         paths = [paths[i] for i in np.linspace(0, len(paths) - 1, MAX_INPUT_IMAGES).astype(int)]
-        note = f" (sub-sampled to {MAX_INPUT_IMAGES} images)"
 
     t_start = time.time()
     workdir = Path(tempfile.mkdtemp(prefix="codecsplat_"))
-    progress(0.05, desc=f"COLMAP on {len(paths)} images")
+    progress(0.05, desc=f"COLMAP: estimating cameras for {len(paths)} images")
+    t_sfm = time.time()
     try:
-        sfm = run_sfm(paths, workdir, max_image_size=1024, matcher=matcher,
+        sfm = run_sfm(paths, workdir, max_image_size=960, matcher=matcher,
                       log=lambda m: print(m, flush=True))
-    except (ValueError, RuntimeError) as exc:
+    except ValueError as exc:
         raise gr.Error(str(exc)) from exc
-    t_sfm = time.time() - t_start
-    img_paths, c2w, K = _scene_arrays(sfm["frames"])
-    _nvidia_fds("before GPU call")
-
-    progress(0.55, desc="CodecSplat: encode -> bitstream -> decode -> render")
-    r = codecsplat_gpu(img_paths, c2w, K, float(lmb))
-
-    progress(0.85, desc="Exporting Gaussians")
-    out = workdir / "out"
-    out.mkdir(exist_ok=True)
-    bin_path = out / f"codecsplat_scene_lmb{int(round(lmb))}.bin"
-    bin_path.write_bytes(r["bitstream"])
-    viewer_ply = cp.export_gaussians(r["gaussians"], r["context_c2w"], out)
-
-    video_path = str(out / "flythrough.mp4")
-    imageio.mimwrite(video_path, r["video"], fps=24, quality=8, macro_block_size=16)
-
-    gallery_ctx = [(im, f"input view {i}") for im, i in zip(r["context_images"], r["ctx_idx"])]
-    gallery_nv = []
-    for rend, gt, i in zip(r["renders"], r["target_images"], r["tgt_idx"]):
-        gallery_nv.append((np.concatenate([rend, gt], axis=1), f"held-out view {i}: render | photo"))
-
-    n_g = int(r["gaussians"]["means"].shape[0])
-    raw_bytes = n_g * (3 + 3 + 4 + 1 + r["gaussians"]["harmonics"][0].numel()) * 4
-    kb = r["num_bytes"] / 1024
+    t_sfm = time.time() - t_sfm
+
+    progress(0.6, desc="CodecSplat: encode -> bitstream -> decode -> render")
+    viewer_ply, video, gallery, bin_path, ply_path, s = codecsplat_run(sfm["scene_dir"], float(lmb))
+
+    ratio = s["ply_bytes"] / max(s["num_bytes"], 1)
     summary = (
-        f"**Compressed scene: {kb:.1f} KB** (features {r['feature_bytes']/1024:.1f} KB + cameras "
-        f"{r['camera_bytes']/1024:.2f} KB) at λ = {lmb:g} — decodes to **{n_g:,} Gaussians** "
-        f"(≈{raw_bytes/2**20:.0f} MB as raw float32, **{raw_bytes/max(r['num_bytes'],1):,.0f}× smaller**).\n\n"
-        f"Held-out novel views: **PSNR {r['psnr']:.2f} dB** · {len(r['ctx_idx'])} input views at "
-        f"{r['hw'][0]}×{r['hw'][1]} · COLMAP registered {sfm['num_registered']}/{sfm['num_input']} "
-        f"images{note} in {t_sfm:.1f}s · CodecSplat encode+decode {r['t_codec']:.2f}s · "
-        f"total {time.time() - t_start:.1f}s"
+        f"### Scene bitstream: **{s['num_bytes'] / 1024:.1f} KB** at λ = {lmb:g}\n"
+        f"(features {s['feature_bytes'] / 1024:.1f} KB + cameras {s['camera_bytes'] / 1024:.2f} KB) → decodes to "
+        f"**{s['num_gaussians']:,} Gaussians** (uncompressed .ply: {s['ply_bytes'] / 2**20:.1f} MB, **{ratio:,.0f}× smaller**).\n\n"
+        f"Held-out views ({s['num_target']}): **PSNR {s['psnr']:.2f} dB** · SSIM {s['ssim']:.3f} at {s['resolution']}  \n"
+        f"COLMAP registered {sfm['num_registered']}/{sfm['num_input']} images in {t_sfm:.1f}s · "
+        f"encode {s['t_encode']:.2f}s · decode {s['t_decode']:.2f}s · {time.time() - t_start:.1f}s total."
     )
-    return viewer_ply, video_path, gallery_nv, gallery_ctx, str(bin_path), summary
+    return viewer_ply, video, gallery, bin_path, ply_path, summary
 
 
 # ---------------------------------------------------------------------------
-# 5. UI
+# 6. UI
 # ---------------------------------------------------------------------------
-EX_DIR = APP_DIR / "examples"
-EXAMPLES = [sorted(str(p) for p in (EX_DIR / s).glob("*.jpg")) for s in ("scene_a", "scene_b")]
+EXAMPLES_DIR = APP_DIR / "examples"
+EXAMPLES = [
+    [sorted(str(p) for p in (EXAMPLES_DIR / s).glob("*.jpg"))]
+    for s in ("kitchen", "garden")
+    if (EXAMPLES_DIR / s).is_dir()
+]
 
-DESCRIPTION = """
-# CodecSplat: Ultra-Compact Latent Coding for Feed-Forward 3D Gaussian Splatting
+CSS = """
+#col-container { max-width: 1200px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
 
-CodecSplat turns a few photos into 3D Gaussians **through a real compressed bitstream**: the 2D
-features used for depth and Gaussian prediction are entropy-coded (tens of KB per scene), then
-decoded into a renderable 3D Gaussian scene. Upload overlapping photos of a static scene —
-camera poses are estimated with COLMAP right here — and pick a rate (λ).
+DESCRIPTION = """# CodecSplat: Ultra-Compact Latent Coding for Feed-Forward 3D Gaussian Splatting
 
-[Paper](https://huggingface.co/papers/2605.25563) · [Code](https://github.com/pengpeng-yu/CodecSplat) · [Weights](https://huggingface.co/pengpeng-yu/CodecSplat)
-"""
+CodecSplat doesn't compress the final 3D Gaussians. It entropy-codes the **2D latent features** a feed-forward
+3DGS model uses to predict them, so a whole scene fits in a **few kilobytes** and decodes back to Gaussians
+in one forward pass. Upload overlapping photos of a scene. COLMAP estimates the cameras, CodecSplat
+compresses 8 views into a real bitstream, then decodes it into 3D Gaussians you can fly through.
 
-CSS = """
-#col-container { max-width: 1280px; margin: 0 auto; }
-.dark .gradio-container { color: var(--body-text-color); }
+[Paper](https://huggingface.co/papers/2605.25563) · [Code](https://github.com/pengpeng-yu/CodecSplat) · [Weights](https://huggingface.co/pengpeng-yu/CodecSplat)
 """
 
 with gr.Blocks(title="CodecSplat") as demo:
@@ -255,48 +465,47 @@ with gr.Blocks(title="CodecSplat") as demo:
         with gr.Row():
             with gr.Column(scale=1):
                 image_files = gr.File(
-                    label="Photos of one static scene (8-30 overlapping shots work best)",
+                    label="Photos of one static scene (8-30, overlapping)",
                     file_count="multiple", file_types=["image"], type="filepath",
                 )
-                lmb = gr.Slider(16, 1024, value=1024, step=1, label="Rate λ",
-                                info="16 = smallest bitstream, 1024 = highest quality")
+                lmb = gr.Slider(16, 1024, value=1024, step=1, label="Rate point λ",
+                                info="Higher λ = more bytes and higher fidelity. The model covers 16-1024.")
                 run_btn = gr.Button("Compress & reconstruct", variant="primary")
                 with gr.Accordion("Advanced options", open=False):
                     matcher = gr.Radio(["exhaustive", "sequential"], value="exhaustive",
                                        label="COLMAP matching",
-                                       info="exhaustive = unordered photos; sequential = video frames")
+                                       info="exhaustive = unordered photos; sequential = video frames.")
                 gr.Markdown(
-                    "8 input views are picked by farthest-point sampling over the COLMAP cameras and "
-                    "resized/cropped to 256×448 (the checkpoint's training resolution); up to 8 other "
-                    "photos are held out to measure novel-view quality.\n\n"
+                    "**Tips**: walk around the scene, keep lots of overlap between neighbouring shots, "
+                    "and avoid moving objects. The 8 context views are picked by farthest-point sampling "
+                    "over the COLMAP cameras. The remaining views are held out for evaluation. "
                     "**Viewer**: drag to orbit, right-drag to pan, scroll to zoom."
                 )
             with gr.Column(scale=2):
-                summary = gr.Markdown()
-                splat = gr.Model3D(label="Decoded 3D Gaussians (interactive)", height=440,
-                                   clear_color=[0.0, 0.0, 0.0, 1.0])
-                video = gr.Video(label="Fly-through rendered from the decoded Gaussians", autoplay=True, loop=True)
-                bin_out = gr.File(label="Compressed scene bitstream (.bin)")
-        novel = gr.Gallery(label="Held-out novel views: render | real photo", columns=2, height=320,
-                           object_fit="contain")
-        inputs_gal = gr.Gallery(label="Selected input views (after crop)", columns=4, height=200,
-                                object_fit="contain")
-
-        outputs = [splat, video, novel, inputs_gal, bin_out, summary]
+                summary_out = gr.Markdown()
+                splat_out = gr.Model3D(label="Decoded 3D Gaussians (interactive)", height=440,
+                                       clear_color=[0.0, 0.0, 0.0, 1.0])
+                video_out = gr.Video(label="Fly-through through the context cameras", autoplay=True, loop=True)
+                gallery_out = gr.Gallery(label="Held-out views: CodecSplat | real photo", columns=2,
+                                         height=260, object_fit="contain")
+                with gr.Row():
+                    bin_out = gr.File(label="Compressed scene bitstream (.bin)")
+                    ply_out = gr.File(label="Decoded 3D Gaussians (.ply)")
+
+        outputs = [splat_out, video_out, gallery_out, bin_out, ply_out, summary_out]
         run_btn.click(reconstruct, inputs=[image_files, lmb, matcher], outputs=outputs,
                       api_name="reconstruct")
 
-        gr.Examples(
-            examples=[[EXAMPLES[0], 1024], [EXAMPLES[1], 64]],
-            inputs=[image_files, lmb],
-            outputs=outputs,
-            fn=reconstruct,
-            cache_examples=True,
-            cache_mode="lazy",
-            label="Frames from the DL3DV demo scene (via the DepthSplat authors)",
-        )
-
-_nvidia_fds("after UI build")
+        if EXAMPLES:
+            gr.Examples(
+                examples=EXAMPLES,
+                inputs=[image_files],
+                outputs=outputs,
+                fn=reconstruct,
+                cache_examples=True,
+                cache_mode="lazy",
+                label="Mip-NeRF 360 scenes (kitchen, garden): 14 frames each",
+            )
 
 if __name__ == "__main__":
     demo.queue(max_size=12).launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True, show_error=True)

colmap_io.py

@@ -0,0 +1,414 @@
+"""COLMAP scene loading + view selection.
+
+Verbatim helpers from the ReSplat reference inference script
+(https://github.com/cvg/resplat, scripts/infer_colmap.py, MIT), which shares
+the DepthSplat camera conventions used by CodecSplat.
+"""
+
+import collections
+import os
+import struct
+
+import numpy as np
+import torch
+
+CameraModel = collections.namedtuple(
+    "CameraModel", ["model_id", "model_name", "num_params"]
+)
+ColmapCamera = collections.namedtuple(
+    "ColmapCamera", ["id", "model", "width", "height", "params"]
+)
+BaseImage = collections.namedtuple(
+    "ColmapImage",
+    ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"],
+)
+
+CAMERA_MODELS = {
+    CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3),
+    CameraModel(model_id=1, model_name="PINHOLE", num_params=4),
+    CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4),
+    CameraModel(model_id=3, model_name="RADIAL", num_params=5),
+    CameraModel(model_id=4, model_name="OPENCV", num_params=8),
+    CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8),
+    CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12),
+    CameraModel(model_id=7, model_name="FOV", num_params=5),
+    CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4),
+    CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5),
+    CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12),
+}
+CAMERA_MODEL_IDS = dict(
+    [(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS]
+)
+
+
+class ColmapImage(BaseImage):
+    def qvec2rotmat(self):
+        return qvec2rotmat(self.qvec)
+
+
+def qvec2rotmat(qvec):
+    return np.array(
+        [
+            [
+                1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2,
+                2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
+                2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2],
+            ],
+            [
+                2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
+                1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2,
+                2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1],
+            ],
+            [
+                2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],
+                2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
+                1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2,
+            ],
+        ]
+    )
+
+
+def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
+    data = fid.read(num_bytes)
+    return struct.unpack(endian_character + format_char_sequence, data)
+
+
+def read_extrinsics_binary(path_to_model_file):
+    images = {}
+    with open(path_to_model_file, "rb") as fid:
+        num_reg_images = read_next_bytes(fid, 8, "Q")[0]
+        for _ in range(num_reg_images):
+            binary_image_properties = read_next_bytes(
+                fid, num_bytes=64, format_char_sequence="idddddddi"
+            )
+            image_id = binary_image_properties[0]
+            qvec = np.array(binary_image_properties[1:5])
+            tvec = np.array(binary_image_properties[5:8])
+            camera_id = binary_image_properties[8]
+            image_name = ""
+            current_char = read_next_bytes(fid, 1, "c")[0]
+            while current_char != b"\x00":
+                image_name += current_char.decode("utf-8")
+                current_char = read_next_bytes(fid, 1, "c")[0]
+            num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[
+                0
+            ]
+            x_y_id_s = read_next_bytes(
+                fid,
+                num_bytes=24 * num_points2D,
+                format_char_sequence="ddq" * num_points2D,
+            )
+            xys = np.column_stack(
+                [
+                    tuple(map(float, x_y_id_s[0::3])),
+                    tuple(map(float, x_y_id_s[1::3])),
+                ]
+            )
+            point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3])))
+            images[image_id] = ColmapImage(
+                id=image_id,
+                qvec=qvec,
+                tvec=tvec,
+                camera_id=camera_id,
+                name=image_name,
+                xys=xys,
+                point3D_ids=point3D_ids,
+            )
+    return images
+
+
+def read_extrinsics_text(path):
+    images = {}
+    with open(path, "r") as fid:
+        while True:
+            line = fid.readline()
+            if not line:
+                break
+            line = line.strip()
+            if len(line) > 0 and line[0] != "#":
+                elems = line.split()
+                image_id = int(elems[0])
+                qvec = np.array(tuple(map(float, elems[1:5])))
+                tvec = np.array(tuple(map(float, elems[5:8])))
+                camera_id = int(elems[8])
+                image_name = elems[9]
+                elems = fid.readline().split()
+                xys = np.column_stack(
+                    [
+                        tuple(map(float, elems[0::3])),
+                        tuple(map(float, elems[1::3])),
+                    ]
+                )
+                point3D_ids = np.array(tuple(map(int, elems[2::3])))
+                images[image_id] = ColmapImage(
+                    id=image_id,
+                    qvec=qvec,
+                    tvec=tvec,
+                    camera_id=camera_id,
+                    name=image_name,
+                    xys=xys,
+                    point3D_ids=point3D_ids,
+                )
+    return images
+
+
+def read_intrinsics_binary(path_to_model_file):
+    cameras = {}
+    with open(path_to_model_file, "rb") as fid:
+        num_cameras = read_next_bytes(fid, 8, "Q")[0]
+        for _ in range(num_cameras):
+            camera_properties = read_next_bytes(
+                fid, num_bytes=24, format_char_sequence="iiQQ"
+            )
+            camera_id = camera_properties[0]
+            model_id = camera_properties[1]
+            model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name
+            width = camera_properties[2]
+            height = camera_properties[3]
+            num_params = CAMERA_MODEL_IDS[model_id].num_params
+            params = read_next_bytes(
+                fid, num_bytes=8 * num_params, format_char_sequence="d" * num_params
+            )
+            cameras[camera_id] = ColmapCamera(
+                id=camera_id,
+                model=model_name,
+                width=width,
+                height=height,
+                params=np.array(params),
+            )
+        assert len(cameras) == num_cameras
+    return cameras
+
+
+def read_intrinsics_text(path):
+    cameras = {}
+    with open(path, "r") as fid:
+        while True:
+            line = fid.readline()
+            if not line:
+                break
+            line = line.strip()
+            if len(line) > 0 and line[0] != "#":
+                elems = line.split()
+                camera_id = int(elems[0])
+                model = elems[1]
+                width = int(elems[2])
+                height = int(elems[3])
+                params = np.array(tuple(map(float, elems[4:])))
+                cameras[camera_id] = ColmapCamera(
+                    id=camera_id,
+                    model=model,
+                    width=width,
+                    height=height,
+                    params=params,
+                )
+    return cameras
+
+
+# =============================================================================
+# Scene Loading
+# =============================================================================
+
+
+def load_colmap_scene(scene_path, sparse_dir="sparse/0", images_dir="images"):
+    """Load COLMAP reconstruction and convert to ReSplat format.
+
+    Returns dict with:
+        image_paths: list[str] - full paths to images
+        image_names: list[str] - image filenames
+        c2w: np.ndarray [N, 4, 4] - camera-to-world matrices (OpenCV convention)
+        intrinsics: np.ndarray [N, 3, 3] - normalized intrinsic matrices
+        image_sizes: list[tuple[int, int]] - (width, height) per image
+    """
+    sparse_path = os.path.join(scene_path, sparse_dir)
+
+    # Try binary format first, fall back to text
+    cameras_bin = os.path.join(sparse_path, "cameras.bin")
+    cameras_txt = os.path.join(sparse_path, "cameras.txt")
+    images_bin = os.path.join(sparse_path, "images.bin")
+    images_txt = os.path.join(sparse_path, "images.txt")
+
+    if os.path.exists(cameras_bin) and os.path.exists(images_bin):
+        cam_intrinsics = read_intrinsics_binary(cameras_bin)
+        cam_extrinsics = read_extrinsics_binary(images_bin)
+        print(f"Loaded COLMAP binary format from {sparse_path}")
+    elif os.path.exists(cameras_txt) and os.path.exists(images_txt):
+        cam_intrinsics = read_intrinsics_text(cameras_txt)
+        cam_extrinsics = read_extrinsics_text(images_txt)
+        print(f"Loaded COLMAP text format from {sparse_path}")
+    else:
+        raise FileNotFoundError(
+            f"No COLMAP reconstruction found in {sparse_path}. "
+            "Expected cameras.bin/txt and images.bin/txt"
+        )
+
+    images_root = os.path.join(scene_path, images_dir)
+
+    # Sort images by name for deterministic ordering
+    sorted_images = sorted(cam_extrinsics.values(), key=lambda x: x.name)
+
+    image_paths = []
+    image_names = []
+    c2w_list = []
+    intrinsics_list = []
+    image_sizes = []
+    point3D_ids_list = []
+
+    for img in sorted_images:
+        # Check image exists
+        img_path = os.path.join(images_root, img.name)
+        if not os.path.exists(img_path):
+            print(f"Warning: Image not found, skipping: {img_path}")
+            continue
+
+        # Get camera intrinsics
+        cam = cam_intrinsics[img.camera_id]
+
+        # Only support PINHOLE and SIMPLE_PINHOLE (undistorted images)
+        if cam.model == "PINHOLE":
+            fx, fy, cx, cy = cam.params[:4]
+        elif cam.model == "SIMPLE_PINHOLE":
+            f, cx, cy = cam.params[:3]
+            fx = fy = f
+        else:
+            print(
+                f"Warning: Unsupported camera model '{cam.model}' for image "
+                f"{img.name}, skipping. Only PINHOLE and SIMPLE_PINHOLE are "
+                f"supported. Please undistort images first using COLMAP."
+            )
+            continue
+
+        # Build normalized intrinsic matrix (fx/width, fy/height, cx/width, cy/height)
+        K = np.eye(3, dtype=np.float32)
+        K[0, 0] = fx / cam.width
+        K[1, 1] = fy / cam.height
+        K[0, 2] = cx / cam.width
+        K[1, 2] = cy / cam.height
+
+        # Build W2C matrix from COLMAP qvec + tvec
+        # COLMAP: qvec is world-to-camera rotation, tvec is translation in camera frame
+        R = qvec2rotmat(img.qvec)
+        w2c = np.eye(4, dtype=np.float32)
+        w2c[:3, :3] = R
+        w2c[:3, 3] = img.tvec
+
+        # Invert to get C2W (camera-to-world)
+        # COLMAP is already in OpenCV convention (Y-down, Z-forward)
+        c2w = np.linalg.inv(w2c).astype(np.float32)
+
+        # Extract valid point3D_ids (filter out -1 which means unmatched)
+        valid_pts = img.point3D_ids[img.point3D_ids >= 0]
+
+        image_paths.append(img_path)
+        image_names.append(img.name)
+        c2w_list.append(c2w)
+        intrinsics_list.append(K)
+        image_sizes.append((cam.width, cam.height))
+        point3D_ids_list.append(set(valid_pts.tolist()))
+
+    if len(image_paths) == 0:
+        raise RuntimeError(f"No valid images found in {images_root}")
+
+    return {
+        "image_paths": image_paths,
+        "image_names": image_names,
+        "c2w": np.stack(c2w_list, axis=0),  # [N, 4, 4]
+        "intrinsics": np.stack(intrinsics_list, axis=0),  # [N, 3, 3]
+        "image_sizes": image_sizes,
+        "point3D_ids": point3D_ids_list,  # list of sets, one per image
+    }
+
+
+# =============================================================================
+# View Selection
+# =============================================================================
+
+
+def farthest_point_sample(xyz, npoint):
+    """Farthest point sampling on camera positions.
+
+    Adapted from src/dataset/view_sampler/view_sampler_bounded_v2.py
+
+    Args:
+        xyz: [B, N, 3] point cloud data
+        npoint: number of samples
+    Returns:
+        centroids: [B, npoint] sampled indices
+    """
+    device = xyz.device
+    B, N, C = xyz.shape
+
+    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
+    distance = torch.ones(B, N).to(device) * 1e10
+
+    batch_indices = torch.arange(B, dtype=torch.long).to(device)
+
+    barycenter = torch.sum(xyz, 1)
+    barycenter = barycenter / xyz.shape[1]
+    barycenter = barycenter.view(B, 1, 3)
+
+    dist = torch.sum((xyz - barycenter) ** 2, -1)
+    farthest = torch.max(dist, 1)[1]
+
+    for i in range(npoint):
+        centroids[:, i] = farthest
+        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
+        dist = torch.sum((xyz - centroid) ** 2, -1)
+        mask = dist < distance
+        distance[mask] = dist[mask]
+        farthest = torch.max(distance, -1)[1]
+
+    return centroids
+
+
+def select_context_views(c2w, num_context, strategy="fps"):
+    """Select context view indices from all available views.
+
+    Args:
+        c2w: [N, 4, 4] camera-to-world matrices
+        num_context: number of context views to select
+        strategy: "fps" | "uniform"
+    Returns:
+        context_indices: sorted numpy array of indices
+    """
+    N = len(c2w)
+    if num_context >= N:
+        return np.arange(N)
+
+    if strategy == "fps":
+        positions = torch.tensor(c2w[:, :3, 3], dtype=torch.float32).unsqueeze(0)
+        indices = farthest_point_sample(positions, num_context)[0].numpy()
+        return np.sort(indices)
+    elif strategy == "uniform":
+        return np.linspace(0, N - 1, num_context, dtype=int)
+    else:
+        raise ValueError(f"Unknown context selection strategy: {strategy}")
+
+
+def select_target_views(num_images, context_indices, target_selection="remaining",
+                        num_target=None):
+    """Select target views for rendering.
+
+    Args:
+        num_images: total number of images
+        context_indices: indices used as context
+        target_selection: "remaining" | "all"
+        num_target: optional max number of target views
+    Returns:
+        target_indices: numpy array of indices
+    """
+    if target_selection == "remaining":
+        all_indices = np.arange(num_images)
+        target_indices = np.setdiff1d(all_indices, context_indices)
+    elif target_selection == "all":
+        target_indices = np.arange(num_images)
+    else:
+        raise ValueError(f"Unknown target selection: {target_selection}")
+
+    if num_target is not None and len(target_indices) > num_target:
+        step = len(target_indices) / num_target
+        selected = [int(i * step) for i in range(num_target)]
+        target_indices = target_indices[selected]
+
+    return target_indices
+

colmap_sfm.py

@@ -7,7 +7,7 @@ to a PINHOLE camera model and writes a standard COLMAP scene layout
     <out>/images/*.jpg
     <out>/sparse/{cameras,images,points3D}.bin
 
-which is exactly what CodecSplat's reference COLMAP loader expects.
+which is exactly what ReSplat's reference COLMAP loader expects.
 """
 
 from __future__ import annotations
@@ -18,11 +18,6 @@ from pathlib import Path
 import pycolmap
 from PIL import Image, ImageOps
 
-# ZeroGPU hosts expose ~192 CPUs. With COLMAP's default num_threads=-1 the incremental
-# mapper then fails to register anything beyond the initial pair (reproduced locally by
-# forcing 192 threads), so cap the thread count explicitly.
-NUM_THREADS = 16
-
 IMAGE_EXT = (".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff")
 
 
@@ -80,12 +75,6 @@ def run_sfm(
             "Upload several overlapping photos of the same scene."
         )
     log(f"[colmap] {len(names)} images prepared at <= {max_image_size}px")
-    try:
-        import os as _os
-        cpu = next((l.split(":", 1)[1].strip() for l in open("/proc/cpuinfo") if l.startswith("model name")), "?")
-        log(f"[colmap] cpu={cpu} ncpu={_os.cpu_count()} threads={NUM_THREADS}")
-    except Exception:
-        pass
 
     database_path = workdir / "colmap.db"
     if database_path.exists():
@@ -97,7 +86,6 @@ def run_sfm(
     extraction_options.sift.domain_size_pooling = False
     extraction_options.use_gpu = False
     extraction_options.max_image_size = int(max_image_size)
-    extraction_options.num_threads = NUM_THREADS
 
     reader_options = pycolmap.ImageReaderOptions()
     reader_options.camera_model = "SIMPLE_RADIAL"
@@ -112,28 +100,23 @@ def run_sfm(
     )
     log("[colmap] features extracted")
 
-    matching_options = pycolmap.FeatureMatchingOptions()
-    matching_options.num_threads = NUM_THREADS
     if matcher == "sequential":
         pairing = pycolmap.SequentialPairingOptions()
         pairing.overlap = 10
         pairing.quadratic_overlap = True
         pycolmap.match_sequential(
             database_path=database_path,
-            matching_options=matching_options,
             pairing_options=pairing,
             device=pycolmap.Device.cpu,
         )
     else:
-        pycolmap.match_exhaustive(database_path=database_path, matching_options=matching_options,
-                                  device=pycolmap.Device.cpu)
+        pycolmap.match_exhaustive(database_path=database_path, device=pycolmap.Device.cpu)
     log(f"[colmap] {matcher} matching done")
 
     sparse_dir = workdir / "sparse_raw"
     sparse_dir.mkdir(parents=True, exist_ok=True)
     options = pycolmap.IncrementalPipelineOptions()
     options.min_num_matches = 15
-    options.num_threads = NUM_THREADS
     options.multiple_models = False
     options.extract_colors = True
     maps = pycolmap.incremental_mapping(
@@ -151,7 +134,7 @@ def run_sfm(
             "rotation / blurry frames, and keep the scene static."
         )
     # Deterministic world scale (same convention COLMAP itself uses), so that the
-    # fixed near/far planes of the CodecSplat checkpoint stay meaningful.
+    # fixed near/far planes of the ReSplat checkpoint stay meaningful.
     rec.normalize(fixed_scale=False, extent=10.0)
     log(f"[colmap] mapped {rec.num_reg_images()}/{len(names)} images, {rec.num_points3D()} points")
 
@@ -171,7 +154,7 @@ def run_sfm(
     )
     log("[colmap] undistorted to PINHOLE")
 
-    # Keep only what CodecSplat reads (the undistorter also writes stereo/ scaffolding).
+    # Keep only what ReSplat reads (the undistorter also writes stereo/ scaffolding).
     for extra in ("stereo",):
         p = scene_dir / extra
         if p.is_dir():
@@ -183,32 +166,7 @@ def run_sfm(
     except Exception:
         pass
 
-    # Read back the undistorted PINHOLE cameras (normalized K, OpenCV c2w), sorted by name.
-    und = pycolmap.Reconstruction(str(scene_dir / "sparse"))
-    frames = []
-    for _, img in und.images.items():
-        cam = und.cameras[img.camera_id]
-        p = list(cam.params)
-        if cam.model.name == "PINHOLE":
-            fx, fy, cx, cy = p[:4]
-        elif cam.model.name == "SIMPLE_PINHOLE":
-            fx = fy = p[0]
-            cx, cy = p[1:3]
-        else:
-            continue
-        cfw = img.cam_from_world() if callable(img.cam_from_world) else img.cam_from_world
-        w2c = cfw.matrix()  # 3x4
-        path = scene_dir / "images" / img.name
-        if path.exists():
-            frames.append({
-                "name": img.name, "path": str(path),
-                "K": [[fx / cam.width, 0, cx / cam.width], [0, fy / cam.height, cy / cam.height], [0, 0, 1]],
-                "w2c": [list(map(float, row)) for row in w2c],
-            })
-    frames.sort(key=lambda f: f["name"])
-
     return {
-        "frames": frames,
         "scene_dir": str(scene_dir),
         "sparse_sub": "sparse",
         "images_sub": "images",
@@ -217,55 +175,3 @@ def run_sfm(
         "num_points3D": int(rec.num_points3D()),
         "camera_model": cam_model,
     }
-
-
-def run_sfm_subprocess(image_paths, workdir, max_image_size=1280, matcher="exhaustive",
-                       timeout=900, log=print) -> dict:
-    """Run `run_sfm` in a fresh Python process that only imports pycolmap.
-
-    Inside the Gradio process torch / CUDA / OpenMP / BLAS libraries are already loaded,
-    which makes pycolmap's bundled OpenBLAS / OpenMP misbehave (crashes in matching and
-    failed image registration). A clean child process avoids all of that.
-    """
-    import json
-    import os
-    import subprocess
-    import sys
-
-    workdir = Path(workdir)
-    workdir.mkdir(parents=True, exist_ok=True)
-    spec = workdir / "sfm_job.json"
-    result = workdir / "sfm_result.json"
-    spec.write_text(json.dumps({
-        "image_paths": [str(p) for p in image_paths], "workdir": str(workdir),
-        "max_image_size": int(max_image_size), "matcher": matcher, "result": str(result),
-    }))
-    env = {k: v for k, v in os.environ.items() if k not in ("LD_PRELOAD",)}
-    env["LD_LIBRARY_PATH"] = os.environ.get("_ORIG_LD_LIBRARY_PATH", "")
-    env["OPENBLAS_NUM_THREADS"] = "1"
-    env.pop("CUDA_VISIBLE_DEVICES", None)
-    proc = subprocess.run([sys.executable, os.path.abspath(__file__), str(spec)], env=env,
-                          capture_output=True, text=True, timeout=timeout)
-    for line in (proc.stdout or "").splitlines():
-        if line.startswith("[colmap]"):
-            log(line)
-    if not result.exists():
-        tail = "\n".join(((proc.stderr or "") + (proc.stdout or "")).splitlines()[-15:])
-        raise RuntimeError(f"COLMAP subprocess failed (exit {proc.returncode}):\n{tail}")
-    data = json.loads(result.read_text())
-    if "error" in data:
-        raise ValueError(data["error"])
-    return data
-
-
-if __name__ == "__main__":
-    import json
-    import sys
-
-    job = json.loads(Path(sys.argv[1]).read_text())
-    try:
-        out = run_sfm(job["image_paths"], Path(job["workdir"]), job["max_image_size"],
-                      job["matcher"], log=lambda m: print(m, flush=True))
-    except ValueError as exc:
-        out = {"error": str(exc)}
-    Path(job["result"]).write_text(json.dumps(out))

requirements.txt

@@ -1,24 +1,27 @@
-torch==2.11.0
-torchvision==0.26.0
-https://huggingface.co/datasets/multimodalart/zerogpu-blackwell-wheels/resolve/main/wheels/pt211-cu130-cp312/diff_gaussian_rasterization-0.0.0-cp312-cp312-linux_x86_64.whl
-compressai==1.2.8
+torchvision
+gsplat==1.5.3
+ninja
+rich
 numpy<2
+compressai==1.2.8
 pycolmap==4.1.1
-einops
-jaxtyping>=0.3.7,<0.3.11
 beartype
-hydra-core
-omegaconf
-dacite
 colorama
 colorspacious
+dacite
 e3nn
+einops
+hydra-core
+omegaconf
+imageio
+imageio-ffmpeg
+jaxtyping>=0.3.7,<0.3.11
 lpips
-plyfile
-scipy
 matplotlib
 opencv-python-headless
 Pillow
+plyfile
+pytorch-lightning
+scipy
 tqdm
-imageio
-imageio-ffmpeg
+wandb

src/main.py

@@ -0,0 +1,339 @@
+import os
+from pathlib import Path
+import copy
+
+import hydra
+import torch
+import wandb
+from colorama import Fore
+from jaxtyping import install_import_hook
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import (
+    LearningRateMonitor,
+    ModelCheckpoint,
+)
+from pytorch_lightning.loggers.wandb import WandbLogger
+
+from pytorch_lightning.plugins.environments import LightningEnvironment
+
+
+# Configure beartype and jaxtyping.
+with install_import_hook(
+    ("src",),
+    ("beartype", "beartype"),
+):
+    from src.config import load_typed_root_config
+    from src.dataset.data_module import DataModule
+    from src.global_cfg import set_cfg
+    from src.loss import get_losses
+    from src.misc.LocalLogger import LocalLogger
+    from src.misc.step_tracker import StepTracker
+    from src.misc.wandb_tools import update_checkpoint_path
+    from src.misc.resume_ckpt import find_latest_ckpt
+    from src.model.decoder import get_decoder
+    from src.model.encoder import get_encoder
+    from src.model.model_wrapper import ModelWrapper
+
+
+def cyan(text: str) -> str:
+    return f"{Fore.CYAN}{text}{Fore.RESET}"
+
+
+def load_state_dict_with_shape_filter(
+    module: torch.nn.Module,
+    state_dict: dict[str, torch.Tensor],
+    strict: bool,
+    label: str,
+) -> None:
+    if strict:
+        module.load_state_dict(state_dict, strict=True)
+        return
+
+    module_state = module.state_dict()
+    filtered_state_dict = {}
+    shape_mismatches = []
+    for key, value in state_dict.items():
+        if key not in module_state:
+            continue
+        if module_state[key].shape != value.shape:
+            shape_mismatches.append(
+                (key, tuple(value.shape), tuple(module_state[key].shape))
+            )
+            continue
+        filtered_state_dict[key] = value
+
+    incompatible = module.load_state_dict(filtered_state_dict, strict=False)
+    missing_keys = sorted(incompatible.missing_keys)
+    unexpected_keys = sorted(incompatible.unexpected_keys)
+
+    print(cyan(f"{label}: strict=False, filtered checkpoint loading summary"))
+    print(f"  loaded tensors: {len(filtered_state_dict)}")
+    print(f"  missing model tensors: {len(missing_keys)}")
+    for key in missing_keys:
+        print(f"    [missing] {key}")
+    print(f"  shape-mismatched checkpoint tensors: {len(shape_mismatches)}")
+    for key, ckpt_shape, model_shape in shape_mismatches:
+        print(
+            f"    [shape-mismatch] {key}: checkpoint {ckpt_shape} vs model {model_shape}"
+        )
+    print(f"  unexpected checkpoint tensors: {len(unexpected_keys)}")
+    for key in unexpected_keys:
+        print(f"    [unexpected] {key}")
+
+
+@hydra.main(
+    version_base=None,
+    config_path="../config",
+    config_name="main",
+)
+def train(cfg_dict: DictConfig):
+    if cfg_dict["mode"] == "train" and cfg_dict["train"]["eval_model_every_n_val"] > 0:
+        eval_cfg_dict = copy.deepcopy(cfg_dict)
+        dataset_dir = str(cfg_dict["dataset"]["roots"]).lower()
+        eval_num_context_views = cfg_dict["dataset"]["view_sampler"][
+            "num_context_views"
+        ]
+        if "re10k" in dataset_dir:
+            if cfg_dict["dataset"]["view_sampler"]["num_context_views"] == 2:
+                eval_path = "assets/evaluation_index_re10k.json"
+            else:
+                raise ValueError("unsupported number of views for re10k")
+        elif "dl3dv" in dataset_dir:
+            if eval_num_context_views == 8:
+                eval_path = (
+                    "assets/dl3dv_evaluation/"
+                    "dl3dv_start_0_distance_40_ctx_8v_tgt_8v.json"
+                )
+            else:
+                raise ValueError("unsupported number of views for dl3dv")
+        else:
+            raise Exception("Fail to load eval index path")
+        eval_cfg_dict["dataset"]["view_sampler"] = {
+            "name": "evaluation",
+            "index_path": eval_path,
+            "num_context_views": eval_num_context_views,
+        }
+        eval_cfg = load_typed_root_config(eval_cfg_dict)
+    else:
+        eval_cfg = None
+
+    cfg = load_typed_root_config(cfg_dict)
+    set_cfg(cfg_dict)
+
+    # Set up the output directory.
+    if cfg_dict.output_dir is None:
+        output_dir = Path(
+            hydra.core.hydra_config.HydraConfig.get()["runtime"]["output_dir"]
+        )
+    else:  # for resuming
+        output_dir = Path(cfg_dict.output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+    print(cyan(f"Saving outputs to {output_dir}."))
+
+    # Set up logging with wandb.
+    callbacks = []
+    if cfg_dict.wandb.mode != "disabled" and cfg.mode == "train":
+        wandb_extra_kwargs = {}
+        if cfg_dict.wandb.id is not None:
+            wandb_extra_kwargs.update({'id': cfg_dict.wandb.id,
+                                       'resume': "must"})
+        logger = WandbLogger(
+            entity=cfg_dict.wandb.entity,
+            project=cfg_dict.wandb.project,
+            mode=cfg_dict.wandb.mode,
+            name=os.path.basename(cfg_dict.output_dir),
+            tags=cfg_dict.wandb.get("tags", None),
+            log_model=False,
+            save_dir=output_dir,
+            config=OmegaConf.to_container(cfg_dict),
+            **wandb_extra_kwargs,
+        )
+        callbacks.append(LearningRateMonitor("step", True))
+
+        if wandb.run is not None:
+            wandb.run.log_code("src")
+    else:
+        logger = LocalLogger(output_dir / "local")
+
+    # Set up checkpointing.
+    callbacks.append(
+        ModelCheckpoint(
+            output_dir / "checkpoints",
+            every_n_train_steps=cfg.checkpointing.every_n_train_steps,
+            save_top_k=cfg.checkpointing.save_top_k,
+            monitor="info/global_step",
+            mode="max",
+        )
+    )
+    for cb in callbacks:
+        cb.CHECKPOINT_EQUALS_CHAR = '_'
+
+    # Prepare the checkpoint for loading.
+    if cfg.checkpointing.resume:
+        if not os.path.exists(output_dir / 'checkpoints'):
+            checkpoint_path = None
+        else:
+            checkpoint_path = find_latest_ckpt(output_dir / 'checkpoints')
+            print(f'resume from {checkpoint_path}')
+    else:
+        checkpoint_path = update_checkpoint_path(cfg.checkpointing.load, cfg.wandb)
+
+    # This allows the current step to be shared with the data loader processes.
+    step_tracker = StepTracker()
+
+    if cfg.mode == "test":
+        trainer_devices = 1
+        trainer_strategy = "auto"
+        trainer_num_nodes = 1
+    else:
+        trainer_devices = torch.cuda.device_count()
+        trainer_strategy = "ddp" if trainer_devices > 1 else "auto"
+        trainer_num_nodes = cfg.trainer.num_nodes
+
+    trainer = Trainer(
+        max_epochs=-1,
+        accelerator="gpu",
+        logger=logger,
+        devices=trainer_devices,
+        strategy=trainer_strategy,
+        callbacks=callbacks,
+        val_check_interval=cfg.trainer.val_check_interval,
+        enable_progress_bar=cfg.mode == "test",
+        gradient_clip_val=cfg.trainer.gradient_clip_val,
+        accumulate_grad_batches=cfg.trainer.accumulate_grad_batches,
+        precision=cfg.trainer.precision,
+        max_steps=cfg.trainer.max_steps,
+        num_sanity_val_steps=cfg.trainer.num_sanity_val_steps,
+        num_nodes=trainer_num_nodes,
+        plugins=LightningEnvironment() if cfg.use_plugins else None,
+    )
+    torch.manual_seed(cfg_dict.seed + trainer.global_rank)
+
+    encoder = get_encoder(cfg.model.encoder)
+
+    model_wrapper = ModelWrapper(
+        cfg.optimizer,
+        cfg.test,
+        cfg.train,
+        encoder,
+        get_decoder(cfg.model.decoder, cfg.dataset),
+        get_losses(cfg.loss),
+        step_tracker,
+        eval_data_cfg=(
+            None if eval_cfg is None else eval_cfg.dataset
+        ),
+    )
+    data_module = DataModule(
+        cfg.dataset,
+        cfg.data_loader,
+        step_tracker,
+        global_rank=trainer.global_rank,
+    )
+
+    if cfg.mode == "train":
+        print("train:", len(data_module.train_dataloader()))
+        print("val:", len(data_module.val_dataloader()))
+        print("test:", len(data_module.test_dataloader()))
+
+    strict_load = not cfg.checkpointing.no_strict_load
+
+    if cfg.mode == "train":
+        # only load monodepth
+        if cfg.checkpointing.pretrained_monodepth is not None:
+            strict_load = False
+            pretrained_model = torch.load(cfg.checkpointing.pretrained_monodepth, map_location='cpu')
+            if 'state_dict' in pretrained_model:
+                pretrained_model = pretrained_model['state_dict']
+
+            model_wrapper.encoder.depth_predictor.load_state_dict(pretrained_model, strict=strict_load)
+            print(
+                cyan(
+                    f"Loaded pretrained monodepth: {cfg.checkpointing.pretrained_monodepth}"
+                )
+            )
+
+        # load pretrained mvdepth
+        if cfg.checkpointing.pretrained_mvdepth is not None:
+            pretrained_model = torch.load(cfg.checkpointing.pretrained_mvdepth, map_location='cpu')['model']
+
+            model_wrapper.encoder.depth_predictor.load_state_dict(pretrained_model, strict=False)
+            print(
+                cyan(
+                    f"Loaded pretrained mvdepth: {cfg.checkpointing.pretrained_mvdepth}"
+                )
+            )
+
+        # load full model
+        if cfg.checkpointing.pretrained_model is not None:
+            pretrained_model = torch.load(cfg.checkpointing.pretrained_model, map_location='cpu')
+            if 'state_dict' in pretrained_model:
+                pretrained_model = pretrained_model['state_dict']
+
+            load_state_dict_with_shape_filter(
+                model_wrapper,
+                pretrained_model,
+                strict=strict_load,
+                label="pretrained_model",
+            )
+            print(
+                cyan(
+                    f"Loaded pretrained weights: {cfg.checkpointing.pretrained_model}"
+                )
+            )
+
+        # load pretrained depth
+        if cfg.checkpointing.pretrained_depth is not None:
+            pretrained_model = torch.load(cfg.checkpointing.pretrained_depth, map_location='cpu')['model']
+
+            strict_load = True
+            model_wrapper.encoder.depth_predictor.load_state_dict(pretrained_model, strict=strict_load)
+            print(
+                cyan(
+                    f"Loaded pretrained depth: {cfg.checkpointing.pretrained_depth}"
+                )
+            )
+
+        trainer.fit(model_wrapper, datamodule=data_module, ckpt_path=checkpoint_path)
+    else:
+        # load full model
+        if cfg.checkpointing.pretrained_model is not None:
+            pretrained_model = torch.load(cfg.checkpointing.pretrained_model, map_location='cpu')
+            if 'state_dict' in pretrained_model:
+                pretrained_model = pretrained_model['state_dict']
+
+            load_state_dict_with_shape_filter(
+                model_wrapper,
+                pretrained_model,
+                strict=strict_load,
+                label="pretrained_model",
+            )
+            print(
+                cyan(
+                    f"Loaded pretrained weights: {cfg.checkpointing.pretrained_model}"
+                )
+            )
+
+        # load pretrained depth model only
+        if cfg.checkpointing.pretrained_depth is not None:
+            pretrained_model = torch.load(cfg.checkpointing.pretrained_depth, map_location='cpu')['model']
+
+            strict_load = True
+            model_wrapper.encoder.depth_predictor.load_state_dict(pretrained_model, strict=strict_load)
+            print(
+                cyan(
+                    f"Loaded pretrained depth: {cfg.checkpointing.pretrained_depth}"
+                )
+            )
+
+        trainer.test(
+            model_wrapper,
+            datamodule=data_module,
+            ckpt_path=checkpoint_path,
+        )
+
+
+if __name__ == "__main__":
+    torch.set_float32_matmul_precision('high')
+
+    train()