app.py

"""CodecSplat: Ultra-Compact Latent Coding for Feed-Forward 3D Gaussian Splatting.

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 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))

_BISECT = os.environ.get("BISECT", "")


def _bisect(stage: str) -> None:
    """Diagnostics: boot only the GPU probe after `stage` (set via the BISECT variable)."""
    if _BISECT == stage:
        print(f"[bisect] stopping after stage {stage}", flush=True)
        exec(open(APP_DIR / "probe_app.py").read(), {"__name__": "probe"})
        raise SystemExit(0)


_bisect("minimal")

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")
# pycolmap's bundled OpenBLAS crashes when many matcher threads each spawn BLAS threads.
os.environ["OPENBLAS_NUM_THREADS"] = "1"

import spaces  # noqa: E402  (must precede any CUDA-touching import)
import torch  # noqa: E402

torch.set_float32_matmul_precision("high")


def _nvidia_fds(stage: str) -> None:
    """Diagnostics: /dev/nvidia* handles in the GPU-less main process mean the CUDA
    driver was touched there, which breaks ZeroGPU's forked workers."""
    fds = set()
    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.add(t)
    libs = set()
    try:
        for line in open("/proc/self/maps"):
            for key in ("libcuda.so", "libnvidia-ml", "libcudart"):
                if key in line:
                    libs.add(line.split()[-1].rsplit("/", 1)[-1])
    except OSError:
        pass
    init = getattr(torch.cuda, "_initialized", None)
    print(f"[diag] {stage}: nvidia fds={sorted(fds)} libs={sorted(libs)} torch.cuda._initialized={init}", flush=True)


_nvidia_fds("after torch import")
_bisect("s1_torch")


# ---------------------------------------------------------------------------
# 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", "")
_bisect("s2_env")

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)
_nvidia_fds("after gsplat")
_bisect("s3_gsplat")

# ---------------------------------------------------------------------------
# 2. Compatibility shims for the vendored reference code
# ---------------------------------------------------------------------------
import matplotlib.cm  # noqa: E402
import matplotlib.pyplot as _plt  # noqa: E402

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

# torchvision's CUDA build (pulled in by `compressai/__init__`, lpips, lightning) breaks
# the ZeroGPU forked worker ("No CUDA GPUs are available" in worker_init), bisected on
# this Space. CodecSplat never uses torchvision at inference time, so:
#  * import compressai's entropy-model / ops submodules without running its package
#    __init__ (which imports models/datasets -> torchvision),
#  * stub the two unused `torchvision` imports in the vendored dataset/vis helpers.
import importlib.machinery  # noqa: E402
import importlib.util  # noqa: E402

_ca_spec = importlib.util.find_spec("compressai")  # does not execute the package
_ca = types.ModuleType("compressai")
_ca.__path__ = list(_ca_spec.submodule_search_locations)
_ca.__spec__ = _ca_spec
_ca.get_entropy_coder = lambda: "ans"
_ca.available_entropy_coders = lambda: ["ans"]
sys.modules["compressai"] = _ca


class _UnusedModule(types.ModuleType):
    def __getattr__(self, name):
        if name.startswith("__"):
            raise AttributeError(name)
        qual = f"{self.__name__}.{name}"

        def _unavailable(*args, **kwargs):
            raise ImportError(f"{qual} is not available in this Space")

        return type(name, (), {"__init__": _unavailable, "__call__": _unavailable})


for _name in ("torchvision", "torchvision.transforms", "torchvision.utils"):
    _m = _UnusedModule(_name)
    _m.__path__ = []
    _m.__spec__ = importlib.machinery.ModuleSpec(_name, None, is_package=True)
    sys.modules[_name] = _m
sys.modules["torchvision"].transforms = sys.modules["torchvision.transforms"]
sys.modules["torchvision"].utils = sys.modules["torchvision.utils"]

# ---------------------------------------------------------------------------
# 3. Model
# ---------------------------------------------------------------------------
import gradio as gr  # noqa: E402
import imageio.v2 as imageio  # noqa: E402
import numpy as np  # 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

from colmap_io import load_colmap_scene, select_context_views, select_target_views  # noqa: E402
_bisect("imp_pre_colmap")
from colmap_sfm import IMAGE_EXT, run_sfm  # noqa: E402
_bisect("imp_colmap")

MODEL_REPO = "pengpeng-yu/CodecSplat"
CHECKPOINT_FILE = "codecsplat-base-dl3dv-256x448-view8-075062d0.pth"
NUM_CONTEXT = 8
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)

    # Typed configs built directly with dacite (same as src.config.load_typed_config),
    # without importing src.config -> model_wrapper -> pytorch_lightning / wandb / lpips.
    from dacite import Config as DaciteConfig, from_dict
    from omegaconf import OmegaConf

    from src.model.decoder import DecoderSplattingGsplatCfg, get_decoder
    from src.model.encoder import EncoderCodecSplatCfg, get_encoder

    dcfg = DaciteConfig(type_hooks={Path: Path})
    enc_cfg = from_dict(EncoderCodecSplatCfg, OmegaConf.to_container(cfg_dict.model.encoder, resolve=True), config=dcfg)
    dec_cfg = from_dict(DecoderSplattingGsplatCfg, OmegaConf.to_container(cfg_dict.model.decoder, resolve=True), config=dcfg)
    dataset_cfg = types.SimpleNamespace(background_color=list(cfg_dict.dataset.background_color))
    encoder = get_encoder(enc_cfg)
    decoder = get_decoder(dec_cfg, dataset_cfg)

    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()


_bisect("s4_imports")
print("[setup] downloading checkpoint...", flush=True)
CHECKPOINT_PATH = hf_hub_download(MODEL_REPO, CHECKPOINT_FILE)
print("[setup] building CodecSplat...", flush=True)
ENCODER, DECODER, DATA_SHIM = build_model(CHECKPOINT_PATH)
_nvidia_fds("after build (cpu)")
_bisect("s5_build")
ENCODER.to("cuda")
DECODER.to("cuda")
print("[setup] model ready", flush=True)
_nvidia_fds("after model build")
_bisect("s6_cuda")

from src.dataset.shims.crop_shim import rescale_and_crop  # 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


# ---------------------------------------------------------------------------
# 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)


# ---------------------------------------------------------------------------
# 5. GPU stage: encode -> bitstream -> decode -> render
# ---------------------------------------------------------------------------
@spaces.GPU(duration=int(os.environ.get("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()
    t_run = time.time() - t0

    gt = batch["target"]["image"][0].float()
    mse = ((gt - rendered.clamp(0, 1)) ** 2).flatten(1).mean(dim=1)
    psnr = float((-10.0 * torch.log10(mse.clamp_min(1e-10))).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,
        "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


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: 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 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 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.")
    if len(paths) > MAX_INPUT_IMAGES:
        paths = [paths[i] for i in np.linspace(0, len(paths) - 1, MAX_INPUT_IMAGES).astype(int)]

    t_start = time.time()
    workdir = Path(tempfile.mkdtemp(prefix="codecsplat_"))
    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=960, matcher=matcher,
                      log=lambda m: print(m, flush=True))
    except ValueError as exc:
        raise gr.Error(str(exc)) from exc
    t_sfm = time.time() - t_sfm

    _nvidia_fds("after COLMAP")
    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"### 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** 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, gallery, bin_path, ply_path, summary


# ---------------------------------------------------------------------------
# 6. UI
# ---------------------------------------------------------------------------
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()
]

CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

DESCRIPTION = """# CodecSplat: Ultra-Compact Latent Coding for Feed-Forward 3D Gaussian Splatting

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.

[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:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(DESCRIPTION)
        with gr.Row():
            with gr.Column(scale=1):
                image_files = gr.File(
                    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 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.")
                gr.Markdown(
                    "**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_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")

        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)