Delete files codecsplat_pipeline.py LICENSE-CodecSplat with huggingface_hub

LICENSE-CodecSplat

@@ -1,21 +0,0 @@
-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.

codecsplat_pipeline.py

@@ -1,338 +0,0 @@
-"""CodecSplat inference helpers for the Space.
-
-Everything here mirrors the authors' evaluation path
-(`scripts/test_dl3dv_view8_256x448.sh` -> `src/main.py` -> `ModelWrapper.test_step`):
-the encoder runs the *complete* learned-codec path (encode context -> entropy-code the
-2D latent features into a real bitstream -> decode the bitstream -> predict Gaussians),
-and the gsplat decoder renders novel views, all under bf16 autocast like
-`trainer.precision=bf16-mixed`.
-"""
-
-from __future__ import annotations
-
-import os
-from pathlib import Path
-from types import SimpleNamespace
-
-import numpy as np
-import torch
-from dacite import Config as DaciteConfig, from_dict
-from einops import rearrange, repeat
-from omegaconf import OmegaConf
-from plyfile import PlyData, PlyElement
-from scipy.spatial.transform import Rotation as R
-from scipy.spatial.transform import Slerp
-
-from src.dataset.shims.crop_shim import apply_crop_shim
-from src.dataset.shims.patch_shim import apply_patch_shim
-from src.misc.sh_rotation import rotate_sh
-from src.model.encoder import get_encoder
-from src.model.encoder.encoder_codecsplat import EncoderCodecSplatCfg
-
-APP_DIR = Path(__file__).parent.resolve()
-
-# ---- authors' DL3DV 8-view 256x448 evaluation preset ---------------------------------
-IMAGE_SHAPE = (256, 448)
-NEAR, FAR = 0.5, 200.0  # config/experiment/dl3dv.yaml
-ENCODER_OVERRIDES = {
-    "num_depth_candidates": 128,
-    "shim_patch_size": 16,
-    "num_scales": 2,
-    "upsample_factor": 4,
-    "lowest_feature_resolution": 8,
-    "monodepth_vit_type": "vitb",
-    "enable_feature_codec": True,
-    "enable_depth_guided_mv_refiner": True,
-    "enable_post_vae_mv_refiner": True,
-    "codec_lmb_range": [16.0, 1024.0],
-    "supervise_intermediate_depth": False,
-}
-DECODER_OVERRIDES = {
-    "use_covariances": True,
-    "rasterize_mode": "antialiased",
-    "eps2d": 0.1,
-    "packed": True,
-    "scale_invariant": False,
-}
-
-
-def build_encoder(checkpoint_path: str):
-    enc_cfg = OmegaConf.load(APP_DIR / "config/model/encoder/codecsplat.yaml")
-    enc_cfg = OmegaConf.merge(enc_cfg, ENCODER_OVERRIDES)
-    enc_cfg = from_dict(
-        EncoderCodecSplatCfg,
-        OmegaConf.to_container(enc_cfg, resolve=True),
-        config=DaciteConfig(cast=[tuple], strict=False),
-    )
-    encoder = get_encoder(enc_cfg)
-
-    ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
-    state = ckpt.get("state_dict", ckpt)
-    state = {k[len("encoder."):]: v for k, v in state.items() if k.startswith("encoder.")}
-    encoder.load_state_dict(state, strict=True)
-    # eval() also runs `compress_mode` -> builds the entropy-coder CDF tables (on CPU).
-    encoder.eval()
-    return encoder
-
-
-def build_decoder():
-    """The authors' `model/decoder=splatting_cuda` decoder (Inria diff-gaussian-rasterization).
-
-    The evaluation scripts use the gsplat decoder; on ZeroGPU we use the repo's other
-    supported decoder because a prebuilt sm_120 wheel exists for it (gsplat would have to
-    be JIT-compiled against a CUDA toolkit at startup).
-    """
-    from src.model.decoder.decoder_splatting_cuda import (
-        DecoderSplattingCUDA,
-        DecoderSplattingCUDACfg,
-    )
-
-    decoder = DecoderSplattingCUDA(
-        DecoderSplattingCUDACfg(name="splatting_cuda"),
-        SimpleNamespace(background_color=[0.0, 0.0, 0.0]),
-    )
-    return decoder.eval()
-
-
-# ---- view selection ------------------------------------------------------------------
-def farthest_point_sample(xyz: np.ndarray, npoint: int) -> np.ndarray:
-    n = len(xyz)
-    if npoint >= n:
-        return np.arange(n)
-    dist = np.full(n, 1e10)
-    farthest = int(np.argmax(((xyz - xyz.mean(0)) ** 2).sum(-1)))
-    out = []
-    for _ in range(npoint):
-        out.append(farthest)
-        d = ((xyz - xyz[farthest]) ** 2).sum(-1)
-        dist = np.minimum(dist, d)
-        farthest = int(np.argmax(dist))
-    return np.sort(np.array(out))
-
-
-def select_views(c2w: np.ndarray, num_context: int = 8, num_target: int = 8):
-    n = len(c2w)
-    ctx = farthest_point_sample(c2w[:, :3, 3], min(num_context, n))
-    rest = np.setdiff1d(np.arange(n), ctx)
-    if len(rest) > num_target:
-        rest = rest[np.linspace(0, len(rest) - 1, num_target).astype(int)]
-    return ctx, rest
-
-
-# ---- batch construction (same conventions as the dataset loaders) --------------------
-def _load_images(paths):
-    from PIL import Image
-
-    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))
-    return torch.stack(imgs)
-
-
-def image_shape_for(h: int, w: int) -> tuple[int, int]:
-    # The checkpoint is trained at 256x448 (landscape). Portrait inputs use 448x256.
-    return IMAGE_SHAPE if w >= h else (IMAGE_SHAPE[1], IMAGE_SHAPE[0])
-
-
-def build_example(image_paths, c2w, K, ctx_idx, tgt_idx):
-    """Build a single (un-batched) example, apply the authors' crop + patch shims."""
-    images = _load_images(image_paths)  # [N, 3, H, W]
-    shape = image_shape_for(images.shape[-2], images.shape[-1])
-    c2w = torch.as_tensor(c2w, dtype=torch.float32)
-    K = torch.as_tensor(K, dtype=torch.float32)
-
-    # Express all poses relative to the middle context camera (rigid change of frame).
-    pivot = c2w[ctx_idx[len(ctx_idx) // 2]]
-    c2w = torch.linalg.inv(pivot)[None] @ c2w
-
-    def views(idx):
-        idx = torch.as_tensor(np.asarray(idx), dtype=torch.long)
-        return {
-            "extrinsics": c2w[idx],
-            "intrinsics": K[idx],
-            "image": images[idx],
-            "near": torch.full((len(idx),), NEAR),
-            "far": torch.full((len(idx),), FAR),
-            "index": idx,
-        }
-
-    example = {"context": views(ctx_idx), "target": views(tgt_idx), "scene": "user_scene"}
-    example = apply_crop_shim(example, shape)
-    batch = {
-        "context": {k: v[None] for k, v in example["context"].items()},
-        "target": {k: v[None] for k, v in example["target"].items()},
-        "scene": ["user_scene"],
-    }
-    # shim_patch_size * downscale_factor = 16 * 4
-    return apply_patch_shim(batch, patch_size=64)
-
-
-def to_device(batch, device):
-    out = {}
-    for k, v in batch.items():
-        if isinstance(v, dict):
-            out[k] = {kk: vv.to(device) for kk, vv in v.items()}
-        else:
-            out[k] = v
-    return out
-
-
-# ---- smooth fly-through trajectory through the context cameras ----------------------
-def order_cameras(c2w: torch.Tensor) -> list[int]:
-    """Greedy nearest-neighbour ordering of cameras (start from an extreme one)."""
-    pos = c2w[:, :3, 3].cpu().numpy()
-    n = len(pos)
-    if n <= 2:
-        return list(range(n))
-    start = int(np.argmax(((pos - pos.mean(0)) ** 2).sum(-1)))
-    order, left = [start], set(range(n)) - {start}
-    while left:
-        last = pos[order[-1]]
-        nxt = min(left, key=lambda i: float(((pos[i] - last) ** 2).sum()))
-        order.append(nxt)
-        left.remove(nxt)
-    return order
-
-
-def interpolate_trajectory(c2w: torch.Tensor, K: torch.Tensor, frames_per_segment: int = 12):
-    order = order_cameras(c2w)
-    c2w = c2w[order].cpu().double().numpy()
-    K = K[order].cpu()
-    poses, intrs = [], []
-    for i in range(len(order) - 1):
-        rots = R.from_matrix(c2w[i : i + 2, :3, :3])
-        slerp = Slerp([0.0, 1.0], rots)
-        for t in np.linspace(0, 1, frames_per_segment, endpoint=False):
-            m = np.eye(4)
-            m[:3, :3] = slerp([t]).as_matrix()[0]
-            m[:3, 3] = (1 - t) * c2w[i, :3, 3] + t * c2w[i + 1, :3, 3]
-            poses.append(m)
-            intrs.append((1 - t) * K[i] + t * K[i + 1])
-    last = np.eye(4)
-    last[:3] = c2w[-1, :3]
-    poses.append(last)
-    intrs.append(K[-1])
-    poses = poses + poses[::-1][1:]  # ping-pong loop
-    intrs = intrs + intrs[::-1][1:]
-    return (
-        torch.tensor(np.stack(poses), dtype=torch.float32),
-        torch.stack([torch.as_tensor(k, dtype=torch.float32) for k in intrs]),
-    )
-
-
-# ---- core inference ------------------------------------------------------------------
-@torch.no_grad()
-def encode_decode(encoder, batch, lmb: float):
-    """Full learned-codec round trip: context -> bitstream -> Gaussians."""
-    encoder.feature_codec.cfg.lmb = float(lmb)
-    with torch.autocast("cuda", dtype=torch.bfloat16, enabled=batch["context"]["image"].is_cuda):
-        out = encoder(batch["context"], 0, visualization_dump={}, return_bitstream=True)
-    return out
-
-
-@torch.no_grad()
-def render(decoder, gaussians, extrinsics, intrinsics, hw, chunk: int = 16):
-    """Render [V] views. `extrinsics`/`intrinsics` are [V,4,4]/[V,3,3] on device."""
-    v = extrinsics.shape[0]
-    near = torch.full((1, v), NEAR, device=extrinsics.device)
-    far = torch.full((1, v), FAR, device=extrinsics.device)
-    colors = []
-    with torch.autocast("cuda", dtype=torch.bfloat16, enabled=extrinsics.is_cuda):
-        for s in range(0, v, chunk):
-            o = decoder.forward(
-                gaussians,
-                extrinsics[None, s : s + chunk],
-                intrinsics[None, s : s + chunk],
-                near[:, s : s + chunk],
-                far[:, s : s + chunk],
-                hw,
-                depth_mode=None,
-            )
-            colors.append(o.color[0].float().clamp(0, 1))
-    return torch.cat(colors, dim=0)
-
-
-def psnr(gt: torch.Tensor, pred: torch.Tensor) -> float:
-    mse = ((gt - pred) ** 2).mean(dim=(1, 2, 3)).clamp_min(1e-10)
-    return float((-10 * torch.log10(mse)).mean())
-
-
-# ---- Gaussian export -----------------------------------------------------------------
-def world_rotations(rotations: torch.Tensor, context_c2w: torch.Tensor) -> torch.Tensor:
-    """Per-pixel Gaussian rotations are predicted in each context camera's frame."""
-    v = context_c2w.shape[0]
-    n = rotations.shape[0]
-    c2w_rot = repeat(context_c2w[:, :3, :3], "v i j -> (v r) i j", r=n // v)
-    local = R.from_quat(rotations.cpu().double().numpy()).as_matrix()
-    world = c2w_rot.cpu().double().numpy() @ local
-    return torch.from_numpy(R.from_matrix(world).as_quat()).float()
-
-
-def write_ply(
-    path,
-    pivot_c2w: torch.Tensor,
-    means: torch.Tensor,
-    scales: torch.Tensor,
-    rotations_xyzw_world: torch.Tensor,
-    harmonics: torch.Tensor,
-    opacities: torch.Tensor,
-):
-    """Same output as the authors' `export_ply(..., align_to_view=True)`, vectorised."""
-    means = means.float().cpu()
-    scales = scales.float().cpu()
-    harmonics = harmonics.float().cpu()
-    opacities = opacities.float().cpu()
-    pivot = pivot_c2w.float().cpu()
-
-    view_rot = pivot[:3, :3].T
-    means = (means - pivot[:3, 3]) @ view_rot.T
-    harmonics = rotate_sh(harmonics, view_rot)
-    rot = R.from_quat(rotations_xyzw_world.double().numpy()).as_matrix()
-    rot = R.from_matrix(view_rot.double().numpy() @ rot).as_quat()  # xyzw
-    rot = np.concatenate([rot[:, 3:4], rot[:, :3]], axis=1)  # wxyz
-
-    dc = harmonics[..., 0].numpy()
-    rest = harmonics[..., 1:].flatten(start_dim=1).numpy()
-    logit_op = torch.logit(opacities.clamp(1e-6, 1 - 1e-6))[:, None].numpy()
-    log_scales = scales.clamp_min(1e-12).log().numpy()
-    m = means.numpy()
-
-    cols = [m, np.zeros_like(m), dc, rest, logit_op, log_scales, rot]
-    data = np.concatenate(cols, axis=1).astype(np.float32)
-    names = ["x", "y", "z", "nx", "ny", "nz", "f_dc_0", "f_dc_1", "f_dc_2"]
-    names += [f"f_rest_{i}" for i in range(rest.shape[1])]
-    names += ["opacity", "scale_0", "scale_1", "scale_2", "rot_0", "rot_1", "rot_2", "rot_3"]
-    arr = np.empty(data.shape[0], dtype=[(n, "f4") for n in names])
-    for i, n in enumerate(names):
-        arr[n] = data[:, i]
-    Path(path).parent.mkdir(parents=True, exist_ok=True)
-    PlyData([PlyElement.describe(arr, "vertex")]).write(str(path))
-    return str(path)
-
-
-def export_gaussians(g: dict, context_c2w: torch.Tensor, outdir: Path,
-                     viewer_max: int = 600_000, viewer_min_opacity: float = 0.02):
-    """Write a compact (DC colour, pruned), re-framed .ply of the Gaussians for the web viewer."""
-    pivot = context_c2w[len(context_c2w) // 2]
-    rot_world = world_rotations(g["rotations"], context_c2w)
-
-    keep = g["opacities"] > viewer_min_opacity
-    if int(keep.sum()) < 1000:
-        keep = torch.ones_like(keep)
-    idx = torch.nonzero(keep).squeeze(1)
-    if idx.numel() > viewer_max:
-        idx = idx[torch.randperm(idx.numel())[:viewer_max]]
-    means, scales = g["means"][idx], g["scales"][idx]
-
-    # Re-centre / re-scale (in the pivot camera frame) so the viewer's orbit camera frames it.
-    view_rot = pivot[:3, :3].T
-    local = (means - pivot[:3, 3]) @ view_rot.T
-    center = local.median(dim=0).values
-    radius = torch.quantile((local - center).norm(dim=-1).float(), 0.7)
-    s = float(3.0 / max(float(radius), 1e-6))
-    # shift in world coordinates so that the pivot-frame centre maps to the origin
-    means = (means - pivot[:3, 3] - center @ view_rot) * s + pivot[:3, 3]
-    viewer_path = write_ply(outdir / "codecsplat_viewer.ply", pivot, means, scales * s,
-                            rot_world[idx], g["harmonics"][idx][..., :1], g["opacities"][idx])
-    return viewer_path