diff --git a/.gitattributes b/.gitattributes index fde9b584b5e21cf82984e8cdc007e7ea0d9b148e..a44fcb2698ffdaeee5ebbb384a2bcd7daa048631 100644 --- a/.gitattributes +++ b/.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 diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..755da3c59a3f5203aa137aac0be747812d3013cf --- /dev/null +++ b/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. diff --git a/README.md b/README.md index c947873a6a88066771b54712de7799c3d51f9374..7ff77c3954dd9be6110aab01d63069be3492ad93 100644 --- a/README.md +++ b/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). diff --git a/app.py b/app.py index aec87001d9bbb2e48c95cf877b9e079cb9a9b8e9..5f471e120f9a6a79a37fa97b51c53eea02549d5c 100644 --- a/app.py +++ b/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) diff --git a/colmap_io.py b/colmap_io.py new file mode 100644 index 0000000000000000000000000000000000000000..3b75508f805df55ff6b1af6204ff09847c229e54 --- /dev/null +++ b/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 + diff --git a/colmap_sfm.py b/colmap_sfm.py index 69ab71eaa6df06e1017f276956ac66801f8f49ee..ffb65a05320dbdb0ac606519cff15e6c4fc8ed7f 100644 --- a/colmap_sfm.py +++ b/colmap_sfm.py @@ -7,7 +7,7 @@ to a PINHOLE camera model and writes a standard COLMAP scene layout /images/*.jpg /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)) diff --git a/examples/counter/counter_00.jpg b/examples/counter/counter_00.jpg new file mode 100644 index 0000000000000000000000000000000000000000..94b1f27dd5f6ff757b6ef6e09c378305f9e03381 --- /dev/null +++ b/examples/counter/counter_00.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:55861fb1ef4c3706f791792f6499dc8362abec51e08e687e431f235d7aae745a +size 155343 diff --git a/examples/counter/counter_01.jpg b/examples/counter/counter_01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..59a91f0c23967bfd96ee856859da0c8f252f9bc3 --- /dev/null +++ b/examples/counter/counter_01.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9cd4f301243ff680ede6bb14d0626297358f343237a6555ec1dd0d9867f18c5f +size 150119 diff --git a/examples/counter/counter_02.jpg b/examples/counter/counter_02.jpg new file mode 100644 index 0000000000000000000000000000000000000000..de5a35eb8b038e268d7e737b36c46a538d0d599a --- /dev/null +++ b/examples/counter/counter_02.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:368b0d29e19e95f8d2e022fa1786d348fcc1c6826bfd6a5809473370edb571a4 +size 152871 diff --git a/examples/counter/counter_03.jpg b/examples/counter/counter_03.jpg new file mode 100644 index 0000000000000000000000000000000000000000..2395424822b7faffd26c451a38fe8ae81067cbba --- /dev/null +++ b/examples/counter/counter_03.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:80e632d5a9131dcbe9979002cac5a148cdf9d502921a9b16cc3ff8adca1a2d5d +size 165327 diff --git a/examples/counter/counter_04.jpg b/examples/counter/counter_04.jpg new file mode 100644 index 0000000000000000000000000000000000000000..bf639ca20b926ca930a230c76c09fe698e15bb69 --- /dev/null +++ b/examples/counter/counter_04.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9b0b1ef3118ae61cff522eaa589d90f0a8ba8f4d6e6c6eed50816003aa43917e +size 158799 diff --git a/examples/counter/counter_05.jpg b/examples/counter/counter_05.jpg new file mode 100644 index 0000000000000000000000000000000000000000..3b2e79ffc4aebd0990057f50a3bca769f163d128 --- /dev/null +++ b/examples/counter/counter_05.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d3d24b2aa67d18d66880c493e893a30be1fc01415e605956790fcdef94ab441a +size 157310 diff --git a/examples/counter/counter_06.jpg b/examples/counter/counter_06.jpg new file mode 100644 index 0000000000000000000000000000000000000000..ec11485091fa4fe4fbd7f0ad3b4c8aa03ec78fae --- /dev/null +++ b/examples/counter/counter_06.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6fa5aa6458400c1f5720ff0b34e4e9f87b016e4adea92bdab089819feb3d8ef3 +size 145093 diff --git a/examples/counter/counter_07.jpg b/examples/counter/counter_07.jpg new file mode 100644 index 0000000000000000000000000000000000000000..6a0b5a7d14abf8456f9247cc39fa7a39177860eb --- /dev/null +++ b/examples/counter/counter_07.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a0270c24dfc52c690b97b8b01bce0ccb24a6792fba871798a8e252df62ab5819 +size 160714 diff --git a/examples/counter/counter_08.jpg b/examples/counter/counter_08.jpg new file mode 100644 index 0000000000000000000000000000000000000000..faa34b9b553bab0be4413ada2d647061d718f940 --- /dev/null +++ b/examples/counter/counter_08.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d8a765d64e817d2fbbb3bf793e23b172ce460237d9c8a86944ee8b130121a50a +size 132544 diff --git a/examples/counter/counter_09.jpg b/examples/counter/counter_09.jpg new file mode 100644 index 0000000000000000000000000000000000000000..9087235e4b9c23a402bc43a9cdad25fd494ca8db --- /dev/null +++ b/examples/counter/counter_09.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:326bce521fe4bae35c6db4974bbd33edf11285531398a5461d94009a7bd897a3 +size 139395 diff --git a/examples/counter/counter_10.jpg b/examples/counter/counter_10.jpg new file mode 100644 index 0000000000000000000000000000000000000000..a837d416e7c004db96bc47fcbe802dffa98fa1f6 --- /dev/null +++ b/examples/counter/counter_10.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d4436bf228ecb79b9f6f065df8dd57a5c0d2af63339d3b07bdab04b77620c51f +size 130991 diff --git a/examples/counter/counter_11.jpg b/examples/counter/counter_11.jpg new file mode 100644 index 0000000000000000000000000000000000000000..652bd81395f06e689757838bcc7d083d0ac1eba8 --- /dev/null +++ b/examples/counter/counter_11.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:77c459e34478d364b486349c269e6c6fc7f52a38d3f92d6e4a9d2c1291d3bf25 +size 133649 diff --git a/examples/counter/counter_12.jpg b/examples/counter/counter_12.jpg new file mode 100644 index 0000000000000000000000000000000000000000..a4ba4e81668585c9708b7118b1e50e8f520a0aa7 --- /dev/null +++ b/examples/counter/counter_12.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3681d14a5f00f93baaa2851474bc8cfe33af5f6df8c283bdbaf908bf9a9c4f2d +size 133966 diff --git a/examples/counter/counter_13.jpg b/examples/counter/counter_13.jpg new file mode 100644 index 0000000000000000000000000000000000000000..91fcbc2491eb31eb46525f6dff75740523749476 --- /dev/null +++ b/examples/counter/counter_13.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d892ac18f9e40600cf0ed7fc3bbaa85038b4989815c1f9322737691aa1d0a78f +size 132173 diff --git a/examples/garden/garden_00.jpg b/examples/garden/garden_00.jpg new file mode 100644 index 0000000000000000000000000000000000000000..98e758c9a4793240b96d02f00019019bfdde9419 --- /dev/null +++ b/examples/garden/garden_00.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e4d991b1ef48375a72cab84679d15224e7e733f55ce6818d279de6d24462b718 +size 262605 diff --git a/examples/garden/garden_01.jpg b/examples/garden/garden_01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..121377a7b818e397b3764059a54127807844f370 --- /dev/null +++ b/examples/garden/garden_01.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ee05e4a6b6f0d97f3acdfc7c217ab4638be72525a0d9c5fe3f094a4356d28efa +size 262742 diff --git a/examples/garden/garden_02.jpg b/examples/garden/garden_02.jpg new file mode 100644 index 0000000000000000000000000000000000000000..fee4c3b37ed60b296292d088b7b7596e3587c58b --- /dev/null +++ b/examples/garden/garden_02.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3283c838528eb146bb81d822072bd51d8a3fbeb7a8c4df3974062eb7a8320e77 +size 256832 diff --git a/examples/garden/garden_03.jpg b/examples/garden/garden_03.jpg new file mode 100644 index 0000000000000000000000000000000000000000..05f191883eef1ade0d4150ba442204a0dbf15fd0 --- /dev/null +++ b/examples/garden/garden_03.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e0cc90c70c0f8bc070c659c6020d96f8c6058a2ec03987a60f2ef4d3cf67da34 +size 244355 diff --git a/examples/garden/garden_04.jpg b/examples/garden/garden_04.jpg new file mode 100644 index 0000000000000000000000000000000000000000..332f44a58184e704a3b6dd5f5555f455849877b8 --- /dev/null +++ b/examples/garden/garden_04.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0a5fe978ab7f45ac68d3cf7bcdc4d3444e2479001ced9ce325858a5be7dfcabd +size 241580 diff --git a/examples/garden/garden_05.jpg b/examples/garden/garden_05.jpg new file mode 100644 index 0000000000000000000000000000000000000000..09f5d1b2feecee23af1579950f66678bf3747241 --- /dev/null +++ b/examples/garden/garden_05.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fada3657684ebc8aadb5db2fbe889423490ba05347609adf8a3454dfe2453204 +size 233851 diff --git a/examples/garden/garden_06.jpg b/examples/garden/garden_06.jpg new file mode 100644 index 0000000000000000000000000000000000000000..7adec1a9c9a9a8aa036d4d953019125848caada7 --- /dev/null +++ b/examples/garden/garden_06.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:40daa01a73b176797d8aec8e212ae4a07100412c231ab96a73408c2eaa0cf3b0 +size 243243 diff --git a/examples/garden/garden_07.jpg b/examples/garden/garden_07.jpg new file mode 100644 index 0000000000000000000000000000000000000000..4ad70a4f53932560b61dbeabe8f275a9e8342ecf --- /dev/null +++ b/examples/garden/garden_07.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:405dc63c34ac8b84892094e75b27940819e9615b0d2e5d1ee0908cae0b4e47cb +size 249204 diff --git a/examples/garden/garden_08.jpg b/examples/garden/garden_08.jpg new file mode 100644 index 0000000000000000000000000000000000000000..fb723eb95c1cfb4970c8e84dd384108e3cc0f47e --- /dev/null +++ b/examples/garden/garden_08.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3219899ccc8ec7f05d90c0e63358330f144566f16200a032df42b316ed14ae9e +size 250688 diff --git a/examples/garden/garden_09.jpg b/examples/garden/garden_09.jpg new file mode 100644 index 0000000000000000000000000000000000000000..97ddb0693e009e40a5e6439a65035c9f98c13561 --- /dev/null +++ b/examples/garden/garden_09.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:14b2ae74378da4561b2b87e132c4366911729d5c3ef718c69c1d591cb04bce7f +size 254304 diff --git a/examples/garden/garden_10.jpg b/examples/garden/garden_10.jpg new file mode 100644 index 0000000000000000000000000000000000000000..37ca69b01f5e86bf5708efab366931d46c365005 --- /dev/null +++ b/examples/garden/garden_10.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:51dd65c626bbb84a831d49d09d20ea1c9b04d4896e0ac939ce2c662a94202330 +size 264429 diff --git a/examples/garden/garden_11.jpg b/examples/garden/garden_11.jpg new file mode 100644 index 0000000000000000000000000000000000000000..afe60213e5070af64429e8df3deba87a7bfb3597 --- /dev/null +++ b/examples/garden/garden_11.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:63d8b21d8d7482641a7f9f46adada683188c527c3efeb046bf8fa00dcc7971c9 +size 257730 diff --git a/examples/garden/garden_12.jpg b/examples/garden/garden_12.jpg new file mode 100644 index 0000000000000000000000000000000000000000..5fc5021c3a694cccaa1813b1990d09a43901bb8b --- /dev/null +++ b/examples/garden/garden_12.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9823dc7364e2348feb9af224c052058dfa4e80c9f1128d0f0a983e33f85d241c +size 255866 diff --git a/examples/garden/garden_13.jpg b/examples/garden/garden_13.jpg new file mode 100644 index 0000000000000000000000000000000000000000..30da74bc28e0204e3636927954dd92a98160d74b --- /dev/null +++ b/examples/garden/garden_13.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e98b2b00200c618d360d7fee41e70c159bbfb57eec9dea45c7f19c09f890c849 +size 243618 diff --git a/examples/kitchen/kitchen_00.jpg b/examples/kitchen/kitchen_00.jpg new file mode 100644 index 0000000000000000000000000000000000000000..1f7b35407ded5daf24669df2bd83a8c5a508f23b --- /dev/null +++ b/examples/kitchen/kitchen_00.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:65983e10ba2c32f8ccd1a1c7161f32da551a3dff29f95e411e6b9152e6fb7e86 +size 174474 diff --git a/examples/kitchen/kitchen_01.jpg b/examples/kitchen/kitchen_01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..6dca9beff381f9a8e23ac8f6073a152980c5e723 --- /dev/null +++ b/examples/kitchen/kitchen_01.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:973d24f802907f706ea270463aabf94346f00442fb6dc89fce0692467183fea8 +size 176014 diff --git a/examples/kitchen/kitchen_02.jpg b/examples/kitchen/kitchen_02.jpg new file mode 100644 index 0000000000000000000000000000000000000000..fc3293b70997f7cb64706b095caf89bff1254251 --- /dev/null +++ b/examples/kitchen/kitchen_02.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:93d0074367a09807bceda9794f9fb1dd6e7123c5a8a75d68daad317ef4dee72d +size 175367 diff --git a/examples/kitchen/kitchen_03.jpg b/examples/kitchen/kitchen_03.jpg new file mode 100644 index 0000000000000000000000000000000000000000..81961f4d3cb32e555d7a10abc895a7df53b62852 --- /dev/null +++ b/examples/kitchen/kitchen_03.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e9658f6d941e27882e9ff7fdad4466661551ce840a7eca1ce142283c0577a8b4 +size 178136 diff --git a/examples/kitchen/kitchen_04.jpg b/examples/kitchen/kitchen_04.jpg new file mode 100644 index 0000000000000000000000000000000000000000..ff7edf38c2a9d3a6ce6f9d0fb9ed3689622fba04 --- /dev/null +++ b/examples/kitchen/kitchen_04.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3233937a8bea606d086e9bd8b3295a869df294c91b110d9a5bd1ec2729421bfb +size 180588 diff --git a/examples/kitchen/kitchen_05.jpg b/examples/kitchen/kitchen_05.jpg new file mode 100644 index 0000000000000000000000000000000000000000..8cd26ddfdb7201b2c4093a599ff7bd7c6e5bfd8d --- /dev/null +++ b/examples/kitchen/kitchen_05.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:63be0d19f9c07da6ea35f1e4cea8e0bcf5a698277e339bfcfb309e3c3923b19b +size 179745 diff --git a/examples/kitchen/kitchen_06.jpg b/examples/kitchen/kitchen_06.jpg new file mode 100644 index 0000000000000000000000000000000000000000..b226fb75c7b391a5fa189ba3d2abd84e8e85c743 --- /dev/null +++ b/examples/kitchen/kitchen_06.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1d7de183db23070776186de17a6af7025f4e943e25b1df08d32664a4e6b7cae4 +size 174108 diff --git a/examples/kitchen/kitchen_07.jpg b/examples/kitchen/kitchen_07.jpg new file mode 100644 index 0000000000000000000000000000000000000000..e4710b3efefa9c75ce97d17a26e0a855399b6683 --- /dev/null +++ b/examples/kitchen/kitchen_07.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:92a92e068cd7789d033119ae0279609fa919ec491224a34e54985c55436d27c0 +size 183491 diff --git a/examples/kitchen/kitchen_08.jpg b/examples/kitchen/kitchen_08.jpg new file mode 100644 index 0000000000000000000000000000000000000000..eaa41ece45125417c9826ca76b888bf7606ed25a --- /dev/null +++ b/examples/kitchen/kitchen_08.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:02afd2f9570a59d916b786fc0e1a60b7b47d7a0b22e36c5f57d2de41b3078a47 +size 203653 diff --git a/examples/kitchen/kitchen_09.jpg b/examples/kitchen/kitchen_09.jpg new file mode 100644 index 0000000000000000000000000000000000000000..6b178505045561cc0e3629e55925c9b2e641bbba --- /dev/null +++ b/examples/kitchen/kitchen_09.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:63207aa638cfe8d1d843869c43ad0ace49d664b2e3f516d2b56acc8ccdeaf841 +size 194643 diff --git a/examples/kitchen/kitchen_10.jpg b/examples/kitchen/kitchen_10.jpg new file mode 100644 index 0000000000000000000000000000000000000000..74b700843d93327aa8675aec414125850a0fdddf --- /dev/null +++ b/examples/kitchen/kitchen_10.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3fe3f4b06105e6b1382c69aa650eb265d9acd8fbfa25f984d600847c0f1ee042 +size 181180 diff --git a/examples/kitchen/kitchen_11.jpg b/examples/kitchen/kitchen_11.jpg new file mode 100644 index 0000000000000000000000000000000000000000..0b1995effd2eb3485781fc888c42318100c5d7ac --- /dev/null +++ b/examples/kitchen/kitchen_11.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:902486cb003afceded1e36c252afe6f9ccc55c5d4d00215a21f33b8c9ba472ef +size 190062 diff --git a/examples/kitchen/kitchen_12.jpg b/examples/kitchen/kitchen_12.jpg new file mode 100644 index 0000000000000000000000000000000000000000..e69ed64a7372cc97008e9c9a264b04134cd0bf44 --- /dev/null +++ b/examples/kitchen/kitchen_12.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c58a7144ef46804658b8170a1b20a91763124f923000726492d21e99114c0f7f +size 178648 diff --git a/examples/kitchen/kitchen_13.jpg b/examples/kitchen/kitchen_13.jpg new file mode 100644 index 0000000000000000000000000000000000000000..7598edd039e5386e8a8dff1263cba86af26e230e --- /dev/null +++ b/examples/kitchen/kitchen_13.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d3b5f2057aae83f4b39dff64afc6d31aec0ca5d5adb0ee5faed52a363fba98f6 +size 180327 diff --git a/requirements.txt b/requirements.txt index 4a4b45a4b17b240a9be8f26ea81d86070131a561..2f69bf9a4f8087ba963cfdcb53fcf98a4f617c22 100644 --- a/requirements.txt +++ b/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 diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000000000000000000000000000000000000..161a9c4e194dfac242f62a3cf73a94f9da0166a3 --- /dev/null +++ b/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() diff --git a/src/model/decoder/__init__.py b/src/model/decoder/__init__.py index 77200dc6035c5d47c5077c9ce990f9ebe1afd076..5dd7b13c326241ca41a01b6c75e609e612fc7d79 100644 --- a/src/model/decoder/__init__.py +++ b/src/model/decoder/__init__.py @@ -16,22 +16,15 @@ except ImportError as exc: else: CUDA_SPLATTING_IMPORT_ERROR = None -try: # gsplat is optional in this Space (we render with diff-gaussian-rasterization) - from .decoder_splatting_gsplat import ( - DecoderSplattingGsplat, - DecoderSplattingGsplatCfg, - ) -except ImportError: - DecoderSplattingGsplat = None - - @dataclass - class DecoderSplattingGsplatCfg: - name: Literal["splatting_gsplat"] - -DECODERS = {} -if DecoderSplattingGsplat is not None: - DECODERS["splatting_gsplat"] = DecoderSplattingGsplat - DECODERS["gsplat"] = DecoderSplattingGsplat +from .decoder_splatting_gsplat import ( + DecoderSplattingGsplat, + DecoderSplattingGsplatCfg, +) + +DECODERS = { + "splatting_gsplat": DecoderSplattingGsplat, + "gsplat": DecoderSplattingGsplat, +} DecoderCfg = DecoderSplattingGsplatCfg | DecoderSplattingCUDACfg