diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..fde9b584b5e21cf82984e8cdc007e7ea0d9b148e 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,53 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +examples/scene_a/00.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/01.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/02.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/03.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/04.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/05.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/06.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/07.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/08.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/09.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/10.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/11.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/12.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/13.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/14.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/15.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/16.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/17.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/18.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/19.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/20.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/21.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/22.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/23.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_a/24.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/00.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/01.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/02.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/03.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/04.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/05.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/06.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/07.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/08.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/09.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/10.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/11.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/12.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/13.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/14.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/15.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/16.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/17.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/18.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/19.jpg filter=lfs diff=lfs merge=lfs -text +examples/scene_b/20.jpg filter=lfs diff=lfs merge=lfs -text +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 diff --git a/LICENSE-CodecSplat b/LICENSE-CodecSplat new file mode 100644 index 0000000000000000000000000000000000000000..755da3c59a3f5203aa137aac0be747812d3013cf --- /dev/null +++ b/LICENSE-CodecSplat @@ -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 6ce5daec57e6969c831dff10eadd491e086591a1..0126c971c122d2fa423d7e30c6ba0b0229f90187 100644 --- a/README.md +++ b/README.md @@ -1,13 +1,50 @@ --- -title: Codecsplat -emoji: 👀 -colorFrom: green -colorTo: purple +title: CodecSplat +emoji: 🗜️ +colorFrom: red +colorTo: yellow sdk: gradio sdk_version: 6.30.0 -python_version: '3.12' app_file: app.py -pinned: false +short_description: Compress photos into a tiny 3D Gaussian splat bitstream +python_version: "3.12" +startup_duration_timeout: 1h +models: + - pengpeng-yu/CodecSplat --- -Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference +# 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' gsplat decoder (antialiased, eps2d=0.1) 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: +* gsplat's CUDA kernels are JIT-compiled at startup against the CUDA toolkit in the ZeroGPU + image, following `hugging-apps/resplat-recurrent-gs`. +* 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. diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..69a6af0b27daae42e2fe4d82a805708b9ff500a8 --- /dev/null +++ b/app.py @@ -0,0 +1,338 @@ +"""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 -> gsplat rendering. +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). +""" + +from __future__ import annotations + +import ctypes +import glob +import os +import sys +import tempfile +import time +from pathlib import Path + +APP_DIR = Path(__file__).parent.resolve() +sys.path.insert(0, str(APP_DIR)) + +os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0") +os.environ.setdefault("MAX_JOBS", str(min(32, (os.cpu_count() or 8)))) + +import spaces # noqa: E402 (must precede any CUDA-touching import) +import torch # noqa: E402 + +torch.set_float32_matmul_precision("high") + + +# --------------------------------------------------------------------------- +# 1. gsplat CUDA backend: JIT-compile once at startup against the CUDA toolkit +# shipped in the ZeroGPU image (same approach as hugging-apps/resplat-recurrent-gs). +# --------------------------------------------------------------------------- +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 gsplat CUDA kernels (one-off, several minutes)...", 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. Model +# --------------------------------------------------------------------------- +import imageio.v2 as imageio # noqa: E402 +import numpy as np # noqa: E402 +import gradio as gr # noqa: E402 +import pycolmap # noqa: E402 +from huggingface_hub import hf_hub_download # noqa: E402 +from PIL import Image # noqa: E402 + +import codecsplat_pipeline as cp # 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 + +print("[setup] downloading checkpoint...", flush=True) +ckpt_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") +print("[setup] model ready", flush=True) + + +def _u8(t: torch.Tensor) -> np.ndarray: + return (t.clamp(0, 1).permute(1, 2, 0).float().cpu().numpy() * 255).round().astype(np.uint8) + + +# --------------------------------------------------------------------------- +# 3. GPU part: 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") + h, w = batch["context"]["image"].shape[-2:] + + 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(), + } + + +# --------------------------------------------------------------------------- +# 4. Full pipeline (CPU COLMAP -> GPU CodecSplat -> CPU export) +# --------------------------------------------------------------------------- +def _read_colmap_scene(scene_dir: str): + rec = pycolmap.Reconstruction(os.path.join(scene_dir, "sparse")) + items = [] + for _, img in rec.images.items(): + cam = rec.cameras[img.camera_id] + p = 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 + K = np.eye(3, dtype=np.float32) + K[0, 0], K[1, 1] = fx / cam.width, fy / cam.height + K[0, 2], K[1, 2] = cx / cam.width, cy / cam.height + w2c = np.eye(4, dtype=np.float64) + w2c[:3, :4] = img.cam_from_world().matrix() if callable(img.cam_from_world) else img.cam_from_world.matrix() + path = os.path.join(scene_dir, "images", img.name) + if os.path.exists(path): + items.append((img.name, path, np.linalg.inv(w2c).astype(np.float32), K)) + items.sort(key=lambda x: x[0]) + if len(items) < 2: + raise gr.Error("COLMAP registered fewer than 2 usable images.") + return [i[1] for i in items], np.stack([i[2] for i in items]), np.stack([i[3] for i in items]) + + +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. + + 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). + + Returns: + Interactive splat (.ply), fly-through video, held-out novel views (render vs. photo), + selected input views, the scene bitstream (.bin) and a summary. + """ + if not images: + raise gr.Error("Please upload at least 3 overlapping photos 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") + try: + sfm = run_sfm(paths, workdir, max_image_size=1024, 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_start + img_paths, c2w, K = _read_colmap_scene(sfm["scene_dir"]) + + 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 + 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" + ) + return viewer_ply, video_path, gallery_nv, gallery_ctx, str(bin_path), summary + + +# --------------------------------------------------------------------------- +# 5. 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")] + +DESCRIPTION = """ +# CodecSplat: Ultra-Compact Latent Coding for Feed-Forward 3D Gaussian Splatting + +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 (λ). + +[Paper](https://huggingface.co/papers/2605.25563) · [Code](https://github.com/pengpeng-yu/CodecSplat) · [Weights](https://huggingface.co/pengpeng-yu/CodecSplat) +""" + +CSS = """ +#col-container { max-width: 1280px; margin: 0 auto; } +.dark .gradio-container { color: var(--body-text-color); } +""" + +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 shots work best)", + 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") + 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( + "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" + "**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] + 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)", + ) + +if __name__ == "__main__": + demo.queue(max_size=12).launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True, show_error=True) diff --git a/codecsplat_pipeline.py b/codecsplat_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..295f2ad8c738500886f36ba2818199f54d5e86e0 --- /dev/null +++ b/codecsplat_pipeline.py @@ -0,0 +1,336 @@ +"""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(): + from src.model.decoder.decoder_splatting_gsplat import ( + DecoderSplattingGsplat, + DecoderSplattingGsplatCfg, + ) + + dec_cfg = OmegaConf.load(APP_DIR / "config/model/decoder/splatting_gsplat.yaml") + dec_cfg = OmegaConf.merge(dec_cfg, DECODER_OVERRIDES) + dec_cfg = from_dict( + DecoderSplattingGsplatCfg, + OmegaConf.to_container(dec_cfg, resolve=True), + config=DaciteConfig(strict=False), + ) + decoder = DecoderSplattingGsplat(dec_cfg, 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 diff --git a/colmap_sfm.py b/colmap_sfm.py new file mode 100644 index 0000000000000000000000000000000000000000..918ffc56390351c691111ce4f33aa4b1093f4d5c --- /dev/null +++ b/colmap_sfm.py @@ -0,0 +1,177 @@ +"""COLMAP structure-from-motion preprocessing (runs inside the Space, on CPU). + +Takes a set of raw multi-view images of a scene, estimates camera intrinsics / +poses / a sparse point cloud with COLMAP (via `pycolmap`), undistorts the images +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. +""" + +from __future__ import annotations + +import shutil +from pathlib import Path + +import pycolmap +from PIL import Image, ImageOps + +IMAGE_EXT = (".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff") + + +def prepare_images(image_paths, dest: Path, max_image_size: int = 1280) -> list[str]: + """Copy the uploaded images into `dest`, EXIF-rotated, RGB, size-capped.""" + dest.mkdir(parents=True, exist_ok=True) + names: list[str] = [] + for i, src in enumerate(str(p) for p in image_paths): # caller decides the order + if not src.lower().endswith(IMAGE_EXT): + continue + try: + img = Image.open(src) + img = ImageOps.exif_transpose(img).convert("RGB") + except Exception as exc: + print(f"[colmap] skipping unreadable image {src}: {exc}", flush=True) + continue + scale = max_image_size / max(img.size) + if scale < 1.0: + img = img.resize( + (max(1, int(round(img.width * scale))), max(1, int(round(img.height * scale)))), + Image.LANCZOS, + ) + name = f"{i:04d}.jpg" + img.save(dest / name, quality=95) + names.append(name) + return names + + +def _largest_reconstruction(maps) -> pycolmap.Reconstruction | None: + recs = list(maps.values()) if isinstance(maps, dict) else list(maps) + if not recs: + return None + return max(recs, key=lambda r: r.num_reg_images()) + + +def run_sfm( + image_paths, + workdir: Path, + max_image_size: int = 1280, + matcher: str = "exhaustive", + max_num_features: int = 8192, + log=print, +) -> dict: + """Run COLMAP SfM on `image_paths`; return info about the written scene. + + Returns a dict with `scene_dir`, `sparse_sub`, `images_sub`, `num_input`, + `num_registered`, `num_points3D`, `camera_model`. + """ + workdir = Path(workdir) + image_dir = workdir / "sfm_images" + names = prepare_images(image_paths, image_dir, max_image_size) + if len(names) < 3: + raise ValueError( + f"COLMAP needs at least 3 usable images, got {len(names)}. " + "Upload several overlapping photos of the same scene." + ) + log(f"[colmap] {len(names)} images prepared at <= {max_image_size}px") + + database_path = workdir / "colmap.db" + if database_path.exists(): + database_path.unlink() + + extraction_options = pycolmap.FeatureExtractionOptions() + extraction_options.sift.max_num_features = int(max_num_features) + extraction_options.sift.estimate_affine_shape = False + extraction_options.sift.domain_size_pooling = False + extraction_options.use_gpu = False + extraction_options.max_image_size = int(max_image_size) + + reader_options = pycolmap.ImageReaderOptions() + reader_options.camera_model = "SIMPLE_RADIAL" + + pycolmap.extract_features( + database_path=database_path, + image_path=image_dir, + camera_mode=pycolmap.CameraMode.SINGLE, + reader_options=reader_options, + extraction_options=extraction_options, + device=pycolmap.Device.cpu, + ) + log("[colmap] features extracted") + + if matcher == "sequential": + pairing = pycolmap.SequentialPairingOptions() + pairing.overlap = 10 + pairing.quadratic_overlap = True + pycolmap.match_sequential( + database_path=database_path, + pairing_options=pairing, + device=pycolmap.Device.cpu, + ) + else: + 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.multiple_models = False + options.extract_colors = True + maps = pycolmap.incremental_mapping( + database_path=database_path, + image_path=image_dir, + output_path=sparse_dir, + options=options, + ) + rec = _largest_reconstruction(maps) + if rec is None or rec.num_reg_images() < 3: + raise ValueError( + "COLMAP could not reconstruct this image set (only " + f"{0 if rec is None else rec.num_reg_images()} of {len(names)} images registered). " + "Use photos with plenty of visual overlap between neighbours, avoid pure " + "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. + rec.normalize(fixed_scale=False, extent=10.0) + log(f"[colmap] mapped {rec.num_reg_images()}/{len(names)} images, {rec.num_points3D()} points") + + model_dir = workdir / "sparse_model" + if model_dir.exists(): + shutil.rmtree(model_dir) + model_dir.mkdir(parents=True) + rec.write(str(model_dir)) + + scene_dir = workdir / "colmap_scene" + if scene_dir.exists(): + shutil.rmtree(scene_dir) + pycolmap.undistort_images( + output_path=scene_dir, + input_path=model_dir, + image_path=image_dir, + ) + log("[colmap] undistorted to PINHOLE") + + # Keep only what CodecSplat reads (the undistorter also writes stereo/ scaffolding). + for extra in ("stereo",): + p = scene_dir / extra + if p.is_dir(): + shutil.rmtree(p, ignore_errors=True) + + cam_model = "PINHOLE" + try: + cam_model = list(rec.cameras.values())[0].model.name + except Exception: + pass + + return { + "scene_dir": str(scene_dir), + "sparse_sub": "sparse", + "images_sub": "images", + "num_input": len(names), + "num_registered": int(rec.num_reg_images()), + "num_points3D": int(rec.num_points3D()), + "camera_model": cam_model, + } diff --git a/config/dataset/dl3dv.yaml b/config/dataset/dl3dv.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6dc4f4324a54edd8900408abdcec1cfde5f1d6a9 --- /dev/null +++ b/config/dataset/dl3dv.yaml @@ -0,0 +1,27 @@ +defaults: + - view_sampler: boundedv2_360 + +name: dl3dv +roots: [datasets/DL3DV-ALL-480P-Processed] +augment: true + +image_shape: [270, 480] +background_color: [0.0, 0.0, 0.0] +cameras_are_circular: false + +max_fov: 100.0 + +skip_bad_shape: true +near: -1. +far: -1. +shuffle_val: true +test_len: -1 +test_chunk_interval: 1 +sort_target_index: true +sort_context_index: true + +train_times_per_scene: 1 +test_times_per_scene: 1 +ori_image_shape: [270, 480] +overfit_max_views: 148 +use_index_to_load_chunk: false diff --git a/config/dataset/re10k.yaml b/config/dataset/re10k.yaml new file mode 100644 index 0000000000000000000000000000000000000000..891116f36d29b9eb8867163288d747783552640f --- /dev/null +++ b/config/dataset/re10k.yaml @@ -0,0 +1,21 @@ +defaults: + - view_sampler: bounded + +name: re10k +roots: [datasets/RE10K-Torch-360p] +augment: true + +image_shape: [180, 320] +background_color: [0.0, 0.0, 0.0] +cameras_are_circular: false + +max_fov: 100.0 + +skip_bad_shape: true +near: -1. +far: -1. +shuffle_val: true +test_len: -1 +test_chunk_interval: 1 + +use_index_to_load_chunk: false diff --git a/config/dataset/view_sampler/all.yaml b/config/dataset/view_sampler/all.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9c49fb4661ff2e30ecd98c7e233e2835c9071014 --- /dev/null +++ b/config/dataset/view_sampler/all.yaml @@ -0,0 +1 @@ +name: all diff --git a/config/dataset/view_sampler/arbitrary.yaml b/config/dataset/view_sampler/arbitrary.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c947c2ce3ddf89fac2136ee8133c6662536af3d5 --- /dev/null +++ b/config/dataset/view_sampler/arbitrary.yaml @@ -0,0 +1,7 @@ +name: arbitrary + +num_target_views: 1 +num_context_views: 2 + +# If you want to hard-code context views, do so here. +context_views: null diff --git a/config/dataset/view_sampler/bounded.yaml b/config/dataset/view_sampler/bounded.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1e7b67d31a27bfe83598a980cf73ab6a2bd581a9 --- /dev/null +++ b/config/dataset/view_sampler/bounded.yaml @@ -0,0 +1,12 @@ +name: bounded + +num_target_views: 1 +num_context_views: 2 + +min_distance_between_context_views: 2 +max_distance_between_context_views: 6 +min_distance_to_context_views: 0 + +warm_up_steps: 0 +initial_min_distance_between_context_views: 2 +initial_max_distance_between_context_views: 6 \ No newline at end of file diff --git a/config/dataset/view_sampler/boundedv2.yaml b/config/dataset/view_sampler/boundedv2.yaml new file mode 100644 index 0000000000000000000000000000000000000000..065e7ea4c7a9eedccb017c1f3bee916da4fc1648 --- /dev/null +++ b/config/dataset/view_sampler/boundedv2.yaml @@ -0,0 +1,15 @@ +name: boundedv2 + +num_target_views: 1 +num_context_views: 2 + +min_distance_between_context_views: 2 +max_distance_between_context_views: 6 +max_distance_to_context_views: 0 + +context_gap_warm_up_steps: 0 +target_gap_warm_up_steps: 0 + +initial_min_distance_between_context_views: 2 +initial_max_distance_between_context_views: 6 +initial_max_distance_to_context_views: 0 diff --git a/config/dataset/view_sampler/boundedv2_360.yaml b/config/dataset/view_sampler/boundedv2_360.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b6c2609aa231ed1b995249c319ab4992cbc3dd93 --- /dev/null +++ b/config/dataset/view_sampler/boundedv2_360.yaml @@ -0,0 +1,17 @@ +name: boundedv2 + +num_target_views: 4 +num_context_views: 4 + +min_distance_between_context_views: 20 +max_distance_between_context_views: 50 +max_distance_to_context_views: 0 + +context_gap_warm_up_steps: 10000 +target_gap_warm_up_steps: 0 + +initial_min_distance_between_context_views: 15 +initial_max_distance_between_context_views: 30 +initial_max_distance_to_context_views: 0 +extra_views_sampling_strategy: farthest_point +target_views_replace_sample: false diff --git a/config/dataset/view_sampler/evaluation.yaml b/config/dataset/view_sampler/evaluation.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3931de80100fc279027bc7285893d11871d546ab --- /dev/null +++ b/config/dataset/view_sampler/evaluation.yaml @@ -0,0 +1,4 @@ +name: evaluation + +index_path: assets/evaluation_index_re10k_video.json +num_context_views: 2 diff --git a/config/dataset/view_sampler_dataset_specific_config/bounded_re10k.yaml b/config/dataset/view_sampler_dataset_specific_config/bounded_re10k.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5c8030734831b35eb63677f8f912cd88369e7edf --- /dev/null +++ b/config/dataset/view_sampler_dataset_specific_config/bounded_re10k.yaml @@ -0,0 +1,11 @@ +# @package _global_ + +dataset: + view_sampler: + min_distance_between_context_views: 45 + max_distance_between_context_views: 135 + min_distance_to_context_views: 0 + warm_up_steps: 30000 + initial_min_distance_between_context_views: 25 + initial_max_distance_between_context_views: 45 + num_target_views: 4 diff --git a/config/dataset/view_sampler_dataset_specific_config/boundedv2_dl3dv.yaml b/config/dataset/view_sampler_dataset_specific_config/boundedv2_dl3dv.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5d52951033d0afaf0d246b95c9727eb8b9745438 --- /dev/null +++ b/config/dataset/view_sampler_dataset_specific_config/boundedv2_dl3dv.yaml @@ -0,0 +1,14 @@ +# @package _global_ + +dataset: + view_sampler: + min_distance_between_context_views: 20 + max_distance_between_context_views: 50 + max_distance_to_context_views: 0 + context_gap_warm_up_steps: 10000 + target_gap_warm_up_steps: 0 + initial_min_distance_between_context_views: 15 + initial_max_distance_between_context_views: 30 + initial_max_distance_to_context_views: 0 + extra_views_sampling_strategy: farthest_point + num_target_views: 4 diff --git a/config/dataset/view_sampler_dataset_specific_config/evaluation_dl3dv.yaml b/config/dataset/view_sampler_dataset_specific_config/evaluation_dl3dv.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b98fb4d9197bafd4baa233cb2ac7248136e34366 --- /dev/null +++ b/config/dataset/view_sampler_dataset_specific_config/evaluation_dl3dv.yaml @@ -0,0 +1,7 @@ +# @package _global_ + +dataset: + image_shape: [256, 448] + view_sampler: + index_path: assets/dl3dv_evaluation/dl3dv_start_0_distance_40_ctx_8v_tgt_8v.json + num_context_views: 8 diff --git a/config/dataset/view_sampler_dataset_specific_config/evaluation_re10k.yaml b/config/dataset/view_sampler_dataset_specific_config/evaluation_re10k.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8c9235712813bac9403dc56d99c7d2ca36a10547 --- /dev/null +++ b/config/dataset/view_sampler_dataset_specific_config/evaluation_re10k.yaml @@ -0,0 +1,5 @@ +# @package _global_ + +dataset: + view_sampler: + index_path: assets/evaluation_index_re10k.json diff --git a/config/experiment/dl3dv.yaml b/config/experiment/dl3dv.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4f732a20e6a29a724c6c1a5d2e2028680886738c --- /dev/null +++ b/config/experiment/dl3dv.yaml @@ -0,0 +1,41 @@ +# @package _global_ + +defaults: + - override /dataset: dl3dv + - override /model/encoder: codecsplat + - override /loss: [mse, lpips] + - override /dataset/view_sampler: boundedv2_360 + +wandb: + name: dl3dv + tags: [dl3dv, 270x480] + +data_loader: + train: + batch_size: 1 + +trainer: + max_steps: 300_001 + num_nodes: 1 + +model: + encoder: + num_depth_candidates: 128 + shim_patch_size: 16 + +# lpips loss +loss: + lpips: + apply_after_step: 0 + weight: 0.05 + +dataset: + near: 0.5 + far: 200. + min_views: 0 + max_views: 0 + highres: false + +test: + eval_time_skip_steps: 0 + compute_scores: true diff --git a/config/experiment/re10k.yaml b/config/experiment/re10k.yaml new file mode 100644 index 0000000000000000000000000000000000000000..57cbdd16e8410de151e34541c691a4de0268c3f6 --- /dev/null +++ b/config/experiment/re10k.yaml @@ -0,0 +1,40 @@ +# @package _global_ + +defaults: + - override /dataset: re10k + - override /model/encoder: codecsplat + - override /loss: [mse, lpips] + +wandb: + name: re10k + tags: [re10k, 256x256] + +data_loader: + train: + batch_size: 14 + +trainer: + max_steps: 300_001 + num_nodes: 1 + +model: + encoder: + num_depth_candidates: 128 + +# lpips loss +loss: + lpips: + apply_after_step: 0 + weight: 0.05 + +dataset: + image_shape: [256, 256] + roots: [datasets/RE10K-Torch-360p] + near: 0.5 + far: 100. + train_times_per_scene: 1 + highres: false + +test: + eval_time_skip_steps: 0 + compute_scores: true diff --git a/config/loss/lpips.yaml b/config/loss/lpips.yaml new file mode 100644 index 0000000000000000000000000000000000000000..af4c548c655ba9bc4b4fe2ae3d8f388aa0357238 --- /dev/null +++ b/config/loss/lpips.yaml @@ -0,0 +1,3 @@ +lpips: + weight: 0.05 + apply_after_step: 150_000 diff --git a/config/loss/mse.yaml b/config/loss/mse.yaml new file mode 100644 index 0000000000000000000000000000000000000000..80cc0be6dc7950661998336bd2bb5cb4ff06ba07 --- /dev/null +++ b/config/loss/mse.yaml @@ -0,0 +1,2 @@ +mse: + weight: 1.0 diff --git a/config/main.yaml b/config/main.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ed2395575070d0ae90e7048cefc1ef68f653aa9f --- /dev/null +++ b/config/main.yaml @@ -0,0 +1,115 @@ +defaults: + - dataset: re10k + - optional dataset/view_sampler_dataset_specific_config: ${dataset/view_sampler}_${dataset} + - model/encoder: codecsplat + - model/decoder: splatting_gsplat + - loss: [mse] + - _self_ + +wandb: + project: codecsplat + entity: null + name: placeholder + mode: online + id: null + +mode: train + +dataset: + overfit_to_scene: null + +data_loader: + train: + num_workers: 10 + persistent_workers: true + batch_size: 4 + seed: 1234 + test: + num_workers: 4 + persistent_workers: false + batch_size: 1 + seed: 2345 + val: + num_workers: 1 + persistent_workers: true + batch_size: 1 + seed: 3456 + +optimizer: + lr: 2.e-4 + lr_monodepth: 2.e-6 + warm_up_steps: 2000 + weight_decay: 0.01 + div_factor: 25. + final_div_factor: 1.e4 + +checkpointing: + load: null + every_n_train_steps: 5000 + save_top_k: 5 + pretrained_model: null + pretrained_monodepth: null + pretrained_mvdepth: null + pretrained_depth: null + no_strict_load: false + resume: false + +train: + depth_mode: null + extended_visualization: false + print_log_every_n_steps: 100 + eval_model_every_n_val: 2 # quantitative evaluation every n val + eval_data_length: 999999 + eval_time_skip_steps: 3 + eval_save_model: true + l1_loss: false + intermediate_loss_weight: 0.9 + no_viz_video: false + train_ignore_large_loss: 0. + +test: + compute_scores: true + eval_time_skip_steps: 0 + save_image: false + save_video: false + save_gt_image: false + save_input_images: false + save_depth: false + save_depth_npy: false + save_depth_concat_img: false + save_gaussian: false + save_bitstream: false + render_chunk_size: null + stabilize_camera: false + stab_camera_kernel: 50 + gaussian_export_all_points: true + gaussian_export_align_to_view: false + +seed: 111123 + +trainer: + max_steps: -1 + val_check_interval: 0.5 + gradient_clip_val: 0.5 + accumulate_grad_batches: 1 + precision: 32-true + num_sanity_val_steps: 2 + num_nodes: 1 + +output_dir: checkpoints/tmp + +use_plugins: false + +hydra: + run: + dir: ${output_dir}/hydra/${now:%Y-%m-%d-%H-%M-%S} + job: + chdir: false + output_subdir: . + job_logging: + handlers: + file: + filename: /dev/null + root: + handlers: + - console diff --git a/config/model/decoder/splatting_cuda.yaml b/config/model/decoder/splatting_cuda.yaml new file mode 100644 index 0000000000000000000000000000000000000000..762a80c3b621a4a64adb13ae2a24e27564ad2d87 --- /dev/null +++ b/config/model/decoder/splatting_cuda.yaml @@ -0,0 +1 @@ +name: splatting_cuda diff --git a/config/model/decoder/splatting_gsplat.yaml b/config/model/decoder/splatting_gsplat.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3e5c8d09aac7177ca89c1f6f0f4421f725551cee --- /dev/null +++ b/config/model/decoder/splatting_gsplat.yaml @@ -0,0 +1,10 @@ +name: splatting_gsplat +render_mode: RGB +rasterize_mode: antialiased +use_covariances: true +packed: true +eps2d: 0.1 +sparse_grad: false +absgrad: false +channel_chunk: 32 +scale_invariant: false diff --git a/config/model/encoder/codecsplat.yaml b/config/model/encoder/codecsplat.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1594663a5b1a451c3873ba85ffc58fde4aa7a58d --- /dev/null +++ b/config/model/encoder/codecsplat.yaml @@ -0,0 +1,48 @@ +name: codecsplat + +num_depth_candidates: 128 +num_surfaces: 1 + +gaussian_adapter: + gaussian_scale_min: 1e-10 + gaussian_scale_max: 3. + sh_degree: 2 + exp_scale_bias: 4. + no_rotate_sh: false + no_sh_mask: false + +# Input alignment +downscale_factor: 4 +shim_patch_size: 4 + +# monodepth +monodepth_vit_type: vits + +# return depth +supervise_intermediate_depth: true +return_depth: true + +# mv_unimatch +num_scales: 1 +upsample_factor: 4 +lowest_feature_resolution: 4 +depth_unet_channels: 128 +grid_sample_disable_cudnn: false +local_mv_match: 2 + +# color branch +init_sh_input_img: false +gaussian_regressor_channels: 64 + +# feature codec path +enable_feature_codec: false +enable_depth_guided_mv_refiner: false +enable_post_vae_mv_refiner: false +codec_loss_weight: 1.0 +codec_depth_loss_weight: 0.1 +codec_lmb_range: [16.0, 1024.0] +freeze_legacy_modules_for_codec: false +codec_estimated_feature_num_bytes: false +deflate_feature_ablation: false +deflate_feature_dtype: float16 +deflate_feature_level: 9 diff --git a/examples/scene_a/00.jpg b/examples/scene_a/00.jpg new file mode 100644 index 0000000000000000000000000000000000000000..817490a33bf1775df1ab210238c07b8ac09797aa --- /dev/null +++ b/examples/scene_a/00.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:864f1efe8677fd03f2e0a825c1b6e36eb9dcdc6c1fd7b5e69f3ee07f5ba62fa0 +size 151310 diff --git a/examples/scene_a/01.jpg b/examples/scene_a/01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..73f1b135e3cf15afa9faf1fb64956d1d8635daa0 --- /dev/null +++ b/examples/scene_a/01.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:414937f97e97b07a47c3b4e81b0325f4627818ea89f638f754fe738caecc2962 +size 151368 diff --git a/examples/scene_a/02.jpg b/examples/scene_a/02.jpg new file mode 100644 index 0000000000000000000000000000000000000000..321c30ce346ab4e56f6857445d2d72fc32fdd84c --- /dev/null +++ b/examples/scene_a/02.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:97312a3a25fc1f01df44d648cdb8681802d5341a3721f6fd12c24afd5df480de +size 150304 diff --git a/examples/scene_a/03.jpg b/examples/scene_a/03.jpg new file mode 100644 index 0000000000000000000000000000000000000000..1af2d44d329209c978587bc8b780c76d40c9deab --- /dev/null +++ b/examples/scene_a/03.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:67054b1d5cf5563913502c2ca32f4c47739af7def7c10185ae840ee4cf812c86 +size 148448 diff --git a/examples/scene_a/04.jpg b/examples/scene_a/04.jpg new file mode 100644 index 0000000000000000000000000000000000000000..e4009454a42f29ea36f08cd86d2b20b82d9db337 --- /dev/null +++ b/examples/scene_a/04.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1d733527d3fa5e7d3c48b35af5db16d15f182ffeb2a656e075879f6983a8364f +size 148568 diff --git a/examples/scene_a/05.jpg b/examples/scene_a/05.jpg new file mode 100644 index 0000000000000000000000000000000000000000..6e1bc51119bccdb9db5c9cb457455cc6d33cb1fb --- /dev/null +++ b/examples/scene_a/05.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3d3891369fd5687128e5e50f5020c323304008b04270ca928698a9f3e93af3fe +size 146970 diff --git a/examples/scene_a/06.jpg b/examples/scene_a/06.jpg new file mode 100644 index 0000000000000000000000000000000000000000..0f3700a53f760699ebbc80d995cc4ffea6bf4219 --- /dev/null +++ b/examples/scene_a/06.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8f2f368c20a98f027e3cbd8cde63830c3db02ed4e6d032a72c7595c6cc708563 +size 142300 diff --git a/examples/scene_a/07.jpg b/examples/scene_a/07.jpg new file mode 100644 index 0000000000000000000000000000000000000000..33b5610a2a0f4a9615e09799af2e5b01728d5726 --- /dev/null +++ b/examples/scene_a/07.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6e1c35fac53bb629bb048395d5d7593bced09a8df310e6805f09bee8cdae9323 +size 146387 diff --git a/examples/scene_a/08.jpg b/examples/scene_a/08.jpg new file mode 100644 index 0000000000000000000000000000000000000000..b3f88b4100c960992b2a59105a00a630dab5ee45 --- /dev/null +++ b/examples/scene_a/08.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f04ed65071df1cbef47180edbdc6a135d4c3e6a40e1d1c93dfe539297d226aa5 +size 145494 diff --git a/examples/scene_a/09.jpg b/examples/scene_a/09.jpg new file mode 100644 index 0000000000000000000000000000000000000000..93aa29b350c769190e83c901f49e9210ee90c0d7 --- /dev/null +++ b/examples/scene_a/09.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d8d92f95fedc7b924046d3ef56a08e084943f5846fe0197269eb6d94d12fcbc0 +size 146737 diff --git a/examples/scene_a/10.jpg b/examples/scene_a/10.jpg new file mode 100644 index 0000000000000000000000000000000000000000..c862fd41d39ff75d7b9c09de82de83671b062076 --- /dev/null +++ b/examples/scene_a/10.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6771ccb3ee35aa5b61df20cc0e38f5423d0a815a53506d08ded513475b1d1ff3 +size 151239 diff --git a/examples/scene_a/11.jpg b/examples/scene_a/11.jpg new file mode 100644 index 0000000000000000000000000000000000000000..c16dcdbc5a3af89c12c46a38c28e4bade6371602 --- /dev/null +++ b/examples/scene_a/11.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f6eb4c48ec8c18dcc44dd0ed140cb335a7705fef25545430d9e4847ac7294f56 +size 150482 diff --git a/examples/scene_a/12.jpg b/examples/scene_a/12.jpg new file mode 100644 index 0000000000000000000000000000000000000000..7d05f8a3e5611d10fa79290177c983e4f360651f --- /dev/null +++ b/examples/scene_a/12.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f0fa1ed39c728f9e1f8464ac6a0ed919361a39b1b1bd21c38b307fd0d4adab29 +size 153629 diff --git a/examples/scene_a/13.jpg b/examples/scene_a/13.jpg new file mode 100644 index 0000000000000000000000000000000000000000..9c454a45d5125cd6e724ee5748a1f13741c42f5c --- /dev/null +++ b/examples/scene_a/13.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:03291cbdef1c6f7d3e00c2ee85787df834cf0e8e15698750cdc1b3ad73d04893 +size 152007 diff --git a/examples/scene_a/14.jpg b/examples/scene_a/14.jpg new file mode 100644 index 0000000000000000000000000000000000000000..1acd34f98dea95433e361b7e65fa8869d5e672af --- /dev/null +++ b/examples/scene_a/14.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b0d2135cb30ef243a02fffdd7d20b61239eca64baab7665158d8854693744c77 +size 145530 diff --git a/examples/scene_a/15.jpg b/examples/scene_a/15.jpg new file mode 100644 index 0000000000000000000000000000000000000000..3943a1694778d75a5076620daa3c0e0a34835e03 --- /dev/null +++ b/examples/scene_a/15.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:295911b689adc5f498451c23b5d1fd1ed668b6dad98284109b59e973756fdf62 +size 153357 diff --git a/examples/scene_a/16.jpg b/examples/scene_a/16.jpg new file mode 100644 index 0000000000000000000000000000000000000000..0e67da20c082d09cef021e92595d86ce00ee7ce2 --- /dev/null +++ b/examples/scene_a/16.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2080b7bce75b0714e37987719a22e116c836d884a02e10435c36d57e2fe27002 +size 147034 diff --git a/examples/scene_a/17.jpg b/examples/scene_a/17.jpg new file mode 100644 index 0000000000000000000000000000000000000000..94466512ba07db02b85290a40c892327d3105f43 --- /dev/null +++ b/examples/scene_a/17.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1606e2337ee00c78b700612e6957d1bca17aeb8edbce4808d595614a3af538f6 +size 147154 diff --git a/examples/scene_a/18.jpg b/examples/scene_a/18.jpg new file mode 100644 index 0000000000000000000000000000000000000000..241c2f222f73d6b675b4d06cbd18f52c01e25954 --- /dev/null +++ b/examples/scene_a/18.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8e6444b113ccafd634cb184b5a878a9e73b3fcfee198c863ac9930879f76daed +size 144352 diff --git a/examples/scene_a/19.jpg b/examples/scene_a/19.jpg new file mode 100644 index 0000000000000000000000000000000000000000..f045a8b2453cd5bc82d9b4d2febdb586b477c8f2 --- /dev/null +++ b/examples/scene_a/19.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:116b1524e488083cc2b0262c115f9c9955840f761f91299caf92220b929481b2 +size 143790 diff --git a/examples/scene_a/20.jpg b/examples/scene_a/20.jpg new file mode 100644 index 0000000000000000000000000000000000000000..8921d9d70ce4f88a6bf3dc7170086bbdc533cc72 --- /dev/null +++ b/examples/scene_a/20.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2656b43536aba2536d8a070b3e2ca3dab1af93e84161f07464a89d1dc4995387 +size 143586 diff --git a/examples/scene_a/21.jpg b/examples/scene_a/21.jpg new file mode 100644 index 0000000000000000000000000000000000000000..042162e2f13027e439e674672aea0cb57f2697ca --- /dev/null +++ b/examples/scene_a/21.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:03bb71d6292638d7c052504eefdba3555164a6c21618c447963d2799cc94c60e +size 143356 diff --git a/examples/scene_a/22.jpg b/examples/scene_a/22.jpg new file mode 100644 index 0000000000000000000000000000000000000000..8f58aadb93692a933bfeaf1816bdad09164cfad4 --- /dev/null +++ b/examples/scene_a/22.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5f3eb3f7f8fa5c942e4c1d3aaabdacbdf21ae54d3c4e20cdc5da82122c47a052 +size 136242 diff --git a/examples/scene_a/23.jpg b/examples/scene_a/23.jpg new file mode 100644 index 0000000000000000000000000000000000000000..8b58289e6148a58ada989f4d6c222b9596db0b5b --- /dev/null +++ b/examples/scene_a/23.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a6b1ff00c86b98bda90b26526626239eab2cf24959954798d74ac9c57d222a17 +size 136475 diff --git a/examples/scene_a/24.jpg b/examples/scene_a/24.jpg new file mode 100644 index 0000000000000000000000000000000000000000..33cf45c14a2b0d2c3f3183567ef4868202a0e40f --- /dev/null +++ b/examples/scene_a/24.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:63a0083d2f1dfd9c01d2fd5faa96e5b103cc844e794399bb3daa55ca2d83e21d +size 134872 diff --git a/examples/scene_b/00.jpg b/examples/scene_b/00.jpg new file mode 100644 index 0000000000000000000000000000000000000000..781f48f151dc48c4699fbea12606cb1842a7d0ca --- /dev/null +++ b/examples/scene_b/00.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:571d06f0aed3d981446cc5d791306ebd86b2fb6d10bab4de0fa846902ad45bb6 +size 114913 diff --git a/examples/scene_b/01.jpg b/examples/scene_b/01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..a8583d54d365942c26f9c4f2af70b1671b30d6a9 --- /dev/null +++ b/examples/scene_b/01.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3f1e5ee3c7f66a3f7e94cd388d07265a77d983f6f3482afcc16ae2d15c290903 +size 113487 diff --git a/examples/scene_b/02.jpg b/examples/scene_b/02.jpg new file mode 100644 index 0000000000000000000000000000000000000000..03b692e82ad7925f394408b985c4daf072b826a7 --- /dev/null +++ b/examples/scene_b/02.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ae982a465ba81bb669ff22572b4acdd7d67b11ab8b1c80f6a0414dd70f94f502 +size 110840 diff --git a/examples/scene_b/03.jpg b/examples/scene_b/03.jpg new file mode 100644 index 0000000000000000000000000000000000000000..1348618df5ed22de549f695ad3e5b56585d251ee --- /dev/null +++ b/examples/scene_b/03.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a31c68906558581490ec997bd95be751e381ae4b4ff972d06f2a3f84e23c4737 +size 122456 diff --git a/examples/scene_b/04.jpg b/examples/scene_b/04.jpg new file mode 100644 index 0000000000000000000000000000000000000000..3c2285c0bc58d81c8db9c14916060537a49e68c0 --- /dev/null +++ b/examples/scene_b/04.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e47238cd70e8a711c38289c8b988db960868afdcf3cddaac5da2484c23891282 +size 113853 diff --git a/examples/scene_b/05.jpg b/examples/scene_b/05.jpg new file mode 100644 index 0000000000000000000000000000000000000000..2e2d7ef346441ab296943c94ebfcc546fdd0101a --- /dev/null +++ b/examples/scene_b/05.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fad67389bd63fc2db0082ebc8ce3fc31ebd711190c5f66da7eaeba45300c25b9 +size 120135 diff --git a/examples/scene_b/06.jpg b/examples/scene_b/06.jpg new file mode 100644 index 0000000000000000000000000000000000000000..a1c16b40547053b2968a1d56b991be279db9ec4f --- /dev/null +++ b/examples/scene_b/06.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1a6840accff5fbf178f927213e92592354297cb6d7c925727318ad0d4288ebfe +size 121284 diff --git a/examples/scene_b/07.jpg b/examples/scene_b/07.jpg new file mode 100644 index 0000000000000000000000000000000000000000..676c16eb8bb5a68ed6cdacbbaec29251c59465cd --- /dev/null +++ b/examples/scene_b/07.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d29b87d9e8414f0c3e937ae4759b3bacb003a62643f53de260859371648955a8 +size 119865 diff --git a/examples/scene_b/08.jpg b/examples/scene_b/08.jpg new file mode 100644 index 0000000000000000000000000000000000000000..4c4dab37c08ab43ef417550157dc0c9ec5505713 --- /dev/null +++ b/examples/scene_b/08.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:66f2f1daac80d15f5738a35189b7feeaf15539e8d63e16a572a28820958ea42f +size 113789 diff --git a/examples/scene_b/09.jpg b/examples/scene_b/09.jpg new file mode 100644 index 0000000000000000000000000000000000000000..589f71c1c827987ffdf184f23c6bbcf7d9f7ad66 --- /dev/null +++ b/examples/scene_b/09.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e4e40c4b709f316132a470b366bcbc3ec9585821f3b4bf5807f278d1205d2e0b +size 118476 diff --git a/examples/scene_b/10.jpg b/examples/scene_b/10.jpg new file mode 100644 index 0000000000000000000000000000000000000000..1ea54261948d5361070791fcc155f15f3fb494ce --- /dev/null +++ b/examples/scene_b/10.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9c45c013f9ac0b5723d6ac496c08cf2603163c10a70fd113fbed241d29ba22b8 +size 122007 diff --git a/examples/scene_b/11.jpg b/examples/scene_b/11.jpg new file mode 100644 index 0000000000000000000000000000000000000000..24127baabb5d13670088ec00f9e8241c6e1fc8cb --- /dev/null +++ b/examples/scene_b/11.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5ac9c91f9321ee750ab9841058c9178e8f83933d146fba67715cacaf76694fa6 +size 117874 diff --git a/examples/scene_b/12.jpg b/examples/scene_b/12.jpg new file mode 100644 index 0000000000000000000000000000000000000000..f9d21b5ebc72892452086b4d5d74ec5cfc7ac432 --- /dev/null +++ b/examples/scene_b/12.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c38ab169eea3c3ef63dd3fea01dca6206e1ed5f42758eefc9ad9473042184947 +size 116113 diff --git a/examples/scene_b/13.jpg b/examples/scene_b/13.jpg new file mode 100644 index 0000000000000000000000000000000000000000..66263bed5258f6703092c2d319865b0f29011ef6 --- /dev/null +++ b/examples/scene_b/13.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a040224ea60e080f866f98e919fb0d001e91fa0518f403cfa262efde90021f6f +size 126699 diff --git a/examples/scene_b/14.jpg b/examples/scene_b/14.jpg new file mode 100644 index 0000000000000000000000000000000000000000..c1a100498fa290893c77be8b129dbe537721f283 --- /dev/null +++ b/examples/scene_b/14.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:81af4c2022e932a2d59702f92527c251d5e459a317f304e6a9ff430dd7f96759 +size 124809 diff --git a/examples/scene_b/15.jpg b/examples/scene_b/15.jpg new file mode 100644 index 0000000000000000000000000000000000000000..01702dadc438f978addea1551be5b901a3a20974 --- /dev/null +++ b/examples/scene_b/15.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a94c9c241c911d79bb420fdddee93c7a782ff13721b936f15bb103f41461971f +size 127789 diff --git a/examples/scene_b/16.jpg b/examples/scene_b/16.jpg new file mode 100644 index 0000000000000000000000000000000000000000..877cdf314a896c21e274c516749bdabfc7d8a752 --- /dev/null +++ b/examples/scene_b/16.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:78ba0db18e1593d58458ca41444271f30ea0a8d5b47207f824282186fc9f1db6 +size 129340 diff --git a/examples/scene_b/17.jpg b/examples/scene_b/17.jpg new file mode 100644 index 0000000000000000000000000000000000000000..37f0670e0ce99bbb3fc411a1b3e33437941eed7a --- /dev/null +++ b/examples/scene_b/17.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8fb6764a501d588e4c0bfd79c844a01ac6bacd40f10614bfc0e5a835bcb45fe5 +size 131565 diff --git a/examples/scene_b/18.jpg b/examples/scene_b/18.jpg new file mode 100644 index 0000000000000000000000000000000000000000..02b118c3aaf5564cc999d7ca711cb1bb830ba1c5 --- /dev/null +++ b/examples/scene_b/18.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:27c52859e3e788dc52c08c19f3a16ef439396fcef83e0d2da87b12b20eb558c5 +size 134459 diff --git a/examples/scene_b/19.jpg b/examples/scene_b/19.jpg new file mode 100644 index 0000000000000000000000000000000000000000..b603daeaff206d342e80859788df9803881d977d --- /dev/null +++ b/examples/scene_b/19.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7824aa0964e4dd708cd783f5e80af8c1c2611f546c00195366fddd76a9073c30 +size 141144 diff --git a/examples/scene_b/20.jpg b/examples/scene_b/20.jpg new file mode 100644 index 0000000000000000000000000000000000000000..b1be93aee351894533e4f5f155aeba4e25fba50d --- /dev/null +++ b/examples/scene_b/20.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:86d5b4c3458dbd119b8fe7bf0fd1a53bdfb54d4c301f904f670d1497efef9b96 +size 149557 diff --git a/examples/scene_b/21.jpg b/examples/scene_b/21.jpg new file mode 100644 index 0000000000000000000000000000000000000000..f75c4681fb8373d00daa542b6f2395f621c5483e --- /dev/null +++ b/examples/scene_b/21.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f2d55c8f8f00f0dede327fbf0ffc86b49eb3b9ae535d270cc1db44ec4b7b88d0 +size 144901 diff --git a/examples/scene_b/22.jpg b/examples/scene_b/22.jpg new file mode 100644 index 0000000000000000000000000000000000000000..60afa8026e936da0a6a0fbeb9ab3d2b01c14dd08 --- /dev/null +++ b/examples/scene_b/22.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:109095f08595c41bbb563c017c2e1c36edd0b19213295311ac3f64e889a0b7c6 +size 149236 diff --git a/examples/scene_b/23.jpg b/examples/scene_b/23.jpg new file mode 100644 index 0000000000000000000000000000000000000000..80f3f6138e90c60b5c6d64dad49795c3aabf0b60 --- /dev/null +++ b/examples/scene_b/23.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:be334c3f7a7cfbbfb3f16779c551727bfcf54673959fa02dc0c83b3e97b00d5d +size 152088 diff --git a/examples/scene_b/24.jpg b/examples/scene_b/24.jpg new file mode 100644 index 0000000000000000000000000000000000000000..94765088263b7bd8ca57732659044e6272d9fb09 --- /dev/null +++ b/examples/scene_b/24.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5444c5823ad7729022b875a5ad389c7c828438a1f0a6e2ac72993abfec56ed49 +size 153433 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..93792bfa38260c574e8a7078be9c4b1b5c71d389 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,24 @@ +torchvision +gsplat==1.5.3 +ninja +compressai==1.2.8 +numpy<2 +pycolmap==4.1.1 +einops +jaxtyping>=0.3.7,<0.3.11 +beartype +hydra-core +omegaconf +dacite +colorama +colorspacious +e3nn +lpips +plyfile +scipy +matplotlib +opencv-python-headless +Pillow +tqdm +imageio +imageio-ffmpeg diff --git a/src/config.py b/src/config.py new file mode 100644 index 0000000000000000000000000000000000000000..2a601907676cb0477fcba321083a05289d93eca9 --- /dev/null +++ b/src/config.py @@ -0,0 +1,99 @@ +from dataclasses import dataclass +from pathlib import Path +from typing import Literal, Optional, Type, TypeVar + +from dacite import Config, from_dict +from omegaconf import DictConfig, OmegaConf + +from .dataset.data_module import DataLoaderCfg, DatasetCfg +from .loss import LossCfgWrapper +from .model.decoder import DecoderCfg +from .model.encoder import EncoderCfg +from .model.model_wrapper import OptimizerCfg, TestCfg, TrainCfg + + +@dataclass +class CheckpointingCfg: + load: Optional[str] # Not a path, since it could be something like wandb://... + every_n_train_steps: int + save_top_k: int + pretrained_model: Optional[str] + pretrained_monodepth: Optional[str] + pretrained_mvdepth: Optional[str] + pretrained_depth: Optional[str] + no_strict_load: bool + resume: bool + + +@dataclass +class ModelCfg: + decoder: DecoderCfg + encoder: EncoderCfg + + +@dataclass +class TrainerCfg: + max_steps: int + val_check_interval: int | float | None + gradient_clip_val: int | float | None + accumulate_grad_batches: int + precision: str + num_sanity_val_steps: int + num_nodes: int + + +@dataclass +class RootCfg: + wandb: dict + mode: Literal["train", "test"] + dataset: DatasetCfg + data_loader: DataLoaderCfg + model: ModelCfg + optimizer: OptimizerCfg + checkpointing: CheckpointingCfg + trainer: TrainerCfg + loss: list[LossCfgWrapper] + test: TestCfg + train: TrainCfg + seed: int + use_plugins: bool + + +TYPE_HOOKS = { + Path: Path, +} + + +T = TypeVar("T") + + +def load_typed_config( + cfg: DictConfig, + data_class: Type[T], + extra_type_hooks: dict = {}, +) -> T: + return from_dict( + data_class, + OmegaConf.to_container(cfg), + config=Config(type_hooks={**TYPE_HOOKS, **extra_type_hooks}), + ) + + +def separate_loss_cfg_wrappers(joined: dict) -> list[LossCfgWrapper]: + # The dummy allows the union to be converted. + @dataclass + class Dummy: + dummy: LossCfgWrapper + + return [ + load_typed_config(DictConfig({"dummy": {k: v}}), Dummy).dummy + for k, v in joined.items() + ] + + +def load_typed_root_config(cfg: DictConfig) -> RootCfg: + return load_typed_config( + cfg, + RootCfg, + {list[LossCfgWrapper]: separate_loss_cfg_wrappers}, + ) diff --git a/src/dataset/__init__.py b/src/dataset/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..05f0be852038505ae2d767553918d71647fd33fe --- /dev/null +++ b/src/dataset/__init__.py @@ -0,0 +1,30 @@ +from torch.utils.data import Dataset + +from ..misc.step_tracker import StepTracker +from .dataset_re10k import DatasetRE10k, DatasetRE10kCfg +from .dataset_dl3dv import DatasetDL3DV, DatasetDL3DVCfg +from .types import Stage +from .view_sampler import get_view_sampler + +DATASETS: dict[str, Dataset] = { + "re10k": DatasetRE10k, + "dl3dv": DatasetDL3DV, +} + + +DatasetCfg = DatasetRE10kCfg | DatasetDL3DVCfg + + +def get_dataset( + cfg: DatasetCfg, + stage: Stage, + step_tracker: StepTracker | None, +) -> Dataset: + view_sampler = get_view_sampler( + cfg.view_sampler, + stage, + cfg.overfit_to_scene is not None, + cfg.cameras_are_circular, + step_tracker, + ) + return DATASETS[cfg.name](cfg, stage, view_sampler) diff --git a/src/dataset/data_module.py b/src/dataset/data_module.py new file mode 100644 index 0000000000000000000000000000000000000000..dcf49010b3ab2fe051db0f25a966b89b37884c16 --- /dev/null +++ b/src/dataset/data_module.py @@ -0,0 +1,130 @@ +import random +from dataclasses import dataclass +from typing import Callable + +import numpy as np +import torch +from pytorch_lightning import LightningDataModule +from torch import Generator, nn +from torch.utils.data import DataLoader, Dataset, IterableDataset + +from ..misc.step_tracker import StepTracker +from . import DatasetCfg, get_dataset +from .types import DataShim, Stage +from .validation_wrapper import ValidationWrapper + + +def get_data_shim(encoder: nn.Module) -> DataShim: + """Get functions that modify the batch. It's sometimes necessary to modify batches + outside the data loader because GPU computations are required to modify the batch or + because the modification depends on something outside the data loader. + """ + + shims: list[DataShim] = [] + if hasattr(encoder, "get_data_shim"): + shims.append(encoder.get_data_shim()) + + def combined_shim(batch): + for shim in shims: + batch = shim(batch) + return batch + + return combined_shim + + +@dataclass +class DataLoaderStageCfg: + batch_size: int + num_workers: int + persistent_workers: bool + seed: int | None + + +@dataclass +class DataLoaderCfg: + train: DataLoaderStageCfg + test: DataLoaderStageCfg + val: DataLoaderStageCfg + + +DatasetShim = Callable[[Dataset, Stage], Dataset] + + +def worker_init_fn(worker_id: int) -> None: + random.seed(int(torch.utils.data.get_worker_info().seed) % (2**32 - 1)) + np.random.seed(int(torch.utils.data.get_worker_info().seed) % (2**32 - 1)) + + +class DataModule(LightningDataModule): + dataset_cfg: DatasetCfg + data_loader_cfg: DataLoaderCfg + step_tracker: StepTracker | None + dataset_shim: DatasetShim + global_rank: int + + def __init__( + self, + dataset_cfg: DatasetCfg, + data_loader_cfg: DataLoaderCfg, + step_tracker: StepTracker | None = None, + dataset_shim: DatasetShim = lambda dataset, _: dataset, + global_rank: int = 0, + ) -> None: + super().__init__() + self.dataset_cfg = dataset_cfg + self.data_loader_cfg = data_loader_cfg + self.step_tracker = step_tracker + self.dataset_shim = dataset_shim + self.global_rank = global_rank + + def get_persistent(self, loader_cfg: DataLoaderStageCfg) -> bool | None: + return None if loader_cfg.num_workers == 0 else loader_cfg.persistent_workers + + def get_generator(self, loader_cfg: DataLoaderStageCfg) -> torch.Generator | None: + if loader_cfg.seed is None: + return None + generator = Generator() + generator.manual_seed(loader_cfg.seed + self.global_rank) + return generator + + def train_dataloader(self): + dataset = get_dataset(self.dataset_cfg, "train", self.step_tracker) + dataset = self.dataset_shim(dataset, "train") + return DataLoader( + dataset, + self.data_loader_cfg.train.batch_size, + shuffle=not isinstance(dataset, IterableDataset), + num_workers=self.data_loader_cfg.train.num_workers, + generator=self.get_generator(self.data_loader_cfg.train), + worker_init_fn=worker_init_fn, + persistent_workers=self.get_persistent(self.data_loader_cfg.train), + ) + + def val_dataloader(self): + dataset = get_dataset(self.dataset_cfg, "val", self.step_tracker) + dataset = self.dataset_shim(dataset, "val") + return DataLoader( + ValidationWrapper(dataset, 1), + self.data_loader_cfg.val.batch_size, + num_workers=self.data_loader_cfg.val.num_workers, + generator=self.get_generator(self.data_loader_cfg.val), + worker_init_fn=worker_init_fn, + persistent_workers=self.get_persistent(self.data_loader_cfg.val), + ) + + def test_dataloader(self, dataset_cfg=None): + dataset = get_dataset( + self.dataset_cfg if dataset_cfg is None else dataset_cfg, + "test", + self.step_tracker, + ) + dataset = self.dataset_shim(dataset, "test") + return DataLoader( + dataset, + self.data_loader_cfg.test.batch_size, + num_workers=self.data_loader_cfg.test.num_workers, + generator=self.get_generator(self.data_loader_cfg.test), + worker_init_fn=worker_init_fn, + persistent_workers=self.get_persistent(self.data_loader_cfg.test), + shuffle=False, + ) diff --git a/src/dataset/dataset.py b/src/dataset/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..d5df79186dc7d165ac67466b73fe51ea02cbec34 --- /dev/null +++ b/src/dataset/dataset.py @@ -0,0 +1,12 @@ +from dataclasses import dataclass + +from .view_sampler import ViewSamplerCfg + + +@dataclass +class DatasetCfgCommon: + image_shape: list[int] + background_color: list[float] + cameras_are_circular: bool + overfit_to_scene: str | None + view_sampler: ViewSamplerCfg diff --git a/src/dataset/dataset_dl3dv.py b/src/dataset/dataset_dl3dv.py new file mode 100644 index 0000000000000000000000000000000000000000..7a9b5922181d582f048012cfa436f056442e49f4 --- /dev/null +++ b/src/dataset/dataset_dl3dv.py @@ -0,0 +1,395 @@ +import json +from dataclasses import dataclass +from functools import cached_property +from io import BytesIO +from pathlib import Path +from typing import Literal, Optional + +import torch +import torchvision.transforms as tf +from einops import rearrange, repeat +from jaxtyping import Float, UInt8 +from PIL import Image +from torch import Tensor +from torch.utils.data import IterableDataset +import numpy as np + +from ..geometry.projection import get_fov +from .dataset import DatasetCfgCommon +from .shims.augmentation_shim import apply_augmentation_shim +from .shims.crop_shim import apply_crop_shim +from .types import Stage +from .view_sampler import ViewSampler + + +@dataclass +class DatasetDL3DVCfg(DatasetCfgCommon): + name: Literal["dl3dv"] + roots: list[Path] + max_fov: float + augment: bool + test_len: int + test_chunk_interval: int + train_times_per_scene: int + test_times_per_scene: int + ori_image_shape: list[int] + skip_bad_shape: bool = True + near: float = -1.0 + far: float = -1.0 + shuffle_val: bool = True + min_views: int = 0 + max_views: int = 0 + highres: bool = False + sort_target_index: Optional[bool] = False + overfit_max_views: Optional[int] = None + sort_context_index: Optional[bool] = False + use_index_to_load_chunk: Optional[bool] = False + + +class DatasetDL3DV(IterableDataset): + cfg: DatasetDL3DVCfg + stage: Stage + view_sampler: ViewSampler + + to_tensor: tf.ToTensor + chunks: list[Path] + near: float = 0.1 + far: float = 1000.0 + + def __init__( + self, + cfg: DatasetDL3DVCfg, + stage: Stage, + view_sampler: ViewSampler, + ) -> None: + super().__init__() + + self.cfg = cfg + self.stage = stage + self.view_sampler = view_sampler + self.to_tensor = tf.ToTensor() + if cfg.near != -1: + self.near = cfg.near + if cfg.far != -1: + self.far = cfg.far + + # Collect chunks. + self.chunks = [] + for i, root in enumerate(cfg.roots): + root = root / self.data_stage + if self.cfg.use_index_to_load_chunk: + with open(root / "index.json", "r") as f: + json_dict = json.load(f) + root_chunks = sorted(list(set(json_dict.values()))) + else: + root_chunks = sorted( + [path for path in root.iterdir() if path.suffix == ".torch"] + ) + + self.chunks.extend(root_chunks) + if self.cfg.overfit_to_scene is not None: + chunk_path = self.index[self.cfg.overfit_to_scene] + self.chunks = [chunk_path] * len(self.chunks) + if self.stage == "test": + # fast testing + self.chunks = self.chunks[:: cfg.test_chunk_interval] + if self.stage == "val": + self.chunks = self.chunks * int(1e6 // len(self.chunks)) + + def shuffle(self, lst: list) -> list: + indices = torch.randperm(len(lst)) + return [lst[x] for x in indices] + + def __iter__(self): + # Chunks must be shuffled here (not inside __init__) for validation to show + # random chunks. + if self.stage in (("train", "val") if self.cfg.shuffle_val else ("train")): + self.chunks = self.shuffle(self.chunks) + + # When testing, the data loaders alternate chunks. + worker_info = torch.utils.data.get_worker_info() + if self.stage == "test" and worker_info is not None: + self.chunks = [ + chunk + for chunk_index, chunk in enumerate(self.chunks) + if chunk_index % worker_info.num_workers == worker_info.id + ] + + for chunk_path in self.chunks: + # Load the chunk. + chunk = torch.load(chunk_path) + + if self.cfg.overfit_to_scene is not None: + item = [x for x in chunk if x["key"] + == self.cfg.overfit_to_scene] + assert len(item) == 1 + if self.stage == "test": + chunk = item + else: + chunk = item * len(chunk) + + if self.stage in (("train", "val") if self.cfg.shuffle_val else ("train")): + chunk = self.shuffle(chunk) + + times_per_scene = ( + self.cfg.test_times_per_scene + if self.stage == "test" + else self.cfg.train_times_per_scene + ) + + for run_idx in range(int(times_per_scene * len(chunk))): + example = chunk[run_idx // times_per_scene] + + extrinsics, intrinsics = self.convert_poses(example["cameras"]) + + scene = example["key"] + + try: + extra_kwargs = {} + if self.cfg.overfit_to_scene is not None and self.stage != "test": + extra_kwargs.update( + { + "max_num_views": ( + 148 + if self.cfg.overfit_max_views is None + else self.cfg.overfit_max_views + ) + } + ) + + out_data = self.view_sampler.sample( + scene, + extrinsics, + intrinsics, + min_context_views=self.cfg.min_views, + max_context_views=self.cfg.max_views, + **extra_kwargs, + ) + if isinstance(out_data, tuple): + context_indices, target_indices = out_data[:2] + c_list = [ + ( + context_indices.sort()[0] + if self.cfg.sort_context_index + else context_indices + ) + ] + t_list = [ + ( + target_indices.sort()[0] + if self.cfg.sort_target_index + else target_indices + ) + ] + if isinstance(out_data, list): + c_list = [ + ( + a.context.sort()[0] + if self.cfg.sort_context_index + else a.context + ) + for a in out_data + ] + t_list = [ + ( + a.target.sort()[0] + if self.cfg.sort_target_index + else a.target + ) + for a in out_data + ] + + except ValueError: + # Skip because the example doesn't have enough frames. + continue + + # Skip the example if the field of view is too wide. + if (get_fov(intrinsics).rad2deg() > self.cfg.max_fov).any(): + continue + + for context_indices, target_indices in zip(c_list, t_list): + # Load the images. + context_images = [ + example["images"][index.item()] for index in context_indices + ] + + try: + context_images = self.convert_images(context_images) + except OSError: + # some data might be corrupted + continue + + target_images = [ + example["images"][index.item()] for index in target_indices + ] + + try: + target_images = self.convert_images(target_images) + except OSError: + # some data might be corrupted + continue + + # Skip the example if the images don't have the right shape. + expected_shape = tuple( + [3, *self.cfg.ori_image_shape] + ) # (3, 270, 480) or (3, 540, 960) + + context_image_invalid = context_images.shape[1:] != expected_shape + target_image_invalid = target_images.shape[1:] != expected_shape + + if self.cfg.skip_bad_shape and ( + context_image_invalid or target_image_invalid + ): + print( + f"Skipped bad example {example['key']}. Context shape was " + f"{context_images.shape}, target shape was " + f"{target_images.shape}, and expected shape was {expected_shape}" + ) + continue + + # check the extrinsics + if any(torch.isnan(torch.det(extrinsics[context_indices][:, :3, :3]))): + # print('invalid extrinsics') + continue + + if any(torch.isnan(torch.det(extrinsics[target_indices][:, :3, :3]))): + # print('invalid extrinsics') + continue + + # check the extrinsics: translation could be very large + # https://github.com/DL3DV-10K/Dataset/issues/34 + if (extrinsics[context_indices][:, :3, 3] > 1e3).any(): + # print('extremely large camera translation') + continue + + if (extrinsics[target_indices][:, :3, 3] > 1e3).any(): + # print('extremely large camera translation') + continue + + if not torch.allclose(torch.det(extrinsics[context_indices][:, :3, :3]), torch.det(extrinsics[context_indices][:, :3, :3]).new_tensor(1)): + # print('invalid extrinsics') + continue + if not torch.allclose(torch.det(extrinsics[target_indices][:, :3, :3]), torch.det(extrinsics[target_indices][:, :3, :3]).new_tensor(1)): + # print('invalid extrinsics') + continue + + nf_scale = 1.0 + example_out = { + "context": { + "extrinsics": extrinsics[context_indices], + "intrinsics": intrinsics[context_indices], + "image": context_images, + "near": self.get_bound("near", len(context_indices)) + / nf_scale, + "far": self.get_bound("far", len(context_indices)) + / nf_scale, + "index": context_indices, + }, + "target": { + "extrinsics": extrinsics[target_indices], + "intrinsics": intrinsics[target_indices], + "image": target_images, + "near": self.get_bound("near", len(target_indices)) + / nf_scale, + "far": self.get_bound("far", len(target_indices)) + / nf_scale, + "index": target_indices, + }, + "scene": scene, + } + + if self.stage == "train" and self.cfg.augment: + example_out = apply_augmentation_shim(example_out) + if self.cfg.image_shape == list(context_images.shape[2:]): + yield example_out + else: + yield apply_crop_shim(example_out, tuple(self.cfg.image_shape)) + + def convert_poses( + self, + poses: Float[Tensor, "batch 18"], + ) -> tuple[ + Float[Tensor, "batch 4 4"], # extrinsics + Float[Tensor, "batch 3 3"], # intrinsics + ]: + b, _ = poses.shape + + # Convert the intrinsics to a 3x3 normalized K matrix. + intrinsics = torch.eye(3, dtype=torch.float32) + intrinsics = repeat(intrinsics, "h w -> b h w", b=b).clone() + fx, fy, cx, cy = poses[:, :4].T + intrinsics[:, 0, 0] = fx + intrinsics[:, 1, 1] = fy + intrinsics[:, 0, 2] = cx + intrinsics[:, 1, 2] = cy + + # Convert the extrinsics to a 4x4 OpenCV-style C2W matrix. + w2c = repeat(torch.eye(4, dtype=torch.float32), + "h w -> b h w", b=b).clone() + w2c[:, :3] = rearrange(poses[:, 6:], "b (h w) -> b h w", h=3, w=4) + return w2c.inverse(), intrinsics + + def convert_images( + self, + images: list[UInt8[Tensor, "..."]], + ) -> Float[Tensor, "batch 3 height width"]: + torch_images = [] + for image in images: + image = Image.open(BytesIO(image.numpy().tobytes())) + torch_images.append(self.to_tensor(image)) + return torch.stack(torch_images) + + def get_bound( + self, + bound: Literal["near", "far"], + num_views: int, + ) -> Float[Tensor, " view"]: + value = torch.tensor(getattr(self, bound), dtype=torch.float32) + return repeat(value, "-> v", v=num_views) + + @property + def data_stage(self) -> Stage: + if self.cfg.overfit_to_scene is not None: + return "test" + if self.stage == "val": + return "test" + return self.stage + + @cached_property + def index(self) -> dict[str, Path]: + merged_index = {} + data_stages = [self.data_stage] + if self.cfg.overfit_to_scene is not None: + data_stages = ("test", "train") + for data_stage in data_stages: + for i, root in enumerate(self.cfg.roots): + if not (root / data_stage).is_dir(): + continue + + # Load the root's index. + with (root / data_stage / "index.json").open("r") as f: + index = json.load(f) + index = {k: Path(root / data_stage / v) + for k, v in index.items()} + + # The constituent datasets should have unique keys. + assert not (set(merged_index.keys()) & set(index.keys())) + + # Merge the root's index into the main index. + merged_index = {**merged_index, **index} + return merged_index + + def __len__(self) -> int: + if self.stage in ['train', 'test']: + return ( + min( + len(self.index.keys()) * self.cfg.test_times_per_scene, + self.cfg.test_len, + ) + if self.stage == "test" and self.cfg.test_len > 0 + else len(self.index.keys()) * self.cfg.train_times_per_scene + ) + else: + # set a very large value here to ensure the validation keep going + # and do not exhaust; it will be wrap to length 1 anyway. + return int(1e10) diff --git a/src/dataset/dataset_re10k.py b/src/dataset/dataset_re10k.py new file mode 100644 index 0000000000000000000000000000000000000000..ad4ec57cf8edfeafbeb511fc71d63f50a19099f4 --- /dev/null +++ b/src/dataset/dataset_re10k.py @@ -0,0 +1,269 @@ +import json +from dataclasses import dataclass +from functools import cached_property +from io import BytesIO +from pathlib import Path +from typing import Literal, Optional + +import numpy as np +import torch +import torchvision.transforms as tf +from einops import rearrange, repeat +from jaxtyping import Float, UInt8 +from PIL import Image +from torch import Tensor +from torch.utils.data import IterableDataset + +from ..geometry.projection import get_fov +from .dataset import DatasetCfgCommon +from .shims.augmentation_shim import apply_augmentation_shim +from .shims.crop_shim import apply_crop_shim +from .types import Stage +from .view_sampler import ViewSampler + + +@dataclass +class DatasetRE10kCfg(DatasetCfgCommon): + name: Literal["re10k"] + roots: list[Path] + max_fov: float + augment: bool + test_len: int + test_chunk_interval: int + skip_bad_shape: bool = True + near: float = -1.0 + far: float = -1.0 + shuffle_val: bool = True + train_times_per_scene: int = 1 + highres: bool = False + use_index_to_load_chunk: Optional[bool] = False + + +class DatasetRE10k(IterableDataset): + cfg: DatasetRE10kCfg + stage: Stage + view_sampler: ViewSampler + + to_tensor: tf.ToTensor + chunks: list[Path] + near: float = 0.1 + far: float = 1000.0 + + def __init__( + self, + cfg: DatasetRE10kCfg, + stage: Stage, + view_sampler: ViewSampler, + ) -> None: + super().__init__() + self.cfg = cfg + self.stage = stage + self.view_sampler = view_sampler + self.to_tensor = tf.ToTensor() + if cfg.near != -1: + self.near = cfg.near + if cfg.far != -1: + self.far = cfg.far + + # Collect chunks. + self.chunks = [] + for i, root in enumerate(cfg.roots): + root = root / self.data_stage + if self.cfg.use_index_to_load_chunk: + with open(root / "index.json", "r") as f: + json_dict = json.load(f) + root_chunks = sorted(list(set(json_dict.values()))) + else: + root_chunks = sorted( + [path for path in root.iterdir() if path.suffix == ".torch"] + ) + + self.chunks.extend(root_chunks) + if self.cfg.overfit_to_scene is not None: + chunk_path = self.index[self.cfg.overfit_to_scene] + self.chunks = [chunk_path] * len(self.chunks) + if self.stage == "test": + # testing on a subset for fast speed + self.chunks = self.chunks[::cfg.test_chunk_interval] + + def shuffle(self, lst: list) -> list: + indices = torch.randperm(len(lst)) + return [lst[x] for x in indices] + + def __iter__(self): + # Chunks must be shuffled here (not inside __init__) for validation to show + # random chunks. + if self.stage in (("train", "val") if self.cfg.shuffle_val else ("train")): + self.chunks = self.shuffle(self.chunks) + + # When testing, the data loaders alternate chunks. + worker_info = torch.utils.data.get_worker_info() + if self.stage == "test" and worker_info is not None: + self.chunks = [ + chunk + for chunk_index, chunk in enumerate(self.chunks) + if chunk_index % worker_info.num_workers == worker_info.id + ] + + for chunk_path in self.chunks: + # Load the chunk. + chunk = torch.load(chunk_path) + + if self.cfg.overfit_to_scene is not None: + item = [x for x in chunk if x["key"] == self.cfg.overfit_to_scene] + assert len(item) == 1 + chunk = item * len(chunk) + + if self.stage in (("train", "val") if self.cfg.shuffle_val else ("train")): + chunk = self.shuffle(chunk) + + times_per_scene = ( + 1 + if self.stage == "test" + else self.cfg.train_times_per_scene + ) + + for run_idx in range(int(times_per_scene * len(chunk))): + example = chunk[run_idx // times_per_scene] + extrinsics, intrinsics = self.convert_poses(example["cameras"]) + scene = example["key"] + + try: + context_indices, target_indices = self.view_sampler.sample( + scene, + extrinsics, + intrinsics, + ) + except ValueError: + # Skip because the example doesn't have enough frames. + continue + + # Skip the example if the field of view is too wide. + if (get_fov(intrinsics).rad2deg() > self.cfg.max_fov).any(): + continue + + # Load the images. + context_images = [ + example["images"][index.item()] for index in context_indices + ] + context_images = self.convert_images(context_images) + target_images = [ + example["images"][index.item()] for index in target_indices + ] + target_images = self.convert_images(target_images) + + # Skip the example if the images don't have the right shape. + if self.cfg.highres: + expected_shape = (3, 720, 1280) + else: + expected_shape = (3, 360, 640) + context_image_invalid = context_images.shape[1:] != expected_shape + target_image_invalid = target_images.shape[1:] != expected_shape + if self.cfg.skip_bad_shape and (context_image_invalid or target_image_invalid): + print( + f"Skipped bad example {example['key']}. Context shape was " + f"{context_images.shape} and target shape was " + f"{target_images.shape}." + ) + continue + + nf_scale = 1.0 + example = { + "context": { + "extrinsics": extrinsics[context_indices], + "intrinsics": intrinsics[context_indices], + "image": context_images, + "near": self.get_bound("near", len(context_indices)) / nf_scale, + "far": self.get_bound("far", len(context_indices)) / nf_scale, + "index": context_indices, + }, + "target": { + "extrinsics": extrinsics[target_indices], + "intrinsics": intrinsics[target_indices], + "image": target_images, + "near": self.get_bound("near", len(target_indices)) / nf_scale, + "far": self.get_bound("far", len(target_indices)) / nf_scale, + "index": target_indices, + }, + "scene": scene, + } + + if self.stage == "train" and self.cfg.augment: + example = apply_augmentation_shim(example) + yield apply_crop_shim(example, tuple(self.cfg.image_shape)) + + def convert_poses( + self, + poses: Float[Tensor, "batch 18"], + ) -> tuple[ + Float[Tensor, "batch 4 4"], # extrinsics + Float[Tensor, "batch 3 3"], # intrinsics + ]: + b, _ = poses.shape + + # Convert the intrinsics to a 3x3 normalized K matrix. + intrinsics = torch.eye(3, dtype=torch.float32) + intrinsics = repeat(intrinsics, "h w -> b h w", b=b).clone() + fx, fy, cx, cy = poses[:, :4].T + intrinsics[:, 0, 0] = fx + intrinsics[:, 1, 1] = fy + intrinsics[:, 0, 2] = cx + intrinsics[:, 1, 2] = cy + + # Convert the extrinsics to a 4x4 OpenCV-style C2W matrix. + w2c = repeat(torch.eye(4, dtype=torch.float32), "h w -> b h w", b=b).clone() + w2c[:, :3] = rearrange(poses[:, 6:], "b (h w) -> b h w", h=3, w=4) + return w2c.inverse(), intrinsics + + def convert_images( + self, + images: list[UInt8[Tensor, "..."]], + ) -> Float[Tensor, "batch 3 height width"]: + torch_images = [] + for image in images: + image = Image.open(BytesIO(image.numpy().tobytes())) + torch_images.append(self.to_tensor(image)) + return torch.stack(torch_images) + + def get_bound( + self, + bound: Literal["near", "far"], + num_views: int, + ) -> Float[Tensor, " view"]: + value = torch.tensor(getattr(self, bound), dtype=torch.float32) + return repeat(value, "-> v", v=num_views) + + @property + def data_stage(self) -> Stage: + if self.cfg.overfit_to_scene is not None: + return "test" + if self.stage == "val": + return "test" + return self.stage + + @cached_property + def index(self) -> dict[str, Path]: + merged_index = {} + data_stages = [self.data_stage] + if self.cfg.overfit_to_scene is not None: + data_stages = ("test", "train") + for data_stage in data_stages: + for i, root in enumerate(self.cfg.roots): + # Load the root's index. + with (root / data_stage / "index.json").open("r") as f: + index = json.load(f) + index = {k: Path(root / data_stage / v) for k, v in index.items()} + + # The constituent datasets should have unique keys. + assert not (set(merged_index.keys()) & set(index.keys())) + + # Merge the root's index into the main index. + merged_index = {**merged_index, **index} + return merged_index + + def __len__(self) -> int: + return ( + min(len(self.index.keys()), self.cfg.test_len) + if self.stage == "test" and self.cfg.test_len > 0 + else len(self.index.keys()) * self.cfg.train_times_per_scene + ) diff --git a/src/dataset/shims/augmentation_shim.py b/src/dataset/shims/augmentation_shim.py new file mode 100644 index 0000000000000000000000000000000000000000..9e56a503ecbfd3a2da115cbc3a5674c33499b744 --- /dev/null +++ b/src/dataset/shims/augmentation_shim.py @@ -0,0 +1,37 @@ +import torch +from jaxtyping import Float +from torch import Tensor + +from ..types import AnyExample, AnyViews + + +def reflect_extrinsics( + extrinsics: Float[Tensor, "*batch 4 4"], +) -> Float[Tensor, "*batch 4 4"]: + reflect = torch.eye(4, dtype=torch.float32, device=extrinsics.device) + reflect[0, 0] = -1 + return reflect @ extrinsics @ reflect + + +def reflect_views(views: AnyViews) -> AnyViews: + return { + **views, + "image": views["image"].flip(-1), + "extrinsics": reflect_extrinsics(views["extrinsics"]), + } + + +def apply_augmentation_shim( + example: AnyExample, + generator: torch.Generator | None = None, +) -> AnyExample: + """Randomly augment the training images.""" + # Do not augment with 50% chance. + if torch.rand(tuple(), generator=generator) < 0.5: + return example + + return { + **example, + "context": reflect_views(example["context"]), + "target": reflect_views(example["target"]), + } diff --git a/src/dataset/shims/crop_shim.py b/src/dataset/shims/crop_shim.py new file mode 100644 index 0000000000000000000000000000000000000000..b84654ec04b4400512bf3cb80a36d8de086dab58 --- /dev/null +++ b/src/dataset/shims/crop_shim.py @@ -0,0 +1,125 @@ +import numpy as np +import torch +from einops import rearrange +from jaxtyping import Float +from PIL import Image +from torch import Tensor +import torch.nn.functional as F + +from ..types import AnyExample, AnyViews + + +def rescale( + image: Float[Tensor, "3 h_in w_in"], + shape: tuple[int, int], +) -> Float[Tensor, "3 h_out w_out"]: + h, w = shape + image_new = (image * 255).clip(min=0, max=255).type(torch.uint8) + image_new = rearrange(image_new, "c h w -> h w c").detach().cpu().numpy() + image_new = Image.fromarray(image_new) + image_new = image_new.resize((w, h), Image.LANCZOS) + image_new = np.array(image_new) / 255 + image_new = torch.tensor(image_new, dtype=image.dtype, device=image.device) + return rearrange(image_new, "h w c -> c h w") + + +def center_crop( + images: Float[Tensor, "*#batch c h w"], + intrinsics: Float[Tensor, "*#batch 3 3"], + shape: tuple[int, int], + depths: None | Float[Tensor, "*#batch h w"], +) -> ( + tuple[ + Float[Tensor, "*#batch c h_out w_out"], # updated images + Float[Tensor, "*#batch 3 3"], # updated intrinsics + ] + | tuple[ + Float[Tensor, "*#batch c h_out w_out"], # updated images + Float[Tensor, "*#batch 3 3"], # updated intrinsics + Float[Tensor, "*#batch h_out w_out"], # updated depths + ] +): + *_, h_in, w_in = images.shape + h_out, w_out = shape + + # Note that odd input dimensions induce half-pixel misalignments. + row = (h_in - h_out) // 2 + col = (w_in - w_out) // 2 + + # Center-crop the image. + images = images[..., :, row : row + h_out, col : col + w_out] + + # Adjust the intrinsics to account for the cropping. + intrinsics = intrinsics.clone() + intrinsics[..., 0, 0] *= w_in / w_out # fx + intrinsics[..., 1, 1] *= h_in / h_out # fy + + if depths is not None: + depths = depths[..., :, row : row + h_out, col : col + w_out] + return images, intrinsics, depths + + return images, intrinsics + + +def rescale_and_crop( + images: Float[Tensor, "*#batch c h w"], + intrinsics: Float[Tensor, "*#batch 3 3"], + shape: tuple[int, int], + depths: None | Float[Tensor, "*#batch h w"], +) -> ( + tuple[ + Float[Tensor, "*#batch c h_out w_out"], # updated images + Float[Tensor, "*#batch 3 3"], # updated intrinsics + ] + | tuple[ + Float[Tensor, "*#batch c h_out w_out"], # updated images + Float[Tensor, "*#batch 3 3"], # updated intrinsics + Float[Tensor, "*#batch h_out w_out"], # updated depths + ] +): + *_, h_in, w_in = images.shape + h_out, w_out = shape + assert h_out <= h_in and w_out <= w_in + + scale_factor = max(h_out / h_in, w_out / w_in) + h_scaled = round(h_in * scale_factor) + w_scaled = round(w_in * scale_factor) + assert h_scaled == h_out or w_scaled == w_out + + # Reshape the images to the correct size. Assume we don't have to worry about + # changing the intrinsics based on how the images are rounded. + *batch, c, h, w = images.shape + images = images.reshape(-1, c, h, w) + images = torch.stack([rescale(image, (h_scaled, w_scaled)) for image in images]) + images = images.reshape(*batch, c, h_scaled, w_scaled) + + # reshape and crop depth as well when available + if depths is not None: + depths = F.interpolate( + depths.unsqueeze(1), + size=(h_scaled, w_scaled), + mode="bilinear", + align_corners=True, + ).squeeze(1) + + return center_crop(images, intrinsics, shape, depths=depths) + + +def apply_crop_shim_to_views(views: AnyViews, shape: tuple[int, int]) -> AnyViews: + images, intrinsics = rescale_and_crop( + views["image"], views["intrinsics"], shape, depths=None + ) + return { + **views, + "image": images, + "intrinsics": intrinsics, + } + + +def apply_crop_shim(example: AnyExample, shape: tuple[int, int]) -> AnyExample: + """Crop images in the example.""" + return { + **example, + "context": apply_crop_shim_to_views(example["context"], shape), + "target": apply_crop_shim_to_views(example["target"], shape), + } diff --git a/src/dataset/shims/patch_shim.py b/src/dataset/shims/patch_shim.py new file mode 100644 index 0000000000000000000000000000000000000000..60dc944995332726b060e7eac9bf914b63c3c840 --- /dev/null +++ b/src/dataset/shims/patch_shim.py @@ -0,0 +1,38 @@ +from ..types import BatchedExample, BatchedViews + + +def apply_patch_shim_to_views(views: BatchedViews, patch_size: int) -> BatchedViews: + _, _, _, h, w = views["image"].shape + + # Image size must be even so that naive center-cropping does not cause misalignment. + assert h % 2 == 0 and w % 2 == 0 + + h_new = (h // patch_size) * patch_size + row = (h - h_new) // 2 + w_new = (w // patch_size) * patch_size + col = (w - w_new) // 2 + + # Center-crop the image. + image = views["image"][:, :, :, row : row + h_new, col : col + w_new] + + # Adjust the intrinsics to account for the cropping. + intrinsics = views["intrinsics"].clone() + intrinsics[:, :, 0, 0] *= w / w_new # fx + intrinsics[:, :, 1, 1] *= h / h_new # fy + + return { + **views, + "image": image, + "intrinsics": intrinsics, + } + + +def apply_patch_shim(batch: BatchedExample, patch_size: int) -> BatchedExample: + """Crop images in the batch so that their dimensions are cleanly divisible by the + specified patch size. + """ + return { + **batch, + "context": apply_patch_shim_to_views(batch["context"], patch_size), + "target": apply_patch_shim_to_views(batch["target"], patch_size), + } diff --git a/src/dataset/types.py b/src/dataset/types.py new file mode 100644 index 0000000000000000000000000000000000000000..03ba3d3894aa91fdfebd7422e6f42fbd3a2e6ddc --- /dev/null +++ b/src/dataset/types.py @@ -0,0 +1,50 @@ +from typing import Callable, Literal, TypedDict + +from jaxtyping import Float, Int64 +from torch import Tensor + +Stage = Literal["train", "val", "test"] + + +# The following types mainly exist to make type-hinted keys show up in VS Code. Some +# dimensions are annotated as "_" because either: +# 1. They're expected to change as part of a function call (e.g., resizing the dataset). +# 2. They're expected to vary within the same function call (e.g., the number of views, +# which differs between context and target BatchedViews). + + +class BatchedViews(TypedDict, total=False): + extrinsics: Float[Tensor, "batch _ 4 4"] # batch view 4 4 + intrinsics: Float[Tensor, "batch _ 3 3"] # batch view 3 3 + image: Float[Tensor, "batch _ _ _ _"] # batch view channel height width + near: Float[Tensor, "batch _"] # batch view + far: Float[Tensor, "batch _"] # batch view + index: Int64[Tensor, "batch _"] # batch view + + +class BatchedExample(TypedDict, total=False): + target: BatchedViews + context: BatchedViews + scene: list[str] + + +class UnbatchedViews(TypedDict, total=False): + extrinsics: Float[Tensor, "_ 4 4"] + intrinsics: Float[Tensor, "_ 3 3"] + image: Float[Tensor, "_ 3 height width"] + near: Float[Tensor, " _"] + far: Float[Tensor, " _"] + index: Int64[Tensor, " _"] + + +class UnbatchedExample(TypedDict, total=False): + target: UnbatchedViews + context: UnbatchedViews + scene: str + + +# A data shim modifies the example after it's been returned from the data loader. +DataShim = Callable[[BatchedExample], BatchedExample] + +AnyExample = BatchedExample | UnbatchedExample +AnyViews = BatchedViews | UnbatchedViews diff --git a/src/dataset/validation_wrapper.py b/src/dataset/validation_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..9771125f0f38059eb43469174c84122542d8c863 --- /dev/null +++ b/src/dataset/validation_wrapper.py @@ -0,0 +1,32 @@ +from typing import Iterator, Optional + +import torch +from torch.utils.data import Dataset, IterableDataset + + +class ValidationWrapper(Dataset): + """Wraps a dataset so that PyTorch Lightning's validation step can be turned into a + visualization step. + """ + + dataset: Dataset + dataset_iterator: Optional[Iterator] + length: int + + def __init__(self, dataset: Dataset, length: int) -> None: + super().__init__() + self.dataset = dataset + self.length = length + self.dataset_iterator = None + + def __len__(self): + return self.length + + def __getitem__(self, index: int): + if isinstance(self.dataset, IterableDataset): + if self.dataset_iterator is None: + self.dataset_iterator = iter(self.dataset) + return next(self.dataset_iterator) + + random_index = torch.randint(0, len(self.dataset), tuple()) + return self.dataset[random_index.item()] diff --git a/src/dataset/view_sampler/__init__.py b/src/dataset/view_sampler/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9c03e5cc29eba8037d697a9722121542c186fea --- /dev/null +++ b/src/dataset/view_sampler/__init__.py @@ -0,0 +1,43 @@ +from typing import Any + +from ...misc.step_tracker import StepTracker +from ..types import Stage +from .view_sampler import ViewSampler +from .view_sampler_all import ViewSamplerAll, ViewSamplerAllCfg +from .view_sampler_arbitrary import ViewSamplerArbitrary, ViewSamplerArbitraryCfg +from .view_sampler_bounded import ViewSamplerBounded, ViewSamplerBoundedCfg +from .view_sampler_evaluation import ViewSamplerEvaluation, ViewSamplerEvaluationCfg +from .view_sampler_bounded_v2 import ViewSamplerBoundedV2, ViewSamplerBoundedV2Cfg + + +VIEW_SAMPLERS: dict[str, ViewSampler[Any]] = { + "all": ViewSamplerAll, + "arbitrary": ViewSamplerArbitrary, + "bounded": ViewSamplerBounded, + "evaluation": ViewSamplerEvaluation, + "boundedv2": ViewSamplerBoundedV2, +} + +ViewSamplerCfg = ( + ViewSamplerArbitraryCfg + | ViewSamplerBoundedCfg + | ViewSamplerEvaluationCfg + | ViewSamplerAllCfg + | ViewSamplerBoundedV2Cfg +) + + +def get_view_sampler( + cfg: ViewSamplerCfg, + stage: Stage, + overfit: bool, + cameras_are_circular: bool, + step_tracker: StepTracker | None, +) -> ViewSampler[Any]: + return VIEW_SAMPLERS[cfg.name]( + cfg, + stage, + overfit, + cameras_are_circular, + step_tracker, + ) diff --git a/src/dataset/view_sampler/view_sampler.py b/src/dataset/view_sampler/view_sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..1caeb29693dd8ecb1ee9bc90dbcd811db315d29b --- /dev/null +++ b/src/dataset/view_sampler/view_sampler.py @@ -0,0 +1,61 @@ +from abc import ABC, abstractmethod +from typing import Generic, TypeVar + +import torch +from jaxtyping import Float, Int64 +from torch import Tensor + +from ...misc.step_tracker import StepTracker +from ..types import Stage + +T = TypeVar("T") + + +class ViewSampler(ABC, Generic[T]): + cfg: T + stage: Stage + is_overfitting: bool + cameras_are_circular: bool + step_tracker: StepTracker | None + + def __init__( + self, + cfg: T, + stage: Stage, + is_overfitting: bool, + cameras_are_circular: bool, + step_tracker: StepTracker | None, + ) -> None: + self.cfg = cfg + self.stage = stage + self.is_overfitting = is_overfitting + self.cameras_are_circular = cameras_are_circular + self.step_tracker = step_tracker + + @abstractmethod + def sample( + self, + scene: str, + extrinsics: Float[Tensor, "view 4 4"], + intrinsics: Float[Tensor, "view 3 3"], + device: torch.device = torch.device("cpu"), + **kwargs, + ) -> tuple[ + Int64[Tensor, " context_view"], # indices for context views + Int64[Tensor, " target_view"], # indices for target views + ]: + pass + + @property + @abstractmethod + def num_target_views(self) -> int: + pass + + @property + @abstractmethod + def num_context_views(self) -> int: + pass + + @property + def global_step(self) -> int: + return 0 if self.step_tracker is None else self.step_tracker.get_step() diff --git a/src/dataset/view_sampler/view_sampler_all.py b/src/dataset/view_sampler/view_sampler_all.py new file mode 100644 index 0000000000000000000000000000000000000000..c5a9e7d637d1ebf8d0b3a6f5c0238cdc13100e58 --- /dev/null +++ b/src/dataset/view_sampler/view_sampler_all.py @@ -0,0 +1,38 @@ +from dataclasses import dataclass +from typing import Literal + +import torch +from jaxtyping import Float, Int64 +from torch import Tensor + +from .view_sampler import ViewSampler + + +@dataclass +class ViewSamplerAllCfg: + name: Literal["all"] + + +class ViewSamplerAll(ViewSampler[ViewSamplerAllCfg]): + def sample( + self, + scene: str, + extrinsics: Float[Tensor, "view 4 4"], + intrinsics: Float[Tensor, "view 3 3"], + device: torch.device = torch.device("cpu"), + **kwargs, + ) -> tuple[ + Int64[Tensor, " context_view"], # indices for context views + Int64[Tensor, " target_view"], # indices for target views + ]: + v, _, _ = extrinsics.shape + all_frames = torch.arange(v, device=device) + return all_frames, all_frames + + @property + def num_context_views(self) -> int: + return 0 + + @property + def num_target_views(self) -> int: + return 0 diff --git a/src/dataset/view_sampler/view_sampler_arbitrary.py b/src/dataset/view_sampler/view_sampler_arbitrary.py new file mode 100644 index 0000000000000000000000000000000000000000..00be8dfb2bb9927794fe71af748e65449fb002a0 --- /dev/null +++ b/src/dataset/view_sampler/view_sampler_arbitrary.py @@ -0,0 +1,71 @@ +from dataclasses import dataclass +from typing import Literal + +import torch +from jaxtyping import Float, Int64 +from torch import Tensor + +from .view_sampler import ViewSampler + + +@dataclass +class ViewSamplerArbitraryCfg: + name: Literal["arbitrary"] + num_context_views: int + num_target_views: int + context_views: list[int] | None + target_views: list[int] | None + + +class ViewSamplerArbitrary(ViewSampler[ViewSamplerArbitraryCfg]): + def sample( + self, + scene: str, + extrinsics: Float[Tensor, "view 4 4"], + intrinsics: Float[Tensor, "view 3 3"], + device: torch.device = torch.device("cpu"), + **kwargs, + ) -> tuple[ + Int64[Tensor, " context_view"], # indices for context views + Int64[Tensor, " target_view"], # indices for target views + ]: + """Arbitrarily sample context and target views.""" + num_views, _, _ = extrinsics.shape + + index_context = torch.randint( + 0, + num_views, + size=(self.cfg.num_context_views,), + device=device, + ) + + # Allow the context views to be fixed. + if self.cfg.context_views is not None: + assert len(self.cfg.context_views) == self.cfg.num_context_views + index_context = torch.tensor( + self.cfg.context_views, dtype=torch.int64, device=device + ) + + index_target = torch.randint( + 0, + num_views, + size=(self.cfg.num_target_views,), + device=device, + ) + + # Allow the target views to be fixed. + if self.cfg.target_views is not None: + assert len(self.cfg.target_views) == self.cfg.num_target_views + index_target = torch.tensor( + self.cfg.target_views, dtype=torch.int64, device=device + ) + + return index_context, index_target + + @property + def num_context_views(self) -> int: + return self.cfg.num_context_views + + @property + def num_target_views(self) -> int: + return self.cfg.num_target_views diff --git a/src/dataset/view_sampler/view_sampler_bounded.py b/src/dataset/view_sampler/view_sampler_bounded.py new file mode 100644 index 0000000000000000000000000000000000000000..d9cc505e019dc8a218526f2b5af0bda57cf9cedd --- /dev/null +++ b/src/dataset/view_sampler/view_sampler_bounded.py @@ -0,0 +1,132 @@ +from dataclasses import dataclass +from typing import Literal + +import torch +from jaxtyping import Float, Int64 +from torch import Tensor + +from .view_sampler import ViewSampler + + +@dataclass +class ViewSamplerBoundedCfg: + name: Literal["bounded"] + num_context_views: int + num_target_views: int + min_distance_between_context_views: int + max_distance_between_context_views: int + min_distance_to_context_views: int + warm_up_steps: int + initial_min_distance_between_context_views: int + initial_max_distance_between_context_views: int + + +class ViewSamplerBounded(ViewSampler[ViewSamplerBoundedCfg]): + def schedule(self, initial: int, final: int) -> int: + fraction = self.global_step / self.cfg.warm_up_steps + return min(initial + int((final - initial) * fraction), final) + + def sample( + self, + scene: str, + extrinsics: Float[Tensor, "view 4 4"], + intrinsics: Float[Tensor, "view 3 3"], + device: torch.device = torch.device("cpu"), + min_view_dist: int | None = None, + max_view_dist: int | None = None, + **kwargs, + ) -> tuple[ + Int64[Tensor, " context_view"], # indices for context views + Int64[Tensor, " target_view"], # indices for target views + ]: + num_views, _, _ = extrinsics.shape + + # Compute the context view spacing based on the current global step. + if self.stage == "test": + # When testing, always use the full gap. + max_gap = self.cfg.max_distance_between_context_views + min_gap = self.cfg.max_distance_between_context_views + elif self.cfg.warm_up_steps > 0: + max_gap = self.schedule( + self.cfg.initial_max_distance_between_context_views, + self.cfg.max_distance_between_context_views, + ) + min_gap = self.schedule( + self.cfg.initial_min_distance_between_context_views, + self.cfg.min_distance_between_context_views, + ) + else: + max_gap = self.cfg.max_distance_between_context_views + min_gap = self.cfg.min_distance_between_context_views + + # Pick the gap between the context views. + if not self.cameras_are_circular: + max_gap = min(num_views - 1, max_gap) + min_gap = max(2 * self.cfg.min_distance_to_context_views, min_gap) + + # overwrite min_gap and max_gap, useful for mixed dataset training + # use different view distance for different dataset + if min_view_dist is not None: + min_gap = min_view_dist + + if max_view_dist is not None: + max_gap = max_view_dist + + if max_gap < min_gap: + raise ValueError("Example does not have enough frames!") + context_gap = torch.randint( + min_gap, + max_gap + 1, + size=tuple(), + device=device, + ).item() + + # Pick the left and right context indices. + index_context_left = torch.randint( + num_views if self.cameras_are_circular else num_views - context_gap, + size=tuple(), + device=device, + ).item() + if self.stage == "test": + index_context_left = index_context_left * 0 + index_context_right = index_context_left + context_gap + + if self.is_overfitting: + index_context_left *= 0 + index_context_right *= 0 + index_context_right += max_gap + + # Pick the target view indices. + if self.stage == "test": + # When testing, pick all. + index_target = torch.arange( + index_context_left, + index_context_right + 1, + device=device, + ) + else: + # When training or validating (visualizing), pick at random. + index_target = torch.randint( + index_context_left + self.cfg.min_distance_to_context_views, + index_context_right + 1 - self.cfg.min_distance_to_context_views, + size=(self.cfg.num_target_views,), + device=device, + ) + + # Apply modulo for circular datasets. + if self.cameras_are_circular: + index_target %= num_views + index_context_right %= num_views + + return ( + torch.tensor((index_context_left, index_context_right)), + index_target, + ) + + @property + def num_context_views(self) -> int: + return 2 + + @property + def num_target_views(self) -> int: + return self.cfg.num_target_views diff --git a/src/dataset/view_sampler/view_sampler_bounded_v2.py b/src/dataset/view_sampler/view_sampler_bounded_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..6aaf7b730d8efbe0fea6dc57bfee9ca3941d6c43 --- /dev/null +++ b/src/dataset/view_sampler/view_sampler_bounded_v2.py @@ -0,0 +1,253 @@ +''' +Modified from latentSplat and pixelSplat to handle extrapolation and more context views +''' + +from dataclasses import dataclass +from typing import Literal, Optional + +import torch +from jaxtyping import Float, Int64 +from torch import Tensor +import random + +from .view_sampler import ViewSampler + + +def farthest_point_sample(xyz, npoint): + """ + Input: + xyz: pointcloud data, [B, N, 3] + npoint: number of samples + Return: + centroids: sampled pointcloud index, [B, npoint] + """ + + 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): + # print("The %d farthest pts %s " % (i, farthest)) + 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 + + +@dataclass +class ViewSamplerBoundedV2Cfg: + name: Literal["boundedv2"] + num_context_views: int + num_target_views: int + min_distance_between_context_views: int + max_distance_between_context_views: int + max_distance_to_context_views: int + context_gap_warm_up_steps: int + target_gap_warm_up_steps: int + initial_min_distance_between_context_views: int + initial_max_distance_between_context_views: int + initial_max_distance_to_context_views: int + extra_views_sampling_strategy: Optional[Literal["random", "farthest_point", "equal"]] = "random" + target_views_replace_sample: Optional[bool] = True + + +class ViewSamplerBoundedV2(ViewSampler[ViewSamplerBoundedV2Cfg]): + def schedule(self, initial: int, final: int, steps: int) -> int: + fraction = self.global_step / steps + return min(initial + int((final - initial) * fraction), final) + + def sample( + self, + scene: str, + extrinsics: Float[Tensor, "view 4 4"], + intrinsics: Float[Tensor, "view 3 3"], + device: torch.device = torch.device("cpu"), + max_num_views: Optional[int] = None, + min_context_views: int = 0, + max_context_views: int = 0, + min_view_dist: int | None = None, + max_view_dist: int | None = None, + ) -> tuple[ + Int64[Tensor, " context_view"], # indices for context views + Int64[Tensor, " target_view"], # indices for target views + ]: + num_views, _, _ = extrinsics.shape + if max_num_views is not None: + num_views = min(num_views, max_num_views) + + if min_context_views > 0 and max_context_views > 0 and self.stage != "test": + random_num_views = random.randint(min_context_views, max_context_views) + else: + random_num_views = None + + # Compute the context view spacing based on the current global step. + if self.stage == "test": + # When testing, always use the full gap. + max_context_gap = self.cfg.max_distance_between_context_views + min_context_gap = self.cfg.max_distance_between_context_views + elif self.cfg.context_gap_warm_up_steps > 0: + max_context_gap = self.schedule( + self.cfg.initial_max_distance_between_context_views, + self.cfg.max_distance_between_context_views, + self.cfg.context_gap_warm_up_steps, + ) + min_context_gap = self.schedule( + self.cfg.initial_min_distance_between_context_views, + self.cfg.min_distance_between_context_views, + self.cfg.context_gap_warm_up_steps, + ) + else: + max_context_gap = self.cfg.max_distance_between_context_views + min_context_gap = self.cfg.min_distance_between_context_views + + if min_view_dist is not None and max_view_dist is not None: + # for mixed dataset training, with different sampling distance + min_context_gap = min_view_dist + max_context_gap = max_view_dist + + if random_num_views is not None: + # smaller context gap accordingly + scale_factor = max(max_context_views // random_num_views, 1) + max_context_gap = max_context_gap // scale_factor + min_context_gap = min_context_gap // scale_factor + + if not self.cameras_are_circular: + max_context_gap = min( + num_views - 1, max_context_gap + ) + + # Compute the margin from context window to target window based on the current global step + if self.stage != "test" and self.cfg.target_gap_warm_up_steps > 0: + max_target_gap = self.schedule( + self.cfg.initial_max_distance_to_context_views, + self.cfg.max_distance_to_context_views, + self.cfg.target_gap_warm_up_steps, + ) + else: + max_target_gap = self.cfg.max_distance_to_context_views + + # Pick the gap between the context views. + if max_context_gap < min_context_gap: + raise ValueError("Example does not have enough frames!") + context_gap = torch.randint( + min_context_gap, + max_context_gap + 1, + size=tuple(), + device=device, + ).item() + + # Pick the left and right context indices. + index_context_left = torch.randint( + low=0, + high=num_views if self.cameras_are_circular else num_views - context_gap, + size=tuple(), + device=device, + ).item() + if self.stage == "test": + index_context_left = index_context_left * 0 + index_context_right = index_context_left + context_gap + + index_target_left = index_context_left - max_target_gap + index_target_right = index_context_right + max_target_gap + + if not self.cameras_are_circular: + index_target_left = max(0, index_target_left) + index_target_right = min(num_views - 1, index_target_right) + + # Pick the target view indices. + if self.stage == "test": + # When testing, pick all. + index_target = torch.arange( + index_target_left, + index_target_right + 1, + device=device, + ) + else: + # When training or validating (visualizing), pick at random. + if self.cfg.target_views_replace_sample: + index_target = torch.randint( + index_target_left, + index_target_right + 1, + size=(self.cfg.num_target_views,), + device=device, + ) + else: # sample without replacement + # similarly, ok to index + index_target_candidates = torch.arange( + index_target_left, + index_target_right + 1, + device=device, + ) + indices = torch.randperm(index_target_right + 1 - index_target_left, device=device)[ + : self.cfg.num_target_views + ] + index_target = index_target_candidates[indices] + + # Apply modulo for circular datasets. + if self.cameras_are_circular: + index_target %= num_views + index_context_right %= num_views + + # If more than two context views are desired, pick extra context views between + # the left and right ones. + if random_num_views is not None: + total_num_views = random_num_views + else: + total_num_views = self.cfg.num_context_views + + if total_num_views > 2: + num_extra_views = total_num_views - 2 + extra_views = [] + if self.cfg.extra_views_sampling_strategy == 'random': + while len(set(extra_views)) != num_extra_views: + extra_views = torch.randint( + index_context_left + 1, + index_context_right, + (num_extra_views,), + ).tolist() + elif self.cfg.extra_views_sampling_strategy == 'farthest_point': + context_bounded_index = torch.arange(index_context_left, index_context_right + 1) + candidate_views_position = extrinsics[context_bounded_index, :3, -1].unsqueeze(0) + index_context_local = farthest_point_sample( + candidate_views_position, total_num_views + ).squeeze(0) + # remap context index back to global scene based index + index_context = context_bounded_index[index_context_local] + index_context_left = index_context[0].item() + index_context_right = index_context[-1].item() + extra_views = index_context[1:-1].tolist() + elif self.cfg.extra_views_sampling_strategy == 'equal': + pass + + # sort the index + extra_views = sorted(extra_views) + else: + extra_views = [] + + return ( + torch.tensor((index_context_left, *extra_views, index_context_right)), + index_target, + ) + + @property + def num_context_views(self) -> int: + return self.cfg.num_context_views + + @property + def num_target_views(self) -> int: + return self.cfg.num_target_views diff --git a/src/dataset/view_sampler/view_sampler_evaluation.py b/src/dataset/view_sampler/view_sampler_evaluation.py new file mode 100644 index 0000000000000000000000000000000000000000..9cf1214abb601259144ec9615ddc1a7fe8dc92f6 --- /dev/null +++ b/src/dataset/view_sampler/view_sampler_evaluation.py @@ -0,0 +1,73 @@ +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +import torch +from dacite import Config, from_dict +from jaxtyping import Float, Int64 +from torch import Tensor + +from ...misc.step_tracker import StepTracker +from ..types import Stage +from .view_sampler import ViewSampler + + +@dataclass +class IndexEntry: + context: tuple[int, ...] + target: tuple[int, ...] + + +@dataclass +class ViewSamplerEvaluationCfg: + name: Literal["evaluation"] + index_path: Path + num_context_views: int + + +class ViewSamplerEvaluation(ViewSampler[ViewSamplerEvaluationCfg]): + index: dict[str, IndexEntry | None] + + def __init__( + self, + cfg: ViewSamplerEvaluationCfg, + stage: Stage, + is_overfitting: bool, + cameras_are_circular: bool, + step_tracker: StepTracker | None, + ) -> None: + super().__init__(cfg, stage, is_overfitting, cameras_are_circular, step_tracker) + + dacite_config = Config(cast=[tuple]) + with cfg.index_path.open("r") as f: + self.index = { + k: None if v is None else from_dict(IndexEntry, v, dacite_config) + for k, v in json.load(f).items() + } + + def sample( + self, + scene: str, + extrinsics: Float[Tensor, "view 4 4"], + intrinsics: Float[Tensor, "view 3 3"], + device: torch.device = torch.device("cpu"), + **kwargs, + ) -> tuple[ + Int64[Tensor, " context_view"], # indices for context views + Int64[Tensor, " target_view"], # indices for target views + ]: + entry = self.index.get(scene) + if entry is None: + raise ValueError(f"No indices available for scene {scene}.") + context_indices = torch.tensor(entry.context, dtype=torch.int64, device=device) + target_indices = torch.tensor(entry.target, dtype=torch.int64, device=device) + return context_indices, target_indices + + @property + def num_context_views(self) -> int: + return 0 + + @property + def num_target_views(self) -> int: + return 0 diff --git a/src/evaluation/metrics.py b/src/evaluation/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..1981b8c6475227797f7945b524559182a8ddb8fb --- /dev/null +++ b/src/evaluation/metrics.py @@ -0,0 +1,129 @@ +from functools import cache + +import torch +import torch.nn.functional as F +from einops import reduce +from jaxtyping import Float +from lpips import LPIPS +from torch import Tensor + + +@torch.no_grad() +def compute_psnr( + ground_truth: Float[Tensor, "batch channel height width"], + predicted: Float[Tensor, "batch channel height width"], +) -> Float[Tensor, " batch"]: + with torch.autocast(device_type=predicted.device.type, enabled=False): + ground_truth = ground_truth.float().clip(min=0, max=1) + predicted = predicted.float().clip(min=0, max=1) + mse = reduce((ground_truth - predicted) ** 2, "b c h w -> b", "mean") + return -10 * mse.log10() + + +@cache +def get_lpips(device: torch.device) -> LPIPS: + model = LPIPS(net="vgg").eval().requires_grad_(False) + return model.to(device=device, dtype=torch.float32) + + +@torch.no_grad() +def compute_lpips( + ground_truth: Float[Tensor, "batch channel height width"], + predicted: Float[Tensor, "batch channel height width"], +) -> Float[Tensor, " batch"]: + with torch.autocast(device_type=predicted.device.type, enabled=False): + value = get_lpips(predicted.device).forward( + ground_truth.float(), + predicted.float(), + normalize=True, + ) + return value[:, 0, 0, 0].float() + + +@torch.no_grad() +def compute_ssim( + ground_truth: Float[Tensor, "batch channel height width"], + predicted: Float[Tensor, "batch channel height width"], + *, + win_size: int = 11, + gaussian_weights: bool = True, + data_range: float = 1.0, + use_sample_covariance: bool = True, + sigma: float = 1.5, + k1: float = 0.01, + k2: float = 0.03, +) -> Float[Tensor, " batch"]: + if ground_truth.shape != predicted.shape or ground_truth.ndim != 4: + raise ValueError("SSIM inputs must have the same BCHW shape.") + if win_size <= 0 or win_size % 2 == 0: + raise ValueError("win_size must be a positive odd integer.") + if min(ground_truth.shape[-2:]) < win_size: + raise ValueError("win_size must not exceed the image dimensions.") + if gaussian_weights and sigma <= 0: + raise ValueError("sigma must be positive when using Gaussian weights.") + if data_range <= 0 or k1 < 0 or k2 < 0: + raise ValueError("data_range must be positive and k1/k2 non-negative.") + + with torch.autocast(device_type=predicted.device.type, enabled=False): + ground_truth = ground_truth.float() + predicted = predicted.float() + + if gaussian_weights: + coordinates = torch.arange( + win_size, + dtype=torch.float32, + device=predicted.device, + ) + coordinates -= (win_size - 1) / 2 + kernel_1d = torch.exp(-(coordinates.square()) / (2 * sigma**2)) + kernel_1d /= kernel_1d.sum() + kernel = torch.outer(kernel_1d, kernel_1d) + else: + kernel = torch.full( + (win_size, win_size), + 1 / win_size**2, + dtype=torch.float32, + device=predicted.device, + ) + + channels = predicted.shape[1] + moments = torch.cat( + ( + ground_truth, + predicted, + ground_truth.square(), + predicted.square(), + ground_truth * predicted, + ), + dim=1, + ) + kernels = kernel[None, None].expand(5 * channels, 1, -1, -1) + ux, uy, uxx, uyy, uxy = F.conv2d( + moments, + kernels, + groups=5 * channels, + ).chunk(5, dim=1) + + covariance_normalization = 1.0 + if use_sample_covariance: + num_pixels = win_size**2 + if num_pixels == 1: + raise ValueError( + "Sample covariance requires a window larger than one pixel." + ) + covariance_normalization = num_pixels / (num_pixels - 1) + + vx = covariance_normalization * (uxx - ux.square()) + vy = covariance_normalization * (uyy - uy.square()) + vxy = covariance_normalization * (uxy - ux * uy) + + c1 = (k1 * data_range) ** 2 + c2 = (k2 * data_range) ** 2 + ssim = ( + (2 * ux * uy + c1) + * (2 * vxy + c2) + / ((ux.square() + uy.square() + c1) * (vx + vy + c2)) + ) + ssim = ssim.mean(dim=(1, 2, 3)) + + return ssim diff --git a/src/geometry/projection.py b/src/geometry/projection.py new file mode 100644 index 0000000000000000000000000000000000000000..eb8f0e3a48e13e658e2f607b0118fc6fed141fb0 --- /dev/null +++ b/src/geometry/projection.py @@ -0,0 +1,247 @@ +from math import prod + +import torch +from einops import einsum, rearrange, reduce, repeat +from jaxtyping import Bool, Float, Int64 +from torch import Tensor + + +def homogenize_points( + points: Float[Tensor, "*batch dim"], +) -> Float[Tensor, "*batch dim+1"]: + """Convert batched points (xyz) to (xyz1).""" + return torch.cat([points, torch.ones_like(points[..., :1])], dim=-1) + + +def homogenize_vectors( + vectors: Float[Tensor, "*batch dim"], +) -> Float[Tensor, "*batch dim+1"]: + """Convert batched vectors (xyz) to (xyz0).""" + return torch.cat([vectors, torch.zeros_like(vectors[..., :1])], dim=-1) + + +def transform_rigid( + homogeneous_coordinates: Float[Tensor, "*#batch dim"], + transformation: Float[Tensor, "*#batch dim dim"], +) -> Float[Tensor, "*batch dim"]: + """Apply a rigid-body transformation to points or vectors.""" + return einsum(transformation, homogeneous_coordinates, "... i j, ... j -> ... i") + + +def transform_cam2world( + homogeneous_coordinates: Float[Tensor, "*#batch dim"], + extrinsics: Float[Tensor, "*#batch dim dim"], +) -> Float[Tensor, "*batch dim"]: + """Transform points from 3D camera coordinates to 3D world coordinates.""" + return transform_rigid(homogeneous_coordinates, extrinsics) + + +def transform_world2cam( + homogeneous_coordinates: Float[Tensor, "*#batch dim"], + extrinsics: Float[Tensor, "*#batch dim dim"], +) -> Float[Tensor, "*batch dim"]: + """Transform points from 3D world coordinates to 3D camera coordinates.""" + return transform_rigid(homogeneous_coordinates, extrinsics.inverse()) + + +def project_camera_space( + points: Float[Tensor, "*#batch dim"], + intrinsics: Float[Tensor, "*#batch dim dim"], + epsilon: float = torch.finfo(torch.float32).eps, + infinity: float = 1e8, +) -> Float[Tensor, "*batch dim-1"]: + points = points / (points[..., -1:] + epsilon) + points = points.nan_to_num(posinf=infinity, neginf=-infinity) + points = einsum(intrinsics, points, "... i j, ... j -> ... i") + return points[..., :-1] + + +def project( + points: Float[Tensor, "*#batch dim"], + extrinsics: Float[Tensor, "*#batch dim+1 dim+1"], + intrinsics: Float[Tensor, "*#batch dim dim"], + epsilon: float = torch.finfo(torch.float32).eps, +) -> tuple[ + Float[Tensor, "*batch dim-1"], # xy coordinates + Bool[Tensor, " *batch"], # whether points are in front of the camera +]: + points = homogenize_points(points) + points = transform_world2cam(points, extrinsics)[..., :-1] + in_front_of_camera = points[..., -1] >= 0 + return project_camera_space(points, intrinsics, epsilon=epsilon), in_front_of_camera + + +def unproject( + coordinates: Float[Tensor, "*#batch dim"], + z: Float[Tensor, "*#batch"], + intrinsics: Float[Tensor, "*#batch dim+1 dim+1"], +) -> Float[Tensor, "*batch dim+1"]: + """Unproject 2D camera coordinates with the given Z values.""" + + # Apply the inverse intrinsics to the coordinates. + coordinates = homogenize_points(coordinates) + ray_directions = einsum( + intrinsics.inverse(), coordinates, "... i j, ... j -> ... i" + ) + + # Apply the supplied depth values. + return ray_directions * z[..., None] + + +def get_world_rays( + coordinates: Float[Tensor, "*#batch dim"], + extrinsics: Float[Tensor, "*#batch dim+2 dim+2"], + intrinsics: Float[Tensor, "*#batch dim+1 dim+1"], +) -> tuple[ + Float[Tensor, "*batch dim+1"], # origins + Float[Tensor, "*batch dim+1"], # directions +]: + # Get camera-space ray directions. + directions = unproject( + coordinates, + torch.ones_like(coordinates[..., 0]), + intrinsics, + ) + directions = directions / directions[..., -1:] + + # Transform ray directions to world coordinates. + directions = homogenize_vectors(directions) + directions = transform_cam2world(directions, extrinsics)[..., :-1] + + # Tile the ray origins to have the same shape as the ray directions. + origins = extrinsics[..., :-1, -1].broadcast_to(directions.shape) + + return origins, directions + + +def sample_image_grid( + shape: tuple[int, ...], + device: torch.device = torch.device("cpu"), +) -> tuple[ + Float[Tensor, "*shape dim"], # float coordinates (xy indexing) + Int64[Tensor, "*shape dim"], # integer indices (ij indexing) +]: + """Get normalized (range 0 to 1) coordinates and integer indices for an image.""" + + # Each entry is a pixel-wise integer coordinate. In the 2D case, each entry is a + # (row, col) coordinate. + indices = [torch.arange(length, device=device) for length in shape] + stacked_indices = torch.stack(torch.meshgrid(*indices, indexing="ij"), dim=-1) + + # Each entry is a floating-point coordinate in the range (0, 1). In the 2D case, + # each entry is an (x, y) coordinate. + coordinates = [(idx + 0.5) / length for idx, length in zip(indices, shape)] + coordinates = reversed(coordinates) + coordinates = torch.stack(torch.meshgrid(*coordinates, indexing="xy"), dim=-1) + + return coordinates, stacked_indices + + +def sample_training_rays( + image: Float[Tensor, "batch view channel ..."], + intrinsics: Float[Tensor, "batch view dim dim"], + extrinsics: Float[Tensor, "batch view dim+1 dim+1"], + num_rays: int, +) -> tuple[ + Float[Tensor, "batch ray dim"], # origins + Float[Tensor, "batch ray dim"], # directions + Float[Tensor, "batch ray 3"], # sampled color +]: + device = extrinsics.device + b, v, _, *grid_shape = image.shape + + # Generate all possible target rays. + xy, _ = sample_image_grid(tuple(grid_shape), device) + origins, directions = get_world_rays( + rearrange(xy, "... d -> ... () () d"), + extrinsics, + intrinsics, + ) + origins = rearrange(origins, "... b v xy -> b (v ...) xy", b=b, v=v) + directions = rearrange(directions, "... b v xy -> b (v ...) xy", b=b, v=v) + pixels = rearrange(image, "b v c ... -> b (v ...) c") + + # Sample random rays. + num_possible_rays = v * prod(grid_shape) + ray_indices = torch.randint(num_possible_rays, (b, num_rays), device=device) + batch_indices = repeat(torch.arange(b, device=device), "b -> b n", n=num_rays) + + return ( + origins[batch_indices, ray_indices], + directions[batch_indices, ray_indices], + pixels[batch_indices, ray_indices], + ) + + +def intersect_rays( + origins_x: Float[Tensor, "*#batch 3"], + directions_x: Float[Tensor, "*#batch 3"], + origins_y: Float[Tensor, "*#batch 3"], + directions_y: Float[Tensor, "*#batch 3"], + eps: float = 1e-5, + inf: float = 1e10, +) -> Float[Tensor, "*batch 3"]: + """Compute the least-squares intersection of rays. Uses the math from here: + https://math.stackexchange.com/a/1762491/286022 + """ + + # Broadcast the rays so their shapes match. + shape = torch.broadcast_shapes( + origins_x.shape, + directions_x.shape, + origins_y.shape, + directions_y.shape, + ) + origins_x = origins_x.broadcast_to(shape) + directions_x = directions_x.broadcast_to(shape) + origins_y = origins_y.broadcast_to(shape) + directions_y = directions_y.broadcast_to(shape) + + # Detect and remove batch elements where the directions are parallel. + parallel = einsum(directions_x, directions_y, "... xyz, ... xyz -> ...") > 1 - eps + origins_x = origins_x[~parallel] + directions_x = directions_x[~parallel] + origins_y = origins_y[~parallel] + directions_y = directions_y[~parallel] + + # Stack the rays into (2, *shape). + origins = torch.stack([origins_x, origins_y], dim=0) + directions = torch.stack([directions_x, directions_y], dim=0) + dtype = origins.dtype + device = origins.device + + # Compute n_i * n_i^T - eye(3) from the equation. + n = einsum(directions, directions, "r b i, r b j -> r b i j") + n = n - torch.eye(3, dtype=dtype, device=device).broadcast_to((2, 1, 3, 3)) + + # Compute the left-hand side of the equation. + lhs = reduce(n, "r b i j -> b i j", "sum") + + # Compute the right-hand side of the equation. + rhs = einsum(n, origins, "r b i j, r b j -> r b i") + rhs = reduce(rhs, "r b i -> b i", "sum") + + # Left-matrix-multiply both sides by the pseudo-inverse of lhs to find p. + result = torch.linalg.lstsq(lhs, rhs).solution + + # Handle the case of parallel lines by setting depth to infinity. + result_all = torch.ones(shape, dtype=dtype, device=device) * inf + result_all[~parallel] = result + return result_all + + +def get_fov(intrinsics: Float[Tensor, "batch 3 3"]) -> Float[Tensor, "batch 2"]: + intrinsics_inv = intrinsics.inverse() + + def process_vector(vector): + vector = torch.tensor(vector, dtype=torch.float32, device=intrinsics.device) + vector = einsum(intrinsics_inv, vector, "b i j, j -> b i") + return vector / vector.norm(dim=-1, keepdim=True) + + left = process_vector([0, 0.5, 1]) + right = process_vector([1, 0.5, 1]) + top = process_vector([0.5, 0, 1]) + bottom = process_vector([0.5, 1, 1]) + fov_x = (left * right).sum(dim=-1).acos() + fov_y = (top * bottom).sum(dim=-1).acos() + return torch.stack((fov_x, fov_y), dim=-1) diff --git a/src/global_cfg.py b/src/global_cfg.py new file mode 100644 index 0000000000000000000000000000000000000000..6796eea86201b8d7b89b67d2f88277a18bb74d16 --- /dev/null +++ b/src/global_cfg.py @@ -0,0 +1,15 @@ +from typing import Optional + +from omegaconf import DictConfig + +cfg: Optional[DictConfig] = None + + +def get_cfg() -> DictConfig: + global cfg + return cfg + + +def set_cfg(new_cfg: DictConfig) -> None: + global cfg + cfg = new_cfg diff --git a/src/loss/__init__.py b/src/loss/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ad5e7e4b9674bb567a99c84ba49d60658bcd9328 --- /dev/null +++ b/src/loss/__init__.py @@ -0,0 +1,14 @@ +from .loss import Loss +from .loss_lpips import LossLpips, LossLpipsCfgWrapper +from .loss_mse import LossMse, LossMseCfgWrapper + +LOSSES = { + LossLpipsCfgWrapper: LossLpips, + LossMseCfgWrapper: LossMse, +} + +LossCfgWrapper = LossLpipsCfgWrapper | LossMseCfgWrapper + + +def get_losses(cfgs: list[LossCfgWrapper]) -> list[Loss]: + return [LOSSES[type(cfg)](cfg) for cfg in cfgs] diff --git a/src/loss/loss.py b/src/loss/loss.py new file mode 100644 index 0000000000000000000000000000000000000000..e331e79d800f8f5bebb94cb1be7df73bd8a0a441 --- /dev/null +++ b/src/loss/loss.py @@ -0,0 +1,36 @@ +from abc import ABC, abstractmethod +from dataclasses import fields +from typing import Generic, TypeVar + +from jaxtyping import Float +from torch import Tensor, nn + +from ..dataset.types import BatchedExample +from ..model.decoder.decoder import DecoderOutput +from ..model.types import Gaussians + +T_cfg = TypeVar("T_cfg") +T_wrapper = TypeVar("T_wrapper") + + +class Loss(nn.Module, ABC, Generic[T_cfg, T_wrapper]): + cfg: T_cfg + name: str + + def __init__(self, cfg: T_wrapper) -> None: + super().__init__() + + # Extract the configuration from the wrapper. + (field,) = fields(type(cfg)) + self.cfg = getattr(cfg, field.name) + self.name = field.name + + @abstractmethod + def forward( + self, + prediction: DecoderOutput, + batch: BatchedExample, + gaussians: Gaussians, + global_step: int, + ) -> Float[Tensor, ""]: + pass diff --git a/src/loss/loss_lpips.py b/src/loss/loss_lpips.py new file mode 100644 index 0000000000000000000000000000000000000000..ab8c2cbb3b37bcdc11ab6e553bf7866164a716d7 --- /dev/null +++ b/src/loss/loss_lpips.py @@ -0,0 +1,79 @@ +from dataclasses import dataclass + +import torch +from einops import rearrange +from jaxtyping import Float +from lpips import LPIPS +from torch import Tensor + +from ..dataset.types import BatchedExample +from ..model.decoder.decoder import DecoderOutput +from ..model.types import Gaussians +from .loss import Loss + + +@dataclass +class LossLpipsCfg: + weight: float + apply_after_step: int + + +@dataclass +class LossLpipsCfgWrapper: + lpips: LossLpipsCfg + + +class LossLpips(Loss[LossLpipsCfg, LossLpipsCfgWrapper]): + lpips: LPIPS + + def __init__(self, cfg: LossLpipsCfgWrapper) -> None: + super().__init__(cfg) + + lpips = LPIPS(net="vgg").eval() + lpips.requires_grad_(False) + self.__dict__["lpips"] = lpips + + def lpips_on_device(self, device: torch.device) -> LPIPS: + lpips = self.lpips + param = next(iter(lpips.parameters())) + if param.device != device or param.dtype != torch.float32: + lpips.to(device=device, dtype=torch.float32) + return lpips + + def forward( + self, + prediction: DecoderOutput, + batch: BatchedExample, + gaussians: Gaussians, + global_step: int, + valid_depth_mask: Tensor | None, + reduction: str = "mean", + ) -> Float[Tensor, "*batch"]: + image = batch["target"]["image"] + + # Before the specified step, don't apply the loss. + if global_step < self.cfg.apply_after_step: + return torch.tensor(0, dtype=torch.float32, device=image.device) + + prediction_color = prediction.color + if valid_depth_mask is not None and valid_depth_mask.max() > 0.5: + prediction_color = prediction_color.clone() + image = image.clone() + prediction_color[valid_depth_mask] = 0 + image[valid_depth_mask] = 0 + + with torch.autocast(device_type=image.device.type, enabled=False): + loss = self.lpips_on_device(image.device).forward( + rearrange(prediction_color.float(), "b v c h w -> (b v) c h w"), + rearrange(image.float(), "b v c h w -> (b v) c h w"), + normalize=True, + ) + batch_size, num_views = image.shape[:2] + sample_loss = loss.reshape(batch_size, num_views, -1).mean(dim=(1, 2)) + + if reduction == "mean": + return self.cfg.weight * sample_loss.mean() + if reduction == "none": + return self.cfg.weight * sample_loss + + raise ValueError(f"Unsupported reduction={reduction}") diff --git a/src/loss/loss_mse.py b/src/loss/loss_mse.py new file mode 100644 index 0000000000000000000000000000000000000000..416f2e7aee36f53fc35fe7d35bc50256f353c105 --- /dev/null +++ b/src/loss/loss_mse.py @@ -0,0 +1,69 @@ +from dataclasses import dataclass + +from jaxtyping import Float +from torch import Tensor +import torch + +from ..dataset.types import BatchedExample +from ..model.decoder.decoder import DecoderOutput +from ..model.types import Gaussians +from .loss import Loss + + +@dataclass +class LossMseCfg: + weight: float + + +@dataclass +class LossMseCfgWrapper: + mse: LossMseCfg + + +class LossMse(Loss[LossMseCfg, LossMseCfgWrapper]): + def forward( + self, + prediction: DecoderOutput, + batch: BatchedExample, + gaussians: Gaussians, + global_step: int, + l1_loss: bool, + clamp_large_error: float, + valid_depth_mask: Tensor | None, + reduction: str = "mean", + ) -> Float[Tensor, "*batch"]: + # Loss accumulation is not a feature transform; keep it in FP32 under + # mixed precision for stable gradients and logging. + delta = prediction.color.float() - batch["target"]["image"].float() + + if l1_loss: + loss_map = delta.abs() + else: + loss_map = delta**2 + + valid_map = torch.ones_like(loss_map, dtype=torch.bool) + if ( + valid_depth_mask is not None + and valid_depth_mask.max() > 0.5 + and valid_depth_mask.min() < 0.5 + ): + valid_map = valid_map & (~valid_depth_mask) + + if clamp_large_error > 0: + valid_map = valid_map & ((delta ** 2) < clamp_large_error) + + if reduction == "mean": + masked_loss = loss_map[valid_map] + if masked_loss.numel() == 0: + return torch.tensor(0.0, device=loss_map.device, dtype=loss_map.dtype) + return self.cfg.weight * masked_loss.mean() + + if reduction == "none": + batch_size = loss_map.shape[0] + flat_loss = loss_map.reshape(batch_size, -1) + flat_valid = valid_map.reshape(batch_size, -1).to(loss_map.dtype) + valid_counts = flat_valid.sum(dim=1).clamp_min(1.0) + sample_loss = (flat_loss * flat_valid).sum(dim=1) / valid_counts + return self.cfg.weight * sample_loss + + raise ValueError(f"Unsupported reduction={reduction}") diff --git a/src/misc/LocalLogger.py b/src/misc/LocalLogger.py new file mode 100644 index 0000000000000000000000000000000000000000..4e11c21c000a62ceafef6f4c81ead5cec47ecdc6 --- /dev/null +++ b/src/misc/LocalLogger.py @@ -0,0 +1,52 @@ +from pathlib import Path +import shutil +from typing import Any, Optional + +import numpy as np +import torch +from PIL import Image +from pytorch_lightning.loggers.logger import Logger +from pytorch_lightning.utilities import rank_zero_only + + +class LocalLogger(Logger): + def __init__(self, log_path: Path | str = Path("outputs/local")) -> None: + super().__init__() + self.experiment = None + self.log_path = Path(log_path) + shutil.rmtree(self.log_path, ignore_errors=True) + + @property + def name(self): + return "LocalLogger" + + @property + def version(self): + return 0 + + @rank_zero_only + def log_hyperparams(self, params): + pass + + @rank_zero_only + def log_metrics(self, metrics, step): + pass + + @rank_zero_only + def log_image( + self, + key: str, + images: list[Any], + step: Optional[int] = None, + **kwargs, + ): + # The function signature is the same as the wandb logger's, but the step is + # actually required. + assert step is not None + for index, image in enumerate(images): + path = self.log_path / f"{key}/{index:0>2}_{step:0>6}.png" + path.parent.mkdir(exist_ok=True, parents=True) + if isinstance(image, torch.Tensor): + Image.fromarray(image.permute(1, 2, 0).numpy().astype(np.uint8)).save(path) + else: + Image.fromarray(image).save(path) diff --git a/src/misc/benchmarker.py b/src/misc/benchmarker.py new file mode 100644 index 0000000000000000000000000000000000000000..4969ef686a8baa2f50148131b522d14e4437d13b --- /dev/null +++ b/src/misc/benchmarker.py @@ -0,0 +1,49 @@ +import json +from collections import defaultdict +from contextlib import contextmanager +from pathlib import Path +from time import time + +import numpy as np +import torch + + +class Benchmarker: + def __init__(self): + self.execution_times = defaultdict(list) + self.memory_stats_reset = False + + def reset_memory_stats(self) -> None: + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + self.memory_stats_reset = True + + @contextmanager + def time(self, tag: str, num_calls: int = 1): + if not self.memory_stats_reset: + self.reset_memory_stats() + try: + start_time = time() + yield + finally: + end_time = time() + for _ in range(num_calls): + self.execution_times[tag].append((end_time - start_time) / num_calls) + + def dump(self, path: Path) -> None: + path.parent.mkdir(exist_ok=True, parents=True) + with path.open("w") as f: + json.dump(dict(self.execution_times), f) + + def dump_memory(self, path: Path) -> None: + path.parent.mkdir(exist_ok=True, parents=True) + with path.open("w") as f: + json.dump(torch.cuda.memory_stats()["allocated_bytes.all.peak"], f) + + def summarize(self) -> None: + for tag, times in self.execution_times.items(): + print(f"{tag}: {len(times)} calls, avg. {np.mean(times)} seconds per call") + + def clear_history(self) -> None: + self.execution_times = defaultdict(list) + self.reset_memory_stats() diff --git a/src/misc/image_io.py b/src/misc/image_io.py new file mode 100644 index 0000000000000000000000000000000000000000..5446f00a5f826e172dd0ba74922982aedf21bad5 --- /dev/null +++ b/src/misc/image_io.py @@ -0,0 +1,96 @@ +import io +from pathlib import Path +from typing import Union + +import imageio.v2 as imageio + +import numpy as np +import torch +import torchvision.transforms as tf +from einops import rearrange, repeat +from jaxtyping import Float, UInt8 +from matplotlib.figure import Figure +from PIL import Image +from torch import Tensor + +FloatImage = Union[ + Float[Tensor, "height width"], + Float[Tensor, "channel height width"], + Float[Tensor, "batch channel height width"], +] + + +def fig_to_image( + fig: Figure, + dpi: int = 100, + device: torch.device = torch.device("cpu"), +) -> Float[Tensor, "3 height width"]: + buffer = io.BytesIO() + fig.savefig(buffer, format="raw", dpi=dpi) + buffer.seek(0) + data = np.frombuffer(buffer.getvalue(), dtype=np.uint8) + h = int(fig.bbox.bounds[3]) + w = int(fig.bbox.bounds[2]) + data = rearrange(data, "(h w c) -> c h w", h=h, w=w, c=4) + buffer.close() + return (torch.tensor(data, device=device, dtype=torch.float32) / 255)[:3] + + +def prep_image(image: FloatImage) -> UInt8[np.ndarray, "height width channel"]: + # Handle batched images. + if image.ndim == 4: + image = rearrange(image, "b c h w -> c h (b w)") + + # Handle single-channel images. + if image.ndim == 2: + image = rearrange(image, "h w -> () h w") + + # Ensure that there are 3 or 4 channels. + channel, _, _ = image.shape + if channel == 1: + image = repeat(image, "() h w -> c h w", c=3) + assert image.shape[0] in (3, 4) + + image = (image.detach().clip(min=0, max=1) * 255).type(torch.uint8) + return rearrange(image, "c h w -> h w c").cpu().numpy() + + +def save_image( + image: FloatImage, + path: Union[Path, str], +) -> None: + """Save an image. Assumed to be in range 0-1.""" + + # Create the parent directory if it doesn't already exist. + path = Path(path) + path.parent.mkdir(exist_ok=True, parents=True) + + # Save the image. + Image.fromarray(prep_image(image)).save(path) + + +def load_image( + path: Union[Path, str], +) -> Float[Tensor, "3 height width"]: + return tf.ToTensor()(Image.open(path))[:3] + + +def save_video( + images: list[FloatImage], + path: Union[Path, str], + fps: None | int = None, +) -> None: + """Save a video whose input frames are assumed to be in range 0-1.""" + + path = Path(path) + path.parent.mkdir(exist_ok=True, parents=True) + + with imageio.get_writer( + str(path), + fps=fps or 30, + codec="libx264", + pixelformat="yuv420p", + output_params=["-crf", "0"], + ) as writer: + for image in images: + writer.append_data(prep_image(image)) diff --git a/src/misc/resume_ckpt.py b/src/misc/resume_ckpt.py new file mode 100644 index 0000000000000000000000000000000000000000..feb1fd6215648f53bd5651e569f6bf50dcda5577 --- /dev/null +++ b/src/misc/resume_ckpt.py @@ -0,0 +1,21 @@ +import os + + +# Function to extract the step number from the filename +def extract_step(file_name): + step_str = file_name.split("-")[1].split("_")[1].replace(".ckpt", "") + return int(step_str) + + +def find_latest_ckpt(ckpt_dir): + # List all files in the directory that end with .ckpt + ckpt_files = [f for f in os.listdir(ckpt_dir) if f.endswith(".ckpt")] + + # Check if there are any .ckpt files in the directory + if not ckpt_files: + raise ValueError(f"No .ckpt files found in {ckpt_dir}.") + else: + # Find the file with the maximum step + latest_ckpt_file = max(ckpt_files, key=extract_step) + + return ckpt_dir / latest_ckpt_file diff --git a/src/misc/sh_rotation.py b/src/misc/sh_rotation.py new file mode 100644 index 0000000000000000000000000000000000000000..4b1a26d8821232ed4f8aab257a7c67c46d905ff0 --- /dev/null +++ b/src/misc/sh_rotation.py @@ -0,0 +1,81 @@ +from math import isqrt + +import torch +from e3nn.o3 import matrix_to_angles, wigner_D +from einops import einsum +from jaxtyping import Float +from torch import Tensor + + +def rotate_sh( + sh_coefficients: Float[Tensor, "*#batch n"], + rotations: Float[Tensor, "*#batch 3 3"], +) -> Float[Tensor, "*batch n"]: + device = sh_coefficients.device + dtype = sh_coefficients.dtype + + *_, n = sh_coefficients.shape + alpha, beta, gamma = matrix_to_angles(rotations) + result = [] + for degree in range(isqrt(n)): + with torch.device(device): + sh_rotations = wigner_D(degree, alpha, beta, gamma).type(dtype) + sh_rotated = einsum( + sh_rotations, + sh_coefficients[..., degree**2 : (degree + 1) ** 2], + "... i j, ... j -> ... i", + ) + result.append(sh_rotated) + + return torch.cat(result, dim=-1) + + +if __name__ == "__main__": + from pathlib import Path + + import matplotlib.pyplot as plt + from e3nn.o3 import spherical_harmonics + from matplotlib import cm + from scipy.spatial.transform.rotation import Rotation as R + + device = torch.device("cuda") + + # Generate random spherical harmonics coefficients. + degree = 4 + coefficients = torch.rand((degree + 1) ** 2, dtype=torch.float32, device=device) + + def plot_sh(sh_coefficients, path: Path) -> None: + phi = torch.linspace(0, torch.pi, 100, device=device) + theta = torch.linspace(0, 2 * torch.pi, 100, device=device) + phi, theta = torch.meshgrid(phi, theta, indexing="xy") + x = torch.sin(phi) * torch.cos(theta) + y = torch.sin(phi) * torch.sin(theta) + z = torch.cos(phi) + xyz = torch.stack([x, y, z], dim=-1) + sh = spherical_harmonics(list(range(degree + 1)), xyz, True) + result = einsum(sh, sh_coefficients, "... n, n -> ...") + result = (result - result.min()) / (result.max() - result.min()) + + # Set the aspect ratio to 1 so our sphere looks spherical + fig = plt.figure(figsize=plt.figaspect(1.0)) + ax = fig.add_subplot(111, projection="3d") + ax.plot_surface( + x.cpu().numpy(), + y.cpu().numpy(), + z.cpu().numpy(), + rstride=1, + cstride=1, + facecolors=cm.seismic(result.cpu().numpy()), + ) + # Turn off the axis planes + ax.set_axis_off() + path.parent.mkdir(exist_ok=True, parents=True) + plt.savefig(path) + + for i, angle in enumerate(torch.linspace(0, 2 * torch.pi, 30)): + rotation = torch.tensor( + R.from_euler("x", angle.item()).as_matrix(), device=device + ) + plot_sh(rotate_sh(coefficients, rotation), Path(f"sh_rotation/{i:0>3}.png")) + + print("Done!") diff --git a/src/misc/stabilize_camera.py b/src/misc/stabilize_camera.py new file mode 100644 index 0000000000000000000000000000000000000000..7389e9f2037546f70356764efac86300b6d86118 --- /dev/null +++ b/src/misc/stabilize_camera.py @@ -0,0 +1,51 @@ +""" +https://github.com/google/dynibar/blob/main/ibrnet/data_loaders/llff_data_utils.py +""" + +import numpy as np +import cv2 + + +def render_stabilization_path(poses, k_size=45): + """Render a stabilized camera path.""" + + # hwf = poses[0, :, 4:5] + num_frames = poses.shape[0] + output_poses = [] + + input_poses = [] + + for i in range(num_frames): + input_poses.append( + np.concatenate( + [poses[i, :3, 0:1], poses[i, :3, 1:2], poses[i, :3, 3:4]], axis=-1 + ) + ) + + input_poses = np.array(input_poses) + + gaussian_kernel = cv2.getGaussianKernel(ksize=k_size, sigma=-1) + output_r1 = cv2.filter2D(input_poses[:, :, 0], -1, gaussian_kernel) + output_r2 = cv2.filter2D(input_poses[:, :, 1], -1, gaussian_kernel) + + output_r1 = output_r1 / np.linalg.norm(output_r1, axis=-1, keepdims=True) + output_r2 = output_r2 / np.linalg.norm(output_r2, axis=-1, keepdims=True) + + output_t = cv2.filter2D(input_poses[:, :, 2], -1, gaussian_kernel) + + for i in range(num_frames): + output_r3 = np.cross(output_r1[i], output_r2[i]) + + render_pose = np.concatenate( + [ + output_r1[i, :, None], + output_r2[i, :, None], + output_r3[:, None], + output_t[i, :, None], + ], + axis=-1, + ) + + output_poses.append(render_pose[:3, :]) + + return output_poses diff --git a/src/misc/step_tracker.py b/src/misc/step_tracker.py new file mode 100644 index 0000000000000000000000000000000000000000..8dd9cd99adb310456970f90e03b69bbe9f934fa3 --- /dev/null +++ b/src/misc/step_tracker.py @@ -0,0 +1,22 @@ +from multiprocessing import RLock + +import torch +from jaxtyping import Int64 +from torch import Tensor + + +class StepTracker: + lock: RLock + step: Int64[Tensor, ""] + + def __init__(self): + self.lock = torch.multiprocessing.RLock() + self.step = torch.tensor(0, dtype=torch.int64).share_memory_() + + def set_step(self, step: int) -> None: + with self.lock: + self.step.fill_(step) + + def get_step(self) -> int: + with self.lock: + return self.step.item() diff --git a/src/misc/wandb_tools.py b/src/misc/wandb_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..72daa77b93ebab5fc1c681949bf9a7943c06306d --- /dev/null +++ b/src/misc/wandb_tools.py @@ -0,0 +1,62 @@ +from pathlib import Path + +import wandb + + +def version_to_int(artifact) -> int: + """Convert versions of the form vX to X. For example, v12 to 12.""" + return int(artifact.version[1:]) + + +def download_checkpoint( + run_id: str, + download_dir: Path, + version: str | None, +) -> Path: + api = wandb.Api() + run = api.run(run_id) + + # Find the latest saved model checkpoint. + chosen = None + for artifact in run.logged_artifacts(): + if artifact.type != "model" or artifact.state != "COMMITTED": + continue + + # If no version is specified, use the latest. + if version is None: + if chosen is None or version_to_int(artifact) > version_to_int(chosen): + chosen = artifact + + # If a specific version is specified, look for it. + elif version == artifact.version: + chosen = artifact + break + + # Download the checkpoint. + download_dir.mkdir(exist_ok=True, parents=True) + root = download_dir / run_id + chosen.download(root=root) + return root / "model.ckpt" + + +def update_checkpoint_path(path: str | None, wandb_cfg: dict) -> Path | None: + if path is None: + return None + + if not str(path).startswith("wandb://"): + return Path(path) + + run_id, *version = path[len("wandb://") :].split(":") + if len(version) == 0: + version = None + elif len(version) == 1: + version = version[0] + else: + raise ValueError("Invalid version specifier!") + + project = wandb_cfg["project"] + return download_checkpoint( + f"{project}/{run_id}", + Path("checkpoints"), + version, + ) diff --git a/src/model/decoder/__init__.py b/src/model/decoder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5dd7b13c326241ca41a01b6c75e609e612fc7d79 --- /dev/null +++ b/src/model/decoder/__init__.py @@ -0,0 +1,43 @@ +from dataclasses import dataclass +from typing import Literal + +from ...dataset import DatasetCfg +from .decoder import Decoder + +try: + from .decoder_splatting_cuda import DecoderSplattingCUDA, DecoderSplattingCUDACfg +except ImportError as exc: + DecoderSplattingCUDA = None + CUDA_SPLATTING_IMPORT_ERROR = exc + + @dataclass + class DecoderSplattingCUDACfg: + name: Literal["splatting_cuda"] +else: + CUDA_SPLATTING_IMPORT_ERROR = None + +from .decoder_splatting_gsplat import ( + DecoderSplattingGsplat, + DecoderSplattingGsplatCfg, +) + +DECODERS = { + "splatting_gsplat": DecoderSplattingGsplat, + "gsplat": DecoderSplattingGsplat, +} + +DecoderCfg = DecoderSplattingGsplatCfg | DecoderSplattingCUDACfg + +if DecoderSplattingCUDA is not None: + DECODERS["splatting_cuda"] = DecoderSplattingCUDA + + +def get_decoder(decoder_cfg: DecoderCfg, dataset_cfg: DatasetCfg) -> Decoder: + decoder_cls = DECODERS.get(decoder_cfg.name) + if decoder_cls is None: + raise ImportError( + "The 'splatting_cuda' decoder was requested, but its CUDA rasterizer " + "dependencies could not be imported. Use model/decoder=splatting_gsplat " + "or install the original CUDA splatting dependencies." + ) from CUDA_SPLATTING_IMPORT_ERROR + return decoder_cls(decoder_cfg, dataset_cfg) diff --git a/src/model/decoder/cuda_splatting.py b/src/model/decoder/cuda_splatting.py new file mode 100644 index 0000000000000000000000000000000000000000..a9c49868edd0410f36702b913bff44d447089395 --- /dev/null +++ b/src/model/decoder/cuda_splatting.py @@ -0,0 +1,190 @@ +from math import isqrt +from typing import Literal + +import torch +from diff_gaussian_rasterization import ( + GaussianRasterizationSettings, + GaussianRasterizer, +) +from einops import einsum, rearrange, repeat +from jaxtyping import Float +from torch import Tensor + +from ...geometry.projection import get_fov, homogenize_points + + +def get_projection_matrix( + near: Float[Tensor, " batch"], + far: Float[Tensor, " batch"], + fov_x: Float[Tensor, " batch"], + fov_y: Float[Tensor, " batch"], +) -> Float[Tensor, "batch 4 4"]: + """Maps points in the viewing frustum to (-1, 1) on the X/Y axes and (0, 1) on the Z + axis. Differs from the OpenGL version in that Z doesn't have range (-1, 1) after + transformation and that Z is flipped. + """ + tan_fov_x = (0.5 * fov_x).tan() + tan_fov_y = (0.5 * fov_y).tan() + + top = tan_fov_y * near + bottom = -top + right = tan_fov_x * near + left = -right + + (b,) = near.shape + result = torch.zeros((b, 4, 4), dtype=torch.float32, device=near.device) + result[:, 0, 0] = 2 * near / (right - left) + result[:, 1, 1] = 2 * near / (top - bottom) + result[:, 0, 2] = (right + left) / (right - left) + result[:, 1, 2] = (top + bottom) / (top - bottom) + result[:, 3, 2] = 1 + result[:, 2, 2] = far / (far - near) + result[:, 2, 3] = -(far * near) / (far - near) + return result + + +@torch.autocast('cuda', enabled=False) +def render_cuda( + extrinsics: Float[Tensor, "batch 4 4"], + intrinsics: Float[Tensor, "batch 3 3"], + near: Float[Tensor, " batch"], + far: Float[Tensor, " batch"], + image_shape: tuple[int, int], + background_color: Float[Tensor, "batch 3"], + gaussian_means: Float[Tensor, "batch gaussian 3"], + gaussian_covariances: Float[Tensor, "batch gaussian 3 3"], + gaussian_sh_coefficients: Float[Tensor, "batch gaussian 3 d_sh"], + gaussian_opacities: Float[Tensor, "batch gaussian"], + scale_invariant: bool = True, + use_sh: bool = True, +) -> Float[Tensor, "batch 3 height width"]: + assert use_sh or gaussian_sh_coefficients.shape[-1] == 1 + + extrinsics = extrinsics.float() + intrinsics = intrinsics.float() + near = near.float() + far = far.float() + background_color = background_color.float() + gaussian_means = gaussian_means.float() + gaussian_covariances = gaussian_covariances.float() + gaussian_sh_coefficients = gaussian_sh_coefficients.float() + gaussian_opacities = gaussian_opacities.float() + + # Make sure everything is in a range where numerical issues don't appear. + if scale_invariant: + scale = 1 / near + extrinsics = extrinsics.clone() + extrinsics[..., :3, 3] = extrinsics[..., :3, 3] * scale[:, None] + gaussian_covariances = gaussian_covariances * (scale[:, None, None, None] ** 2) + gaussian_means = gaussian_means * scale[:, None, None] + near = near * scale + far = far * scale + + _, _, _, n = gaussian_sh_coefficients.shape + degree = isqrt(n) - 1 + shs = rearrange(gaussian_sh_coefficients, "b g xyz n -> b g n xyz").contiguous() + + b, _, _ = extrinsics.shape + h, w = image_shape + + fov_x, fov_y = get_fov(intrinsics).unbind(dim=-1) + tan_fov_x = (0.5 * fov_x).tan() + tan_fov_y = (0.5 * fov_y).tan() + + projection_matrix = get_projection_matrix(near, far, fov_x, fov_y) + projection_matrix = rearrange(projection_matrix, "b i j -> b j i") + view_matrix = rearrange(extrinsics.inverse(), "b i j -> b j i") + full_projection = view_matrix @ projection_matrix + + all_images = [] + all_radii = [] + for i in range(b): + # Set up a tensor for the gradients of the screen-space means. + mean_gradients = torch.zeros_like(gaussian_means[i], requires_grad=True) + try: + mean_gradients.retain_grad() + except Exception: + pass + + settings = GaussianRasterizationSettings( + image_height=h, + image_width=w, + tanfovx=tan_fov_x[i].item(), + tanfovy=tan_fov_y[i].item(), + bg=background_color[i], + scale_modifier=1.0, + viewmatrix=view_matrix[i], + projmatrix=full_projection[i], + sh_degree=degree, + campos=extrinsics[i, :3, 3], + prefiltered=False, # This matches the original usage. + debug=False, + ) + rasterizer = GaussianRasterizer(settings) + + row, col = torch.triu_indices(3, 3) + + image, radii = rasterizer( + means3D=gaussian_means[i], + means2D=mean_gradients, + shs=shs[i] if use_sh else None, + colors_precomp=None if use_sh else shs[i, :, 0, :], + opacities=gaussian_opacities[i, ..., None], + cov3D_precomp=gaussian_covariances[i, :, row, col], + ) + all_images.append(image) + all_radii.append(radii) + return torch.stack(all_images) + +DepthRenderingMode = Literal["depth", "disparity", "relative_disparity", "log"] + + +@torch.autocast('cuda', enabled=False) +def render_depth_cuda( + extrinsics: Float[Tensor, "batch 4 4"], + intrinsics: Float[Tensor, "batch 3 3"], + near: Float[Tensor, " batch"], + far: Float[Tensor, " batch"], + image_shape: tuple[int, int], + gaussian_means: Float[Tensor, "batch gaussian 3"], + gaussian_covariances: Float[Tensor, "batch gaussian 3 3"], + gaussian_opacities: Float[Tensor, "batch gaussian"], + scale_invariant: bool = True, + mode: DepthRenderingMode = "depth", +) -> Float[Tensor, "batch height width"]: + extrinsics = extrinsics.float() + intrinsics = intrinsics.float() + near = near.float() + far = far.float() + gaussian_means = gaussian_means.float() + gaussian_covariances = gaussian_covariances.float() + gaussian_opacities = gaussian_opacities.float() + + # Specify colors according to Gaussian depths. + camera_space_gaussians = einsum( + extrinsics.inverse(), homogenize_points(gaussian_means), "b i j, b g j -> b g i" + ) + fake_color = camera_space_gaussians[..., 2] + + if mode == "disparity": + fake_color = 1 / fake_color + elif mode == "log": + fake_color = fake_color.minimum(near[:, None]).maximum(far[:, None]).log() + + # Render using depth as color. + b, _ = fake_color.shape + result = render_cuda( + extrinsics, + intrinsics, + near, + far, + image_shape, + torch.zeros((b, 3), dtype=fake_color.dtype, device=fake_color.device), + gaussian_means, + gaussian_covariances, + repeat(fake_color, "b g -> b g c ()", c=3), + gaussian_opacities, + scale_invariant=scale_invariant, + use_sh=False, + ) + return result.mean(dim=1) diff --git a/src/model/decoder/decoder.py b/src/model/decoder/decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..ddc784b9b79cfb1f9ab6d08453e75c16abc9c2aa --- /dev/null +++ b/src/model/decoder/decoder.py @@ -0,0 +1,49 @@ +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Generic, Literal, TypeVar + +from jaxtyping import Float +from torch import Tensor, nn + +from ...dataset import DatasetCfg +from ..types import Gaussians + +DepthRenderingMode = Literal[ + "depth", + "log", + "disparity", + "relative_disparity", +] + + +@dataclass +class DecoderOutput: + color: Float[Tensor, "batch view 3 height width"] + depth: Float[Tensor, "batch view height width"] | None + accumulated_alpha: Float[Tensor, "batch view height width"] | None = None + + +T = TypeVar("T") + + +class Decoder(nn.Module, ABC, Generic[T]): + cfg: T + dataset_cfg: DatasetCfg + + def __init__(self, cfg: T, dataset_cfg: DatasetCfg) -> None: + super().__init__() + self.cfg = cfg + self.dataset_cfg = dataset_cfg + + @abstractmethod + def forward( + self, + gaussians: Gaussians, + extrinsics: Float[Tensor, "batch view 4 4"], + intrinsics: Float[Tensor, "batch view 3 3"], + near: Float[Tensor, "batch view"], + far: Float[Tensor, "batch view"], + image_shape: tuple[int, int], + depth_mode: DepthRenderingMode | None = None, + ) -> DecoderOutput: + pass diff --git a/src/model/decoder/decoder_splatting_cuda.py b/src/model/decoder/decoder_splatting_cuda.py new file mode 100644 index 0000000000000000000000000000000000000000..b77b29442bc8fdb53e9aa08def6341d6166cc8e8 --- /dev/null +++ b/src/model/decoder/decoder_splatting_cuda.py @@ -0,0 +1,91 @@ +from dataclasses import dataclass +from typing import Literal + +import torch +from einops import rearrange, repeat +from jaxtyping import Float +from torch import Tensor + +from ...dataset import DatasetCfg +from ..types import Gaussians +from .cuda_splatting import DepthRenderingMode, render_cuda, render_depth_cuda +from .decoder import Decoder, DecoderOutput + + +@dataclass +class DecoderSplattingCUDACfg: + name: Literal["splatting_cuda"] + + +class DecoderSplattingCUDA(Decoder[DecoderSplattingCUDACfg]): + background_color: Float[Tensor, "3"] + + def __init__( + self, + cfg: DecoderSplattingCUDACfg, + dataset_cfg: DatasetCfg, + ) -> None: + super().__init__(cfg, dataset_cfg) + self.register_buffer( + "background_color", + torch.tensor(dataset_cfg.background_color, dtype=torch.float32), + persistent=False, + ) + + def forward( + self, + gaussians: Gaussians, + extrinsics: Float[Tensor, "batch view 4 4"], + intrinsics: Float[Tensor, "batch view 3 3"], + near: Float[Tensor, "batch view"], + far: Float[Tensor, "batch view"], + image_shape: tuple[int, int], + depth_mode: DepthRenderingMode | None = None, + ) -> DecoderOutput: + b, v, _, _ = extrinsics.shape + color = render_cuda( + rearrange(extrinsics, "b v i j -> (b v) i j"), + rearrange(intrinsics, "b v i j -> (b v) i j"), + rearrange(near, "b v -> (b v)"), + rearrange(far, "b v -> (b v)"), + image_shape, + repeat(self.background_color, "c -> (b v) c", b=b, v=v), + repeat(gaussians.means, "b g xyz -> (b v) g xyz", v=v), + repeat(gaussians.covariances, "b g i j -> (b v) g i j", v=v), + repeat(gaussians.harmonics, "b g c d_sh -> (b v) g c d_sh", v=v), + repeat(gaussians.opacities, "b g -> (b v) g", v=v), + ) + color = rearrange(color, "(b v) c h w -> b v c h w", b=b, v=v) + + return DecoderOutput( + color, + None + if depth_mode is None + else self.render_depth( + gaussians, extrinsics, intrinsics, near, far, image_shape, depth_mode + ), + ) + + def render_depth( + self, + gaussians: Gaussians, + extrinsics: Float[Tensor, "batch view 4 4"], + intrinsics: Float[Tensor, "batch view 3 3"], + near: Float[Tensor, "batch view"], + far: Float[Tensor, "batch view"], + image_shape: tuple[int, int], + mode: DepthRenderingMode = "depth", + ) -> Float[Tensor, "batch view height width"]: + b, v, _, _ = extrinsics.shape + result = render_depth_cuda( + rearrange(extrinsics, "b v i j -> (b v) i j"), + rearrange(intrinsics, "b v i j -> (b v) i j"), + rearrange(near, "b v -> (b v)"), + rearrange(far, "b v -> (b v)"), + image_shape, + repeat(gaussians.means, "b g xyz -> (b v) g xyz", v=v), + repeat(gaussians.covariances, "b g i j -> (b v) g i j", v=v), + repeat(gaussians.opacities, "b g -> (b v) g", v=v), + mode=mode, + ) + return rearrange(result, "(b v) h w -> b v h w", b=b, v=v) diff --git a/src/model/decoder/decoder_splatting_gsplat.py b/src/model/decoder/decoder_splatting_gsplat.py new file mode 100644 index 0000000000000000000000000000000000000000..2110df0ff77892012ec0fd2276d0242fa3729951 --- /dev/null +++ b/src/model/decoder/decoder_splatting_gsplat.py @@ -0,0 +1,144 @@ +from dataclasses import dataclass +from typing import Literal + +import torch +from jaxtyping import Float +from torch import Tensor + +from ...dataset import DatasetCfg +from ..types import Gaussians +from .decoder import Decoder, DecoderOutput +from .gsplat_splatting import ( + GsplatRasterizeMode, + GsplatRenderMode, + render_depth_gsplat, + render_gsplat, +) + + +@dataclass +class DecoderSplattingGsplatCfg: + name: Literal["splatting_gsplat", "gsplat"] + render_mode: GsplatRenderMode = "RGB" + rasterize_mode: GsplatRasterizeMode = "antialiased" + use_covariances: bool = True + packed: bool = True + eps2d: float = 0.1 + sparse_grad: bool = False + absgrad: bool = False + channel_chunk: int = 32 + scale_invariant: bool = False + + +class DecoderSplattingGsplat(Decoder[DecoderSplattingGsplatCfg]): + background_color: Float[Tensor, "3"] + + def __init__( + self, + cfg: DecoderSplattingGsplatCfg, + dataset_cfg: DatasetCfg, + ) -> None: + super().__init__(cfg, dataset_cfg) + self.register_buffer( + "background_color", + torch.tensor(dataset_cfg.background_color, dtype=torch.float32), + persistent=False, + ) + + def forward( + self, + gaussians: Gaussians, + extrinsics: Float[Tensor, "batch view 4 4"], + intrinsics: Float[Tensor, "batch view 3 3"], + near: Float[Tensor, "batch view"], + far: Float[Tensor, "batch view"], + image_shape: tuple[int, int], + depth_mode=None, + ) -> DecoderOutput: + render_mode = self.cfg.render_mode + if depth_mode is not None: + render_mode = "RGB+ED" if depth_mode == "relative_disparity" else "RGB+D" + + color, depth, accumulated_alpha = render_gsplat( + extrinsics, + intrinsics, + near, + far, + image_shape, + self.background_color[None].expand(extrinsics.shape[0], -1), + gaussians.means, + gaussians.covariances, + gaussians.harmonics, + gaussians.opacities, + gaussian_scales=getattr(gaussians, "scales", None), + gaussian_rotations=getattr( + gaussians, + "rotations_unnorm", + getattr(gaussians, "rotations", None), + ), + render_mode=render_mode, + use_covariances=self.cfg.use_covariances, + eps2d=self.cfg.eps2d, + packed=self.cfg.packed, + sparse_grad=self.cfg.sparse_grad, + absgrad=self.cfg.absgrad, + rasterize_mode=self.cfg.rasterize_mode, + channel_chunk=self.cfg.channel_chunk, + scale_invariant=self.cfg.scale_invariant, + ) + assert color is not None + + if depth_mode in ("disparity", "relative_disparity"): + assert depth is not None + depth = 1.0 / depth.clamp_min(1e-8) + elif depth_mode == "log": + assert depth is not None + depth = depth.clamp_min(near[..., None, None]).clamp_max( + far[..., None, None] + ).log() + elif depth_mode is None and "D" not in render_mode: + depth = None + + return DecoderOutput( + color, + depth, + accumulated_alpha=accumulated_alpha, + ) + + def render_depth( + self, + gaussians: Gaussians, + extrinsics: Float[Tensor, "batch view 4 4"], + intrinsics: Float[Tensor, "batch view 3 3"], + near: Float[Tensor, "batch view"], + far: Float[Tensor, "batch view"], + image_shape: tuple[int, int], + mode="depth", + ) -> Float[Tensor, "batch view height width"]: + return render_depth_gsplat( + extrinsics, + intrinsics, + near, + far, + image_shape, + self.background_color[None].expand(extrinsics.shape[0], -1), + gaussians.means, + gaussians.covariances, + gaussians.harmonics, + gaussians.opacities, + gaussian_scales=getattr(gaussians, "scales", None), + gaussian_rotations=getattr( + gaussians, + "rotations_unnorm", + getattr(gaussians, "rotations", None), + ), + mode=mode, + use_covariances=self.cfg.use_covariances, + eps2d=self.cfg.eps2d, + packed=self.cfg.packed, + sparse_grad=self.cfg.sparse_grad, + absgrad=self.cfg.absgrad, + rasterize_mode=self.cfg.rasterize_mode, + channel_chunk=self.cfg.channel_chunk, + scale_invariant=self.cfg.scale_invariant, + ) diff --git a/src/model/decoder/gsplat_splatting.py b/src/model/decoder/gsplat_splatting.py new file mode 100644 index 0000000000000000000000000000000000000000..9d1bd19bf5429219c2c62f7ad1465924f8026349 --- /dev/null +++ b/src/model/decoder/gsplat_splatting.py @@ -0,0 +1,228 @@ +from math import isqrt +from typing import Literal + +import torch +from einops import rearrange, repeat +from jaxtyping import Float +from torch import Tensor + +from gsplat.rendering import rasterization + + +GsplatRenderMode = Literal["RGB", "RGB+D", "RGB+ED"] +GsplatRasterizeMode = Literal["classic", "antialiased"] + + +def normalized_intrinsics_to_pixel_K( + intrinsics: Float[Tensor, "batch view 3 3"], + image_shape: tuple[int, int], +) -> Float[Tensor, "batch view 3 3"]: + h, w = image_shape + K = intrinsics.clone().float() + K[..., 0, 0] *= w + K[..., 1, 1] *= h + K[..., 0, 2] *= w + K[..., 1, 2] *= h + return K + + +def _identity_quats_like( + means: Float[Tensor, "batch gaussian 3"], +) -> Float[Tensor, "batch gaussian 4"]: + quats = torch.zeros((*means.shape[:-1], 4), device=means.device, dtype=means.dtype) + # gsplat uses wxyz quaternions. + quats[..., 0] = 1.0 + return quats + + +def _xyzw_to_wxyz( + quats: Float[Tensor, "batch gaussian 4"], +) -> Float[Tensor, "batch gaussian 4"]: + return torch.cat((quats[..., 3:4], quats[..., :3]), dim=-1) + + +@torch.autocast("cuda", enabled=False) +def render_gsplat( + extrinsics: Float[Tensor, "batch view 4 4"], + intrinsics: Float[Tensor, "batch view 3 3"], + near: Float[Tensor, "batch view"], + far: Float[Tensor, "batch view"], + image_shape: tuple[int, int], + background_color: Float[Tensor, "batch 3"], + gaussian_means: Float[Tensor, "batch gaussian 3"], + gaussian_covariances: Float[Tensor, "batch gaussian 3 3"] | None, + gaussian_sh_coefficients: Float[Tensor, "batch gaussian 3 d_sh"], + gaussian_opacities: Float[Tensor, "batch gaussian"], + gaussian_scales: Float[Tensor, "batch gaussian 3"] | None = None, + gaussian_rotations: Float[Tensor, "batch gaussian 4"] | None = None, + render_mode: GsplatRenderMode = "RGB", + use_covariances: bool = True, + eps2d: float = 0.1, + packed: bool = True, + sparse_grad: bool = False, + absgrad: bool = False, + rasterize_mode: GsplatRasterizeMode = "antialiased", + channel_chunk: int = 32, + scale_invariant: bool = False, +) -> tuple[ + Float[Tensor, "batch view 3 height width"] | None, + Float[Tensor, "batch view height width"] | None, + Float[Tensor, "batch view height width"], +]: + b, v, _, _ = extrinsics.shape + h, w = image_shape + + extrinsics = extrinsics.float() + intrinsics = intrinsics.float() + near = near.float() + far = far.float() + background_color = background_color.float() + gaussian_means = gaussian_means.float() + gaussian_sh_coefficients = gaussian_sh_coefficients.float() + gaussian_opacities = gaussian_opacities.float() + + covars = gaussian_covariances.float() if gaussian_covariances is not None else None + scales = gaussian_scales.float() if gaussian_scales is not None else None + # The adapter stores quaternions as xyzw, while gsplat expects wxyz. + quats = ( + _xyzw_to_wxyz(gaussian_rotations.float()) + if gaussian_rotations is not None + else None + ) + + if use_covariances: + if covars is None: + raise ValueError("gsplat covariance mode requires gaussians.covariances.") + if scales is None: + scales = torch.ones_like(gaussian_means) + if quats is None: + quats = _identity_quats_like(gaussian_means) + else: + covars = None + if scales is None or quats is None: + raise ValueError( + "gsplat scale/quaternion mode requires gaussians.scales and " + "gaussians.rotations. Prefer use_covariances=true." + ) + + if scale_invariant: + scale = 1.0 / near[:, :1].clamp_min(1e-8) + extrinsics = extrinsics.clone() + extrinsics[..., :3, 3] = extrinsics[..., :3, 3] * scale[..., None] + gaussian_means = gaussian_means * scale[:, :, None] + near = near * scale + far = far * scale + if covars is not None: + covars = covars * (scale[:, :, None, None] ** 2) + elif scales is not None: + scales = scales * scale[:, :, None] + + Ks = normalized_intrinsics_to_pixel_K(intrinsics, image_shape) + viewmats = torch.linalg.inv(extrinsics) + sh_degree = isqrt(gaussian_sh_coefficients.shape[-1]) - 1 + shs = rearrange( + gaussian_sh_coefficients, "b g c n -> b g n c" + ).contiguous() + + render_rgb = "RGB" in render_mode + render_depth = "D" in render_mode + + # gsplat's packed path has limited background broadcasting across render + # modes. Our datasets use black backgrounds, so omitting the background keeps + # the memory-efficient packed path available without changing semantics. + backgrounds = ( + None + if packed + else repeat(background_color, "b c -> b v c", v=v).to(Ks) + ) + + render_out, alpha, _ = rasterization( + means=gaussian_means, + quats=quats, + scales=scales, + opacities=gaussian_opacities, + colors=shs, + viewmats=viewmats, + Ks=Ks, + width=w, + height=h, + near_plane=float(near.detach().min().item()), + far_plane=float(far.detach().max().item()), + eps2d=eps2d, + sh_degree=sh_degree, + packed=packed, + backgrounds=backgrounds, + render_mode=render_mode, + sparse_grad=sparse_grad, + absgrad=absgrad, + rasterize_mode=rasterize_mode, + channel_chunk=channel_chunk, + covars=covars, + ) + + color = ( + rearrange(render_out[..., :3], "b v h w c -> b v c h w") + if render_rgb + else None + ) + depth = render_out[..., -1] if render_depth else None + return color, depth, alpha.squeeze(-1) + + +def render_depth_gsplat( + extrinsics: Float[Tensor, "batch view 4 4"], + intrinsics: Float[Tensor, "batch view 3 3"], + near: Float[Tensor, "batch view"], + far: Float[Tensor, "batch view"], + image_shape: tuple[int, int], + background_color: Float[Tensor, "batch 3"], + gaussian_means: Float[Tensor, "batch gaussian 3"], + gaussian_covariances: Float[Tensor, "batch gaussian 3 3"] | None, + gaussian_sh_coefficients: Float[Tensor, "batch gaussian 3 d_sh"], + gaussian_opacities: Float[Tensor, "batch gaussian"], + gaussian_scales: Float[Tensor, "batch gaussian 3"] | None = None, + gaussian_rotations: Float[Tensor, "batch gaussian 4"] | None = None, + mode: Literal["depth", "disparity", "relative_disparity", "log"] = "depth", + use_covariances: bool = True, + eps2d: float = 0.1, + packed: bool = True, + sparse_grad: bool = False, + absgrad: bool = False, + rasterize_mode: GsplatRasterizeMode = "antialiased", + channel_chunk: int = 32, + scale_invariant: bool = False, +) -> Float[Tensor, "batch view height width"]: + render_mode = "RGB+ED" if mode == "relative_disparity" else "RGB+D" + _, depth, _ = render_gsplat( + extrinsics, + intrinsics, + near, + far, + image_shape, + background_color, + gaussian_means, + gaussian_covariances, + gaussian_sh_coefficients, + gaussian_opacities, + gaussian_scales=gaussian_scales, + gaussian_rotations=gaussian_rotations, + render_mode=render_mode, + use_covariances=use_covariances, + eps2d=eps2d, + packed=packed, + sparse_grad=sparse_grad, + absgrad=absgrad, + rasterize_mode=rasterize_mode, + channel_chunk=channel_chunk, + scale_invariant=scale_invariant, + ) + assert depth is not None + + if mode in ("disparity", "relative_disparity"): + depth = 1.0 / depth.clamp_min(1e-8) + elif mode == "log": + depth = depth.clamp_min(near[..., None, None]).clamp_max( + far[..., None, None] + ).log() + + return depth diff --git a/src/model/encoder/__init__.py b/src/model/encoder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3381972472295b55348a29d174c5f2d07bc2b962 --- /dev/null +++ b/src/model/encoder/__init__.py @@ -0,0 +1,12 @@ +from .encoder import Encoder +from .encoder_codecsplat import EncoderCodecSplat, EncoderCodecSplatCfg + +ENCODERS = { + "codecsplat": EncoderCodecSplat, +} + +EncoderCfg = EncoderCodecSplatCfg + + +def get_encoder(cfg: EncoderCfg) -> Encoder: + return ENCODERS[cfg.name](cfg) diff --git a/src/model/encoder/common/depth_guided_mv_refiner.py b/src/model/encoder/common/depth_guided_mv_refiner.py new file mode 100644 index 0000000000000000000000000000000000000000..7689f56414bc42853a0285dfd231e2cc6fcb26dd --- /dev/null +++ b/src/model/encoder/common/depth_guided_mv_refiner.py @@ -0,0 +1,266 @@ +import torch +import torch.nn.functional as F +from einops import rearrange, repeat +from torch import nn + +from ....geometry.projection import ( + homogenize_points, + project, + sample_image_grid, + transform_cam2world, + transform_world2cam, + unproject, +) + +class DepthGuidedMVRefiner(nn.Module): + def __init__( + self, + feature_channels: int, + hidden_channels: int | None = None, + concat_ref_depth: bool = False, + ) -> None: + super().__init__() + hidden_channels = feature_channels if hidden_channels is None else hidden_channels + depth_channels = max(hidden_channels // 4, 16) + self.concat_ref_depth = concat_ref_depth + + self.depth_encoder = nn.Sequential( + nn.Conv2d( + 1, + depth_channels, + 1, + 1, + 0, + ), + nn.GELU(), + nn.Conv2d( + depth_channels, + depth_channels, + 1, + 1, + 0, + ), + nn.GELU(), + ) + self.feature_encoder = nn.Sequential( + nn.Conv2d( + feature_channels * 2 + + depth_channels * (2 if self.concat_ref_depth else 1), + hidden_channels, + 3, + 1, + 1, + padding_mode="replicate", + ), + nn.GELU(), + nn.Conv2d( + hidden_channels, + hidden_channels, + 3, + 1, + 1, + padding_mode="replicate", + ), + nn.GELU(), + ) + self.post_fuse = nn.Sequential( + nn.Conv2d( + feature_channels + hidden_channels + depth_channels * int(self.concat_ref_depth), + hidden_channels, + 3, + 1, + 1, + padding_mode="replicate", + ), + nn.GELU(), + nn.Conv2d( + hidden_channels, + feature_channels, + 3, + 1, + 1, + padding_mode="replicate", + ), + ) + nn.init.zeros_(self.post_fuse[-1].weight) + nn.init.zeros_(self.post_fuse[-1].bias) + + @staticmethod + def _normalized_coords_to_grid( + xy: torch.Tensor, + h: int, + w: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + x_pixel = xy[..., 0] * w - 0.5 + y_pixel = xy[..., 1] * h - 0.5 + + if w > 1: + x_grid = 2 * x_pixel / (w - 1) - 1 + else: + x_grid = torch.zeros_like(x_pixel) + + if h > 1: + y_grid = 2 * y_pixel / (h - 1) - 1 + else: + y_grid = torch.zeros_like(y_pixel) + + grid = torch.stack((x_grid, y_grid), dim=-1) + in_bounds = ( + (x_pixel >= 0) + & (x_pixel <= (w - 1)) + & (y_pixel >= 0) + & (y_pixel <= (h - 1)) + ) + return grid, in_bounds + + def forward( + self, + features: torch.Tensor, + depth: torch.Tensor, + intrinsics: torch.Tensor, + extrinsics: torch.Tensor, + source_view_indices: torch.Tensor | None = None, + ) -> torch.Tensor: + b, v, c, h, w = features.shape + if v <= 1: + return features + + depth = depth.detach() + intrinsics = intrinsics.detach() + extrinsics = extrinsics.detach() + feature_dtype = features.dtype + + xy_ref, _ = sample_image_grid((h, w), device=features.device) + xy_ref = repeat(xy_ref, "h w xy -> b h w xy", b=b) + + refined_features = [] + geom_eps = torch.finfo(torch.float32).eps + count_eps = torch.finfo(features.dtype).eps + batch_indices = torch.arange(b, device=features.device) + if source_view_indices is not None: + if source_view_indices.shape[:2] != (b, v): + raise ValueError( + "source_view_indices must have shape [B, V, M], got " + f"{tuple(source_view_indices.shape)} for features " + f"{tuple(features.shape)}" + ) + source_view_indices = source_view_indices.to( + device=features.device, + dtype=torch.long, + ) + + for ref_idx in range(v): + ref_features = features[:, ref_idx] + ref_depth = depth[:, ref_idx].float() + ref_intrinsics = intrinsics[:, ref_idx].float() + ref_extrinsics = extrinsics[:, ref_idx].float() + ref_depth_repr = None + if self.concat_ref_depth: + ref_depth_input = ref_depth.unsqueeze(1).clamp_min(geom_eps).reciprocal() + ref_depth_repr = self.depth_encoder(ref_depth_input.to(feature_dtype)) + + with torch.no_grad(), torch.autocast( + device_type=features.device.type, enabled=False + ): + ref_points_cam = unproject( + xy_ref, + ref_depth, + ref_intrinsics[:, None, None], + ) + ref_points_world = transform_cam2world( + homogenize_points(ref_points_cam), + ref_extrinsics[:, None, None], + )[..., :-1] + + fused_sum = None + valid_count = None + + if source_view_indices is None: + selected_src_indices = [ + src_idx for src_idx in range(v) if src_idx != ref_idx + ] + else: + selected_src_indices = source_view_indices[:, ref_idx].unbind(dim=-1) + + for src_idx in selected_src_indices: + if isinstance(src_idx, int): + src_features = features[:, src_idx] + src_extrinsics = extrinsics[:, src_idx].float()[:, None, None] + src_intrinsics = intrinsics[:, src_idx].float()[:, None, None] + else: + src_features = features[batch_indices, src_idx] + src_extrinsics = ( + extrinsics[batch_indices, src_idx].float()[:, None, None] + ) + src_intrinsics = ( + intrinsics[batch_indices, src_idx].float()[:, None, None] + ) + + with torch.no_grad(), torch.autocast( + device_type=features.device.type, enabled=False + ): + src_xy, src_in_front = project( + ref_points_world, + src_extrinsics, + src_intrinsics, + ) + src_grid, src_in_bounds = self._normalized_coords_to_grid( + src_xy, h, w + ) + valid_mask_fp32 = (src_in_front & src_in_bounds).unsqueeze(1).float() + + with torch.autocast(device_type=features.device.type, enabled=False): + warped_features = F.grid_sample( + src_features.float(), + src_grid, + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) + valid_mask = valid_mask_fp32.to(feature_dtype) + warped_features = warped_features.to(feature_dtype) + warped_features = warped_features * valid_mask + + with torch.no_grad(), torch.autocast( + device_type=features.device.type, enabled=False + ): + src_points_cam = transform_world2cam( + homogenize_points(ref_points_world), + src_extrinsics, + )[..., :-1] + warped_depth = rearrange( + src_points_cam[..., -1], "b h w -> b 1 h w" + ) + warped_depth_input = torch.where( + valid_mask_fp32 > 0, + warped_depth.clamp_min(geom_eps).reciprocal(), + torch.zeros_like(warped_depth), + ) + + depth_repr = self.depth_encoder(warped_depth_input.to(feature_dtype)) + feature_inputs = [ref_features, warped_features] + if ref_depth_repr is not None: + feature_inputs.append(ref_depth_repr) + feature_inputs.append(depth_repr) + value = self.feature_encoder(torch.cat(feature_inputs, dim=1)) * valid_mask + + if fused_sum is None: + fused_sum = value + valid_count = valid_mask + else: + fused_sum = fused_sum + value + valid_count = valid_count + valid_mask + + if fused_sum is None or valid_count is None: + refined_features.append(ref_features) + continue + + fused = fused_sum / valid_count.clamp_min(count_eps) + + post_fuse_inputs = [ref_features, fused] + if ref_depth_repr is not None: + post_fuse_inputs.append(ref_depth_repr) + residual = self.post_fuse(torch.cat(post_fuse_inputs, dim=1)) + refined_features.append(ref_features + residual) + + return torch.stack(refined_features, dim=1) diff --git a/src/model/encoder/common/gaussian_adapter.py b/src/model/encoder/common/gaussian_adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..eb132614d0a34eeb42cf0189a40c2eb7479df223 --- /dev/null +++ b/src/model/encoder/common/gaussian_adapter.py @@ -0,0 +1,144 @@ +from dataclasses import dataclass + +import torch +from einops import rearrange +from jaxtyping import Float +from torch import Tensor, nn +import torch.nn.functional as F + +from ....geometry.projection import get_world_rays +from ....misc.sh_rotation import rotate_sh +from .gaussians import build_covariance + + +@dataclass +class Gaussians: + means: Float[Tensor, "*batch 3"] + covariances: Float[Tensor, "*batch 3 3"] + scales: Float[Tensor, "*batch 3"] + rotations: Float[Tensor, "*batch 4"] + harmonics: Float[Tensor, "*batch 3 _"] + opacities: Float[Tensor, " *batch"] + rotations_unnorm: Float[Tensor, "*batch 4"] + + +@dataclass +class GaussianAdapterCfg: + gaussian_scale_min: float + gaussian_scale_max: float + sh_degree: int + exp_scale_bias: float + no_rotate_sh: bool + no_sh_mask: bool + + +class GaussianAdapter(nn.Module): + cfg: GaussianAdapterCfg + + def __init__(self, cfg: GaussianAdapterCfg): + super().__init__() + self.cfg = cfg + + # Create a mask for the spherical harmonics coefficients. This ensures that at + # initialization, the coefficients are biased towards having a large DC + # component and small view-dependent components. + self.register_buffer( + "sh_mask", + torch.ones((self.d_sh,), dtype=torch.float32), + persistent=False, + ) + for degree in range(1, self.cfg.sh_degree + 1): + self.sh_mask[degree**2 : (degree + 1) ** 2] = 0.1 * 0.25**degree + + def forward( + self, + extrinsics: Float[Tensor, "*#batch 4 4"], + intrinsics: Float[Tensor, "*#batch 3 3"] | None, + coordinates: Float[Tensor, "*#batch 2"], + depths: Float[Tensor, "*#batch"] | None, + opacities: Float[Tensor, "*#batch"], + raw_gaussians: Float[Tensor, "*#batch _"], + image_shape: tuple[int, int], + eps: float = 1e-8, + point_cloud: Float[Tensor, "*#batch 3"] | None = None, + input_images: Tensor | None = None, + ) -> Gaussians: + # Gaussian geometry is consumed by CUDA/gsplat rasterizers, so keep the + # position, covariance, and SH rotation math in FP32 under mixed precision. + with torch.autocast(device_type=raw_gaussians.device.type, enabled=False): + extrinsics = extrinsics.float() + intrinsics = intrinsics.float() if intrinsics is not None else None + coordinates = coordinates.float() + depths = depths.float() if depths is not None else None + opacities = opacities.float() + raw_gaussians = raw_gaussians.float() + + scales, rotations_unnorm, sh = raw_gaussians.split( + (3, 4, 3 * self.d_sh), dim=-1 + ) + + scales = torch.clamp( + F.softplus(scales - self.cfg.exp_scale_bias), + min=self.cfg.gaussian_scale_min, + max=self.cfg.gaussian_scale_max, + ) + + # Normalize the quaternion features to yield a valid quaternion. + rotations = rotations_unnorm / ( + rotations_unnorm.norm(dim=-1, keepdim=True) + eps + ) + + # Split RGB spherical-harmonic coefficients into the expected layout. + sh = rearrange(sh, "... (xyz d_sh) -> ... xyz d_sh", xyz=3) + sh = sh.broadcast_to((*opacities.shape, 3, self.d_sh)) + if self.cfg.no_sh_mask: + sh = sh.clone() + else: + sh = sh * self.sh_mask + + if input_images is not None: + # [B, V, H*W, 1, 1, 3] + imgs = rearrange( + input_images.float(), "b v c h w -> b v (h w) () () c" + ) + # init sh with input images + sh[..., 0] = sh[..., 0] + RGB2SH(imgs) + + # Create world-space covariance matrices. + covariances = build_covariance(scales, rotations) + c2w_rotations = extrinsics[..., :3, :3] + covariances = ( + c2w_rotations @ covariances @ c2w_rotations.transpose(-1, -2) + ) + + # Compute Gaussian means. + origins, directions = get_world_rays(coordinates, extrinsics, intrinsics) + means = origins + directions * depths[..., None] + harmonics = ( + sh + if self.cfg.no_rotate_sh + else rotate_sh(sh, c2w_rotations[..., None, :, :]) + ) + + return Gaussians( + means=means, + covariances=covariances, + harmonics=harmonics, + opacities=opacities, + scales=scales, + rotations=rotations.broadcast_to((*scales.shape[:-1], 4)), + rotations_unnorm=rotations_unnorm.broadcast_to((*scales.shape[:-1], 4)), + ) + + @property + def d_sh(self) -> int: + return (self.cfg.sh_degree + 1) ** 2 + + @property + def d_in(self) -> int: + return 7 + 3 * self.d_sh + + +def RGB2SH(rgb): + C0 = 0.28209479177387814 + return (rgb - 0.5) / C0 diff --git a/src/model/encoder/common/gaussians.py b/src/model/encoder/common/gaussians.py new file mode 100644 index 0000000000000000000000000000000000000000..1bff0519fa2d8c7e93d66ede9da9e5166e554a5e --- /dev/null +++ b/src/model/encoder/common/gaussians.py @@ -0,0 +1,44 @@ +import torch +from einops import rearrange +from jaxtyping import Float +from torch import Tensor + + +# https://github.com/facebookresearch/pytorch3d/blob/main/pytorch3d/transforms/rotation_conversions.py +def quaternion_to_matrix( + quaternions: Float[Tensor, "*batch 4"], + eps: float = 1e-8, +) -> Float[Tensor, "*batch 3 3"]: + # Order changed to match scipy format! + i, j, k, r = torch.unbind(quaternions, dim=-1) + two_s = 2 / ((quaternions * quaternions).sum(dim=-1) + eps) + + o = torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ), + -1, + ) + return rearrange(o, "... (i j) -> ... i j", i=3, j=3) + + +def build_covariance( + scale: Float[Tensor, "*#batch 3"], + rotation_xyzw: Float[Tensor, "*#batch 4"], +) -> Float[Tensor, "*batch 3 3"]: + scale = scale.diag_embed() + rotation = quaternion_to_matrix(rotation_xyzw) + return ( + rotation + @ scale + @ rearrange(scale, "... i j -> ... j i") + @ rearrange(rotation, "... i j -> ... j i") + ) diff --git a/src/model/encoder/common/gs_decoder.py b/src/model/encoder/common/gs_decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..118eb273d178fe827cc9dcedfef809b888e1ceb6 --- /dev/null +++ b/src/model/encoder/common/gs_decoder.py @@ -0,0 +1,103 @@ +import torch +from torch import nn + + +class GaussianLatentDecoder(nn.Module): + def __init__( + self, + latent_channels: int, + hidden_channels: int, + num_gaussian_parameters: int, + init_sh_input_img: bool = True, + ) -> None: + super().__init__() + in_channels = latent_channels + 1 # out_hat + depth + self.decoder = nn.Sequential( + nn.Conv2d(in_channels, hidden_channels, 3, 1, 1, padding_mode="replicate"), + nn.GELU(), + nn.Conv2d( + hidden_channels, + num_gaussian_parameters, + 3, + 1, + 1, + padding_mode="replicate", + ), + nn.GELU(), + nn.Conv2d( + num_gaussian_parameters, + num_gaussian_parameters, + 3, + 1, + 1, + padding_mode="replicate", + ), + ) + self._init_output_layer(init_sh_input_img) + + def _init_output_layer(self, init_sh_input_img: bool) -> None: + # Keep the Gaussian head initially close to identity-like residuals. + # Channel order: opacity, offset_xy, scales, quaternion, SH residuals. + output_layer = self.decoder[-1] + if init_sh_input_img: + nn.init.zeros_(output_layer.weight[10:]) + nn.init.zeros_(output_layer.bias[10:]) + + nn.init.zeros_(output_layer.weight[3:6]) + nn.init.zeros_(output_layer.bias[3:6]) + + def forward( + self, + latent: torch.Tensor, + depth: torch.Tensor, + ) -> torch.Tensor: + x = torch.cat((latent, depth), dim=1) + return self.decoder(x) + + +class GaussianLatentDepthDecoder(nn.Module): + def __init__( + self, + latent_channels: int, + hidden_channels: int, + ) -> None: + super().__init__() + self.inv_depth_logits = nn.Sequential( + nn.Conv2d(latent_channels, hidden_channels, 3, 1, 1, padding_mode="replicate"), + nn.GELU(), + nn.Conv2d(hidden_channels, hidden_channels, 3, 1, 1, padding_mode="replicate"), + nn.GELU(), + nn.Conv2d(hidden_channels, 1, 3, 1, 1, padding_mode="replicate"), + ) + # Predict the midpoint of the valid inverse-depth interval at initialization. + nn.init.zeros_(self.inv_depth_logits[-1].weight) + nn.init.zeros_(self.inv_depth_logits[-1].bias) + + def forward( + self, + latent: torch.Tensor, + near: torch.Tensor, + far: torch.Tensor, + return_inverse: bool = False, + return_normalized_inverse: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: + with torch.autocast(device_type=latent.device.type, enabled=False): + inv_depth_logits = self.inv_depth_logits(latent.float()) + near = near.float().clamp_min(torch.finfo(torch.float32).eps) + far = far.float().clamp_min(torch.finfo(torch.float32).eps) + inv_min = far.reciprocal() + inv_max = near.reciprocal() + inv_range = (inv_max - inv_min).clamp_min(torch.finfo(torch.float32).eps) + + normalized_inv_depth = torch.sigmoid(inv_depth_logits.float()) + inv_depth = inv_min + normalized_inv_depth * inv_range + depth = inv_depth.clamp_min(torch.finfo(torch.float32).eps).reciprocal() + + outputs: list[torch.Tensor] = [depth] + if return_inverse: + outputs.append(inv_depth) + if return_normalized_inverse: + outputs.append(normalized_inv_depth) + if len(outputs) == 1: + return outputs[0] + return tuple(outputs) diff --git a/src/model/encoder/encoder.py b/src/model/encoder/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..c22402be65b42f835b6f8ea27d253d541b645b23 --- /dev/null +++ b/src/model/encoder/encoder.py @@ -0,0 +1,29 @@ +from abc import ABC, abstractmethod +from typing import Generic, TypeVar + +from torch import nn + +from ...dataset.types import BatchedViews, DataShim +from ..types import Gaussians + +T = TypeVar("T") + + +class Encoder(nn.Module, ABC, Generic[T]): + cfg: T + + def __init__(self, cfg: T) -> None: + super().__init__() + self.cfg = cfg + + @abstractmethod + def forward( + self, + context: BatchedViews, + global_step: int, + ) -> Gaussians: + pass + + def get_data_shim(self) -> DataShim: + """The default shim doesn't modify the batch.""" + return lambda x: x diff --git a/src/model/encoder/encoder_codecsplat.py b/src/model/encoder/encoder_codecsplat.py new file mode 100644 index 0000000000000000000000000000000000000000..598fd476c8417981de4243531e1bc192d8b513de --- /dev/null +++ b/src/model/encoder/encoder_codecsplat.py @@ -0,0 +1,954 @@ +import io +import time +import zipfile +import zlib +from dataclasses import dataclass, field +from typing import Literal, Optional, List + +import numpy as np +import torch +from einops import rearrange, repeat +from torch import nn + +from ...dataset.shims.patch_shim import apply_patch_shim +from ...dataset.types import BatchedExample, DataShim +from ...geometry.projection import sample_image_grid +from ..types import Gaussians +from .common.depth_guided_mv_refiner import DepthGuidedMVRefiner +from .common.gaussian_adapter import GaussianAdapter, GaussianAdapterCfg +from .common.gs_decoder import GaussianLatentDecoder, GaussianLatentDepthDecoder +from .encoder import Encoder +from .vae import coding +from .vae import VariableRateLossyVAE +from .vae.config import Config as VaeConfig + +import torch.nn.functional as F +from numpy.lib.format import write_array + +from .unimatch.mv_unimatch import MultiViewUniMatch +from .unimatch.dpt_head import DPTHead + + +@dataclass +class EncoderCodecSplatCfg: + name: Literal["codecsplat"] + num_depth_candidates: int + num_surfaces: int + gaussian_adapter: GaussianAdapterCfg + downscale_factor: int + shim_patch_size: int + + # mv_unimatch + num_scales: int + upsample_factor: int + lowest_feature_resolution: int + depth_unet_channels: int + grid_sample_disable_cudnn: bool + + # color branch + init_sh_input_img: bool + gaussian_regressor_channels: int + + # loss config + supervise_intermediate_depth: bool + return_depth: bool + + # monodepth config + monodepth_vit_type: str + + # multi-view matching + local_mv_match: int + + # feature codec path + enable_feature_codec: bool = False + enable_depth_guided_mv_refiner: bool = False + enable_post_vae_mv_refiner: bool = False + codec_loss_weight: float = 1.0 + codec_depth_loss_weight: float = 0.1 + codec_lmb_range: List[float] = field(default_factory=lambda: [16.0, 1024.0]) + freeze_legacy_modules_for_codec: bool = False + codec_estimated_feature_num_bytes: bool = False + + # Test-only ablation: entropy-code the stage-1 latent feature/depth tensors + # with plain deflate instead of using the learned VAE codec. + deflate_feature_ablation: bool = False + deflate_feature_dtype: Literal["float16", "float32"] = "float16" + deflate_feature_level: int = 9 + + +class EncoderCodecSplat(Encoder[EncoderCodecSplatCfg]): + def __init__(self, cfg: EncoderCodecSplatCfg) -> None: + super().__init__(cfg) + if cfg.enable_feature_codec: + assert not cfg.init_sh_input_img, ( + "Feature codec mode requires init_sh_input_img=False because " + "decoder-side reconstruction cannot access original input images." + ) + if cfg.deflate_feature_ablation: + assert not cfg.enable_feature_codec, ( + "deflate_feature_ablation is a non-learned baseline and should be " + "run with enable_feature_codec=False." + ) + + self.depth_predictor = MultiViewUniMatch( + num_scales=cfg.num_scales, + upsample_factor=cfg.upsample_factor, + lowest_feature_resolution=cfg.lowest_feature_resolution, + num_depth_candidates=cfg.num_depth_candidates, + vit_type=cfg.monodepth_vit_type, + unet_channels=cfg.depth_unet_channels, + grid_sample_disable_cudnn=cfg.grid_sample_disable_cudnn, + ) + + self.mv_refiner = None + self.post_mv_refiner = None + self.feature_codec = None + self.gs_decoder = None + self.codec_depth_head = None + + # upsample features to the original resolution + model_configs = { + 'vits': {'in_channels': 384, 'features': 64, 'out_channels': [48, 96, 192, 384]}, + 'vitb': {'in_channels': 768, 'features': 96, 'out_channels': [96, 192, 384, 768]}, + 'vitl': {'in_channels': 1024, 'features': 128, 'out_channels': [128, 256, 512, 1024]}, + } + + self.feature_upsampler = DPTHead(**model_configs[cfg.monodepth_vit_type], + downsample_factor=cfg.upsample_factor, + return_feature=True, + num_scales=cfg.num_scales, + ) + feature_upsampler_channels = model_configs[cfg.monodepth_vit_type]["features"] + + # gaussians adapter + self.gaussian_adapter = GaussianAdapter(cfg.gaussian_adapter) + + # If a pre-refiner is enabled, let it handle explicit depth conditioning and keep + # the gaussian_regressor input focused on image / match / feature cues. + self.gaussian_regressor_uses_depth = not cfg.enable_depth_guided_mv_refiner + gaussian_regressor_in_ch = 3 + 1 + feature_upsampler_channels + if self.gaussian_regressor_uses_depth: + gaussian_regressor_in_ch += 1 + channels = self.cfg.gaussian_regressor_channels + + self.gaussian_regressor = nn.Sequential( + nn.Conv2d(gaussian_regressor_in_ch, channels, 3, 1, 1), + nn.GELU(), + nn.Conv2d(channels, channels, 3, 1, 1), + ) + + # predict gaussian parameters: scale, q, sh, offset, opacity + num_gaussian_parameters = self.gaussian_adapter.d_in + 2 + 1 + self.num_gaussian_parameters = num_gaussian_parameters + self.gaussian_regressor_out_channels = channels + + if self.cfg.enable_depth_guided_mv_refiner: + self.mv_refiner = DepthGuidedMVRefiner( + feature_channels=channels, + hidden_channels=channels, + concat_ref_depth=True, + ) + if self.cfg.enable_feature_codec: + codec_lmb_range = tuple(float(x) for x in self.cfg.codec_lmb_range) + self.feature_codec = VariableRateLossyVAE( + channels, + VaeConfig( + lmb_range=codec_lmb_range, + lmb=float(codec_lmb_range[-1]), + ), + ) + self.codec_depth_head = GaussianLatentDepthDecoder( + latent_channels=channels, + hidden_channels=channels, + ) + if self.cfg.enable_post_vae_mv_refiner: + self.post_mv_refiner = DepthGuidedMVRefiner( + feature_channels=channels, + hidden_channels=channels, + ) + self.gs_decoder = GaussianLatentDecoder( + latent_channels=channels, + hidden_channels=channels, + num_gaussian_parameters=num_gaussian_parameters, + init_sh_input_img=self.cfg.init_sh_input_img, + ) + if self.cfg.enable_feature_codec: + if self.cfg.freeze_legacy_modules_for_codec: + self.freeze_legacy_modules_for_codec() + + def freeze_legacy_modules_for_codec(self) -> None: + for param in self.parameters(): + param.requires_grad_(False) + + modules_to_train = [ + self.gaussian_regressor, + self.mv_refiner, + self.feature_codec, + self.codec_depth_head, + self.post_mv_refiner, + self.gs_decoder, + ] + + for module in modules_to_train: + if module is None: + continue + for param in module.parameters(): + param.requires_grad_(True) + + def forward( + self, + context: dict, + global_step: int, + visualization_dump: Optional[dict] = None, + scene_names: Optional[list] = None, + return_bitstream: bool = False, + ): + if self.training: + return self.normal_forward( + context, + global_step, + visualization_dump=visualization_dump, + scene_names=scene_names, + ) + if self.cfg.deflate_feature_ablation: + return self.deflate_feature_forward( + context, + global_step, + visualization_dump=visualization_dump, + scene_names=scene_names, + ) + if self.feature_codec is not None: + return self.codec_forward( + context, + global_step, + visualization_dump=visualization_dump, + scene_names=scene_names, + return_bitstream=return_bitstream, + ) + return self.normal_forward( + context, + global_step, + visualization_dump=visualization_dump, + scene_names=scene_names, + ) + + def compute_view_neighbor_indices(self, extrinsics: torch.Tensor): + v = extrinsics.shape[1] + if v <= 3: + return None + + with torch.no_grad(): + xyzs = extrinsics[:, :, :3, -1].detach().float() + cameras_dist_matrix = torch.cdist(xyzs, xyzs, p=2) + cameras_dist_index = torch.argsort(cameras_dist_matrix) + return cameras_dist_index[:, :, :(self.cfg.local_mv_match + 1)] + + def encode_camera_stream(self, context: dict) -> bytes: + _, v, _, h, w = context["image"].shape + camera_arrays = { + "image_shape": np.array([h, w], dtype=np.int32), + "num_context_views": np.array([v], dtype=np.int32), + "intrinsics": context["intrinsics"][0].detach().cpu().float().numpy(), + "extrinsics": context["extrinsics"][0].detach().cpu().float().numpy(), + "near": context["near"][0].detach().cpu().float().numpy(), + "far": context["far"][0].detach().cpu().float().numpy(), + } + camera_buffer = io.BytesIO() + with zipfile.ZipFile( + camera_buffer, + mode="w", + compression=zipfile.ZIP_DEFLATED, + compresslevel=9, + ) as camera_zip: + for name, array in camera_arrays.items(): + array_buffer = io.BytesIO() + write_array(array_buffer, np.asanyarray(array), allow_pickle=False) + camera_zip.writestr(f"{name}.npy", array_buffer.getvalue()) + camera_stream = camera_buffer.getvalue() + camera_buffer.close() + return camera_stream + + @torch.no_grad() + def compress(self, context: dict, features: torch.Tensor) -> tuple[bytes, dict]: + assert self.feature_codec is not None + b, v, _, h, w = context["image"].shape + assert b == 1, "Scene bitstream export currently supports batch size 1." + assert features.shape[0] == v, ( + f"Expected one feature tensor per view, got {features.shape[0]=}, {v=}." + ) + assert v > 1, "Scene feature stream packing expects at least two context views." + + if features.is_cuda: + torch.cuda.synchronize(features.device) + start = time.perf_counter() + + camera_stream = self.encode_camera_stream(context) + feature_streams = [ + self.feature_codec.compress(features[i : i + 1]) + for i in range(features.shape[0]) + ] + feature_stream = coding.concat_bytes_list(feature_streams) + + camera_num_bytes = len(camera_stream) + assert camera_num_bytes < (1 << 24), ( + f"Camera stream is too large for the 3-byte length header: {camera_num_bytes}" + ) + bytestream = ( + camera_num_bytes.to_bytes(3, byteorder="big", signed=False) + + camera_stream + + feature_stream + ) + + if features.is_cuda: + torch.cuda.synchronize(features.device) + compress_seconds = time.perf_counter() - start + + return bytestream, { + "codec_compress_seconds": compress_seconds, + "codec_camera_num_bytes": float(camera_num_bytes), + "codec_feature_num_bytes": float(len(feature_stream)), + "codec_num_bytes": float(len(bytestream)), + } + + @torch.no_grad() + def run_deflate_feature_ablation( + self, + context: dict, + features: torch.Tensor, + depth: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, dict]: + assert not self.training, "deflate_feature_ablation is intended for evaluation." + + device = features.device + if features.is_cuda: + torch.cuda.synchronize(device) + compress_start = time.perf_counter() + + dtype = self.cfg.deflate_feature_dtype + np_dtype = np.float16 if dtype == "float16" else np.float32 + level = int(self.cfg.deflate_feature_level) + + def tensor_to_npy_bytes(tensor: torch.Tensor) -> bytes: + tensor_np = tensor.detach().cpu().to(torch.float32).numpy().astype( + np_dtype, + copy=False, + ) + tensor_buffer = io.BytesIO() + write_array(tensor_buffer, np.asanyarray(tensor_np), allow_pickle=False) + stream = tensor_buffer.getvalue() + tensor_buffer.close() + return stream + + camera_stream = self.encode_camera_stream(context) + feature_raw_stream = tensor_to_npy_bytes(features) + depth_raw_stream = tensor_to_npy_bytes(depth) + feature_stream = zlib.compress(feature_raw_stream, level=level) + depth_stream = zlib.compress(depth_raw_stream, level=level) + + if features.is_cuda: + torch.cuda.synchronize(device) + compress_seconds = time.perf_counter() - compress_start + + if features.is_cuda: + torch.cuda.synchronize(device) + decompress_start = time.perf_counter() + + features_hat = torch.from_numpy( + np.load(io.BytesIO(zlib.decompress(feature_stream)), allow_pickle=False) + ).to(device=device, dtype=features.dtype) + depth_hat = torch.from_numpy( + np.load(io.BytesIO(zlib.decompress(depth_stream)), allow_pickle=False) + ).to(device=device, dtype=depth.dtype) + + if features.is_cuda: + torch.cuda.synchronize(device) + decompress_seconds = time.perf_counter() - decompress_start + + header_num_bytes = 24 # three uint64 lengths for camera/feature/depth streams + codec_num_bytes = ( + header_num_bytes + + len(camera_stream) + + len(feature_stream) + + len(depth_stream) + ) + return features_hat, depth_hat, { + "codec_compress_seconds": compress_seconds, + "codec_decompress_seconds": decompress_seconds, + "codec_camera_num_bytes": float(len(camera_stream)), + "codec_feature_num_bytes": float(len(feature_stream) + len(depth_stream)), + "codec_latent_feature_num_bytes": float(len(feature_stream)), + "codec_depth_num_bytes": float(len(depth_stream)), + "codec_feature_raw_num_bytes": float(len(feature_raw_stream)), + "codec_depth_raw_num_bytes": float(len(depth_raw_stream)), + "codec_num_bytes": float(codec_num_bytes), + } + + @torch.no_grad() + def decompress( + self, + bytestream: bytes, + visualization_dump: Optional[dict] = None, + ) -> tuple[Gaussians, Optional[torch.Tensor], dict, dict]: + assert self.feature_codec is not None + param = next(self.parameters()) + device = param.device + + if param.is_cuda: + torch.cuda.synchronize(device) + start = time.perf_counter() + + camera_num_bytes = int.from_bytes( + bytestream[:3], byteorder="big", signed=False + ) + camera_stream = bytestream[3 : 3 + camera_num_bytes] + feature_stream = bytestream[3 + camera_num_bytes :] + + with np.load(io.BytesIO(camera_stream)) as camera_data: + image_shape = tuple(int(x) for x in camera_data["image_shape"]) + num_context_views = int(camera_data["num_context_views"][0]) + intrinsics_np = camera_data["intrinsics"] + extrinsics_np = camera_data["extrinsics"] + near_np = camera_data["near"] + far_np = camera_data["far"] + + feature_streams = coding.split_bytes_list(feature_stream, num_context_views) + features_hat = torch.cat( + [self.feature_codec.decompress(stream) for stream in feature_streams], + dim=0, + ) + + decoded_context = {} + decoded_context["image_shape"] = image_shape + decoded_context["intrinsics"] = torch.from_numpy(intrinsics_np).to( + device=device, + dtype=torch.float32, + ).unsqueeze(0) + decoded_context["extrinsics"] = torch.from_numpy(extrinsics_np).to( + device=device, + dtype=torch.float32, + ).unsqueeze(0) + decoded_context["near"] = torch.from_numpy(near_np).to( + device=device, + dtype=torch.float32, + ).unsqueeze(0) + decoded_context["far"] = torch.from_numpy(far_np).to( + device=device, + dtype=torch.float32, + ).unsqueeze(0) + + neighbor_indices = self.compute_view_neighbor_indices( + decoded_context["extrinsics"] + ) + decode_state = { + "device": device, + "b": 1, + "v": num_context_views, + "h": image_shape[0], + "w": image_shape[1], + "refiner_source_indices": ( + neighbor_indices[..., 1:] if neighbor_indices is not None else None + ), + } + gaussians_raw, depth, _, _ = self.decode_codec_features( + features_hat, + decoded_context, + decode_state, + codec_depth_targets=None, + ) + decode_state["depth_preds"] = [depth] + gaussians, depths = self.adapt_gaussians( + decoded_context, + decoded_context, + gaussians_raw, + depth, + decode_state, + visualization_dump, + ) + + if param.is_cuda: + torch.cuda.synchronize(device) + decompress_seconds = time.perf_counter() - start + + return gaussians, depths, decoded_context, { + "codec_decompress_seconds": decompress_seconds, + } + + def encode_context_features(self, context: dict): + device = context["image"].device + b, v, _, h, w = context["image"].shape + + cameras_dist_index = self.compute_view_neighbor_indices(context["extrinsics"]) + refiner_source_indices = ( + cameras_dist_index[..., 1:] if cameras_dist_index is not None else None + ) + + results_dict = self.depth_predictor( + context["image"], + attn_splits_list=[2], + min_depth=1. / context["far"], + max_depth=1. / context["near"], + intrinsics=context["intrinsics"], + extrinsics=context["extrinsics"], + nn_matrix=cameras_dist_index, + supervise_intermediate_depth=self.cfg.supervise_intermediate_depth, + ) + + depth_preds = results_dict['depth_preds'] + depth = depth_preds[-1] + + features = self.feature_upsampler( + results_dict["features_mono_intermediate"], + cnn_features=results_dict["features_cnn_all_scales"][::-1], + mv_features=results_dict["features_mv"][0] + if self.cfg.num_scales == 1 + else results_dict["features_mv"][::-1], + ) + + match_prob = results_dict['match_probs'][-1] + match_prob = torch.max(match_prob, dim=1, keepdim=True)[0] + match_prob = F.interpolate(match_prob, size=depth.shape[-2:], mode='nearest') + + regressor_inputs = [ + rearrange(context["image"], "b v c h w -> (b v) c h w"), + match_prob, + features, + ] + if self.mv_refiner is None: + regressor_inputs.insert(1, rearrange(depth, "b v h w -> (b v) () h w")) + + out = self.gaussian_regressor(torch.cat(regressor_inputs, dim=1)) + + if self.mv_refiner is not None: + out = rearrange(out, "(b v) c h w -> b v c h w", b=b, v=v) + out = self.mv_refiner( + out, + depth, + context["intrinsics"], + context["extrinsics"], + source_view_indices=refiner_source_indices, + ) + out = rearrange(out, "b v c h w -> (b v) c h w") + + return { + "device": device, + "b": b, + "v": v, + "h": h, + "w": w, + "depth_preds": depth_preds, + "depth": depth, + "teacher_depth": depth, + "features": features, + "match_prob": match_prob, + "out": out, + "refiner_source_indices": refiner_source_indices, + } + + def decode_codec_features( + self, + out_hat: torch.Tensor, + decode_context: dict, + state: dict, + codec_depth_targets: Optional[torch.Tensor] = None, + ): + device = state["device"] + b = state["b"] + v = state["v"] + + near_flat = rearrange(decode_context["near"], "b v -> (b v) () () ()") + far_flat = rearrange(decode_context["far"], "b v -> (b v) () () ()") + ( + codec_depth_flat, + _codec_inv_depth_flat, + codec_norm_inv_depth_flat, + ) = self.codec_depth_head( + out_hat, + near_flat, + far_flat, + return_inverse=True, + return_normalized_inverse=True, + ) + codec_depth = rearrange(codec_depth_flat, "(b v) () h w -> b v h w", b=b, v=v) + codec_depth_loss_pred = None + codec_depth_loss_target = None + + if codec_depth_targets is not None: + codec_norm_inv_depth = rearrange( + codec_norm_inv_depth_flat, "(b v) () h w -> b v h w", b=b, v=v + ) + with torch.autocast(device_type=device.type, enabled=False): + near = rearrange(decode_context["near"].float(), "b v -> b v () ()") + far = rearrange(decode_context["far"].float(), "b v -> b v () ()") + inv_min = far.clamp_min(torch.finfo(torch.float32).eps).reciprocal() + inv_max = near.clamp_min(torch.finfo(torch.float32).eps).reciprocal() + inv_range = (inv_max - inv_min).clamp_min(torch.finfo(torch.float32).eps) + teacher_inv_depth = codec_depth_targets.float().clamp_min( + torch.finfo(torch.float32).eps + ).reciprocal() + teacher_inv_depth = teacher_inv_depth.clamp(min=inv_min, max=inv_max) + codec_depth_loss_target = ((teacher_inv_depth - inv_min) / inv_range).clamp(0.0, 1.0) + codec_depth_loss_pred = codec_norm_inv_depth + + if self.post_mv_refiner is not None: + out_hat = rearrange(out_hat, "(b v) c h w -> b v c h w", b=b, v=v) + out_hat = self.post_mv_refiner( + out_hat, + codec_depth, + decode_context["intrinsics"], + decode_context["extrinsics"], + source_view_indices=state["refiner_source_indices"], + ) + out_hat = rearrange(out_hat, "b v c h w -> (b v) c h w") + + depth_flat = rearrange(codec_depth, "b v h w -> (b v) () h w") + assert self.gs_decoder is not None + gaussians = self.gs_decoder(out_hat, depth_flat) + return gaussians, codec_depth, codec_depth_loss_pred, codec_depth_loss_target + + def adapt_gaussians( + self, + context: dict, + adapter_context: dict, + gaussians: torch.Tensor, + depth: torch.Tensor, + state: dict, + visualization_dump: Optional[dict], + ): + device = state["device"] + b = state["b"] + v = state["v"] + h = state["h"] + w = state["w"] + depth_preds = state["depth_preds"] + + gaussians = rearrange(gaussians, "(b v) c h w -> b v c h w", b=b, v=v) + depths = rearrange(depth, "b v h w -> b v (h w) () ()") + raw_gaussians = rearrange(gaussians, "b v c h w -> b v (h w) c") + + if self.cfg.supervise_intermediate_depth and len(depth_preds) > 1: + num_depths = len(depth_preds) + intermediate_depths = torch.cat(depth_preds[:(num_depths - 1)], dim=0) + intermediate_depths = rearrange( + intermediate_depths, "b v h w -> b v (h w) () ()" + ) + depths = torch.cat((intermediate_depths, depths), dim=0) + raw_gaussians = torch.cat([raw_gaussians] * num_depths, dim=0) + b *= num_depths + + with torch.autocast(device_type=device.type, enabled=False): + raw_gaussians = raw_gaussians.float() + depths = depths.float() + opacities = raw_gaussians[..., :1].sigmoid().unsqueeze(-1) + raw_gaussians = raw_gaussians[..., 1:] + + xy_ray, _ = sample_image_grid((h, w), device) + xy_ray = rearrange(xy_ray, "h w xy -> (h w) () xy") + gaussians = rearrange( + raw_gaussians, + "... (srf c) -> ... srf c", + srf=self.cfg.num_surfaces, + ) + offset_xy = gaussians[..., :2].sigmoid() + pixel_size = 1 / torch.tensor((w, h), dtype=torch.float32, device=device) + xy_ray = xy_ray + (offset_xy - 0.5) * pixel_size + + sh_input_images = context.get("image") + if self.cfg.init_sh_input_img: + assert sh_input_images is not None, ( + "init_sh_input_img=True requires input images, but decoded " + "bitstreams intentionally do not contain original RGB images." + ) + if self.cfg.supervise_intermediate_depth and len(depth_preds) > 1: + context_extrinsics = torch.cat( + [adapter_context["extrinsics"]] * len(depth_preds), dim=0 + ) + context_intrinsics = torch.cat( + [adapter_context["intrinsics"]] * len(depth_preds), dim=0 + ) + gaussians = self.gaussian_adapter.forward( + rearrange(context_extrinsics, "b v i j -> b v () () () i j"), + rearrange(context_intrinsics, "b v i j -> b v () () () i j"), + rearrange(xy_ray, "b v r srf xy -> b v r srf () xy"), + depths, + opacities, + rearrange(gaussians[..., 2:], "b v r srf c -> b v r srf () c"), + (h, w), + input_images=sh_input_images.repeat(len(depth_preds), 1, 1, 1, 1) + if self.cfg.init_sh_input_img + else None, + ) + else: + gaussians = self.gaussian_adapter.forward( + rearrange(adapter_context["extrinsics"], "b v i j -> b v () () () i j"), + rearrange(adapter_context["intrinsics"], "b v i j -> b v () () () i j"), + rearrange(xy_ray, "b v r srf xy -> b v r srf () xy"), + depths, + opacities, + rearrange(gaussians[..., 2:], "b v r srf c -> b v r srf () c"), + (h, w), + input_images=sh_input_images if self.cfg.init_sh_input_img else None, + ) + + if visualization_dump is not None: + visualization_dump["depth"] = rearrange( + depths, "b v (h w) srf s -> b v h w srf s", h=h, w=w + ) + visualization_dump["scales"] = rearrange( + gaussians.scales, "b v r srf spp xyz -> b (v r srf spp) xyz" + ) + visualization_dump["rotations"] = rearrange( + gaussians.rotations, "b v r srf spp xyzw -> b (v r srf spp) xyzw" + ) + + gaussians = Gaussians( + means=rearrange(gaussians.means, "b v r srf spp xyz -> b (v r srf spp) xyz"), + covariances=rearrange( + gaussians.covariances, "b v r srf spp i j -> b (v r srf spp) i j" + ), + harmonics=rearrange( + gaussians.harmonics, "b v r srf spp c d_sh -> b (v r srf spp) c d_sh" + ), + opacities=rearrange(gaussians.opacities, "b v r srf spp -> b (v r srf spp)"), + scales=rearrange(gaussians.scales, "b v r srf spp xyz -> b (v r srf spp) xyz"), + rotations=rearrange( + gaussians.rotations, "b v r srf spp xyzw -> b (v r srf spp) xyzw" + ), + rotations_unnorm=rearrange( + gaussians.rotations_unnorm, "b v r srf spp xyzw -> b (v r srf spp) xyzw" + ), + ) + if self.cfg.return_depth: + depths = rearrange( + depths, "b v (h w) srf s -> b v h w srf s", h=h, w=w + ).squeeze(-1).squeeze(-1) + return gaussians, depths + + @torch.no_grad() + def codec_forward( + self, + context: dict, + global_step: int, + visualization_dump: Optional[dict] = None, + scene_names: Optional[list] = None, + return_bitstream: bool = False, + ): + b = context["image"].shape[0] + assert b == 1, "Codec bitstream evaluation currently supports batch size 1." + assert self.feature_codec is not None + assert not self.training + + input_device = context["image"].device + if context["image"].is_cuda: + torch.cuda.synchronize(input_device) + encode_context_start = time.perf_counter() + state = self.encode_context_features(context) + if context["image"].is_cuda: + torch.cuda.synchronize(input_device) + encode_context_seconds = time.perf_counter() - encode_context_start + assert "depth_only_result" not in state, "Codec bitstream evaluation needs feature latents." + + estimated_feature_num_bytes = None + if self.cfg.codec_estimated_feature_num_bytes: + estimated_feature_bpp, _, _ = self.feature_codec.train_forward( + state["out"], + lmb=self.feature_codec.cfg.lmb, + ) + feature_height, feature_width = state["out"].shape[-2:] + estimated_feature_num_bytes = ( + estimated_feature_bpp.float().sum() + * float(feature_height * feature_width) + / 8.0 + ) + + bytestream, compress_stats = self.compress(context, state["out"]) + gaussians, depths, codec_decode_context, decompress_stats = self.decompress( + bytestream, + visualization_dump=visualization_dump, + ) + assert codec_decode_context["image_shape"] == (state["h"], state["w"]), ( + f"Decoded image_shape={codec_decode_context['image_shape']} " + f"does not match context image_shape={(state['h'], state['w'])}." + ) + assert codec_decode_context["intrinsics"].shape[1] == state["v"], ( + "Decoded camera stream context view count does not match " + f"the current context: {codec_decode_context['intrinsics'].shape[1]} vs {state['v']}." + ) + + device = state["device"] + codec_num_bytes = torch.tensor( + compress_stats["codec_num_bytes"], device=device, dtype=torch.float32 + ) + total_pixels_per_scene = state["out"].shape[-2] * state["out"].shape[-1] + result = { + "gaussians": gaussians, + "codec_num_bytes": codec_num_bytes, + "codec_bpp": codec_num_bytes * 8.0 / float(total_pixels_per_scene), + "codec_encode_context_seconds": torch.tensor( + encode_context_seconds, device=device, dtype=torch.float32 + ), + "codec_compress_seconds": torch.tensor( + compress_stats["codec_compress_seconds"], device=device, dtype=torch.float32 + ), + "codec_decompress_seconds": torch.tensor( + decompress_stats["codec_decompress_seconds"], device=device, dtype=torch.float32 + ), + "codec_camera_num_bytes": torch.tensor( + compress_stats["codec_camera_num_bytes"], device=device, dtype=torch.float32 + ), + "codec_feature_num_bytes": torch.tensor( + compress_stats["codec_feature_num_bytes"], device=device, dtype=torch.float32 + ), + } + if self.cfg.return_depth: + result["depths"] = depths + if estimated_feature_num_bytes is not None: + result["codec_estimated_feature_num_bytes"] = ( + estimated_feature_num_bytes + ) + if return_bitstream: + result["bitstream"] = bytestream + return result + + @torch.no_grad() + def deflate_feature_forward( + self, + context: dict, + global_step: int, + visualization_dump: Optional[dict] = None, + scene_names: Optional[list] = None, + ): + assert not self.training + assert self.feature_codec is None + assert context["image"].shape[0] == 1, ( + "Deflate feature evaluation currently supports batch size 1." + ) + + state = self.encode_context_features(context) + assert "depth_only_result" not in state, ( + "Deflate feature evaluation needs Gaussian-generation features." + ) + + device = state["device"] + out, depth, deflate_stats = self.run_deflate_feature_ablation( + context, + state["out"], + state["depth"], + ) + state["depth"] = depth + state["depth_preds"] = [*state["depth_preds"][:-1], depth] + + depth_flat = rearrange(depth, "b v h w -> (b v) () h w") + assert self.gs_decoder is not None + gaussians_raw = self.gs_decoder(out, depth_flat) + gaussians, depths = self.adapt_gaussians( + context, + context, + gaussians_raw, + depth, + state, + visualization_dump, + ) + + codec_num_bytes = torch.tensor( + deflate_stats["codec_num_bytes"], + device=device, + dtype=torch.float32, + ) + total_pixels_per_scene = out.shape[-2] * out.shape[-1] + result = { + "gaussians": gaussians, + "codec_num_bytes": codec_num_bytes, + "codec_bpp": codec_num_bytes * 8.0 / float(total_pixels_per_scene), + } + result.update( + { + name: torch.tensor(value, device=device, dtype=torch.float32) + for name, value in deflate_stats.items() + if name != "codec_num_bytes" + } + ) + if self.cfg.return_depth: + result["depths"] = depths + return result + + def normal_forward( + self, + context: dict, + global_step: int, + visualization_dump: Optional[dict] = None, + scene_names: Optional[list] = None, + ): + state = self.encode_context_features(context) + b = state["b"] + v = state["v"] + out = state["out"] + depth = state["depth"] + + codec_loss = None + codec_lmb = None + codec_depth_targets = None + codec_depth_loss_pred = None + codec_depth_loss_target = None + + if self.feature_codec is not None: + assert self.training, ( + "Learned codec evaluation must use the complete bitstream path." + ) + codec_depth_targets = state["teacher_depth"].detach() + scene_lmb = self.feature_codec.sample_lmb(b, out.device) + view_lmb = repeat(scene_lmb, "b -> (b v)", v=v) + codec_loss, out_hat, _ = self.feature_codec(out, lmb=view_lmb) + codec_loss = rearrange(codec_loss, "(b v) -> b v", b=b, v=v) + codec_loss = codec_loss.sum(dim=1) + codec_lmb = scene_lmb + + gaussians_raw, depth, codec_depth_loss_pred, codec_depth_loss_target = ( + self.decode_codec_features( + out_hat, + context, + state, + codec_depth_targets=codec_depth_targets, + ) + ) + adapter_context = context + else: + depth_flat = rearrange(depth, "b v h w -> (b v) () h w") + assert self.gs_decoder is not None + gaussians_raw = self.gs_decoder(out, depth_flat) + adapter_context = context + + gaussians, depths = self.adapt_gaussians( + context, + adapter_context, + gaussians_raw, + depth, + state, + visualization_dump, + ) + if self.cfg.return_depth: + result = {"gaussians": gaussians, "depths": depths} + if codec_depth_targets is not None: + result["codec_depth_targets"] = codec_depth_targets + if codec_depth_loss_pred is not None: + result["codec_depth_loss_pred"] = codec_depth_loss_pred + if codec_depth_loss_target is not None: + result["codec_depth_loss_target"] = codec_depth_loss_target + if codec_loss is not None: + return result, codec_loss, codec_lmb + return result + + if codec_loss is not None: + return gaussians, codec_loss, codec_lmb + return gaussians + + def get_data_shim(self) -> DataShim: + def data_shim(batch: BatchedExample) -> BatchedExample: + batch = apply_patch_shim( + batch, + patch_size=self.cfg.shim_patch_size + * self.cfg.downscale_factor, + ) + + return batch + + return data_shim diff --git a/src/model/encoder/unimatch/backbone.py b/src/model/encoder/unimatch/backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..8473b09529dbbfa0efc854e7f97766b74d0b375b --- /dev/null +++ b/src/model/encoder/unimatch/backbone.py @@ -0,0 +1,170 @@ +import torch.nn as nn + + +class ResidualBlock(nn.Module): + def __init__( + self, + in_planes, + planes, + norm_layer=nn.InstanceNorm2d, + stride=1, + dilation=1, + ): + super(ResidualBlock, self).__init__() + + self.conv1 = nn.Conv2d( + in_planes, + planes, + kernel_size=3, + dilation=dilation, + padding=dilation, + stride=stride, + bias=False, + ) + self.conv2 = nn.Conv2d( + planes, + planes, + kernel_size=3, + dilation=dilation, + padding=dilation, + bias=False, + ) + self.relu = nn.ReLU(inplace=True) + + self.norm1 = norm_layer(planes) + self.norm2 = norm_layer(planes) + if not stride == 1 or in_planes != planes: + self.norm3 = norm_layer(planes) + + if stride == 1 and in_planes == planes: + self.downsample = None + else: + self.downsample = nn.Sequential( + nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3 + ) + + def forward(self, x): + y = x + y = self.relu(self.norm1(self.conv1(y))) + y = self.relu(self.norm2(self.conv2(y))) + + if self.downsample is not None: + x = self.downsample(x) + + return self.relu(x + y) + + +class CNNEncoder(nn.Module): + def __init__( + self, + output_dim=128, + norm_layer=nn.InstanceNorm2d, + num_output_scales=1, + return_quarter=False, # return 1/4 resolution feature + lowest_scale=8, # lowest resolution, 1/8 or 1/4 + return_all_scales=False, + **kwargs, + ): + super(CNNEncoder, self).__init__() + self.num_scales = num_output_scales + self.return_quarter = return_quarter + self.lowest_scale = lowest_scale + self.return_all_scales = return_all_scales + + feature_dims = [64, 96, 128] + + self.conv1 = nn.Conv2d( + 3, feature_dims[0], kernel_size=7, stride=2, padding=3, bias=False + ) # 1/2 + self.norm1 = norm_layer(feature_dims[0]) + self.relu1 = nn.ReLU(inplace=True) + + self.in_planes = feature_dims[0] + self.layer1 = self._make_layer( + feature_dims[0], stride=1, norm_layer=norm_layer + ) # 1/2 + + if self.lowest_scale == 4: + stride = 1 + else: + stride = 2 + self.layer2 = self._make_layer( + feature_dims[1], stride=stride, norm_layer=norm_layer + ) # 1/2 or 1/4 + + # lowest resolution 1/4 or 1/8 + self.layer3 = self._make_layer( + feature_dims[2], + stride=2, + norm_layer=norm_layer, + ) # 1/4 or 1/8 + + self.conv2 = nn.Conv2d(feature_dims[2], output_dim, 1, 1, 0) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") + elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): + if m.weight is not None: + nn.init.constant_(m.weight, 1) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def _make_layer(self, dim, stride=1, dilation=1, norm_layer=nn.InstanceNorm2d): + layer1 = ResidualBlock( + self.in_planes, dim, norm_layer=norm_layer, stride=stride, dilation=dilation + ) + layer2 = ResidualBlock( + dim, dim, norm_layer=norm_layer, stride=1, dilation=dilation + ) + + layers = (layer1, layer2) + + self.in_planes = dim + return nn.Sequential(*layers) + + def forward(self, x): + output_all_scales = [] + output = [] + x = self.conv1(x) + x = self.norm1(x) + x = self.relu1(x) + + x = self.layer1(x) # 1/2 + + if self.return_all_scales: + output_all_scales.append(x) + + if self.num_scales >= 3: + output.append(x) + + x = self.layer2(x) # 1/2 or 1/4 + if self.return_quarter: + output.append(x) + + if self.return_all_scales: + output_all_scales.append(x) + + if self.num_scales >= 2: + output.append(x) + + x = self.layer3(x) # 1/4 or 1/8 + x = self.conv2(x) + + if self.return_all_scales: + output_all_scales.append(x) + + if self.return_all_scales: + return output_all_scales + + if self.return_quarter: + output.append(x) + return output + + if self.num_scales >= 1: + output.append(x) + return output + + out = [x] + + return out diff --git a/src/model/encoder/unimatch/dinov2/dinov2.py b/src/model/encoder/unimatch/dinov2/dinov2.py new file mode 100644 index 0000000000000000000000000000000000000000..fd2b668158f1c4a30ee8c69bb21fe2749e32b431 --- /dev/null +++ b/src/model/encoder/unimatch/dinov2/dinov2.py @@ -0,0 +1,435 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py + +from functools import partial +import math +import logging +from typing import Sequence, Tuple, Union, Callable + +import torch +import torch.nn as nn +import torch.utils.checkpoint +from torch.nn.init import trunc_normal_ + +from .layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block + + +logger = logging.getLogger("dinov2") + + +def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module: + if not depth_first and include_root: + fn(module=module, name=name) + for child_name, child_module in module.named_children(): + child_name = ".".join((name, child_name)) if name else child_name + named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True) + if depth_first and include_root: + fn(module=module, name=name) + return module + + +class BlockChunk(nn.ModuleList): + def forward(self, x): + for b in self: + x = b(x) + return x + + +class DinoVisionTransformer(nn.Module): + def __init__( + self, + img_size=224, + patch_size=16, + in_chans=3, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4.0, + qkv_bias=True, + ffn_bias=True, + proj_bias=True, + drop_path_rate=0.0, + drop_path_uniform=False, + init_values=None, # for layerscale: None or 0 => no layerscale + embed_layer=PatchEmbed, + act_layer=nn.GELU, + block_fn=Block, + ffn_layer="mlp", + block_chunks=1, + num_register_tokens=0, + interpolate_antialias=False, + interpolate_offset=0.1, + use_checkpointing=False, + ): + """ + Args: + img_size (int, tuple): input image size + patch_size (int, tuple): patch size + in_chans (int): number of input channels + embed_dim (int): embedding dimension + depth (int): depth of transformer + num_heads (int): number of attention heads + mlp_ratio (int): ratio of mlp hidden dim to embedding dim + qkv_bias (bool): enable bias for qkv if True + proj_bias (bool): enable bias for proj in attn if True + ffn_bias (bool): enable bias for ffn if True + drop_path_rate (float): stochastic depth rate + drop_path_uniform (bool): apply uniform drop rate across blocks + weight_init (str): weight init scheme + init_values (float): layer-scale init values + embed_layer (nn.Module): patch embedding layer + act_layer (nn.Module): MLP activation layer + block_fn (nn.Module): transformer block class + ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity" + block_chunks: (int) split block sequence into block_chunks units for FSDP wrap + num_register_tokens: (int) number of extra cls tokens (so-called "registers") + interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings + interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings + """ + super().__init__() + norm_layer = partial(nn.LayerNorm, eps=1e-6) + + self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models + self.num_tokens = 1 + self.n_blocks = depth + self.num_heads = num_heads + self.patch_size = patch_size + self.num_register_tokens = num_register_tokens + self.interpolate_antialias = interpolate_antialias + self.interpolate_offset = interpolate_offset + self.use_checkpointing = use_checkpointing + + self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) + num_patches = self.patch_embed.num_patches + + self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim)) + assert num_register_tokens >= 0 + self.register_tokens = ( + nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None + ) + + if drop_path_uniform is True: + dpr = [drop_path_rate] * depth + else: + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule + + if ffn_layer == "mlp": + logger.info("using MLP layer as FFN") + ffn_layer = Mlp + elif ffn_layer == "swiglufused" or ffn_layer == "swiglu": + logger.info("using SwiGLU layer as FFN") + ffn_layer = SwiGLUFFNFused + elif ffn_layer == "identity": + logger.info("using Identity layer as FFN") + + def f(*args, **kwargs): + return nn.Identity() + + ffn_layer = f + else: + raise NotImplementedError + + blocks_list = [ + block_fn( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + ffn_bias=ffn_bias, + drop_path=dpr[i], + norm_layer=norm_layer, + act_layer=act_layer, + ffn_layer=ffn_layer, + init_values=init_values, + ) + for i in range(depth) + ] + if block_chunks > 0: + self.chunked_blocks = True + chunked_blocks = [] + chunksize = depth // block_chunks + for i in range(0, depth, chunksize): + # this is to keep the block index consistent if we chunk the block list + chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize]) + self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks]) + else: + self.chunked_blocks = False + self.blocks = nn.ModuleList(blocks_list) + + self.norm = norm_layer(embed_dim) + self.head = nn.Identity() + + self.mask_token = nn.Parameter(torch.zeros(1, embed_dim)) + + self.init_weights() + + def init_weights(self): + trunc_normal_(self.pos_embed, std=0.02) + nn.init.normal_(self.cls_token, std=1e-6) + if self.register_tokens is not None: + nn.init.normal_(self.register_tokens, std=1e-6) + named_apply(init_weights_vit_timm, self) + + def interpolate_pos_encoding(self, x, w, h): + previous_dtype = x.dtype + npatch = x.shape[1] - 1 + N = self.pos_embed.shape[1] - 1 + if npatch == N and w == h: + return self.pos_embed + pos_embed = self.pos_embed.float() + class_pos_embed = pos_embed[:, 0] + patch_pos_embed = pos_embed[:, 1:] + dim = x.shape[-1] + w0 = w // self.patch_size + h0 = h // self.patch_size + M = int(math.sqrt(N)) # Recover the number of patches in each dimension + assert N == M * M + kwargs = {} + if self.interpolate_offset: + # Historical kludge: add a small number to avoid floating point error in the interpolation, see https://github.com/facebookresearch/dino/issues/8 + # Note: still needed for backward-compatibility, the underlying operators are using both output size and scale factors + sx = float(w0 + self.interpolate_offset) / M + sy = float(h0 + self.interpolate_offset) / M + kwargs["scale_factor"] = (sx, sy) + else: + # Simply specify an output size instead of a scale factor + kwargs["size"] = (w0, h0) + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed.reshape(1, M, M, dim).permute(0, 3, 1, 2), + mode="bicubic", + antialias=self.interpolate_antialias, + **kwargs, + ) + assert (w0, h0) == patch_pos_embed.shape[-2:] + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype) + + def prepare_tokens_with_masks(self, x, masks=None): + B, nc, w, h = x.shape + x = self.patch_embed(x) + if masks is not None: + x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x) + + x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1) + x = x + self.interpolate_pos_encoding(x, w, h) + + if self.register_tokens is not None: + x = torch.cat( + ( + x[:, :1], + self.register_tokens.expand(x.shape[0], -1, -1), + x[:, 1:], + ), + dim=1, + ) + + return x + + def forward_features_list(self, x_list, masks_list): + x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)] + for blk in self.blocks: + x = blk(x) + + all_x = x + output = [] + for x, masks in zip(all_x, masks_list): + x_norm = self.norm(x) + output.append( + { + "x_norm_clstoken": x_norm[:, 0], + "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1], + "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :], + "x_prenorm": x, + "masks": masks, + } + ) + return output + + def forward_features(self, x, masks=None): + if isinstance(x, list): + return self.forward_features_list(x, masks) + + x = self.prepare_tokens_with_masks(x, masks) + + for blk in self.blocks: + if self.use_checkpointing: + x = torch.utils.checkpoint.checkpoint(blk, x) + else: + x = blk(x) + + x_norm = self.norm(x) + return { + "x_norm_clstoken": x_norm[:, 0], + "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1], + "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :], + "x_prenorm": x, + "masks": masks, + } + + def _get_intermediate_layers_not_chunked(self, x, n=1): + x = self.prepare_tokens_with_masks(x) + # If n is an int, take the n last blocks. If it's a list, take them + output, total_block_len = [], len(self.blocks) + blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n + for i, blk in enumerate(self.blocks): + if self.use_checkpointing: + x = torch.utils.checkpoint.checkpoint(blk, x) + else: + x = blk(x) + if i in blocks_to_take: + output.append(x) + assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found" + return output + + def _get_intermediate_layers_chunked(self, x, n=1): + x = self.prepare_tokens_with_masks(x) + output, i, total_block_len = [], 0, len(self.blocks[-1]) + # If n is an int, take the n last blocks. If it's a list, take them + blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n + for block_chunk in self.blocks: + for blk in block_chunk[i:]: # Passing the nn.Identity() + if self.use_checkpointing: + x = torch.utils.checkpoint.checkpoint(blk, x) + else: + x = blk(x) + if i in blocks_to_take: + output.append(x) + i += 1 + assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found" + return output + + def get_intermediate_layers( + self, + x: torch.Tensor, + n: Union[int, Sequence] = 1, # Layers or n last layers to take + reshape: bool = False, + return_class_token: bool = False, + norm=True, + ): + if self.chunked_blocks: + outputs = self._get_intermediate_layers_chunked(x, n) + else: + outputs = self._get_intermediate_layers_not_chunked(x, n) + if norm: + outputs = [self.norm(out) for out in outputs] + class_tokens = [out[:, 0] for out in outputs] + outputs = [out[:, 1 + self.num_register_tokens :] for out in outputs] + if reshape: + B, _, w, h = x.shape + outputs = [ + out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous() + for out in outputs + ] + if return_class_token: + return tuple(zip(outputs, class_tokens)) + return tuple(outputs) + + def forward(self, *args, is_training=False, **kwargs): + ret = self.forward_features(*args, **kwargs) + if is_training: + return ret + else: + return self.head(ret["x_norm_clstoken"]) + + +def init_weights_vit_timm(module: nn.Module, name: str = ""): + """ViT weight initialization, original timm impl (for reproducibility)""" + if isinstance(module, nn.Linear): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + + +def vit_small(patch_size=16, num_register_tokens=0, use_checkpointing=False, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=384, + depth=12, + num_heads=6, + mlp_ratio=4., + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + use_checkpointing=use_checkpointing, + **kwargs, + ) + return model + + +def vit_base(patch_size=16, num_register_tokens=0, use_checkpointing=False, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4., + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + use_checkpointing=use_checkpointing, + **kwargs, + ) + return model + + +def vit_large(patch_size=16, num_register_tokens=0, use_checkpointing=False, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=1024, + depth=24, + num_heads=16, + mlp_ratio=4., + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + use_checkpointing=use_checkpointing, + **kwargs, + ) + return model + + +def vit_giant2(patch_size=16, num_register_tokens=0, use_checkpointing=False, **kwargs): + """ + Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64 + """ + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=1536, + depth=40, + num_heads=24, + mlp_ratio=4., + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + use_checkpointing=use_checkpointing, + **kwargs, + ) + return model + + +def DINOv2(model_name, + use_checkpointing=False, + ): + model_zoo = { + "vits": vit_small, + "vitb": vit_base, + "vitl": vit_large, + "vitg": vit_giant2 + } + + return model_zoo[model_name]( + img_size=518, + patch_size=14, + init_values=1.0, + ffn_layer="mlp" if model_name != "vitg" else "swiglufused", + block_chunks=0, + num_register_tokens=0, + interpolate_antialias=False, + interpolate_offset=0.1, + use_checkpointing=use_checkpointing, + ) + \ No newline at end of file diff --git a/src/model/encoder/unimatch/dinov2/layers/__init__.py b/src/model/encoder/unimatch/dinov2/layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..05a0b61868e43abb821ca05a813bab2b8b43629e --- /dev/null +++ b/src/model/encoder/unimatch/dinov2/layers/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +from .dino_head import DINOHead +from .mlp import Mlp +from .patch_embed import PatchEmbed +from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused +from .block import NestedTensorBlock +from .attention import MemEffAttention diff --git a/src/model/encoder/unimatch/dinov2/layers/attention.py b/src/model/encoder/unimatch/dinov2/layers/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..517d9f50198199d027dc8b9e0d1f035efe2af96e --- /dev/null +++ b/src/model/encoder/unimatch/dinov2/layers/attention.py @@ -0,0 +1,124 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py + +import torch + +from torch import Tensor +from torch import nn + +import torch.nn.functional as F + + +class Attention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = False, + proj_bias: bool = True, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ) -> None: + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim**-0.5 + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim, bias=proj_bias) + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, x: Tensor) -> Tensor: + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + + q, k, v = qkv[0], qkv[1], qkv[2] + out = F.scaled_dot_product_attention(q, k, v) + x = out.permute(0, 2, 1, 3).reshape(B, N, C) + + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class MemEffAttention(Attention): + def forward(self, x: Tensor, attn_bias=None, y=None) -> Tensor: + if attn_bias is not None: + raise NotImplementedError("attn_bias is not supported without xformers") + + if y is not None: + B, Nq, C = x.shape + context = x if y is None else y + Nk = context.shape[1] + + if not hasattr(self, 'q_proj') or self.q_proj is None: + self.q_proj, self.kv_proj = convert_qkv_to_q_and_kv_proj(self.qkv) + + q = self.q_proj(x).reshape(B, Nq, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) + kv = self.kv_proj(context).reshape(B, Nk, 2, self.num_heads, C // self.num_heads) + k, v = kv.unbind(dim=2) + k = k.permute(0, 2, 1, 3) + v = v.permute(0, 2, 1, 3) + + out = F.scaled_dot_product_attention(q, k, v) + out = out.permute(0, 2, 1, 3).reshape(B, Nq, C) + out = self.proj(out) + out = self.proj_drop(out) + return out + + else: + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) + + q, k, v = torch.unbind(qkv, 2) + q = q.permute(0, 2, 1, 3) + k = k.permute(0, 2, 1, 3) + v = v.permute(0, 2, 1, 3) + + x = F.scaled_dot_product_attention(q, k, v) + x = x.permute(0, 2, 1, 3).reshape(B, N, C) + + x = self.proj(x) + x = self.proj_drop(x) + return x + + +def convert_qkv_to_q_and_kv_proj(qkv_layer: nn.Linear): + """ + Convert a self-attention qkv projection layer (dim -> 3*dim) into + separate q_proj (dim -> dim) and kv_proj (dim -> 2*dim) layers. + + Returns: + q_proj (nn.Linear): projection for query + kv_proj (nn.Linear): projection for key and value + """ + assert isinstance(qkv_layer, nn.Linear), "Expected nn.Linear for qkv_layer" + in_features = qkv_layer.in_features + out_features = qkv_layer.out_features + assert out_features % 3 == 0, "Output features must be divisible by 3" + + dim = out_features // 3 + device = qkv_layer.weight.device + dtype = qkv_layer.weight.dtype + + q_proj = nn.Linear(in_features, dim, bias=qkv_layer.bias is not None).to(device=device, dtype=dtype) + kv_proj = nn.Linear(in_features, dim * 2, bias=qkv_layer.bias is not None).to(device=device, dtype=dtype) + + # Split weights and biases + q_weight, k_weight, v_weight = qkv_layer.weight.chunk(3, dim=0) + q_proj.weight.data.copy_(q_weight) + kv_proj.weight.data.copy_(torch.cat([k_weight, v_weight], dim=0)) + + if qkv_layer.bias is not None: + q_bias, k_bias, v_bias = qkv_layer.bias.chunk(3, dim=0) + q_proj.bias.data.copy_(q_bias) + kv_proj.bias.data.copy_(torch.cat([k_bias, v_bias], dim=0)) + + return q_proj, kv_proj diff --git a/src/model/encoder/unimatch/dinov2/layers/block.py b/src/model/encoder/unimatch/dinov2/layers/block.py new file mode 100644 index 0000000000000000000000000000000000000000..7de35b436da88eade27a8499e06e2642243ff996 --- /dev/null +++ b/src/model/encoder/unimatch/dinov2/layers/block.py @@ -0,0 +1,128 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py + +from typing import Callable, List + +import torch +from torch import nn, Tensor + +from .attention import Attention, MemEffAttention +from .drop_path import DropPath +from .layer_scale import LayerScale +from .mlp import Mlp + + +class Block(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = False, + proj_bias: bool = True, + ffn_bias: bool = True, + drop: float = 0.0, + attn_drop: float = 0.0, + init_values=None, + drop_path: float = 0.0, + act_layer: Callable[..., nn.Module] = nn.GELU, + norm_layer: Callable[..., nn.Module] = nn.LayerNorm, + attn_class: Callable[..., nn.Module] = MemEffAttention, + ffn_layer: Callable[..., nn.Module] = Mlp, + no_attn: bool = False + ) -> None: + super().__init__() + self.norm1 = norm_layer(dim) + if no_attn: + self.attn = nn.Identity() + else: + self.attn = attn_class( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + attn_drop=attn_drop, + proj_drop=drop, + ) + self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = ffn_layer( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + bias=ffn_bias, + ) + self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + self.sample_drop_ratio = drop_path + + def forward(self, x: Tensor, y: Tensor | None = None) -> Tensor: + def attn_residual_func(x: Tensor) -> Tensor: + return self.ls1(self.attn(self.norm1(x), y=y)) + + def ffn_residual_func(x: Tensor) -> Tensor: + return self.ls2(self.mlp(self.norm2(x))) + + if self.training and self.sample_drop_ratio > 0.1: + x = drop_add_residual_stochastic_depth( + x, + residual_func=attn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + ) + x = drop_add_residual_stochastic_depth( + x, + residual_func=ffn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + ) + elif self.training and self.sample_drop_ratio > 0.0: + x = x + self.drop_path1(attn_residual_func(x)) + x = x + self.drop_path2(ffn_residual_func(x)) + else: + x = x + attn_residual_func(x) + x = x + ffn_residual_func(x) + return x + + +def drop_add_residual_stochastic_depth( + x: Tensor, + residual_func: Callable[[Tensor], Tensor], + sample_drop_ratio: float = 0.0, +) -> Tensor: + b, n, d = x.shape + sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1) + brange = (torch.randperm(b, device=x.device))[:sample_subset_size] + x_subset = x[brange] + + residual = residual_func(x_subset) + + x_flat = x.flatten(1) + residual = residual.flatten(1) + + residual_scale_factor = b / sample_subset_size + + x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor) + return x_plus_residual.view_as(x) + + +class NestedTensorBlock(Block): + def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]: + raise NotImplementedError("Nested tensor support requires xformers, which has been removed") + + def forward(self, x_or_x_list): + if isinstance(x_or_x_list, Tensor): + return super().forward(x_or_x_list) + elif isinstance(x_or_x_list, list): + raise NotImplementedError("Nested tensor support requires xformers, which has been removed") + else: + raise AssertionError diff --git a/src/model/encoder/unimatch/dinov2/layers/dino_head.py b/src/model/encoder/unimatch/dinov2/layers/dino_head.py new file mode 100644 index 0000000000000000000000000000000000000000..0ace8ffd6297a1dd480b19db407b662a6ea0f565 --- /dev/null +++ b/src/model/encoder/unimatch/dinov2/layers/dino_head.py @@ -0,0 +1,58 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +import torch +import torch.nn as nn +from torch.nn.init import trunc_normal_ +from torch.nn.utils import weight_norm + + +class DINOHead(nn.Module): + def __init__( + self, + in_dim, + out_dim, + use_bn=False, + nlayers=3, + hidden_dim=2048, + bottleneck_dim=256, + mlp_bias=True, + ): + super().__init__() + nlayers = max(nlayers, 1) + self.mlp = _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=hidden_dim, use_bn=use_bn, bias=mlp_bias) + self.apply(self._init_weights) + self.last_layer = weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False)) + self.last_layer.weight_g.data.fill_(1) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=0.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + + def forward(self, x): + x = self.mlp(x) + eps = 1e-6 if x.dtype == torch.float16 else 1e-12 + x = nn.functional.normalize(x, dim=-1, p=2, eps=eps) + x = self.last_layer(x) + return x + + +def _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=None, use_bn=False, bias=True): + if nlayers == 1: + return nn.Linear(in_dim, bottleneck_dim, bias=bias) + else: + layers = [nn.Linear(in_dim, hidden_dim, bias=bias)] + if use_bn: + layers.append(nn.BatchNorm1d(hidden_dim)) + layers.append(nn.GELU()) + for _ in range(nlayers - 2): + layers.append(nn.Linear(hidden_dim, hidden_dim, bias=bias)) + if use_bn: + layers.append(nn.BatchNorm1d(hidden_dim)) + layers.append(nn.GELU()) + layers.append(nn.Linear(hidden_dim, bottleneck_dim, bias=bias)) + return nn.Sequential(*layers) diff --git a/src/model/encoder/unimatch/dinov2/layers/drop_path.py b/src/model/encoder/unimatch/dinov2/layers/drop_path.py new file mode 100644 index 0000000000000000000000000000000000000000..1d640e0b969b8dcba96260243473700b4e5b24b5 --- /dev/null +++ b/src/model/encoder/unimatch/dinov2/layers/drop_path.py @@ -0,0 +1,34 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py + + +from torch import nn + + +def drop_path(x, drop_prob: float = 0.0, training: bool = False): + if drop_prob == 0.0 or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0: + random_tensor.div_(keep_prob) + output = x * random_tensor + return output + + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob=None): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training) diff --git a/src/model/encoder/unimatch/dinov2/layers/layer_scale.py b/src/model/encoder/unimatch/dinov2/layers/layer_scale.py new file mode 100644 index 0000000000000000000000000000000000000000..51df0d7ce61f2b41fa9e6369f52391dd7fe7d386 --- /dev/null +++ b/src/model/encoder/unimatch/dinov2/layers/layer_scale.py @@ -0,0 +1,27 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110 + +from typing import Union + +import torch +from torch import Tensor +from torch import nn + + +class LayerScale(nn.Module): + def __init__( + self, + dim: int, + init_values: Union[float, Tensor] = 1e-5, + inplace: bool = False, + ) -> None: + super().__init__() + self.inplace = inplace + self.gamma = nn.Parameter(init_values * torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + return x.mul_(self.gamma) if self.inplace else x * self.gamma diff --git a/src/model/encoder/unimatch/dinov2/layers/mlp.py b/src/model/encoder/unimatch/dinov2/layers/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..bbf9432aae9258612caeae910a7bde17999e328e --- /dev/null +++ b/src/model/encoder/unimatch/dinov2/layers/mlp.py @@ -0,0 +1,40 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py + + +from typing import Callable, Optional + +from torch import Tensor, nn + + +class Mlp(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = nn.GELU, + drop: float = 0.0, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features, bias=bias) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features, bias=bias) + self.drop = nn.Dropout(drop) + + def forward(self, x: Tensor) -> Tensor: + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x diff --git a/src/model/encoder/unimatch/dinov2/layers/patch_embed.py b/src/model/encoder/unimatch/dinov2/layers/patch_embed.py new file mode 100644 index 0000000000000000000000000000000000000000..8b7c0804784a42cf80c0297d110dcc68cc85b339 --- /dev/null +++ b/src/model/encoder/unimatch/dinov2/layers/patch_embed.py @@ -0,0 +1,88 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py + +from typing import Callable, Optional, Tuple, Union + +from torch import Tensor +import torch.nn as nn + + +def make_2tuple(x): + if isinstance(x, tuple): + assert len(x) == 2 + return x + + assert isinstance(x, int) + return (x, x) + + +class PatchEmbed(nn.Module): + """ + 2D image to patch embedding: (B,C,H,W) -> (B,N,D) + + Args: + img_size: Image size. + patch_size: Patch token size. + in_chans: Number of input image channels. + embed_dim: Number of linear projection output channels. + norm_layer: Normalization layer. + """ + + def __init__( + self, + img_size: Union[int, Tuple[int, int]] = 224, + patch_size: Union[int, Tuple[int, int]] = 16, + in_chans: int = 3, + embed_dim: int = 768, + norm_layer: Optional[Callable] = None, + flatten_embedding: bool = True, + ) -> None: + super().__init__() + + image_HW = make_2tuple(img_size) + patch_HW = make_2tuple(patch_size) + patch_grid_size = ( + image_HW[0] // patch_HW[0], + image_HW[1] // patch_HW[1], + ) + + self.img_size = image_HW + self.patch_size = patch_HW + self.patches_resolution = patch_grid_size + self.num_patches = patch_grid_size[0] * patch_grid_size[1] + + self.in_chans = in_chans + self.embed_dim = embed_dim + + self.flatten_embedding = flatten_embedding + + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW) + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x: Tensor) -> Tensor: + _, _, H, W = x.shape + patch_H, patch_W = self.patch_size + + assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}" + assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}" + + x = self.proj(x) # B C H W + H, W = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) # B HW C + x = self.norm(x) + if not self.flatten_embedding: + x = x.reshape(-1, H, W, self.embed_dim) # B H W C + return x + + def flops(self) -> float: + Ho, Wo = self.patches_resolution + flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1]) + if self.norm is not None: + flops += Ho * Wo * self.embed_dim + return flops diff --git a/src/model/encoder/unimatch/dinov2/layers/swiglu_ffn.py b/src/model/encoder/unimatch/dinov2/layers/swiglu_ffn.py new file mode 100644 index 0000000000000000000000000000000000000000..a0eba35994b4335e2c996d848c81ec651953bb79 --- /dev/null +++ b/src/model/encoder/unimatch/dinov2/layers/swiglu_ffn.py @@ -0,0 +1,56 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +from typing import Callable, Optional + +from torch import Tensor, nn +import torch.nn.functional as F + + +class SwiGLUFFN(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = None, + drop: float = 0.0, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias) + self.w3 = nn.Linear(hidden_features, out_features, bias=bias) + + def forward(self, x: Tensor) -> Tensor: + x12 = self.w12(x) + x1, x2 = x12.chunk(2, dim=-1) + hidden = F.silu(x1) * x2 + return self.w3(hidden) + + +SwiGLU = SwiGLUFFN + + +class SwiGLUFFNFused(SwiGLU): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = None, + drop: float = 0.0, + bias: bool = True, + ) -> None: + out_features = out_features or in_features + hidden_features = hidden_features or in_features + hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8 + super().__init__( + in_features=in_features, + hidden_features=hidden_features, + out_features=out_features, + bias=bias, + ) diff --git a/src/model/encoder/unimatch/dpt_head.py b/src/model/encoder/unimatch/dpt_head.py new file mode 100644 index 0000000000000000000000000000000000000000..72f558e9a44a1f8aa4cf19776f0f727eb904cf8e --- /dev/null +++ b/src/model/encoder/unimatch/dpt_head.py @@ -0,0 +1,571 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def _make_scratch(in_shape, out_shape, groups=1, expand=False): + scratch = nn.Module() + + out_shape1 = out_shape + out_shape2 = out_shape + out_shape3 = out_shape + if len(in_shape) >= 4: + out_shape4 = out_shape + + if expand: + out_shape1 = out_shape + out_shape2 = out_shape * 2 + out_shape3 = out_shape * 4 + if len(in_shape) >= 4: + out_shape4 = out_shape * 8 + + scratch.layer1_rn = nn.Conv2d( + in_shape[0], + out_shape1, + kernel_size=3, + stride=1, + padding=1, + bias=False, + groups=groups, + ) + scratch.layer2_rn = nn.Conv2d( + in_shape[1], + out_shape2, + kernel_size=3, + stride=1, + padding=1, + bias=False, + groups=groups, + ) + scratch.layer3_rn = nn.Conv2d( + in_shape[2], + out_shape3, + kernel_size=3, + stride=1, + padding=1, + bias=False, + groups=groups, + ) + if len(in_shape) >= 4: + scratch.layer4_rn = nn.Conv2d( + in_shape[3], + out_shape4, + kernel_size=3, + stride=1, + padding=1, + bias=False, + groups=groups, + ) + + return scratch + + +class ResidualConvUnit(nn.Module): + """Residual convolution module.""" + + def __init__(self, features, activation, bn): + """Init. + + Args: + features (int): number of features + """ + super().__init__() + + self.bn = bn + + self.groups = 1 + + self.conv1 = nn.Conv2d( + features, + features, + kernel_size=3, + stride=1, + padding=1, + bias=True, + groups=self.groups, + ) + + self.conv2 = nn.Conv2d( + features, + features, + kernel_size=3, + stride=1, + padding=1, + bias=True, + groups=self.groups, + ) + + if self.bn == True: + self.bn1 = nn.BatchNorm2d(features) + self.bn2 = nn.BatchNorm2d(features) + + self.activation = activation + + self.skip_add = nn.quantized.FloatFunctional() + + def forward(self, x): + """Forward pass. + + Args: + x (tensor): input + + Returns: + tensor: output + """ + + out = self.activation(x) + out = self.conv1(out) + if self.bn == True: + out = self.bn1(out) + + out = self.activation(out) + out = self.conv2(out) + if self.bn == True: + out = self.bn2(out) + + if self.groups > 1: + out = self.conv_merge(out) + + return self.skip_add.add(out, x) + + +class FeatureFusionBlock(nn.Module): + """Feature fusion block.""" + + def __init__( + self, + features, + activation, + deconv=False, + bn=False, + expand=False, + align_corners=True, + size=None, + ): + """Init. + + Args: + features (int): number of features + """ + super(FeatureFusionBlock, self).__init__() + + self.deconv = deconv + self.align_corners = align_corners + + self.groups = 1 + + self.expand = expand + out_features = features + if self.expand == True: + out_features = features // 2 + + self.out_conv = nn.Conv2d( + features, + out_features, + kernel_size=1, + stride=1, + padding=0, + bias=True, + groups=1, + ) + + self.resConfUnit1 = ResidualConvUnit(features, activation, bn) + self.resConfUnit2 = ResidualConvUnit(features, activation, bn) + + self.skip_add = nn.quantized.FloatFunctional() + + self.size = size + + def forward(self, *xs, size=None): + """Forward pass. + + Returns: + tensor: output + """ + output = xs[0] + + if len(xs) == 2: + res = self.resConfUnit1(xs[1]) + output = self.skip_add.add(output, res) + + output = self.resConfUnit2(output) + + if (size is None) and (self.size is None): + modifier = {"scale_factor": 2} + elif size is None: + modifier = {"size": self.size} + else: + modifier = {"size": size} + + output = nn.functional.interpolate( + output, **modifier, mode="bilinear", align_corners=self.align_corners + ) + + output = self.out_conv(output) + + return output + + +def _make_fusion_block(features, use_bn, size=None): + return FeatureFusionBlock( + features, + nn.ReLU(False), + deconv=False, + bn=use_bn, + expand=False, + align_corners=True, + size=size, + ) + + +class DPTHead(nn.Module): + def __init__( + self, + in_channels, + features=256, + use_bn=False, + out_channels=[256, 512, 1024, 1024], + use_clstoken=False, + concat_cnn_features=True, + concat_mv_features=True, + cnn_feature_channels=[64, 96, 128], + concat_features=True, + downsample_factor=8, + return_feature=False, + num_scales=1, + ): + super(DPTHead, self).__init__() + + self.use_clstoken = use_clstoken + + self.concat_cnn_features = concat_cnn_features + self.concat_mv_features = concat_mv_features + self.concat_features = concat_features + self.downsample_factor = downsample_factor + self.return_feature = return_feature + self.num_scales = num_scales + + if self.concat_features: + if self.downsample_factor == 4 and num_scales == 2: + depth_channel = 0 if self.return_feature else 1 + self.concat_projects = nn.ModuleList( + [ + nn.Conv2d( + cnn_feature_channels[0] + out_channels[0], + out_channels[0], + 1, + ), + nn.Conv2d( + cnn_feature_channels[1] + + out_channels[1] + + 64 + + depth_channel, + out_channels[1], + 1, + ), # 1/4 concat(cnn, mono, mv, depth) + nn.Conv2d( + cnn_feature_channels[2] + out_channels[2] + 128, + out_channels[2], + 1, + ), # 1/8 concat(cnn, mono, mv) + ] + ) + elif self.downsample_factor == 2 and num_scales == 2: + depth_channel = 0 if self.return_feature else 1 + self.concat_projects = nn.ModuleList( + [ + nn.Conv2d( + cnn_feature_channels[0] + + cnn_feature_channels[1] + + out_channels[0] + + 64 + + depth_channel, + out_channels[0], + 1, + ), # 1/2 + nn.Conv2d( + cnn_feature_channels[2] + out_channels[1] + 128, + out_channels[1], + 1, + ), # 1/4 concat(cnn, mono, mv, depth) + nn.Conv2d(out_channels[2], out_channels[2], 1), # 1/8 mono + ] + ) + elif self.downsample_factor == 4 and num_scales == 1: + depth_channel = 0 if self.return_feature else 1 + self.concat_projects = nn.ModuleList( + [ + nn.Conv2d( + cnn_feature_channels[0] + + cnn_feature_channels[1] + + out_channels[0], + out_channels[0], + 1, + ), + nn.Conv2d( + cnn_feature_channels[2] + + out_channels[1] + + 128 + + depth_channel, + out_channels[1], + 1, + ), + nn.Conv2d(out_channels[2], out_channels[2], 1), # 1/8 mono + ] + ) + else: + depth_channel = 0 if self.return_feature else 1 + self.concat_projects = nn.ModuleList( + [ + nn.Conv2d( + cnn_feature_channels[0] + out_channels[0], + out_channels[0], + 1, + ), + nn.Conv2d( + cnn_feature_channels[1] + out_channels[1], + out_channels[1], + 1, + ), + nn.Conv2d( + cnn_feature_channels[2] + + out_channels[2] + + 128 + + depth_channel, + out_channels[2], + 1, + ), # 1/8 concat(cnn, mono, mv, depth) + ] + ) + else: + if self.concat_cnn_features: + self.cnn_projects = nn.ModuleList( + [ + nn.Conv2d(cnn_feature_channels[i], out_channels[i], 1) + for i in range(len(cnn_feature_channels)) + ] + ) + + if self.concat_mv_features: + self.mv_projects = nn.Conv2d(128, out_channels[2], 1) + + self.projects = nn.ModuleList( + [ + nn.Conv2d( + in_channels=in_channels, + out_channels=out_channel, + kernel_size=1, + stride=1, + padding=0, + ) + for out_channel in out_channels + ] + ) + + self.resize_layers = nn.ModuleList( + [ + nn.ConvTranspose2d( + in_channels=out_channels[0], + out_channels=out_channels[0], + kernel_size=4, + stride=4, + padding=0, + ), + nn.ConvTranspose2d( + in_channels=out_channels[1], + out_channels=out_channels[1], + kernel_size=2, + stride=2, + padding=0, + ), + nn.Identity(), + nn.Conv2d( + in_channels=out_channels[3], + out_channels=out_channels[3], + kernel_size=3, + stride=2, + padding=1, + ), + ] + ) + + if use_clstoken: + self.readout_projects = nn.ModuleList() + for _ in range(len(self.projects)): + self.readout_projects.append( + nn.Sequential(nn.Linear(2 * in_channels, in_channels), nn.GELU()) + ) + + self.scratch = _make_scratch( + out_channels, + features, + groups=1, + expand=False, + ) + + self.scratch.stem_transpose = None + + self.scratch.refinenet1 = _make_fusion_block(features, use_bn) + self.scratch.refinenet2 = _make_fusion_block(features, use_bn) + self.scratch.refinenet3 = _make_fusion_block(features, use_bn) + self.scratch.refinenet4 = _make_fusion_block(features, use_bn) + + # not used + del self.scratch.refinenet4.resConfUnit1 + + head_features_1 = features + head_features_2 = 16 + + if not self.return_feature: + self.scratch.output_conv = nn.Sequential( + nn.Conv2d( + head_features_1, + head_features_1 // 2, + 3, + 1, + 1, + padding_mode="replicate", + ), + nn.GELU(), + nn.Conv2d( + head_features_1 // 2, + head_features_2, + kernel_size=3, + stride=1, + padding=1, + padding_mode="replicate", + ), + nn.GELU(), + nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0), + ) + + # init delta depth as zero + nn.init.zeros_(self.scratch.output_conv[-1].weight) + nn.init.zeros_(self.scratch.output_conv[-1].bias) + + def forward( + self, + out_features, + downsample_factor=8, + cnn_features=None, + mv_features=None, + depth=None, + ): + out = [] + for i, x in enumerate(out_features): + x = self.projects[i](x) + x = self.resize_layers[i](x) + + out.append(x) + + # 1/2, 1/4, 1/8, 1/16 + layer_1, layer_2, layer_3, layer_4 = out + + if self.concat_features: + if not self.return_feature: + assert depth is not None + + if self.downsample_factor == 4 and self.num_scales == 1: + concat1 = torch.cat((cnn_features[0], cnn_features[1], layer_1), dim=1) + elif self.downsample_factor == 2 and self.num_scales == 2: + if self.return_feature: + concat1 = torch.cat( + (cnn_features[0], cnn_features[1], mv_features[0], layer_1), + dim=1, + ) + else: + concat1 = torch.cat( + ( + cnn_features[0], + cnn_features[1], + mv_features[0], + depth, + layer_1, + ), + dim=1, + ) + else: + concat1 = torch.cat((cnn_features[0], layer_1), dim=1) + layer_1 = self.concat_projects[0](concat1) # 1/2 + + if self.downsample_factor == 4 and self.num_scales == 2: + assert isinstance(mv_features, list) + if self.return_feature: + concat2 = torch.cat( + (cnn_features[1], layer_2, mv_features[0]), dim=1 + ) + else: + concat2 = torch.cat( + (cnn_features[1], layer_2, mv_features[0], depth), dim=1 + ) + layer_2 = self.concat_projects[1](concat2) # 1/4 + + concat3 = torch.cat((cnn_features[2], layer_3, mv_features[1]), dim=1) + layer_3 = self.concat_projects[2](concat3) # 1/8 + elif self.downsample_factor == 2 and self.num_scales == 2: + assert isinstance(mv_features, list) + concat2 = torch.cat((cnn_features[2], layer_2, mv_features[1]), dim=1) + layer_2 = self.concat_projects[1](concat2) # 1/4 + + concat3 = layer_3 + layer_3 = self.concat_projects[2](concat3) # 1/8 + elif self.downsample_factor == 4 and self.num_scales == 1: + if self.return_feature: + concat2 = torch.cat((cnn_features[2], layer_2, mv_features), dim=1) + else: + concat2 = torch.cat( + (cnn_features[2], layer_2, mv_features, depth), dim=1 + ) + layer_2 = self.concat_projects[1](concat2) # 1/4 + + concat3 = layer_3 + layer_3 = self.concat_projects[2](concat3) # 1/8 + else: + concat2 = torch.cat((cnn_features[1], layer_2), dim=1) + layer_2 = self.concat_projects[1](concat2) # 1/4 + + if self.return_feature: + concat3 = torch.cat((cnn_features[2], layer_3, mv_features), dim=1) + else: + concat3 = torch.cat( + (cnn_features[2], layer_3, mv_features, depth), dim=1 + ) + layer_3 = self.concat_projects[2](concat3) # 1/8 + else: + if self.concat_cnn_features: + assert cnn_features is not None + assert len(cnn_features) == 3 # 1/2, 1/4, 1/8 + cnn_features = [ + self.cnn_projects[i](f) for i, f in enumerate(cnn_features) + ] + + layer_1 = layer_1 + cnn_features[0] # 1/2 + layer_2 = layer_2 + cnn_features[1] # 1/4 + layer_3 = layer_3 + cnn_features[2] # 1/8 + + if self.concat_mv_features: + # 1/8 + mv_features = self.mv_projects(mv_features) + + layer_3 = layer_3 + mv_features # 1/8 + + layer_1_rn = self.scratch.layer1_rn(layer_1) + layer_2_rn = self.scratch.layer2_rn(layer_2) + layer_3_rn = self.scratch.layer3_rn(layer_3) + layer_4_rn = self.scratch.layer4_rn(layer_4) + + path_4 = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:]) # 1/8 + path_3 = self.scratch.refinenet3( + path_4, layer_3_rn, size=layer_2_rn.shape[2:] + ) # 1/4 + path_2 = self.scratch.refinenet2( + path_3, layer_2_rn, size=layer_1_rn.shape[2:] + ) # 1/2 + path_1 = self.scratch.refinenet1(path_2, layer_1_rn) # 1 + + if self.return_feature: + return path_1 + + out = self.scratch.output_conv(path_1) + + return out diff --git a/src/model/encoder/unimatch/ldm_unet/__init__.py b/src/model/encoder/unimatch/ldm_unet/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/model/encoder/unimatch/ldm_unet/attention.py b/src/model/encoder/unimatch/ldm_unet/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..d8765df2ba0fc027b7665cdcc8d586998ace2fd2 --- /dev/null +++ b/src/model/encoder/unimatch/ldm_unet/attention.py @@ -0,0 +1,129 @@ +import torch +import torch.nn.functional as F +from torch import nn + + +class GEGLU(nn.Module): + def __init__(self, dim_in, dim_out): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out * 2) + + def forward(self, x): + x, gate = self.proj(x).chunk(2, dim=-1) + return x * F.gelu(gate) + + +class FeedForward(nn.Module): + def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0): + super().__init__() + inner_dim = int(dim * mult) + dim_out = dim if dim_out is None else dim_out + project_in = ( + nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU()) + if not glu + else GEGLU(dim, inner_dim) + ) + self.net = nn.Sequential( + project_in, + nn.Dropout(dropout), + nn.Linear(inner_dim, dim_out), + ) + + def forward(self, x): + return self.net(x) + + +def zero_module(module): + for p in module.parameters(): + p.detach().zero_() + return module + + +def Normalize(in_channels): + return nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + + +class CrossAttention(nn.Module): + def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0): + super().__init__() + inner_dim = dim_head * heads + context_dim = query_dim if context_dim is None else context_dim + + self.heads = heads + self.to_q = nn.Linear(query_dim, inner_dim, bias=False) + self.to_k = nn.Linear(context_dim, inner_dim, bias=False) + self.to_v = nn.Linear(context_dim, inner_dim, bias=False) + self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)) + + def forward(self, x, context=None, mask=None): + b, n, _ = x.shape + h = self.heads + context = x if context is None else context + + q = self.to_q(x).reshape(b, n, h, -1).transpose(1, 2) + k = self.to_k(context).reshape(b, context.shape[1], h, -1).transpose(1, 2) + v = self.to_v(context).reshape(b, context.shape[1], h, -1).transpose(1, 2) + + attn_mask = None + if mask is not None: + attn_mask = mask.reshape(b, -1).to(torch.bool)[:, None, None, :] + + out = F.scaled_dot_product_attention( + q, k, v, attn_mask=attn_mask, dropout_p=0.0 + ) + out = out.transpose(1, 2).reshape(b, n, -1) + return self.to_out(out) + + +class BasicTransformerBlock(nn.Module): + def __init__(self, dim, n_heads, d_head, dropout=0.0, context_dim=None, gated_ff=True): + super().__init__() + self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout) + self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff) + self.attn2 = CrossAttention( + query_dim=dim, + context_dim=context_dim, + heads=n_heads, + dim_head=d_head, + dropout=dropout, + ) + self.norm1 = nn.LayerNorm(dim) + self.norm2 = nn.LayerNorm(dim) + self.norm3 = nn.LayerNorm(dim) + + def forward(self, x, context=None): + x = self.attn1(self.norm1(x)) + x + x = self.attn2(self.norm2(x), context=context) + x + x = self.ff(self.norm3(x)) + x + return x + + +class SpatialTransformer(nn.Module): + def __init__(self, in_channels, n_heads, d_head, depth=1, dropout=0.0, context_dim=None): + super().__init__() + inner_dim = n_heads * d_head + self.norm = Normalize(in_channels) + self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1) + self.transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + inner_dim, + n_heads, + d_head, + dropout=dropout, + context_dim=context_dim, + ) + for _ in range(depth) + ] + ) + self.proj_out = zero_module(nn.Conv2d(inner_dim, in_channels, kernel_size=1)) + + def forward(self, x, context=None): + b, _, h, w = x.shape + residual = x + x = self.proj_in(self.norm(x)) + x = x.flatten(2).transpose(1, 2) + for block in self.transformer_blocks: + x = block(x, context=context) + x = x.transpose(1, 2).reshape(b, -1, h, w) + return self.proj_out(x) + residual diff --git a/src/model/encoder/unimatch/ldm_unet/cross_attention.py b/src/model/encoder/unimatch/ldm_unet/cross_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..3f7864c656c6390463f8be58f4c7bac3fc22a0ed --- /dev/null +++ b/src/model/encoder/unimatch/ldm_unet/cross_attention.py @@ -0,0 +1,147 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class CrossAttention(nn.Module): + def __init__( + self, + in_dim1, + in_dim2, + dim=128, + out_dim=None, + num_heads=4, + qkv_bias=False, + proj_bias=False, + ): + super().__init__() + + if out_dim is None: + out_dim = in_dim1 + + self.num_heads = num_heads + self.dim = dim + self.q = nn.Linear(in_dim1, dim, bias=qkv_bias) + self.kv = nn.Linear(in_dim2, dim * 2, bias=qkv_bias) + self.proj = nn.Linear(dim, out_dim, bias=proj_bias) + + def forward(self, x, y): + c = self.dim + b, n1, c1 = x.shape + n2, c2 = y.shape[1:] + + q = self.q(x).reshape(b, n1, self.num_heads, c // self.num_heads) + kv = self.kv(y).reshape(b, n2, 2, self.num_heads, c // self.num_heads) + k, v = kv.unbind(2) + + x = F.scaled_dot_product_attention( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + dropout_p=0.0, + ) + x = x.transpose(1, 2).reshape(b, n1, c) + + x = self.proj(x) + + return x + + +class UNetCrossAttentionBlock(nn.Module): + def __init__(self, + in_dim1, + in_dim2, + dim=128, + out_dim=None, + num_heads=4, + qkv_bias=False, + proj_bias=False, + with_ffn=False, + concat_cross_attn=False, + concat_output=False, + no_cross_attn=False, + with_norm=False, + concat_conv3x3=False, + ): + super().__init__() + + out_dim = out_dim or in_dim1 + + self.no_cross_attn = no_cross_attn + self.with_norm = with_norm + + if no_cross_attn: + if concat_conv3x3: + self.proj = nn.Conv2d(in_dim1 + in_dim2, out_dim, 3, 1, 1) + else: + self.proj = nn.Conv2d(in_dim1 + in_dim2, out_dim, 1) + else: + self.with_ffn = with_ffn + self.concat_cross_attn = concat_cross_attn + self.concat_output = concat_output + + self.cross_attn = CrossAttention( + in_dim1=in_dim1, + in_dim2=in_dim2, + dim=dim, + out_dim=out_dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + ) + + if with_norm: + self.norm1 = nn.LayerNorm(out_dim) + else: + self.norm1 = nn.Identity() + + if with_ffn: + in_channels = out_dim + in_dim1 if concat_cross_attn else in_dim1 + ffn_dim_expansion = 4 + self.mlp = nn.Sequential( + nn.Linear(in_channels, in_channels * ffn_dim_expansion, bias=False), + nn.GELU(), + nn.Linear(in_channels * ffn_dim_expansion, in_dim1, bias=False), + ) + + if with_norm: + self.norm2 = nn.LayerNorm(in_dim1) + else: + self.norm2 = nn.Identity() + + if self.concat_output: + self.out = nn.Linear(out_dim + in_dim1, in_dim1) + + def forward(self, x, y): + # x: [B, C, H, W] + # y: [B, N, C] or [B, C, H, W] + + if self.no_cross_attn: + assert x.dim() == 4 and y.dim() == 4 + if y.shape[2:] != x.shape[2:]: + y = F.interpolate(y, x.shape[2:], mode='bilinear', align_corners=True) + return self.proj(torch.cat((x, y), dim=1)) + + identity = x + + b, c, h, w = x.size() + x = x.view(b, c, -1).permute(0, 2, 1) + + cross_attn = self.norm1(self.cross_attn(x, y)) + + if self.with_ffn: + if self.concat_cross_attn: + concat = torch.cat((x, cross_attn), dim=-1) + else: + concat = x + cross_attn + + cross_attn = self.norm2(self.mlp(concat)) + + if self.concat_output: + return self.out(torch.cat((x, cross_attn), dim=-1)) + + # reshape back + cross_attn = cross_attn.view(b, h, w, c).permute(0, 3, 1, 2) # [B, C, H, W] + + return identity + cross_attn + diff --git a/src/model/encoder/unimatch/ldm_unet/unet.py b/src/model/encoder/unimatch/ldm_unet/unet.py new file mode 100644 index 0000000000000000000000000000000000000000..a00a3cc43771bf4c2e8c02a59f04cd31af3e3b30 --- /dev/null +++ b/src/model/encoder/unimatch/ldm_unet/unet.py @@ -0,0 +1,961 @@ +from abc import abstractmethod +import torch.nn as nn +import torch.nn.functional as F +import torch + +from .util import ( + conv_nd, + avg_pool_nd, + zero_module, + normalization, +) +from .attention import SpatialTransformer + +from .cross_attention import UNetCrossAttentionBlock + + +class TimestepBlock(nn.Module): + """ + Any module where forward() takes timestep embeddings as a second argument. + """ + + @abstractmethod + def forward(self, x, emb): + """ + Apply the module to `x` given `emb` timestep embeddings. + """ + + +class TimestepEmbedSequential(nn.Sequential, TimestepBlock): + """ + A sequential module that passes timestep embeddings to the children that + support it as an extra input. + """ + + def forward(self, x, emb, context=None): + for layer in self: + if isinstance(layer, TimestepBlock): + x = layer(x, emb) + elif isinstance(layer, SpatialTransformer): + x = layer(x, context) + else: + x = layer(x) + return x + + +class Upsample(nn.Module): + """ + An upsampling layer with an optional convolution. + :param channels: channels in the inputs and outputs. + :param use_conv: a bool determining if a convolution is applied. + :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then + upsampling occurs in the inner-two dimensions. + """ + + def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1, + downsample_3ddim=False, # downsample all 3d dims instead of only spatial dims + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.dims = dims + if use_conv: + self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=padding) + + self.downsample_3ddim = downsample_3ddim + + def forward(self, x, y=None): + assert x.shape[1] == self.channels + if self.dims == 3 and not self.downsample_3ddim: + x = F.interpolate( + x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest" + ) + else: + x = F.interpolate(x, scale_factor=2, mode="nearest") + if self.use_conv: + x = self.conv(x) + return x + + +class Downsample(nn.Module): + """ + A downsampling layer with an optional convolution. + :param channels: channels in the inputs and outputs. + :param use_conv: a bool determining if a convolution is applied. + :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then + downsampling occurs in the inner-two dimensions. + """ + + def __init__(self, channels, use_conv, dims=2, out_channels=None,padding=1, + downsample_3ddim=False, # downsample all 3d dims instead of only spatial dims + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.dims = dims + stride = 2 if dims != 3 else (1, 2, 2) + + if downsample_3ddim: + assert dims == 3 + stride = 2 + + if use_conv: + self.op = conv_nd( + dims, self.channels, self.out_channels, 3, stride=stride, padding=padding + ) + else: + assert self.channels == self.out_channels + self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride) + + def forward(self, x, y=None): + assert x.shape[1] == self.channels + return self.op(x) + + +class ResBlock(TimestepBlock): + """ + A residual block that can optionally change the number of channels. + :param channels: the number of input channels. + :param emb_channels: the number of timestep embedding channels. + :param dropout: the rate of dropout. + :param out_channels: if specified, the number of out channels. + :param use_conv: if True and out_channels is specified, use a spatial + convolution instead of a smaller 1x1 convolution to change the + channels in the skip connection. + :param dims: determines if the signal is 1D, 2D, or 3D. + :param up: if True, use this block for upsampling. + :param down: if True, use this block for downsampling. + """ + + def __init__( + self, + channels, + emb_channels, + dropout, + out_channels=None, + use_conv=False, + use_scale_shift_norm=False, + dims=2, + up=False, + down=False, + postnorm=False, + channels_per_group=None, + kernel_size=3, + ): + super().__init__() + self.channels = channels + self.emb_channels = emb_channels + self.dropout = dropout + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_scale_shift_norm = use_scale_shift_norm + + if postnorm: + self.in_layers = nn.Sequential( + conv_nd(dims, channels, self.out_channels, kernel_size, padding=(kernel_size - 1) // 2), + normalization(self.out_channels, channels_per_group=channels_per_group), + nn.SiLU(), + ) + else: + self.in_layers = nn.Sequential( + normalization(channels, channels_per_group=channels_per_group), + nn.SiLU(), + conv_nd(dims, channels, self.out_channels, kernel_size, padding=(kernel_size - 1) // 2), + ) + + self.updown = up or down + + if up: + self.h_upd = Upsample(channels, False, dims) + self.x_upd = Upsample(channels, False, dims) + elif down: + self.h_upd = Downsample(channels, False, dims) + self.x_upd = Downsample(channels, False, dims) + else: + self.h_upd = self.x_upd = nn.Identity() + + if postnorm: + self.out_layers = nn.Sequential( + conv_nd(dims, self.out_channels, self.out_channels, kernel_size, padding=(kernel_size - 1) // 2), + zero_module( + normalization(self.out_channels, channels_per_group=channels_per_group), + ), + nn.SiLU(), + ) + else: + self.out_layers = nn.Sequential( + normalization(self.out_channels, channels_per_group=channels_per_group), + nn.SiLU(), + nn.Dropout(p=dropout), + zero_module( + conv_nd(dims, self.out_channels, self.out_channels, kernel_size, padding=(kernel_size - 1) // 2) + ), + ) + + if self.out_channels == channels: + self.skip_connection = nn.Identity() + elif use_conv: + self.skip_connection = conv_nd( + dims, channels, self.out_channels, kernel_size, padding=(kernel_size - 1) // 2 + ) + else: + self.skip_connection = conv_nd(dims, channels, self.out_channels, 1) + + def forward(self, x, emb=None): + """ + Apply the block to a Tensor, conditioned on a timestep embedding. + :param x: an [N x C x ...] Tensor of features. + :param emb: an [N x emb_channels] Tensor of timestep embeddings. + :return: an [N x C x ...] Tensor of outputs. + """ + return self._forward(x, emb) + + def _forward(self, x, emb=None): + if self.updown: + in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] + h = in_rest(x) + h = self.h_upd(h) + x = self.x_upd(x) + h = in_conv(h) + else: + h = self.in_layers(x) + h = self.out_layers(h) + return self.skip_connection(x) + h + + +class AttentionBlock(nn.Module): + """ + An attention block that allows spatial positions to attend to each other. + Originally ported from here, but adapted to the N-d case. + https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66. + """ + + def __init__( + self, + channels, + num_heads=1, + num_head_channels=-1, + use_new_attention_order=False, + postnorm=False, + channels_per_group=None, + num_frames=2, + use_cross_view_self_attn=False, + ): + super().__init__() + + self.channels = channels + if num_head_channels == -1: + self.num_heads = num_heads + else: + assert ( + channels % num_head_channels == 0 + ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}" + self.num_heads = channels // num_head_channels + self.qkv = conv_nd(1, channels, channels * 3, 1) + if use_new_attention_order: + # split qkv before split heads + self.attention = QKVAttention(self.num_heads) + else: + # split heads before split qkv + self.attention = QKVAttentionLegacy( + self.num_heads, + n_frames=num_frames, + use_cross_view_self_attn=use_cross_view_self_attn, + ) + + if postnorm: + self.proj_out = conv_nd(1, channels, channels, 1) + self.norm = zero_module(normalization(channels, channels_per_group=channels_per_group)) + else: + self.norm = normalization(channels, channels_per_group=channels_per_group) + self.proj_out = zero_module(conv_nd(1, channels, channels, 1)) + + self.postnorm = postnorm + + def forward(self, x): + return self._forward(x) + + def _forward(self, x): + b, c, *spatial = x.shape + x = x.reshape(b, c, -1) + + if self.postnorm: + qkv = self.qkv(x) + h = self.attention(qkv) + h = self.proj_out(h) + h = self.norm(h) + else: + qkv = self.qkv(self.norm(x)) + h = self.attention(qkv) + h = self.proj_out(h) + + return (x + h).reshape(b, c, *spatial) + + +class CrossAttentionBlock(nn.Module): + """Cross-attention conditioning for 3D UNet features.""" + + def __init__( + self, + channels, + condition_channels, + num_heads=8, + proj_channels=512, + num_views=3, + num_head_channels=-1, + use_new_attention_order=False, + channels_per_group=None, + with_norm=False, + tanh_gating=False, + ffn_after_cross_attn=False, + ): + super().__init__() + + self.channels = channels + self.num_head = num_heads + self.num_views = num_views + self.proj_channels = proj_channels + self.with_norm = with_norm + self.tanh_gating = tanh_gating + self.ffn_after_cross_attn = ffn_after_cross_attn + + self.q_proj = nn.Linear(channels, proj_channels) + self.k_proj = nn.Linear(condition_channels, proj_channels) + self.v_proj = nn.Linear(condition_channels, proj_channels) + + if self.tanh_gating: + self.out_proj = conv_nd(3, proj_channels, channels, 1) + self.attn_gate = nn.Parameter(torch.tensor([0.])) + else: + if self.with_norm: + self.out_proj = conv_nd(3, proj_channels, channels, 1) + self.norm = zero_module(normalization(channels, channels_per_group=channels_per_group)) + else: + self.out_proj = zero_module(conv_nd(3, proj_channels, channels, 1)) + + if self.ffn_after_cross_attn: + self.ffn_gate = nn.Parameter(torch.tensor([0.])) + self.ffn = nn.Sequential( + nn.Conv3d(channels, channels * 4, 1, 1, 0), + normalization(channels * 4, channels_per_group=channels_per_group), + nn.GELU(), + nn.Conv3d(channels * 4, channels, 1, 1, 0), + ) + + def forward(self, x, y=None): + return self._forward(x, y) + + def _forward(self, x, y=None): + # x: [B, C, D, H, W], feature + # y: [B, H*W, D, C] + assert x.dim() == 5 and y.dim() == 4 + + b, c1, d, h, w = x.size() + lx = h * w + ly = y.size(1) + c2 = y.size(-1) + + identity = x + + x = x.permute(0, 2, 3, 4, 1).reshape(b * d, lx, c1) + y = y.permute(0, 2, 1, 3).reshape(b * d, ly, c2) + + c = self.proj_channels + + assert c % self.num_head == 0 + head_dim = c // self.num_head + q = self.q_proj(x).reshape(b * d, lx, self.num_head, head_dim) + k = self.k_proj(y).reshape(b * d, ly, self.num_head, head_dim) + v = self.v_proj(y).reshape(b * d, ly, self.num_head, head_dim) + + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0) + out = out.transpose(1, 2).reshape(b * d, lx, c) + out = out.reshape(b, d, h, w, c).permute(0, 4, 1, 2, 3) + + if self.tanh_gating: + out = self.attn_gate.tanh() * self.out_proj(out) + else: + if self.with_norm: + out = self.out_proj(out) + out = self.norm(out) + else: + out = self.out_proj(out) + + out = identity + out + + if self.ffn_after_cross_attn: + out = out + self.ffn_gate.tanh() * self.ffn(out) + + return out + + + +class QKVAttentionLegacy(nn.Module): + """ + SDPA implementation of the original legacy QKV packing order. + """ + + def __init__(self, n_heads, n_frames=2, use_cross_view_self_attn=False): + super().__init__() + self.n_heads = n_heads + self.n_frames = n_frames + self.use_cross_view_self_attn = use_cross_view_self_attn + + def forward(self, qkv, num_views=None): + """ + Apply attention with the legacy [Q,K,V] packed per head. + :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs. + :return: an [N x (H * C) x T] tensor after attention. + """ + + # Move the view dimension into T for cross-view self-attention. + # Input batch is flattened in '(b v)' order. + if self.use_cross_view_self_attn: + n_views = self.n_frames if num_views is None else num_views + bv, width, length = qkv.shape + b = bv // n_views + qkv = ( + qkv.reshape(b, n_views, width, length) + .permute(0, 2, 1, 3) + .reshape(b, width, n_views * length) + ) + + bs, width, length = qkv.shape + assert width % (3 * self.n_heads) == 0 + ch = width // (3 * self.n_heads) + q, k, v = qkv.reshape(bs, self.n_heads, ch * 3, length).split(ch, dim=2) + q = q.transpose(-2, -1) + k = k.transpose(-2, -1) + v = v.transpose(-2, -1) + a = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0) + a = a.transpose(-2, -1).reshape(bs, -1, length) + + # Move the view dimension back to the original flattened '(b v)' order. + if self.use_cross_view_self_attn: + t = length // n_views + a = ( + a.reshape(bs, -1, n_views, t) + .permute(0, 2, 1, 3) + .reshape(bs * n_views, -1, t) + ) + + return a + + +class QKVAttention(nn.Module): + """ + SDPA implementation for the newer [Q,K,V] packing order. + """ + + def __init__(self, n_heads): + super().__init__() + self.n_heads = n_heads + + def forward(self, qkv): + """ + Apply attention with the newer [Q,K,V] packed-before-heads layout. + :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs. + :return: an [N x (H * C) x T] tensor after attention. + """ + bs, width, length = qkv.shape + assert width % (3 * self.n_heads) == 0 + ch = width // (3 * self.n_heads) + q, k, v = qkv.reshape(bs, 3, self.n_heads, ch, length).unbind(dim=1) + q = q.transpose(-2, -1) + k = k.transpose(-2, -1) + v = v.transpose(-2, -1) + a = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0) + return a.transpose(-2, -1).reshape(bs, -1, length) + + +class UNetModel(nn.Module): + """ + The full UNet model with attention and timestep embedding. + :param in_channels: channels in the input Tensor. + :param model_channels: base channel count for the model. + :param out_channels: channels in the output Tensor. + :param num_res_blocks: number of residual blocks per downsample. + :param attention_resolutions: a collection of downsample rates at which + attention will take place. May be a set, list, or tuple. + For example, if this contains 4, then at 4x downsampling, attention + will be used. + :param dropout: the dropout probability. + :param channel_mult: channel multiplier for each level of the UNet. + :param conv_resample: if True, use learned convolutions for upsampling and + downsampling. + :param dims: determines if the signal is 1D, 2D, or 3D. + :param num_classes: if specified (as an int), then this model will be + class-conditional with `num_classes` classes. + :param num_heads: the number of attention heads in each attention layer. + :param num_heads_channels: if specified, ignore num_heads and instead use + a fixed channel width per attention head. + :param num_heads_upsample: works with num_heads to set a different number + of heads for upsampling. Deprecated. + :param use_scale_shift_norm: use a FiLM-like conditioning mechanism. + :param resblock_updown: use residual blocks for up/downsampling. + :param use_new_attention_order: use a different attention pattern for potentially + increased efficiency. + """ + + def __init__( + self, + image_size, + in_channels, + model_channels, + out_channels, + num_res_blocks, + attention_resolutions, + dropout=0, + channel_mult=(1, 2, 4, 8), + conv_resample=True, + dims=2, + middle_block_attn=False, # use attn in middle block + middle_block_no_identity=False, # some previous models are trained without the identity layer + postnorm=False, # default prenorm doesn't converge + attn_prenorm=False, # try postnorm for resblock and prenorm for attn + downsample_3ddim=False, # downsample all 3d dims instead of only spatial dims + zero_final_layer=False, # init zero final output layer + channels_per_group=None, + num_classes=None, + num_heads=-1, + num_head_channels=-1, + num_heads_upsample=-1, + use_scale_shift_norm=False, + resblock_updown=False, + use_new_attention_order=False, + use_spatial_transformer=False, # custom transformer support + transformer_depth=1, # custom transformer support + context_dim=None, # custom transformer support + n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model + legacy=True, + cross_attn_condition=False, + tanh_gating=False, + ffn_after_cross_attn=False, + cross_attn_with_norm=False, + condition_channels=384, + condition_num_views=3, + no_self_attn=False, + conv_kernel_size=3, + concat_condition=False, + concat_conv3x3=False, + num_frames=2, + use_cross_view_self_attn=False, + downsample_factor=None, + ): + super().__init__() + if use_spatial_transformer: + assert context_dim is not None, 'context_dim is required when using spatial transformer cross-attention.' + + if context_dim is not None: + assert use_spatial_transformer, 'cross-attention conditioning requires use_spatial_transformer=True.' + from omegaconf.listconfig import ListConfig + if type(context_dim) == ListConfig: + context_dim = list(context_dim) + + if num_heads_upsample == -1: + num_heads_upsample = num_heads + + if num_heads == -1: + assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set' + + if num_head_channels == -1: + assert num_heads != -1, 'Either num_heads or num_head_channels has to be set' + + self.image_size = image_size + self.in_channels = in_channels + self.model_channels = model_channels + self.out_channels = out_channels + self.num_res_blocks = num_res_blocks + self.attention_resolutions = attention_resolutions + self.dropout = dropout + self.channel_mult = channel_mult + self.conv_resample = conv_resample + self.num_classes = num_classes + self.num_heads = num_heads + self.num_head_channels = num_head_channels + self.num_heads_upsample = num_heads_upsample + self.predict_codebook_ids = n_embed is not None + + self.middle_block_attn = middle_block_attn + + self.middle_block_no_identity = middle_block_no_identity + + # output lower resolution gaussians + self.downsample_factor = downsample_factor + + time_embed_dim = model_channels * 4 + self.cross_attn_condition = cross_attn_condition + + self.input_blocks = nn.ModuleList( + [ + nn.Sequential( + conv_nd(dims, in_channels, model_channels, 3, padding=1) + ) + ] + ) + self._feature_size = model_channels + input_block_chans = [model_channels] + ch = model_channels + ds = 1 + for level, mult in enumerate(channel_mult): + for _ in range(num_res_blocks): + layers = [ + ResBlock( + ch, + time_embed_dim, + dropout, + out_channels=mult * model_channels, + dims=dims, + use_scale_shift_norm=use_scale_shift_norm, + postnorm=postnorm, + channels_per_group=channels_per_group, + kernel_size=conv_kernel_size, + ) + ] + ch = mult * model_channels + if ds in attention_resolutions: + if num_head_channels == -1: + dim_head = ch // num_heads + else: + num_heads = ch // num_head_channels + dim_head = num_head_channels + if legacy: + dim_head = ch // num_heads if use_spatial_transformer else num_head_channels + + if not no_self_attn: # only cross attn, without self attn + layers.append( + AttentionBlock( + ch, + num_heads=num_heads, + num_head_channels=dim_head, + use_new_attention_order=use_new_attention_order, + postnorm=False if attn_prenorm else postnorm, + channels_per_group=channels_per_group, + num_frames=num_frames, + use_cross_view_self_attn=use_cross_view_self_attn, + ) if not use_spatial_transformer else SpatialTransformer( + ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim + ) + ) + + if cross_attn_condition: + layers.append( + UNetCrossAttentionBlock(ch, + condition_channels, + dim=256, + no_cross_attn=concat_condition, + with_norm=cross_attn_with_norm, + concat_conv3x3=concat_conv3x3, + ) + ) + + self.input_blocks.append(nn.Sequential(*layers)) + self._feature_size += ch + input_block_chans.append(ch) + if level != len(channel_mult) - 1: + out_ch = ch + self.input_blocks.append( + nn.Sequential( + ResBlock( + ch, + time_embed_dim, + dropout, + out_channels=out_ch, + dims=dims, + use_scale_shift_norm=use_scale_shift_norm, + down=True, + postnorm=postnorm, + channels_per_group=channels_per_group, + kernel_size=conv_kernel_size, + ) + if resblock_updown + else Downsample( + ch, conv_resample, dims=dims, out_channels=out_ch, + downsample_3ddim=downsample_3ddim, + ) + ) + ) + ch = out_ch + input_block_chans.append(ch) + ds *= 2 + self._feature_size += ch + + if num_head_channels == -1: + dim_head = ch // num_heads + else: + num_heads = ch // num_head_channels + dim_head = num_head_channels + if legacy: + dim_head = ch // num_heads if use_spatial_transformer else num_head_channels + + if self.middle_block_attn: + self.middle_block = nn.Sequential( + ResBlock( + ch, + time_embed_dim, + dropout, + dims=dims, + use_scale_shift_norm=use_scale_shift_norm, + postnorm=postnorm, + channels_per_group=channels_per_group, + kernel_size=conv_kernel_size, + ), + AttentionBlock( + ch, + num_heads=num_heads, + num_head_channels=dim_head, + use_new_attention_order=use_new_attention_order, + postnorm=False if attn_prenorm else postnorm, + channels_per_group=channels_per_group, + num_frames=num_frames, + use_cross_view_self_attn=use_cross_view_self_attn, + ) if not use_spatial_transformer else SpatialTransformer( + ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim + ), + # cross attention condition + UNetCrossAttentionBlock(ch, + condition_channels, + dim=256, + no_cross_attn=concat_condition, + with_norm=cross_attn_with_norm, + concat_conv3x3=concat_conv3x3, + ) if cross_attn_condition else nn.Identity(), + ResBlock( + ch, + time_embed_dim, + dropout, + dims=dims, + use_scale_shift_norm=use_scale_shift_norm, + postnorm=postnorm, + channels_per_group=channels_per_group, + kernel_size=conv_kernel_size, + ), + ) + else: + if self.middle_block_no_identity: + self.middle_block = nn.Sequential( + ResBlock( + ch, + time_embed_dim, + dropout, + dims=dims, + use_scale_shift_norm=use_scale_shift_norm, + postnorm=postnorm, + channels_per_group=channels_per_group, + kernel_size=conv_kernel_size, + ), + ResBlock( + ch, + time_embed_dim, + dropout, + dims=dims, + use_scale_shift_norm=use_scale_shift_norm, + postnorm=postnorm, + channels_per_group=channels_per_group, + kernel_size=conv_kernel_size, + ), + ) + else: + self.middle_block = nn.Sequential( + ResBlock( + ch, + time_embed_dim, + dropout, + dims=dims, + use_scale_shift_norm=use_scale_shift_norm, + postnorm=postnorm, + channels_per_group=channels_per_group, + kernel_size=conv_kernel_size, + ), + UNetCrossAttentionBlock(ch, + condition_channels, + dim=256, + no_cross_attn=concat_condition, + with_norm=cross_attn_with_norm, + concat_conv3x3=concat_conv3x3, + ) if cross_attn_condition else nn.Identity(), + ResBlock( + ch, + time_embed_dim, + dropout, + dims=dims, + use_scale_shift_norm=use_scale_shift_norm, + postnorm=postnorm, + channels_per_group=channels_per_group, + kernel_size=conv_kernel_size, + ), + ) + self._feature_size += ch + + self.output_blocks = nn.ModuleList([]) + for level, mult in list(enumerate(channel_mult))[::-1]: + for i in range(num_res_blocks + 1): + ich = input_block_chans.pop() + layers = [ + ResBlock( + ch + ich, + time_embed_dim, + dropout, + out_channels=model_channels * mult, + dims=dims, + use_scale_shift_norm=use_scale_shift_norm, + postnorm=postnorm, + channels_per_group=channels_per_group, + kernel_size=conv_kernel_size, + ) + ] + ch = model_channels * mult + if ds in attention_resolutions: + if num_head_channels == -1: + dim_head = ch // num_heads + else: + num_heads = ch // num_head_channels + dim_head = num_head_channels + if legacy: + dim_head = ch // num_heads if use_spatial_transformer else num_head_channels + + if not no_self_attn: # only cross attn, without self attn + layers.append( + AttentionBlock( + ch, + num_heads=num_heads_upsample, + num_head_channels=dim_head, + use_new_attention_order=use_new_attention_order, + postnorm=False if attn_prenorm else postnorm, + channels_per_group=channels_per_group, + num_frames=num_frames, + use_cross_view_self_attn=use_cross_view_self_attn, + ) if not use_spatial_transformer else SpatialTransformer( + ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim + ) + ) + + if cross_attn_condition: + layers.append( + UNetCrossAttentionBlock(ch, + condition_channels, + dim=256, + no_cross_attn=concat_condition, + with_norm=cross_attn_with_norm, + concat_conv3x3=concat_conv3x3, + ) + ) + + if level and i == num_res_blocks: + out_ch = ch + layers.append( + ResBlock( + ch, + time_embed_dim, + dropout, + out_channels=out_ch, + dims=dims, + use_scale_shift_norm=use_scale_shift_norm, + up=True, + postnorm=postnorm, + channels_per_group=channels_per_group, + kernel_size=conv_kernel_size, + ) + if resblock_updown + else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch, + downsample_3ddim=downsample_3ddim, + ) + ) + ds //= 2 + self.output_blocks.append(nn.Sequential(*layers)) + self._feature_size += ch + + if self.downsample_factor is not None: + if self.downsample_factor == 2: + del self.output_blocks[-3:] + elif self.downsample_factor == 4: + del self.output_blocks[-5:] + else: + raise NotImplementedError + + if postnorm: + self.out = nn.Sequential( + conv_nd(dims, model_channels, out_channels, 3, padding=1), + normalization(out_channels, channels_per_group=channels_per_group) if not zero_final_layer else zero_module(normalization(out_channels, channels_per_group=channels_per_group)), + nn.SiLU(), + ) + else: + if self.downsample_factor is not None: + in_channels = self.model_channels * self.channel_mult[self.downsample_factor // 2] + model_channels = model_channels * self.channel_mult[self.downsample_factor // 2] + out_channels = model_channels + else: + in_channels = ch + self.out = nn.Sequential( + normalization(in_channels, channels_per_group=channels_per_group), + nn.SiLU(), + zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)), + ) + + + if self.predict_codebook_ids: + self.id_predictor = nn.Sequential( + normalization(ch, channels_per_group=channels_per_group), + conv_nd(dims, model_channels, n_embed, 1), + ) + + def forward(self, x, num_views=None, timesteps=None, context=None, y=None,**kwargs): + """ + Apply the model to an input batch. + :param x: an [N x C x ...] Tensor of inputs. + :param timesteps: a 1-D batch of timesteps. + :param context: conditioning plugged in via crossattn + :param y: an [N] Tensor of labels, if class-conditional. + :return: an [N x C x ...] Tensor of outputs. + """ + assert (y is not None) == ( + self.num_classes is not None + ), "must specify y if and only if the model is class-conditional" + hs = [] + emb = None + + if self.num_classes is not None: + assert y.shape == (x.shape[0],) + emb = emb + self.label_emb(y) + + # Let autocast control FP32/BF16 execution. Geometry/probability code + # outside this UNet is responsible for opting out of autocast when needed. + h = x + for module in self.input_blocks: + + if self.cross_attn_condition: + for submodule in module: + if 'UNetCrossAttentionBlock' == submodule.__class__.__name__: + h = submodule(h, context) + else: + h = submodule(h) + else: + h = module(h) + hs.append(h) + + for module in self.middle_block: + if 'UNetCrossAttentionBlock' == module.__class__.__name__: + h = module(h, context) + else: + h = module(h) + + + for module in self.output_blocks: + h = torch.cat([h, hs.pop()], dim=1) + if self.cross_attn_condition: + for submodule in module: + if 'UNetCrossAttentionBlock' == submodule.__class__.__name__: + h = submodule(h, context) + else: + h = submodule(h) + else: + h = module(h) + if self.predict_codebook_ids: + return self.id_predictor(h) + else: + return self.out(h) diff --git a/src/model/encoder/unimatch/ldm_unet/util.py b/src/model/encoder/unimatch/ldm_unet/util.py new file mode 100644 index 0000000000000000000000000000000000000000..2c6c9d1198560781993ba9ceb07810a41db0942b --- /dev/null +++ b/src/model/encoder/unimatch/ldm_unet/util.py @@ -0,0 +1,41 @@ +import torch.nn as nn + + +def zero_module(module): + """Zero out a module's parameters and return it.""" + for param in module.parameters(): + param.detach().zero_() + return module + + +def normalization(channels, channels_per_group=None): + if channels_per_group is not None: + return Float32GroupNorm(channels // channels_per_group, channels) + + num_groups = 4 if channels % 8 != 0 else 8 + return Float32GroupNorm(num_groups, channels) + + +class Float32GroupNorm(nn.GroupNorm): + def forward(self, x): + return super().forward(x.float()).to(dtype=x.dtype) + + +def conv_nd(dims, *args, **kwargs): + if dims == 1: + return nn.Conv1d(*args, **kwargs) + if dims == 2: + return nn.Conv2d(*args, **kwargs) + if dims == 3: + return nn.Conv3d(*args, **kwargs) + raise ValueError(f"unsupported dimensions: {dims}") + + +def avg_pool_nd(dims, *args, **kwargs): + if dims == 1: + return nn.AvgPool1d(*args, **kwargs) + if dims == 2: + return nn.AvgPool2d(*args, **kwargs) + if dims == 3: + return nn.AvgPool3d(*args, **kwargs) + raise ValueError(f"unsupported dimensions: {dims}") diff --git a/src/model/encoder/unimatch/matching.py b/src/model/encoder/unimatch/matching.py new file mode 100644 index 0000000000000000000000000000000000000000..48efebe90b8ec9539fa840fe6d8340fe0f4f0e92 --- /dev/null +++ b/src/model/encoder/unimatch/matching.py @@ -0,0 +1,104 @@ +import torch +import torch.nn.functional as F + + +def coords_grid(b, h, w, homogeneous=False, device=None): + y, x = torch.meshgrid( + torch.arange(h), torch.arange(w), indexing="ij" + ) # [H, W] + + stacks = [x, y] + + if homogeneous: + ones = torch.ones_like(x) # [H, W] + stacks.append(ones) + + grid = torch.stack(stacks, dim=0).float() # [2, H, W] or [3, H, W] + + grid = grid[None].expand(b, -1, -1, -1) # [B, 2, H, W] or [B, 3, H, W] + + if device is not None: + grid = grid.to(device) + + return grid + + +def warp_with_pose_depth_candidates( + feature1, + ref_intrinsics, + tgt_intrinsics, + pose, + depth, + clamp_min_depth=1e-3, + grid_sample_disable_cudnn=False, +): + """ + feature1: [B, C, H, W] + ref_intrinsics: [B, 3, 3], used to backproject reference pixels + tgt_intrinsics: [B, 3, 3], used to reproject into target-view pixels + pose: [B, 4, 4] + depth: [B, D, H, W] + """ + + assert ref_intrinsics.size(1) == ref_intrinsics.size(2) == 3 + assert tgt_intrinsics.size(1) == tgt_intrinsics.size(2) == 3 + assert pose.size(1) == pose.size(2) == 4 + assert depth.dim() == 4 + + b, d, h, w = depth.size() + c = feature1.size(1) + feature_dtype = feature1.dtype + + with torch.no_grad(), torch.autocast( + device_type=depth.device.type, enabled=False + ): + ref_intrinsics = ref_intrinsics.float() + tgt_intrinsics = tgt_intrinsics.float() + pose = pose.float() + depth = depth.float() + # pixel coordinates + grid = coords_grid( + b, h, w, homogeneous=True, device=depth.device + ) # [B, 3, H, W] + # back project to 3D and transform viewpoint + points = torch.linalg.inv(ref_intrinsics).bmm( + grid.reshape(b, 3, -1) + ) # [B, 3, H*W] + points = torch.bmm(pose[:, :3, :3], points).unsqueeze(2) * depth.reshape( + b, 1, d, h * w + ) # [B, 3, D, H*W] + points = points + pose[:, :3, -1:].unsqueeze(-1) # [B, 3, D, H*W] + # reproject to 2D image plane + points = torch.bmm(tgt_intrinsics, points.view(b, 3, -1)).view( + b, 3, d, h * w + ) # [B, 3, D, H*W] + pixel_coords = points[:, :2] / points[:, -1:].clamp( + min=clamp_min_depth + ) # [B, 2, D, H*W] + + # normalize to [-1, 1] + x_grid = 2 * pixel_coords[:, 0] / (w - 1) - 1 + y_grid = 2 * pixel_coords[:, 1] / (h - 1) - 1 + + grid = torch.stack([x_grid, y_grid], dim=-1) # [B, D, H*W, 2] + + # Disable cuDNN for large grid_sample inputs; this avoids a known PyTorch + # issue while keeping the default path for smaller inputs. + # Ref: https://github.com/pytorch/pytorch/issues/88380 + if feature1.numel() > 1000000: + grid_sample_disable_cudnn = True + with torch.backends.cudnn.flags(enabled=not grid_sample_disable_cudnn): + # grid_sample requires the feature and grid dtypes to match. Compute the + # warp in FP32 for sub-pixel grid accuracy, then return to the feature dtype. + with torch.autocast(device_type=feature1.device.type, enabled=False): + warped_feature = F.grid_sample( + feature1.float(), + grid.view(b, d * h, w, 2), + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ).view( + b, c, d, h, w + ) # [B, C, D, H, W] + + return warped_feature.to(feature_dtype) diff --git a/src/model/encoder/unimatch/mv_transformer.py b/src/model/encoder/unimatch/mv_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..493d8491c9d2aead3a360b3e58d4e2499da1dda2 --- /dev/null +++ b/src/model/encoder/unimatch/mv_transformer.py @@ -0,0 +1,694 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .utils import split_feature, merge_splits + + +def gather_view_tensor(tensor, view_indices): + view_indices = view_indices.to(device=tensor.device, dtype=torch.long) + batch_indices = torch.arange(tensor.size(0), device=tensor.device) + batch_indices = batch_indices.reshape( + tensor.size(0), *([1] * (view_indices.dim() - 1)) + ) + return tensor[batch_indices, view_indices] + + +def non_self_view_indices(num_views, device): + view_ids = torch.arange(num_views, device=device) + return torch.stack( + [torch.cat((view_ids[:i], view_ids[i + 1 :])) for i in range(num_views)], + dim=0, + ) + + +def target_view_indices(batch_size, num_views, nn_matrix, device): + if nn_matrix is not None: + return nn_matrix[..., 1:].to(device=device, dtype=torch.long) + return non_self_view_indices(num_views, device).unsqueeze(0).expand( + batch_size, -1, -1 + ) + + +def flatten_batch_view(tensor): + # Batch-major flatten: [B, V, ...] -> [B*V, ...]. + return tensor.flatten(0, 1) + + +def flatten_view_batch(tensor): + # View-major flatten: [B, V, ...] -> [V*B, ...]. + return tensor.transpose(0, 1).reshape( + tensor.size(0) * tensor.size(1), *tensor.shape[2:] + ) + + +def unflatten_view_batch(tensor, batch_size, num_views): + # Inverse of flatten_view_batch. + return tensor.reshape(num_views, batch_size, *tensor.shape[1:]).transpose(0, 1) + + +def single_head_full_attention(q, k, v): + # q, k, v: [B, L, C] + assert q.dim() == k.dim() == v.dim() == 3 + + return F.scaled_dot_product_attention( + q.unsqueeze(1), k.unsqueeze(1), v.unsqueeze(1) + ).squeeze(1) + + +def generate_shift_window_attn_mask( + input_resolution, + window_size_h, + window_size_w, + shift_size_h, + shift_size_w, + device=torch.device("cuda"), +): + # Ref: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py + # Build the shifted-window attention mask used by Swin-style local attention. + h, w = input_resolution + img_mask = torch.zeros((1, h, w, 1)).to(device) # 1 H W 1 + h_slices = ( + slice(0, -window_size_h), + slice(-window_size_h, -shift_size_h), + slice(-shift_size_h, None), + ) + w_slices = ( + slice(0, -window_size_w), + slice(-window_size_w, -shift_size_w), + slice(-shift_size_w, None), + ) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = split_feature( + img_mask, num_splits=input_resolution[-1] // window_size_w, channel_last=True + ) + + mask_windows = mask_windows.view(-1, window_size_h * window_size_w) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill( + attn_mask == 0, float(0.0) + ) + + return attn_mask + + +def single_head_split_window_attention( + q, + k, + v, + num_splits=1, + with_shift=False, + h=None, + w=None, + attn_mask=None, +): + # Ref: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py + # q, k, v: [B, L, C] for self-attention or pairwise cross-attention. + # For multi-view cross-attention, q: [B, L, C], k/v: [B, M, L, C]. + + # Multi-view cross-attention, where M is the number of target views per query. + if not (q.dim() == k.dim() == v.dim() == 3): + assert k.dim() == v.dim() == 4 + assert h is not None and w is not None + assert q.size(1) == h * w + + m = k.size(1) + + b, _, c = q.size() + + b_new = b * num_splits * num_splits + + window_size_h = h // num_splits + window_size_w = w // num_splits + + q = q.view(b, h, w, c) # [B, H, W, C] + k = k.view(b, m, h, w, c) # [B, M, H, W, C] + v = v.view(b, m, h, w, c) + + if with_shift: + assert attn_mask is not None # compute once + shift_size_h = window_size_h // 2 + shift_size_w = window_size_w // 2 + + q = torch.roll(q, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) + k = torch.roll(k, shifts=(-shift_size_h, -shift_size_w), dims=(2, 3)) + v = torch.roll(v, shifts=(-shift_size_h, -shift_size_w), dims=(2, 3)) + + q = split_feature( + q, num_splits=num_splits, channel_last=True + ) # [B*K*K, H/K, W/K, C] + k = split_feature( + k.permute(0, 2, 3, 4, 1).reshape(b, h, w, -1), + num_splits=num_splits, + channel_last=True, + ) # [B*K*K, H/K, W/K, C*M] + v = split_feature( + v.permute(0, 2, 3, 4, 1).reshape(b, h, w, -1), + num_splits=num_splits, + channel_last=True, + ) # [B*K*K, H/K, W/K, C*M] + + k = ( + k.view(b_new, h // num_splits, w // num_splits, c, m) + .permute(0, 1, 2, 4, 3) + .reshape(b_new, -1, c) + ) # [B*K*K, H/K*W/K*M, C] + v = ( + v.view(b_new, h // num_splits, w // num_splits, c, m) + .permute(0, 1, 2, 4, 3) + .reshape(b_new, -1, c) + ) # [B*K*K, H/K*W/K*M, C] + + attn_bias = None + if with_shift: + attn_bias = attn_mask.repeat(b, 1, m).unsqueeze(1).to(q.dtype) + + out = F.scaled_dot_product_attention( + q.view(b_new, -1, c).unsqueeze(1), + k.unsqueeze(1), + v.unsqueeze(1), + attn_mask=attn_bias, + ).squeeze(1) + + out = merge_splits( + out.view(b_new, h // num_splits, w // num_splits, c), + num_splits=num_splits, + channel_last=True, + ) # [B, H, W, C] + + # shift back + if with_shift: + out = torch.roll(out, shifts=(shift_size_h, shift_size_w), dims=(1, 2)) + + out = out.view(b, -1, c) + else: + # Self-attention or pairwise cross-attention. + assert q.dim() == k.dim() == v.dim() == 3 + + assert h is not None and w is not None + assert q.size(1) == h * w + + b, _, c = q.size() + + b_new = b * num_splits * num_splits + + window_size_h = h // num_splits + window_size_w = w // num_splits + + q = q.view(b, h, w, c) # [B, H, W, C] + k = k.view(b, h, w, c) + v = v.view(b, h, w, c) + + if with_shift: + assert attn_mask is not None # compute once + shift_size_h = window_size_h // 2 + shift_size_w = window_size_w // 2 + + q = torch.roll(q, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) + k = torch.roll(k, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) + v = torch.roll(v, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) + + q = split_feature( + q, num_splits=num_splits, channel_last=True + ) # [B*K*K, H/K, W/K, C] + k = split_feature(k, num_splits=num_splits, channel_last=True) + v = split_feature(v, num_splits=num_splits, channel_last=True) + + attn_bias = None + if with_shift: + attn_bias = attn_mask.repeat(b, 1, 1).unsqueeze(1).to(q.dtype) + + out = F.scaled_dot_product_attention( + q.view(b_new, -1, c).unsqueeze(1), + k.view(b_new, -1, c).unsqueeze(1), + v.view(b_new, -1, c).unsqueeze(1), + attn_mask=attn_bias, + ).squeeze(1) + + out = merge_splits( + out.view(b_new, h // num_splits, w // num_splits, c), + num_splits=num_splits, + channel_last=True, + ) # [B, H, W, C] + + # shift back + if with_shift: + out = torch.roll(out, shifts=(shift_size_h, shift_size_w), dims=(1, 2)) + + out = out.view(b, -1, c) + + return out + + +def multi_head_split_window_attention( + q, + k, + v, + num_splits=1, + with_shift=False, + h=None, + w=None, + attn_mask=None, + num_head=1, +): + """Windowed multi-head scaled dot-product attention. + + Args: + q: [B, L, C] + k: [B, S, C] + v: [B, S, C] + Returns: + out: [B, L, C] + """ + + assert h is not None and w is not None + assert q.size(1) == h * w + + b, _, c = q.size() + + b_new = b * num_splits * num_splits + + window_size_h = h // num_splits + window_size_w = w // num_splits + + q = q.view(b, h, w, c) # [B, H, W, C] + k = k.view(b, h, w, c) + v = v.view(b, h, w, c) + + assert c % num_head == 0 + + if with_shift: + assert attn_mask is not None # compute once + shift_size_h = window_size_h // 2 + shift_size_w = window_size_w // 2 + + q = torch.roll(q, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) + k = torch.roll(k, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) + v = torch.roll(v, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) + + q = split_feature(q, num_splits=num_splits) # [B*K*K, H/K, W/K, C] + k = split_feature(k, num_splits=num_splits) + v = split_feature(v, num_splits=num_splits) + + # Split channels into attention heads for PyTorch SDPA. + head_dim = c // num_head + q = q.view(b_new, -1, num_head, head_dim).transpose(1, 2) + k = k.view(b_new, -1, num_head, head_dim).transpose(1, 2) + v = v.view(b_new, -1, num_head, head_dim).transpose(1, 2) + + attn_bias = None + if with_shift: + attn_bias = attn_mask.unsqueeze(1).repeat(b, num_head, 1, 1).to(q.dtype) + + out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_bias) + + out = merge_splits( + out.transpose(1, 2).reshape(b_new, h // num_splits, w // num_splits, c), + num_splits=num_splits, + ) # [B, H, W, C] + + # shift back + if with_shift: + out = torch.roll(out, shifts=(shift_size_h, shift_size_w), dims=(1, 2)) + + out = out.view(b, -1, c) + + return out + + +class TransformerLayer(nn.Module): + def __init__( + self, + d_model=256, + nhead=1, + attention_type="swin", + no_ffn=False, + ffn_dim_expansion=4, + with_shift=False, + add_per_view_attn=False, + **kwargs, + ): + super(TransformerLayer, self).__init__() + + self.dim = d_model + self.nhead = nhead + self.attention_type = attention_type + self.no_ffn = no_ffn + self.add_per_view_attn = add_per_view_attn + + self.with_shift = with_shift + + # Linear projections used by both single-head and multi-head attention. + self.q_proj = nn.Linear(d_model, d_model, bias=False) + self.k_proj = nn.Linear(d_model, d_model, bias=False) + self.v_proj = nn.Linear(d_model, d_model, bias=False) + + self.merge = nn.Linear(d_model, d_model, bias=False) + + self.norm1 = nn.LayerNorm(d_model) + + # no ffn after self-attn, with ffn after cross-attn + if not self.no_ffn: + in_channels = d_model * 2 + self.mlp = nn.Sequential( + nn.Linear(in_channels, in_channels * ffn_dim_expansion, bias=False), + nn.GELU(), + nn.Linear(in_channels * ffn_dim_expansion, d_model, bias=False), + ) + + self.norm2 = nn.LayerNorm(d_model) + + def forward( + self, + source, + target, + height=None, + width=None, + shifted_window_attn_mask=None, + attn_num_splits=None, + **kwargs, + ): + if "attn_type" in kwargs: + attn_type = kwargs["attn_type"] + else: + attn_type = self.attention_type + + # source: [B, L, C] + # target: [B, L, C] for pairwise attention or [B, M, L, C] for multi-view. + query, key, value = source, target, target + + # Project tokens before dispatching to local/global attention. + query = self.q_proj(query) # [B, L, C] + key = self.k_proj(key) # [B, L, C] or [B, M, L, C] + value = self.v_proj(value) # [B, L, C] or [B, M, L, C] + + if attn_type == "swin" and attn_num_splits > 1: + if self.nhead > 1: + message = multi_head_split_window_attention( + query, + key, + value, + num_splits=attn_num_splits, + with_shift=self.with_shift, + h=height, + w=width, + attn_mask=shifted_window_attn_mask, + num_head=self.nhead, + ) + else: + if self.add_per_view_attn: + assert query.dim() == 3 and key.dim() == 4 and value.dim() == 4 + b, l, c = query.size() + query = query.unsqueeze(1).expand( + -1, key.size(1), -1, -1 + ).reshape(-1, l, c) # [B*M, L, C] + key = key.reshape(-1, l, c) + value = value.reshape(-1, l, c) + message = single_head_split_window_attention( + query, + key, + value, + num_splits=attn_num_splits, + with_shift=self.with_shift, + h=height, + w=width, + attn_mask=shifted_window_attn_mask, + ) + # Sum per-target-view messages back to [B, L, C]. + message = message.view(b, -1, l, c).sum(1) + else: + message = single_head_split_window_attention( + query, + key, + value, + num_splits=attn_num_splits, + with_shift=self.with_shift, + h=height, + w=width, + attn_mask=shifted_window_attn_mask, + ) + else: + message = single_head_full_attention(query, key, value) # [B, L, C] + + message = self.merge(message) # [B, L, C] + message = self.norm1(message) + + if not self.no_ffn: + message = self.mlp(torch.cat([source, message], dim=-1)) + message = self.norm2(message) + + return source + message + + +class TransformerBlock(nn.Module): + """self attention + cross attention + FFN""" + + def __init__( + self, + d_model=256, + nhead=1, + attention_type="swin", + ffn_dim_expansion=4, + with_shift=False, + add_per_view_attn=False, + no_cross_attn=False, + **kwargs, + ): + super(TransformerBlock, self).__init__() + + self.no_cross_attn = no_cross_attn + + if no_cross_attn: + self.self_attn = TransformerLayer( + d_model=d_model, + nhead=nhead, + attention_type=attention_type, + ffn_dim_expansion=ffn_dim_expansion, + with_shift=with_shift, + add_per_view_attn=add_per_view_attn, + ) + else: + self.self_attn = TransformerLayer( + d_model=d_model, + nhead=nhead, + attention_type=attention_type, + no_ffn=True, + ffn_dim_expansion=ffn_dim_expansion, + with_shift=with_shift, + ) + + self.cross_attn_ffn = TransformerLayer( + d_model=d_model, + nhead=nhead, + attention_type=attention_type, + ffn_dim_expansion=ffn_dim_expansion, + with_shift=with_shift, + add_per_view_attn=add_per_view_attn, + ) + + def forward( + self, + source, + target, + height=None, + width=None, + shifted_window_attn_mask=None, + attn_num_splits=None, + **kwargs, + ): + # source, target: [B, L, C] + # self attention + source = self.self_attn( + source, + source, + height=height, + width=width, + shifted_window_attn_mask=shifted_window_attn_mask, + attn_num_splits=attn_num_splits, + **kwargs, + ) + + if self.no_cross_attn: + return source + + # cross attention and ffn + source = self.cross_attn_ffn( + source, + target, + height=height, + width=width, + shifted_window_attn_mask=shifted_window_attn_mask, + attn_num_splits=attn_num_splits, + **kwargs, + ) + + return source + + +def batch_features(features, nn_matrix=None): + # Construct transformer inputs from tensor views. + # features: [B, V, C, H, W] or [B, V, H*W, C]. + # Returns view-major query/target tensors: [V*B, ...] and [V*B, M, ...]. + assert features.dim() in (4, 5) + batch_size, num_views = features.shape[:2] + indices = target_view_indices( + batch_size, num_views, nn_matrix, features.device + ) + + q = flatten_view_batch(features) + kv = gather_view_tensor(features, indices) + kv = flatten_view_batch(kv) + + return q, kv + + +class MultiViewFeatureTransformer(nn.Module): + def __init__( + self, + num_layers=6, + d_model=128, + nhead=1, + attention_type="swin", + ffn_dim_expansion=4, + add_per_view_attn=False, + no_cross_attn=False, + **kwargs, + ): + super(MultiViewFeatureTransformer, self).__init__() + + self.attention_type = attention_type + + self.d_model = d_model + self.nhead = nhead + + self.layers = nn.ModuleList( + [ + TransformerBlock( + d_model=d_model, + nhead=nhead, + attention_type=attention_type, + ffn_dim_expansion=ffn_dim_expansion, + with_shift=( + True if attention_type == "swin" and i % 2 == 1 else False + ), + add_per_view_attn=add_per_view_attn, + no_cross_attn=no_cross_attn, + ) + for i in range(num_layers) + ] + ) + + for p in self.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + + # zero init layers beyond 6 + if num_layers > 6: + for i in range(6, num_layers): + self.layers[i].self_attn.norm1.weight.data.zero_() + self.layers[i].self_attn.norm1.bias.data.zero_() + self.layers[i].cross_attn_ffn.norm2.weight.data.zero_() + self.layers[i].cross_attn_ffn.norm2.bias.data.zero_() + + def forward( + self, + multi_view_features, + attn_num_splits=None, + **kwargs, + ): + nn_matrix = kwargs.pop("nn_matrix", None) + + # multi_view_features: [B, V, C, H, W]. + assert multi_view_features.dim() == 5 + b, num_views, c, h, w = multi_view_features.shape + assert self.d_model == c + + if self.attention_type == "swin" and attn_num_splits > 1: + # global and refine use different number of splits + window_size_h = h // attn_num_splits + window_size_w = w // attn_num_splits + + # compute attn mask once + shifted_window_attn_mask = generate_shift_window_attn_mask( + input_resolution=(h, w), + window_size_h=window_size_h, + window_size_w=window_size_w, + shift_size_h=window_size_h // 2, + shift_size_w=window_size_w // 2, + device=multi_view_features.device, + ) # [K*K, H/K*W/K, H/K*W/K] + else: + shifted_window_attn_mask = None + + # View-major query/target tensors: [V*B, C, H, W], [V*B, M, C, H, W]. + concat0, concat1 = batch_features(multi_view_features, nn_matrix=nn_matrix) + concat0 = concat0.reshape(num_views * b, c, -1).permute( + 0, 2, 1 + ) # [V*B, H*W, C] + c1_v = num_views - 1 if nn_matrix is None else nn_matrix.shape[-1] - 1 + concat1 = concat1.reshape(num_views * b, c1_v, c, -1).permute( + 0, 1, 3, 2 + ) # [V*B, M, H*W, C] + + for i, layer in enumerate(self.layers): + concat0 = layer( + concat0, + concat1, + height=h, + width=w, + shifted_window_attn_mask=shifted_window_attn_mask, + attn_num_splits=attn_num_splits, + ) + + if i < len(self.layers) - 1: + # Rebuild target-view groups after each layer. + features = unflatten_view_batch(concat0, b, num_views) + concat0, concat1 = batch_features(features, nn_matrix=nn_matrix) + + features = unflatten_view_batch(concat0, b, num_views) + features = features.reshape(b, num_views, h, w, c) + features = features.permute(0, 1, 4, 2, 3).contiguous() + + return features + + +def batch_features_camera_parameters( + features, + intrinsics, + extrinsics, + nn_matrix=None, +): + # Construct batch-major inputs for plane-sweep warping. + # features: [B, V, C, H, W]. + # intrinsics/extrinsics: [B, V, 3, 3] and [B, V, 4, 4]. + # Returns [B*V, ...] query tensors and [B*V, M, ...] target tensors. + assert features.dim() == 5 + assert intrinsics.shape[-2:] == (3, 3) + assert extrinsics.shape[-2:] == (4, 4) + + batch_size, num_views, c, h, w = features.shape + indices = target_view_indices( + batch_size, num_views, nn_matrix, features.device + ) + num_selected_views = indices.size(-1) + + q = features + q_intrinsics = intrinsics + q_extrinsics = extrinsics + kv = gather_view_tensor(features, indices) + kv_intrinsics = gather_view_tensor(intrinsics, indices) + kv_extrinsics = gather_view_tensor(extrinsics, indices) + + q = flatten_batch_view(q) # [BV, C, H, W] + q_intrinsics = flatten_batch_view(q_intrinsics) # [BV, 3, 3] + q_extrinsics = flatten_batch_view(q_extrinsics) # [BV, 4, 4] + kv = kv.reshape(-1, num_selected_views, c, h, w) # [BV, V-1, C, H, W] + kv_intrinsics = kv_intrinsics.reshape(-1, num_selected_views, 3, 3) + kv_extrinsics = kv_extrinsics.reshape(-1, num_selected_views, 4, 4) + + return q, q_intrinsics, q_extrinsics, kv, kv_intrinsics, kv_extrinsics diff --git a/src/model/encoder/unimatch/mv_unimatch.py b/src/model/encoder/unimatch/mv_unimatch.py new file mode 100644 index 0000000000000000000000000000000000000000..467b6bce39821a4f22c49a9a3cc917e45423530b --- /dev/null +++ b/src/model/encoder/unimatch/mv_unimatch.py @@ -0,0 +1,586 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .backbone import CNNEncoder +from .vit_fpn import ViTFeaturePyramid +from .mv_transformer import ( + MultiViewFeatureTransformer, + batch_features_camera_parameters, +) +from .matching import warp_with_pose_depth_candidates +from .utils import mv_feature_add_position +from .dpt_head import DPTHead +from .ldm_unet.unet import UNetModel, AttentionBlock +from .ldm_unet.util import normalization +from .dinov2.dinov2 import DINOv2 + + +def flatten_views(tensor): + return tensor.flatten(0, 1) + + +def unflatten_views(tensor, batch_size, num_views): + return tensor.reshape(batch_size, num_views, *tensor.shape[1:]) + + +class MultiViewUniMatch(nn.Module): + def __init__( + self, + num_scales=1, + feature_channels=128, + upsample_factor=8, + lowest_feature_resolution=8, + num_head=1, + ffn_dim_expansion=4, + num_transformer_layers=6, + num_depth_candidates=128, + vit_type="vits", + unet_channels=128, + unet_channel_mult=[1, 1, 1], + unet_num_res_blocks=1, + unet_attn_resolutions=[4], + grid_sample_disable_cudnn=False, + **kwargs, + ): + super(MultiViewUniMatch, self).__init__() + + # CNN + self.feature_channels = feature_channels + self.num_scales = num_scales + self.lowest_feature_resolution = lowest_feature_resolution + self.upsample_factor = upsample_factor + + # monocular backbones: final + self.vit_type = vit_type + + # cost volume + self.num_depth_candidates = num_depth_candidates + + # upsampler + vit_feature_channel_dict = {"vits": 384, "vitb": 768, "vitl": 1024} + + vit_feature_channel = vit_feature_channel_dict[vit_type] + + # CNN + self.backbone = CNNEncoder( + output_dim=feature_channels, + num_output_scales=num_scales, + downsample_factor=upsample_factor, + lowest_scale=lowest_feature_resolution, + return_all_scales=True, + ) + + # Transformer + self.transformer = MultiViewFeatureTransformer( + num_layers=num_transformer_layers, + d_model=feature_channels, + nhead=num_head, + ffn_dim_expansion=ffn_dim_expansion, + ) + + if self.num_scales > 1: + # generate multi-scale features + self.mv_pyramid = ViTFeaturePyramid( + in_channels=128, scale_factors=[2**i for i in range(self.num_scales)] + ) + + # monodepth + encoder = vit_type # can also be 'vitb' or 'vitl' + self.pretrained = DINOv2( + encoder, + use_checkpointing=kwargs.get("use_checkpointing", False), + ) + + del self.pretrained.mask_token # unused + + if self.num_scales > 1: + # generate multi-scale features + self.mono_pyramid = ViTFeaturePyramid( + in_channels=vit_feature_channel, + scale_factors=[2**i for i in range(self.num_scales)], + ) + + # UNet regressor + self.regressor = nn.ModuleList() + self.regressor_residual = nn.ModuleList() + self.depth_head = nn.ModuleList() + + for i in range(self.num_scales): + curr_depth_candidates = num_depth_candidates // (4**i) + cnn_feature_channels = 128 - (32 * i) + mv_transformer_feature_channels = 128 // (2**i) + + mono_feature_channels = vit_feature_channel // (2**i) + + # concat(cost volume, cnn feature, mv feature, mono feature) + in_channels = ( + curr_depth_candidates + + cnn_feature_channels + + mv_transformer_feature_channels + + mono_feature_channels + ) + + # unet channels + channels = unet_channels // (2**i) + + # unet channel mult & unet_attn_resolutions + if i > 0: + unet_channel_mult = unet_channel_mult + [1] + unet_attn_resolutions = [x * 2 for x in unet_attn_resolutions] + + # unet + modules = [ + nn.Conv2d(in_channels, channels, 3, 1, 1), + normalization(channels), + nn.GELU(), + ] + + modules.append( + UNetModel( + image_size=None, + in_channels=channels, + model_channels=channels, + out_channels=channels, + num_res_blocks=unet_num_res_blocks, + attention_resolutions=unet_attn_resolutions, + channel_mult=unet_channel_mult, + num_head_channels=32, + dims=2, + postnorm=False, + num_frames=2, + use_cross_view_self_attn=True, + ) + ) + + modules.append(nn.Conv2d(channels, channels, 3, 1, 1)) + + self.regressor.append(nn.Sequential(*modules)) + + # regressor residual + self.regressor_residual.append(nn.Conv2d(in_channels, channels, 1)) + + # depth head + self.depth_head.append( + nn.Sequential( + nn.Conv2d( + channels, channels * 2, 3, 1, 1, padding_mode="replicate" + ), + nn.GELU(), + nn.Conv2d( + channels * 2, + curr_depth_candidates, + 3, + 1, + 1, + padding_mode="replicate", + ), + ) + ) + + # upsampler + # concat(lowres_depth, cnn feature, mv feature, mono feature) + in_channels = ( + 1 + + cnn_feature_channels + + mv_transformer_feature_channels + + mono_feature_channels + ) + + model_configs = { + "vits": { + "in_channels": 384, + "features": 32, + "out_channels": [48, 96, 192, 384], + }, + "vitb": { + "in_channels": 768, + "features": 48, + "out_channels": [96, 192, 384, 768], + }, + "vitl": { + "in_channels": 1024, + "features": 64, + "out_channels": [128, 256, 512, 1024], + }, + } + + self.upsampler = DPTHead( + **model_configs[vit_type], + downsample_factor=upsample_factor, + num_scales=num_scales, + ) + + self.grid_sample_disable_cudnn = grid_sample_disable_cudnn + + def normalize_images(self, images): + """Normalize image to match the pretrained UniMatch model. + images: (B, V, C, H, W) + """ + shape = [*[1] * (images.dim() - 3), 3, 1, 1] + mean = torch.tensor([0.485, 0.456, 0.406]).reshape(*shape).to(images.device) + std = torch.tensor([0.229, 0.224, 0.225]).reshape(*shape).to(images.device) + + return (images - mean) / std + + def extract_feature(self, images): + # images: [B, V, C, H, W] + concat = flatten_views(images) + # list of [BV, C, H, W], resolution from high to low + features = self.backbone(concat) + # reverse: resolution from low to high + features = features[::-1] + + return features + + def forward( + self, + images, + attn_splits_list=None, + intrinsics=None, + min_depth=1.0 / 0.5, # inverse depth range + max_depth=1.0 / 100, + num_depth_candidates=128, + extrinsics=None, + nn_matrix=None, + supervise_intermediate_depth=False, + **kwargs, + ): + + results_dict = {} + depth_preds = [] + match_probs = [] + + # first normalize images + images = self.normalize_images(images) + b, v, _, ori_h, ori_w = images.shape + + # update the num_views in unet attention, useful for random input views + set_num_views(self.regressor, num_views=v) + + # NOTE: in this codebase, intrinsics are normalized by image width and height + # in unimatch's codebase: https://github.com/autonomousvision/unimatch, no normalization + intrinsics = intrinsics.clone() + intrinsics[:, :, 0] *= ori_w + intrinsics[:, :, 1] *= ori_h + + # Batch-major flatten: [B, V] -> [B*V], matching flatten_views(). + max_depth = max_depth.view(-1) + min_depth = min_depth.view(-1) + + # Scale list of CNN features, each [B*V, C, H, W], resolution low to high. + features_list_cnn = self.extract_feature(images) + features_list_cnn_all_scales = features_list_cnn + features_list_cnn = features_list_cnn[: self.num_scales] + results_dict.update({"features_cnn_all_scales": features_list_cnn_all_scales}) + # features_cnn is only used internally; returning it keeps extra feature maps alive. + # results_dict.update({"features_cnn": features_list_cnn}) + + # Add positional encoding on flattened features, then run the transformer + # in explicit [B, V, C, H, W] layout. + attn_splits = attn_splits_list[0] + + # [BV, C, H, W] + features_cnn_pos = mv_feature_add_position( + features_list_cnn[0], attn_splits, self.feature_channels + ) + + features_mv = self.transformer( + unflatten_views(features_cnn_pos, b, v), + attn_num_splits=attn_splits, + nn_matrix=nn_matrix, + ) + features_mv = flatten_views(features_mv) # [B*V, C, H, W] + + if self.num_scales > 1: + # Scale list of MV features, each [B*V, C, H, W], resolution low to high. + features_list_mv = self.mv_pyramid(features_mv) + else: + features_list_mv = [features_mv] + + results_dict.update({"features_mv": features_list_mv}) + + # mono feature + ori_h, ori_w = images.shape[-2:] + resize_h, resize_w = ori_h // 14 * 14, ori_w // 14 * 14 + concat = flatten_views(images) + concat = F.interpolate( + concat, (resize_h, resize_w), mode="bilinear", align_corners=True + ) + + # get intermediate features + intermediate_layer_idx = { + "vits": [2, 5, 8, 11], + "vitb": [2, 5, 8, 11], + "vitl": [4, 11, 17, 23], + } + + mono_intermediate_features = list( + self.pretrained.get_intermediate_layers( + concat, intermediate_layer_idx[self.vit_type], return_class_token=False + ) + ) + + for i in range(len(mono_intermediate_features)): + curr_features = ( + mono_intermediate_features[i] + .reshape(concat.shape[0], resize_h // 14, resize_w // 14, -1) + .permute(0, 3, 1, 2) + .contiguous() + ) + # resize to 1/8 resolution + curr_features = F.interpolate( + curr_features, + (ori_h // 8, ori_w // 8), + mode="bilinear", + align_corners=True, + ) + mono_intermediate_features[i] = curr_features + + results_dict.update({"features_mono_intermediate": mono_intermediate_features}) + + # last mono feature + mono_features = mono_intermediate_features[-1] + + if self.lowest_feature_resolution == 4: + mono_features = F.interpolate( + mono_features, scale_factor=2, mode="bilinear", align_corners=True + ) + + if self.num_scales > 1: + # Scale list of mono features, each [B*V, C, H, W], resolution low to high. + features_list_mono = self.mono_pyramid(mono_features) + else: + features_list_mono = [mono_features] + + # features_mono is only used internally; returning it keeps extra feature maps alive. + # results_dict.update({"features_mono": features_list_mono}) + + depth = None + + for scale_idx in range(self.num_scales): + downsample_factor = self.upsample_factor * ( + 2 ** (self.num_scales - 1 - scale_idx) + ) + + # scale intrinsics + intrinsics_curr = intrinsics.clone() # [B, V, 3, 3] + intrinsics_curr[:, :, :2] = intrinsics_curr[:, :, :2] / downsample_factor + + # build cost volume + features_mv = features_list_mv[scale_idx] # [BV, C, H, W] + + # batch_features_camera_parameters returns batch-major query tensors + # and M target views per query. + # ref: [B*V, C, H, W], [B*V, 3, 3], [B*V, 4, 4] + # tgt: [B*V, M, C, H, W], [B*V, M, 3, 3], [B*V, M, 4, 4] + ( + ref_features, + ref_intrinsics, + ref_extrinsics, + tgt_features, + tgt_intrinsics, + tgt_extrinsics, + ) = batch_features_camera_parameters( + unflatten_views(features_mv, b, v), + intrinsics_curr, + extrinsics, + nn_matrix=nn_matrix, + ) + + b_new, _, c, h, w = tgt_features.size() + + # relative pose + # extrinsics: c2w + with torch.autocast(device_type=tgt_extrinsics.device.type, enabled=False): + pose_curr = torch.matmul( + torch.linalg.inv(tgt_extrinsics.float()), + ref_extrinsics.float().unsqueeze(1), + ) # [B*V, M, 4, 4] + + if scale_idx > 0: + # 2x upsample depth + assert depth is not None + depth = F.interpolate( + depth, scale_factor=2, mode="bilinear", align_corners=True + ).detach() + + num_depth_candidates = self.num_depth_candidates // (4**scale_idx) + + # generate depth candidates + with torch.autocast(device_type=features_mv.device.type, enabled=False): + min_depth_fp32 = min_depth.float() + max_depth_fp32 = max_depth.float() + if scale_idx == 0: + # min_depth, max_depth: [B*V] + depth_interval = (max_depth_fp32 - min_depth_fp32) / ( + self.num_depth_candidates - 1 + ) # [B*V] + + linear_space = torch.linspace( + 0, + 1, + num_depth_candidates, + device=features_mv.device, + dtype=torch.float32, + ).view(1, num_depth_candidates, 1, 1) # [1, D, 1, 1] + + depth_candidates = ( + min_depth_fp32.view(-1, 1, 1, 1) + + linear_space + * (max_depth_fp32 - min_depth_fp32).view(-1, 1, 1, 1) + ) # [B*V, D, 1, 1] + else: + # half interval each scale + depth_interval = ( + (max_depth_fp32 - min_depth_fp32) + / (self.num_depth_candidates - 1) + / (2**scale_idx) + ) # [B*V] + # [B*V, 1, 1, 1] + depth_interval = depth_interval.view(-1, 1, 1, 1) + + # [B*V, 1, H, W] + depth_fp32 = depth.float() + depth_range_min = ( + depth_fp32 - depth_interval * (num_depth_candidates // 2) + ).clamp(min=min_depth_fp32.view(-1, 1, 1, 1)) + depth_range_max = ( + depth_fp32 + + depth_interval * (num_depth_candidates // 2 - 1) + ).clamp(max=max_depth_fp32.view(-1, 1, 1, 1)) + + linear_space = torch.linspace( + 0, + 1, + num_depth_candidates, + device=features_mv.device, + dtype=torch.float32, + ).view(1, num_depth_candidates, 1, 1) # [1, D, 1, 1] + depth_candidates = depth_range_min + linear_space * ( + depth_range_max - depth_range_min + ) # [B*V, D, H, W] + + num_target_views = tgt_features.size(1) + if scale_idx == 0: + depth_candidates_full = depth_candidates.expand(-1, -1, h, w) + else: + depth_candidates_full = depth_candidates + depth_candidates_curr = ( + depth_candidates_full.unsqueeze(1) + .expand(-1, num_target_views, -1, -1, -1) + .reshape(-1, num_depth_candidates, h, w) + ) # [B*V*M, D, H, W] + + ref_intrinsics_input = ref_intrinsics.unsqueeze(1).expand( + -1, num_target_views, -1, -1 + ) # [B*V, M, 3, 3] + + warped_tgt_features = warp_with_pose_depth_candidates( + flatten_views(tgt_features), + flatten_views(ref_intrinsics_input), + flatten_views(tgt_intrinsics), + flatten_views(pose_curr), + 1.0 / depth_candidates_curr, # convert inverse depth to depth + grid_sample_disable_cudnn=self.grid_sample_disable_cudnn, + ) # [B*V*M, C, D, H, W] + + # ref: [B*V, C, H, W] + # warped: [B*V*M, C, D, H, W] -> [B*V, M, C, D, H, W] + warped_tgt_features = unflatten_views( + warped_tgt_features, b_new, num_target_views + ) + # [B*V, M, D, H, W] -> [B*V, D, H, W] + # Average over target views. + cost_volume = ( + (ref_features.unsqueeze(-3).unsqueeze(1) * warped_tgt_features).sum(2) + / (c**0.5) + ).mean(1) + + # regressor + features_cnn = features_list_cnn[scale_idx] # [B*V, C, H, W] + + features_mono = features_list_mono[scale_idx] # [B*V, C, H, W] + + concat = torch.cat( + (cost_volume, features_cnn, features_mv, features_mono), dim=1 + ) + + out = self.regressor[scale_idx](concat) + self.regressor_residual[ + scale_idx + ](concat) + + # Depth distribution and expectation are distance/probability + # operations, so keep them in FP32 under mixed precision. + with torch.autocast(device_type=out.device.type, enabled=False): + depth_logits = self.depth_head[scale_idx](out.float()) + match_prob = F.softmax(depth_logits, dim=1) # [B*V, D, H, W] + match_probs.append(match_prob) + + if scale_idx == 0: + # [B*V, D, H, W] + depth_candidates = depth_candidates.expand(-1, -1, h, w) + depth = (match_prob * depth_candidates).sum( + dim=1, keepdim=True + ) # [B*V, 1, H, W] + + # Keep intermediate predictions only when they are actually supervised. + if ( + supervise_intermediate_depth + and self.training + and scale_idx < self.num_scales - 1 + ): + depth_bilinear = F.interpolate( + depth, + scale_factor=downsample_factor, + mode="bilinear", + align_corners=True, + ) + depth_preds.append(depth_bilinear) + + # final output, learned upsampler + if scale_idx == self.num_scales - 1: + residual_depth = self.upsampler( + mono_intermediate_features, + # resolution high to low + cnn_features=features_list_cnn_all_scales[::-1], + mv_features=( + features_mv if self.num_scales == 1 else features_list_mv[::-1] + ), + depth=depth, + ) + + depth_bilinear = F.interpolate( + depth, + scale_factor=self.upsample_factor, + mode="bilinear", + align_corners=True, + ) + depth = (depth_bilinear + residual_depth).clamp( + min=min_depth.float().view(-1, 1, 1, 1), + max=max_depth.float().view(-1, 1, 1, 1), + ) + + depth_preds.append(depth) + + # convert inverse depth to depth + for i in range(len(depth_preds)): + depth_pred = 1.0 / depth_preds[i].squeeze(1) # [B*V, H, W] + depth_preds[i] = unflatten_views(depth_pred, b, v) # [B, V, H, W] + + results_dict.update({"depth_preds": depth_preds}) + results_dict.update({"match_probs": match_probs}) + + return results_dict + + +def set_num_views(module, num_views): + if isinstance(module, AttentionBlock): + module.attention.n_frames = num_views + elif ( + isinstance(module, nn.ModuleList) + or isinstance(module, nn.Sequential) + or isinstance(module, nn.Module) + ): + for submodule in module.children(): + set_num_views(submodule, num_views) diff --git a/src/model/encoder/unimatch/position.py b/src/model/encoder/unimatch/position.py new file mode 100644 index 0000000000000000000000000000000000000000..e51da0746bc7aef3eb7b0af97f9251b7e04d33ca --- /dev/null +++ b/src/model/encoder/unimatch/position.py @@ -0,0 +1,50 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +# https://github.com/facebookresearch/detr/blob/main/models/position_encoding.py + +import torch +import torch.nn as nn +import math + + +class PositionEmbeddingSine(nn.Module): + """ + This is a more standard version of the position embedding, very similar to the one + used by the Attention is all you need paper, generalized to work on images. + """ + + def __init__(self, num_pos_feats=64, temperature=10000, normalize=True, scale=None): + super().__init__() + self.num_pos_feats = num_pos_feats + self.temperature = temperature + self.normalize = normalize + if scale is not None and normalize is False: + raise ValueError("normalize should be True if scale is passed") + if scale is None: + scale = 2 * math.pi + self.scale = scale + + def forward(self, x): + # x = tensor_list.tensors # [B, C, H, W] + # mask = tensor_list.mask # [B, H, W], input with padding, valid as 0 + b, c, h, w = x.size() + mask = torch.ones((b, h, w), device=x.device) # [B, H, W] + y_embed = mask.cumsum(1, dtype=torch.float32) + x_embed = mask.cumsum(2, dtype=torch.float32) + if self.normalize: + eps = 1e-6 + y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale + x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale + + dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) + dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) + + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + pos_x = torch.stack( + (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos_y = torch.stack( + (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + return pos diff --git a/src/model/encoder/unimatch/utils.py b/src/model/encoder/unimatch/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..42a8f88713b4ed3974f2870a3f9ee6bad87b22f1 --- /dev/null +++ b/src/model/encoder/unimatch/utils.py @@ -0,0 +1,81 @@ +import torch +from .position import PositionEmbeddingSine + + +def split_feature( + feature, + num_splits=2, + channel_last=False, +): + if channel_last: # [B, H, W, C] + b, h, w, c = feature.size() + assert h % num_splits == 0 and w % num_splits == 0 + + b_new = b * num_splits * num_splits + h_new = h // num_splits + w_new = w // num_splits + + feature = ( + feature.view(b, num_splits, h // num_splits, num_splits, w // num_splits, c) + .permute(0, 1, 3, 2, 4, 5) + .reshape(b_new, h_new, w_new, c) + ) # [B*K*K, H/K, W/K, C] + else: # [B, C, H, W] + b, c, h, w = feature.size() + assert h % num_splits == 0 and w % num_splits == 0 + + b_new = b * num_splits * num_splits + h_new = h // num_splits + w_new = w // num_splits + + feature = ( + feature.view(b, c, num_splits, h // num_splits, num_splits, w // num_splits) + .permute(0, 2, 4, 1, 3, 5) + .reshape(b_new, c, h_new, w_new) + ) # [B*K*K, C, H/K, W/K] + + return feature + + +def merge_splits( + splits, + num_splits=2, + channel_last=False, +): + if channel_last: # [B*K*K, H/K, W/K, C] + b, h, w, c = splits.size() + new_b = b // num_splits // num_splits + + splits = splits.view(new_b, num_splits, num_splits, h, w, c) + return ( + splits.permute(0, 1, 3, 2, 4, 5) + .contiguous() + .view(new_b, num_splits * h, num_splits * w, c) + ) # [B, H, W, C] + else: # [B*K*K, C, H/K, W/K] + b, c, h, w = splits.size() + new_b = b // num_splits // num_splits + + splits = splits.view(new_b, num_splits, num_splits, c, h, w) + return ( + splits.permute(0, 3, 1, 4, 2, 5) + .contiguous() + .view(new_b, c, num_splits * h, num_splits * w) + ) # [B, C, H, W] + + +def mv_feature_add_position(features, attn_splits, feature_channels): + pos_enc = PositionEmbeddingSine(num_pos_feats=feature_channels // 2) + + assert features.dim() == 4 # [B*V, C, H, W] + + if attn_splits > 1: # Add position inside split windows. + features_splits = split_feature(features, num_splits=attn_splits) + position = pos_enc(features_splits) + features_splits = features_splits + position + features = merge_splits(features_splits, num_splits=attn_splits) + else: + position = pos_enc(features) + features = features + position + + return features diff --git a/src/model/encoder/unimatch/vit_fpn.py b/src/model/encoder/unimatch/vit_fpn.py new file mode 100644 index 0000000000000000000000000000000000000000..d4c16402f6042829afbc4b70494f78d8cbd9cc0b --- /dev/null +++ b/src/model/encoder/unimatch/vit_fpn.py @@ -0,0 +1,66 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +# Ref: https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/vit.py#L363 + + +class ViTFeaturePyramid(nn.Module): + """ + This module implements SimpleFeaturePyramid in :paper:`vitdet`. + It creates pyramid features built on top of the input feature map. + """ + + def __init__( + self, + in_channels, + scale_factors, + ): + """ + Args: + scale_factors (list[float]): list of scaling factors to upsample or downsample + the input features for creating pyramid features. + """ + super(ViTFeaturePyramid, self).__init__() + + self.scale_factors = scale_factors + + out_dim = dim = in_channels + self.stages = nn.ModuleList() + for idx, scale in enumerate(scale_factors): + if scale == 4.0: + layers = [ + nn.ConvTranspose2d(dim, dim // 2, kernel_size=2, stride=2), + nn.GELU(), + nn.ConvTranspose2d(dim // 2, dim // 4, kernel_size=2, stride=2), + ] + out_dim = dim // 4 + elif scale == 2.0: + layers = [nn.ConvTranspose2d(dim, dim // 2, kernel_size=2, stride=2)] + out_dim = dim // 2 + elif scale == 1.0: + layers = [] + elif scale == 0.5: + layers = [nn.MaxPool2d(kernel_size=2, stride=2)] + else: + raise NotImplementedError(f"scale_factor={scale} is not supported yet.") + + if scale != 1.0: + layers.extend( + [ + nn.GELU(), + nn.Conv2d(out_dim, out_dim, 3, 1, 1), + ] + ) + layers = nn.Sequential(*layers) + + self.stages.append(layers) + + def forward(self, x): + results = [] + + for stage in self.stages: + results.append(stage(x)) + + return results diff --git a/src/model/encoder/vae/__init__.py b/src/model/encoder/vae/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..77f2fb5445d9e8b9ec9c48404596f09ca59ce062 --- /dev/null +++ b/src/model/encoder/vae/__init__.py @@ -0,0 +1,3 @@ +from .vae import VariableRateLossyVAE + +__all__ = ["VariableRateLossyVAE"] diff --git a/src/model/encoder/vae/coding.py b/src/model/encoder/vae/coding.py new file mode 100644 index 0000000000000000000000000000000000000000..96d5e797cf9740980468bbe1f80cfde5337363eb --- /dev/null +++ b/src/model/encoder/vae/coding.py @@ -0,0 +1,89 @@ +import io +import math +from typing import List, Optional + +import torch +import torch.nn.functional as F + + +def concat_bytes_list(bytes_list: List[bytes], bs_io: io.BytesIO = None) -> Optional[bytes]: + """ + assert 0 <= bytes_list[i] <= 16777216 and len(bytes_list) > 1 + """ + assert len(bytes_list) > 1 + + head_bits_num = 1 + bytes_len_bytes_list = [] + bytes_len_bytes_len_list = [] + for _ in bytes_list: + bytes_len = len(_) + bytes_len_bytes_len = math.ceil(bytes_len.bit_length() / 8) or 1 + head_bits_num = max(head_bits_num, (bytes_len_bytes_len - 1).bit_length()) + bytes_len_bytes_list.append(bytes_len.to_bytes(bytes_len_bytes_len, 'little', signed=False)) + bytes_len_bytes_len_list.append(bytes_len_bytes_len) + assert head_bits_num == 1 or head_bits_num == 2, head_bits_num + + head_bytes = int( + '1' + ''.join((format(_ - 1, f'0{head_bits_num}b') for _ in bytes_len_bytes_len_list)), 2 + ).to_bytes(math.ceil(len(bytes_list) / (8 // head_bits_num) + 0.25), 'big', signed=False) + if head_bits_num == 2: + head_bytes = bytes([head_bytes[0] | 0x80]) + head_bytes[1:] + + if bs_io is None: + return_bytes = True + bs_io = io.BytesIO() + else: + return_bytes = False + bs_io.write(head_bytes) + for _ in bytes_len_bytes_list: + bs_io.write(_) + for _ in bytes_list: + bs_io.write(_) + if return_bytes: + ret = bs_io.getvalue() + bs_io.close() + return ret + + +def split_bytes_list(concat_bytes: Optional[bytes], bytes_list_len: int, + bs_io: io.BytesIO = None) -> List[bytes]: + if bs_io is None: + bytes_given = True + else: + bytes_given = False + assert concat_bytes is None + + if bytes_given: + bs_io = io.BytesIO(concat_bytes) + first_byte = bs_io.read(1)[0] + head_bits_num = 2 if bool(first_byte & 0x80) else 1 + head_bytes_len = math.ceil(bytes_list_len / (8 // head_bits_num) + 0.25) + head_bits = f"{int.from_bytes(bytes([first_byte & 0x7f]) + bs_io.read(head_bytes_len - 1), 'big'):b}"[1:] + + bytes_len_bytes_len_list = [] + for idx in range(0, bytes_list_len * head_bits_num, head_bits_num): + bytes_len_bytes_len_list.append(int(head_bits[idx: idx + head_bits_num], 2) + 1) + + bytes_len_list = [] + for bytes_len_bytes_len in bytes_len_bytes_len_list: + bytes_len_list.append(int.from_bytes(bs_io.read(bytes_len_bytes_len), 'little')) + + bytes_list = [] + for bytes_len in bytes_len_list: + bytes_list.append(bs_io.read(bytes_len)) + + if bytes_given: + bs_io.close() + return bytes_list + + +@torch.compiler.disable +def pad_tensor_divisible_by(img, div=64): + h_old, w_old = img.shape[2:] + if (h_old % div == 0) and (w_old % div == 0): + return img + h_tgt = div * math.ceil(h_old / div) + w_tgt = div * math.ceil(w_old / div) + # left, top, right, bottom + padded = F.pad(img, (0, w_tgt - w_old, 0, h_tgt - h_old), 'replicate') + return padded diff --git a/src/model/encoder/vae/common.py b/src/model/encoder/vae/common.py new file mode 100644 index 0000000000000000000000000000000000000000..44aa0749fb15e7faa92159eff68580747b452227 --- /dev/null +++ b/src/model/encoder/vae/common.py @@ -0,0 +1,175 @@ +from functools import partial +from itertools import repeat +import collections.abc + +import torch +import torch.nn as nn + + +def _ntuple(n): + def parse(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + return tuple(x) + return tuple(repeat(x, n)) + return parse + + +to_1tuple = _ntuple(1) +to_2tuple = _ntuple(2) +to_3tuple = _ntuple(3) +to_4tuple = _ntuple(4) +to_ntuple = _ntuple + + +def get_conv(in_ch, out_ch, kernel_size, stride, padding, zero_bias=True, zero_weights=False): + conv = nn.Conv2d(in_ch, out_ch, kernel_size, stride, padding) + if zero_bias: + conv.bias.data.mul_(0.0) + if zero_weights: + conv.weight.data.mul_(0.0) + return conv + + +def conv_k1s1(in_ch, out_ch, zero_bias=True, zero_weights=False): + return get_conv(in_ch, out_ch, 1, 1, 0, zero_bias, zero_weights) + + +def conv_k3s1(in_ch, out_ch, zero_bias=True, zero_weights=False): + return get_conv(in_ch, out_ch, 3, 1, 1, zero_bias, zero_weights) + + +def patch_downsample(in_ch, out_ch, rate=2): + return get_conv(in_ch, out_ch, kernel_size=rate, stride=rate, padding=0) + + +def patch_upsample(in_ch, out_ch, rate=2): + conv = nn.Sequential( + get_conv(in_ch, out_ch * (rate ** 2), kernel_size=1, stride=1, padding=0), + nn.PixelShuffle(rate) + ) + return conv + + +class SetKey(nn.Module): + """ A dummy layer that is used to mark the position of a layer in the network. + """ + def __init__(self, key): + super().__init__() + self.key = key + + def forward(self, x): + return x + + +class CompressionStopFlag(nn.Module): + """ A dummy layer that is used to mark the stop position of encoding bits. + """ + def __init__(self): + super().__init__() + + def forward(self, x): + return x + + +class FeatureExtractorWithEmbedding(nn.Module): + def __init__(self, blocks): + super().__init__() + self.enc_blocks = nn.ModuleList(blocks) + + def forward(self, x, emb=None): + features = {} + for i, block in enumerate(self.enc_blocks): + if isinstance(block, SetKey): + features[block.key] = x + elif getattr(block, 'requires_embedding', False): + x = block(x, emb) + else: + x = block(x) + return x, features + + +class Mlp(nn.Module): + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + act_layer=nn.GELU, + norm_layer=None, + bias=True, + use_conv=False, + ): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + bias = to_2tuple(bias) + linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear + + self.fc1 = linear_layer(in_features, hidden_features, bias=bias[0]) + self.act = act_layer() + self.norm = norm_layer(hidden_features) if norm_layer is not None else nn.Identity() + self.fc2 = linear_layer(hidden_features, out_features, bias=bias[1]) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.norm(x) + x = self.fc2(x) + return x + + +class ConvNeXtBlockAdaLN(nn.Module): + def __init__(self, dim, embed_dim=256, out_dim=None, kernel_size=7, mlp_ratio: float = 2, + residual=True, ls_init_value=1e-6): + super().__init__() + # depthwise conv + pad = (kernel_size - 1) // 2 + self.conv_dw = nn.Conv2d(dim, dim, kernel_size=kernel_size, padding=pad, groups=dim) + # layer norm + self.norm = nn.LayerNorm(dim, eps=1e-6, elementwise_affine=False) + self.norm.affine = False # for FLOPs computing + # AdaLN + self.embedding_layer = nn.Sequential( + nn.Sequential(), + nn.Linear(embed_dim, 2*dim), + nn.Unflatten(1, unflattened_size=(1, 1, 2*dim)) + ) + # MLP + hidden = int(mlp_ratio * dim) + out_dim = out_dim or dim + self.mlp = Mlp(dim, hidden_features=hidden, out_features=out_dim, act_layer=nn.GELU) + # layer scaling + if ls_init_value >= 0: + self.gamma = nn.Parameter(torch.full(size=(1, out_dim, 1, 1), fill_value=1e-6)) + else: + self.gamma = None + + self.residual = residual + self.requires_embedding = True + + for m in self.modules(): + if isinstance(m, (nn.Conv2d, nn.Linear)): + nn.init.trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def forward(self, x, emb): + shortcut = x + # depthwise conv + x = self.conv_dw(x) + # layer norm + x = x.permute(0, 2, 3, 1).contiguous() + x = self.norm(x) + # AdaLN + embedding = self.embedding_layer(emb) + shift, scale = torch.chunk(embedding, chunks=2, dim=-1) + x = x * (1 + scale) + shift + # MLP + x = self.mlp(x) + # scaling + if self.gamma is not None: + x = x.mul(self.gamma.reshape(-1)) + x = x.permute(0, 3, 1, 2).contiguous() + if self.residual: + x = x + shortcut + return x diff --git a/src/model/encoder/vae/config.py b/src/model/encoder/vae/config.py new file mode 100644 index 0000000000000000000000000000000000000000..8ec84c51f33d02f90f2d2b6112a0847851d08a13 --- /dev/null +++ b/src/model/encoder/vae/config.py @@ -0,0 +1,19 @@ +from dataclasses import dataclass +from typing import Tuple + + +@dataclass +class Config: + max_stride: int = 64 + lmb_range: Tuple[float, float] = (16.0, 1024.0) + lmb_limit: float = 8192.0 + lmb_embed_dim: Tuple[int, ...] = (256, 256) + sin_period: int = 64 + enc_channels: Tuple[int, ...] = (192, 384, 512, 512, 512) + dec_channels: Tuple[int, ...] = (512, 512, 384, 256, 128) + # Latent channels for stride-64/32/16/8 blocks. + # Original default: (32, 32, 96, 8). + z_dec_channels: Tuple[int, ...] = (32, 32, 64, 16) + + # Test phase + lmb: float = 1024.0 diff --git a/src/model/encoder/vae/entropy_coding.py b/src/model/encoder/vae/entropy_coding.py new file mode 100644 index 0000000000000000000000000000000000000000..c75710d64cbe0897202ba5bd87ebff85e51cfae7 --- /dev/null +++ b/src/model/encoder/vae/entropy_coding.py @@ -0,0 +1,82 @@ +import math +import scipy.stats +import torch +import torch.distributions as td + +from compressai.ops import LowerBound +from compressai.entropy_models import GaussianConditional + + +def _to_float32(*args): + if len(args) == 1: + return args[0].to(dtype=torch.float) + else: + return tuple([t.to(dtype=torch.float) for t in args]) + + +@torch.autocast('cuda', enabled=False) # disable mixed precision +def _safe_log_prob_mass(distribution, x, bin_size=1, prob_clamp=1e-6): + prob_mass = distribution.cdf(x + 0.5*bin_size) - distribution.cdf(x - 0.5*bin_size) + log_prob = torch.where( + prob_mass > prob_clamp, + torch.log(prob_mass.clamp(min=1e-8)), + distribution.log_prob(x) + math.log(bin_size) + ) + return log_prob + + +def _sanity_check_scale_table(scale_table): + assert isinstance(scale_table, torch.Tensor) + assert (scale_table.dim() == 1) and (scale_table.shape[0] >= 1) and (scale_table.min() > 0) + assert torch.equal(scale_table, torch.sort(scale_table)[0]) + + +def gaussian_log_prob_mass(mean, scale, x, bin_size=1.0, prob_clamp=1e-6): + """ Compute log(P) of a "quantized" Normal(`mean`, `scale`) distribution evaluated at `x`, + where P = cdf(`x` + 0.5*bin_size) - cdf(`x` - 0.5*bin_size). + + Args: + mean (Tensor): mean of the Gaussian + scale (Tensor): scale (standard deviation) of the Gaussian + x (Tensor): the quantized Gaussian is evaluated at `x` + bin_size (float): quantization bin size + prob_clamp (float): when prob < prob_clamp, use approximation \ + to improve numerical stability. + """ + mean, scale, x = _to_float32(mean, scale, x) + assert scale.min() > 0, f'invalid scale value = {scale.min()}' + log_prob = _safe_log_prob_mass(td.Normal(mean, scale), x, bin_size, prob_clamp) + return log_prob + + +class DiscretizedGaussian(GaussianConditional): + """ Custom discretized gaussian. + """ + def __init__(self, scale_table=None): + """ + Args: + scale_table (torch.Tensor, optional): a 1-D tensor. + """ + super(GaussianConditional, self).__init__() + + if scale_table is None: + scale_table = self._get_default_scale_table() + _sanity_check_scale_table(scale_table) + self.register_buffer("scale_table", scale_table, persistent=False) + + self.tail_mass = float(1e-9) + self.lower_bound_scale = LowerBound(scale_table[0]) + + self.standard_gaussian = td.Normal(loc=0, scale=1) + + @staticmethod + def _get_default_scale_table(): + scale_table = torch.exp(torch.linspace(math.log(0.11), math.log(20.0), steps=64)) + return scale_table + + @staticmethod + def _standardized_quantile(quantile): + return scipy.stats.norm.ppf(quantile) + + def _standardized_cumulative(self, inputs: torch.Tensor): + return self.standard_gaussian.cdf(inputs) diff --git a/src/model/encoder/vae/vae.py b/src/model/encoder/vae/vae.py new file mode 100644 index 0000000000000000000000000000000000000000..dea67d0fe194b92569e74ad3d7d03440c8ce4f14 --- /dev/null +++ b/src/model/encoder/vae/vae.py @@ -0,0 +1,360 @@ +import math +import struct + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from . import common +from . import coding +from .entropy_coding import DiscretizedGaussian + +from .config import Config + + +one_div_sqrt2 = 1 / math.sqrt(2) +half_log2pi = 0.5 * math.log(2 * math.pi) + + +@torch.autocast('cuda', enabled=False) +def gaussian_log_prob_mass(mean: torch.Tensor, scale: torch.Tensor, x: torch.Tensor, prob_clamp: float = 1e-6): + mean = mean.float() + scale = scale.float() + x = x.float() + scale_sqrt2_r = scale.reciprocal() * one_div_sqrt2 + mean_sub_x = mean - x + prob_mass = 0.5 * ( + (mean_sub_x - 0.5).mul_(scale_sqrt2_r).erfc_() - + (mean_sub_x + 0.5).mul_(scale_sqrt2_r).erfc_()) + mask = prob_mass > prob_clamp + log_prob_mass = prob_mass.clamp_min(1e-8).log_() + normal_log_prob = -(mean_sub_x * scale_sqrt2_r).square_() - scale.log() - half_log2pi + log_prob = torch.where(mask, log_prob_mass, normal_log_prob) + return log_prob + + +class VRLVBlockBase(nn.Module): + def __init__(self, width, z_dim, enc_key, enc_width, embed_dim=256, kernel_size=7, mlp_ratio: float = 2.0): + super().__init__() + self.in_channels = width + self.out_channels = width + self.enc_key = enc_key + + block = common.ConvNeXtBlockAdaLN + self.resnet_front = block(width, embed_dim, kernel_size=kernel_size, mlp_ratio=mlp_ratio) + self.resnet_end = block(width, embed_dim, kernel_size=kernel_size, mlp_ratio=mlp_ratio) + self.posterior0 = block(enc_width, embed_dim, kernel_size=kernel_size) + self.posterior1 = block(width, embed_dim, kernel_size=kernel_size) + self.posterior2 = block(width, embed_dim, kernel_size=kernel_size) + self.post_merge = common.conv_k1s1(width + enc_width, width) + self.posterior = common.conv_k3s1(width, z_dim) + self.z_proj = common.conv_k1s1(z_dim, width) + self.prior = common.conv_k1s1(width, z_dim * 2) + + self.discrete_gaussian = DiscretizedGaussian() + self.is_latent_block = True + + def transform_prior(self, feature, lmb_embedding): + feature = self.resnet_front(feature, lmb_embedding) + pm, plogv = self.prior(feature).chunk(2, dim=1) + with torch.autocast(device_type=plogv.device.type, enabled=False): + # Keep probability scale construction in FP32 under mixed precision. + pv = F.softplus(plogv.float() + 2.3).sub(2.3).clip(max=4.6).exp() + return feature, pm, pv + + def transform_posterior(self, feature, enc_feature, lmb_embedding): + enc_feature = self.posterior0(enc_feature, lmb_embedding) + feature = self.posterior1(feature, lmb_embedding) + merged = torch.cat([feature, enc_feature], dim=1) + merged = self.post_merge(merged) + merged = self.posterior2(merged, lmb_embedding) + qm = self.posterior(merged) + return qm + + def fuse_feature_and_z(self, feature, z): + # add the new information carried by z to the feature + return feature + self.z_proj(z) + + def forward(self, feature, lmb_embedding, enc_feature=None, mode='trainval', + get_latent=False, strings=None): + feature, pm, pv = self.transform_prior(feature, lmb_embedding) + + additional = dict() + if mode == 'trainval': + qm = self.transform_posterior(feature, enc_feature, lmb_embedding) + if self.training: + z = qm + (torch.rand_like(qm) - 0.5) + kl = -gaussian_log_prob_mass(pm, pv, x=z) + else: + z, probs = self.discrete_gaussian(qm, scales=pv, means=pm) + kl = probs.log().neg_() + additional['kl'] = kl + elif mode == 'compress': + qm = self.transform_posterior(feature, enc_feature, lmb_embedding) + indexes = self.discrete_gaussian.build_indexes(pv) + strings = self.discrete_gaussian.compress(qm, indexes, means=pm)[0] + z = self.discrete_gaussian.quantize(qm, mode='dequantize', means=pm) + additional['strings'] = strings + elif mode == 'decompress': + assert strings is not None + indexes = self.discrete_gaussian.build_indexes(pv) + z = self.discrete_gaussian.decompress([strings], indexes, means=pm) + else: + raise ValueError(f'Unknown mode={mode}') + + feature = self.fuse_feature_and_z(feature, z) + feature = self.resnet_end(feature, lmb_embedding) + if get_latent: + additional['z'] = z.detach() + return feature, additional + + def update(self): + self.discrete_gaussian.update() + + +class VariableRateLossyVAE(nn.Module): + log2_e = math.log2(math.e) + + def __init__(self, in_ch, cfg: Config): + super().__init__() + self.cfg = cfg + + enc_dims = cfg.enc_channels + res_block = common.ConvNeXtBlockAdaLN + embed_dim = cfg.lmb_embed_dim[1] + enc_blocks = [ + # 64x64 + common.patch_downsample(in_ch, enc_dims[0], rate=4), + # 16x16 + *[res_block(enc_dims[0], embed_dim, kernel_size=7) for _ in range(3)], # original: 6 + res_block(enc_dims[0]), + common.patch_downsample(enc_dims[0], enc_dims[1]), + # 8x8 + *[res_block(enc_dims[1], embed_dim, kernel_size=7) for _ in range(3)], # original: 6 + common.SetKey('enc_s8'), + res_block(enc_dims[1], embed_dim), + common.patch_downsample(enc_dims[1], enc_dims[2]), + # 4x4 + *[res_block(enc_dims[2], embed_dim, kernel_size=5) for _ in range(3)], # original: 6 + common.SetKey('enc_s16'), + res_block(enc_dims[2], embed_dim), + common.patch_downsample(enc_dims[2], enc_dims[3]), + # 2x2 + *[res_block(enc_dims[3], embed_dim, kernel_size=3) for _ in range(2)], # original: 4 + common.SetKey('enc_s32'), + res_block(enc_dims[3], embed_dim), + common.patch_downsample(enc_dims[3], enc_dims[4]), + # 1x1 + *[res_block(enc_dims[4], embed_dim, kernel_size=1) for _ in range(2)], # original: 4 + common.SetKey('enc_s64'), + ] + + dec_dims = cfg.dec_channels + z_dims = cfg.z_dec_channels + dec_blocks = [ + # 1x1 + *[VRLVBlockBase(dec_dims[0], z_dims[0], enc_key='enc_s64', enc_width=enc_dims[-1], kernel_size=1, + mlp_ratio=4.0) for _ in range(1)], + res_block(dec_dims[0], embed_dim, kernel_size=1, mlp_ratio=4.0), + common.patch_upsample(dec_dims[0], dec_dims[1], rate=2), + # 2x2 + res_block(dec_dims[1], embed_dim, kernel_size=3, mlp_ratio=3.0), + *[VRLVBlockBase(dec_dims[1], z_dims[1], enc_key='enc_s32', enc_width=enc_dims[-2], kernel_size=3, + mlp_ratio=3.0) for _ in range(1)], # original: 2 + res_block(dec_dims[1], embed_dim, kernel_size=3, mlp_ratio=3.0), + common.patch_upsample(dec_dims[1], dec_dims[2], rate=2), + # 4x4 + res_block(dec_dims[2], embed_dim, kernel_size=5, mlp_ratio=2.0), + *[VRLVBlockBase(dec_dims[2], z_dims[2], enc_key='enc_s16', enc_width=enc_dims[-3], kernel_size=5, + mlp_ratio=2.0) for _ in range(2)], # original: 3 + res_block(dec_dims[2], embed_dim, kernel_size=5, mlp_ratio=2.0), + common.patch_upsample(dec_dims[2], dec_dims[3], rate=2), + # 8x8 + res_block(dec_dims[3], kernel_size=7, mlp_ratio=1.75), + *[VRLVBlockBase(dec_dims[3], z_dims[3], enc_key='enc_s8', enc_width=enc_dims[-4], kernel_size=7, + mlp_ratio=1.75) for _ in range(2)], # original: 3 + common.CompressionStopFlag(), # no need to execute remaining blocks when compressing + res_block(dec_dims[3], kernel_size=7, mlp_ratio=1.75), + common.patch_upsample(dec_dims[3], dec_dims[4], rate=2), + # 16x16 + *[res_block(dec_dims[4], kernel_size=7, mlp_ratio=1.5) for _ in range(4)], # original: 8 + common.patch_upsample(dec_dims[4], in_ch, rate=4) + ] + self.encoder = common.FeatureExtractorWithEmbedding(enc_blocks) + self.dec_blocks = nn.ModuleList(dec_blocks) + self.bias = nn.Parameter(torch.zeros(1, self.dec_blocks[0].in_channels, 1, 1)) + self.lmb_embedding = nn.Sequential( + nn.Linear(cfg.lmb_embed_dim[0], cfg.lmb_embed_dim[1]), + nn.GELU(), + nn.Linear(cfg.lmb_embed_dim[1], cfg.lmb_embed_dim[1]), + nn.GELU() + ) + + self.register_buffer( + 'sinusoidal_freqs', + torch.pow(self.cfg.sin_period, 1 - torch.linspace(0, 1, steps=(cfg.lmb_embed_dim[0] // 2))).reshape(1, -1), + persistent=False) + self.sinusoidal_freqs: torch.Tensor + self.num_latent_blocks = sum( + 1 for block in self.dec_blocks if getattr(block, "is_latent_block", False) + ) + + self.compress_mode(True) + + @torch.no_grad() + def sample_lmb(self, n, device): + low, high = self.cfg.lmb_range + if low == high: + return torch.full((n,), float(low), device=device) + p = 3.0 + low, high = math.pow(low, 1/p), math.pow(high, 1/p) # transformed space + lmb = low + (high-low) * torch.rand(n, device=device) + lmb.pow_(p) + return lmb + + def expand_to_tensor(self, input_, n, device): + if isinstance(input_, (float, int)): + input_ = torch.full(size=(n,), fill_value=float(input_), device=device) + assert input_.shape == (n,), f'{input_=}, {input_.shape=}' + return input_ + + def get_lmb_embedding(self, lmb, n, device): + lmb = self.expand_to_tensor(lmb, n, device) + lmb_limit = self.cfg.lmb_limit + scaled = torch.log(lmb) / math.log(lmb_limit) + args = scaled.reshape(-1, 1) * self.sinusoidal_freqs + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + embedding = self.lmb_embedding(embedding) + return embedding + + def get_bias(self, bhw_repeat=(1, 1, 1)): + nB, nH, nW = bhw_repeat + feature = self.bias.expand(nB, -1, nH, nW) + return feature + + def forward_end2end(self, features: torch.Tensor, lmb, mode='trainval', get_latent=False): + x = features + # ================ get lambda embedding ================ + emb = self.get_lmb_embedding(lmb, features.shape[0], x.device) + # ================ Forward pass ================ + enc_features = self.encoder(x, emb)[1] + batch_size, _, height, width = x.shape + feature = self.get_bias(bhw_repeat=(batch_size, height//self.cfg.max_stride, width//self.cfg.max_stride)) + lv_block_results = [] # all latent variable block results + for i, block in enumerate(self.dec_blocks): + if getattr(block, 'is_latent_block', False): + f_enc = enc_features[block.enc_key] + feature, stats = block(feature, emb, enc_feature=f_enc, mode=mode, get_latent=get_latent) + lv_block_results.append(stats) + elif getattr(block, 'requires_embedding', False): + feature = block(feature, emb) + elif isinstance(block, common.CompressionStopFlag) and (mode == 'compress'): + # no need to execute remaining blocks when compressing + return lv_block_results + else: + feature = block(feature) + return feature, lv_block_results + + def forward(self, features, lmb=None): + if self.training: + return self.train_forward(features, lmb=lmb) + return self.test_forward(features) + + def train_forward(self, features, lmb=None): + height, width = features.shape[2:] + + if lmb is None: + lmb = self.sample_lmb(features.shape[0], features.device) + else: + lmb = self.expand_to_tensor(lmb, features.shape[0], features.device) + x_hat, stats_all = self.forward_end2end(features, lmb) + + kl_loss = sum([stat['kl'].sum(dim=(1, 2, 3)) for stat in stats_all]) * ( + 1 / (height * width) + ) + + return kl_loss * self.log2_e, x_hat, lmb + + @torch.compiler.disable + @torch.no_grad() + def test_forward(self, features): + bitstreams = [] + reconstructions = [] + num_bytes = [] + + for i in range(features.shape[0]): + feature_i = features[i : i + 1] + strings = self.compress(feature_i) + feature_hat = self.decompress(strings) + bitstreams.append(strings) + reconstructions.append(feature_hat) + num_bytes.append(len(strings)) + + features_hat = torch.cat(reconstructions, dim=0) + num_bytes_tensor = torch.tensor( + num_bytes, device=features.device, dtype=torch.float32 + ) + return bitstreams, features_hat, num_bytes_tensor + + def train(self, mode: bool = True): + if mode is False: + self.compress_mode(True) + return super(VariableRateLossyVAE, self).train(mode) + + def compress_mode(self, mode=True): + if mode: + for block in self.dec_blocks: + if hasattr(block, 'update'): + block.update() + + def compress_once(self, features, padded_features, lmb): + lv_block_results = self.forward_end2end( + padded_features, lmb, mode='compress', get_latent=True + ) + all_lv_strings = [res['strings'] for res in lv_block_results] + strings = struct.pack('f', torch.tensor(lmb, dtype=torch.float32).item()) + \ + struct.pack('2H', features.shape[2], features.shape[3]) + \ + coding.concat_bytes_list(all_lv_strings) + return strings, lv_block_results + + @torch.no_grad() + def compress(self, features): + assert ( + features.shape[0] == 1 + ), "Feature codec bitstream export currently supports batch size 1." + padded_features = coding.pad_tensor_divisible_by(features, div=self.cfg.max_stride) + lmb = self.cfg.lmb + strings, _ = self.compress_once(features, padded_features, lmb) + return strings + + @torch.no_grad() + def decompress(self, strings): + (lmb,), strings = struct.unpack("f", strings[:4]), strings[4:] + (height, width), strings = struct.unpack("2H", strings[:4]), strings[4:] + downsampled_height = math.ceil(height / self.cfg.max_stride) + downsampled_width = math.ceil(width / self.cfg.max_stride) + all_lv_strings = coding.split_bytes_list(strings, self.num_latent_blocks) + + device = next(self.parameters()).device + lmb_embedding = self.get_lmb_embedding(lmb, 1, device) + feature = self.get_bias( + bhw_repeat=(1, downsampled_height, downsampled_width) + ) + + str_i = 0 + for block in self.dec_blocks: + if getattr(block, "is_latent_block", False): + feature, _ = block( + feature, + lmb_embedding, + mode="decompress", + strings=all_lv_strings[str_i], + ) + str_i += 1 + elif getattr(block, "requires_embedding", False): + feature = block(feature, lmb_embedding) + else: + feature = block(feature) + + assert str_i == len(all_lv_strings), f"str_i={str_i}, len={len(all_lv_strings)}" + return feature[:, :, :height, :width] diff --git a/src/model/model_wrapper.py b/src/model/model_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..42c42c62ec64065d0b015e5c25d88965b679f014 --- /dev/null +++ b/src/model/model_wrapper.py @@ -0,0 +1,1276 @@ +from dataclasses import dataclass +from pathlib import Path +from typing import Optional, Protocol, runtime_checkable + +import torch +import wandb +from einops import rearrange, repeat, einsum +from jaxtyping import Float +from pytorch_lightning import LightningModule +from pytorch_lightning.loggers.wandb import WandbLogger +from pytorch_lightning.utilities import rank_zero_only +from torch import Tensor, nn, optim +import numpy as np +import json +import os +import time +from tqdm import tqdm +import torch.nn.functional as F +import math + +from ..dataset.data_module import get_data_shim +from ..dataset.types import BatchedExample +from ..dataset import DatasetCfg +from ..evaluation.metrics import compute_lpips, compute_psnr, compute_ssim +from ..global_cfg import get_cfg +from ..loss import Loss, LossLpips +from ..misc.benchmarker import Benchmarker +from ..misc.image_io import prep_image, save_image, save_video +from ..misc.LocalLogger import LocalLogger +from ..misc.step_tracker import StepTracker +from ..visualization.annotation import add_label +from ..visualization.camera_trajectory.interpolation import ( + interpolate_extrinsics, + interpolate_intrinsics, +) +from ..visualization.camera_trajectory.wobble import ( + generate_wobble, + generate_wobble_transformation, +) +from ..visualization.color_map import apply_color_map_to_image +from ..visualization.layout import add_border, hcat, vcat +from .decoder.decoder import Decoder, DepthRenderingMode, DecoderOutput +from .encoder import Encoder +from .types import Gaussians +from src.visualization.vis_depth import viz_depth_tensor +from PIL import Image +from ..misc.stabilize_camera import render_stabilization_path +from .ply_export import save_gaussian_ply + + +@dataclass +class OptimizerCfg: + lr: float + warm_up_steps: int + lr_monodepth: float + weight_decay: float + div_factor: float + final_div_factor: float + + +@dataclass +class TestCfg: + compute_scores: bool + save_image: bool + save_video: bool + eval_time_skip_steps: int + save_gt_image: bool + save_input_images: bool + save_depth: bool + save_depth_concat_img: bool + save_depth_npy: bool + save_gaussian: bool + save_bitstream: bool + render_chunk_size: int | None + stabilize_camera: bool + stab_camera_kernel: int + gaussian_export_all_points: bool + gaussian_export_align_to_view: bool + + +@dataclass +class TrainCfg: + depth_mode: DepthRenderingMode | None + extended_visualization: bool + print_log_every_n_steps: int + eval_model_every_n_val: int + eval_data_length: int + eval_time_skip_steps: int + eval_save_model: bool + l1_loss: bool + intermediate_loss_weight: float + no_viz_video: bool + train_ignore_large_loss: float + + +@runtime_checkable +class TrajectoryFn(Protocol): + def __call__( + self, + t: Float[Tensor, " t"], + ) -> tuple[ + Float[Tensor, "batch view 4 4"], # extrinsics + Float[Tensor, "batch view 3 3"], # intrinsics + ]: + pass + + +class ModelWrapper(LightningModule): + logger: Optional[WandbLogger] + encoder: nn.Module + decoder: Decoder + losses: nn.ModuleList + optimizer_cfg: OptimizerCfg + test_cfg: TestCfg + train_cfg: TrainCfg + step_tracker: StepTracker | None + eval_data_cfg: Optional[DatasetCfg | None] + + def __init__( + self, + optimizer_cfg: OptimizerCfg, + test_cfg: TestCfg, + train_cfg: TrainCfg, + encoder: Encoder, + decoder: Decoder, + losses: list[Loss], + step_tracker: StepTracker | None, + eval_data_cfg: Optional[DatasetCfg | None] = None, + ) -> None: + super().__init__() + self.optimizer_cfg = optimizer_cfg + self.test_cfg = test_cfg + self.train_cfg = train_cfg + self.step_tracker = step_tracker + self.eval_data_cfg = eval_data_cfg + + # Set up the model. + self.encoder = encoder + self.decoder = decoder + self.data_shim = get_data_shim(self.encoder) + self.losses = nn.ModuleList(losses) + + # This is used for testing. + self.benchmarker = Benchmarker() + self.eval_cnt = 0 + + if self.test_cfg.compute_scores: + self.test_step_outputs = {} + self.time_skip_steps_dict = {"encoder": 0, "decoder": 0} + + @torch.no_grad() + def compute_lpips( + self, + ground_truth: Float[Tensor, "batch channel height width"], + predicted: Float[Tensor, "batch channel height width"], + ) -> Float[Tensor, " batch"]: + for loss_fn in self.losses: + if isinstance(loss_fn, LossLpips): + with torch.autocast( + device_type=predicted.device.type, + enabled=False, + ): + value = loss_fn.lpips_on_device(predicted.device).forward( + ground_truth.float(), + predicted.float(), + normalize=True, + ) + return value[:, 0, 0, 0] + + return compute_lpips(ground_truth, predicted) + + def _codec_loss_weight(self) -> float: + return float(getattr(getattr(self.encoder, "cfg", None), "codec_loss_weight", 1.0)) + + def _codec_depth_loss_weight(self) -> float: + return float( + getattr(getattr(self.encoder, "cfg", None), "codec_depth_loss_weight", 0.0) + ) + + def _codec_lmb_weight(self, codec_lmb: Tensor) -> Tensor: + cfg = getattr(self.encoder, "cfg", None) + codec_lmb_range = getattr(cfg, "codec_lmb_range", None) + if codec_lmb_range is None or len(codec_lmb_range) == 0: + return torch.ones_like(codec_lmb) + + max_lmb = float(codec_lmb_range[-1]) + if max_lmb <= 0: + return torch.ones_like(codec_lmb) + + return (codec_lmb / max_lmb).clamp_min(0.0) + + def _unpack_encoder_output( + self, + encoder_output, + ) -> tuple[ + Gaussians | None, + Tensor | None, + Tensor | None, + Tensor | None, + Tensor | None, + Tensor | None, + Tensor | None, + Tensor | None, + Tensor | None, + ]: + codec_loss = None + codec_lmb = None + codec_num_bytes = None + codec_bpp = None + codec_depth_targets = None + codec_depth_loss_pred = None + codec_depth_loss_target = None + if isinstance(encoder_output, tuple): + if len(encoder_output) != 3: + raise ValueError( + f"Unexpected encoder output tuple length: {len(encoder_output)}" + ) + encoder_output, codec_loss, codec_lmb = encoder_output + + pred_depths = None + gaussians = encoder_output + if isinstance(encoder_output, dict): + pred_depths = encoder_output.get("depths") + gaussians = encoder_output["gaussians"] + codec_num_bytes = encoder_output.get("codec_num_bytes") + codec_bpp = encoder_output.get("codec_bpp") + codec_depth_targets = encoder_output.get("codec_depth_targets") + codec_depth_loss_pred = encoder_output.get("codec_depth_loss_pred") + codec_depth_loss_target = encoder_output.get("codec_depth_loss_target") + + return ( + gaussians, + pred_depths, + codec_depth_targets, + codec_depth_loss_pred, + codec_depth_loss_target, + codec_loss, + codec_lmb, + codec_num_bytes, + codec_bpp, + ) + + def training_step(self, batch, batch_idx): + batch: BatchedExample = self.data_shim(batch) + batch_size, _, _, h, w = batch["target"]["image"].shape + + # Run the model. + encoder_output = self.encoder( + batch["context"], self.global_step, scene_names=batch["scene"] + ) + ( + gaussians, + pred_depths, + codec_depth_targets, + codec_depth_loss_pred, + codec_depth_loss_target, + codec_loss, + codec_lmb, + _, + _, + ) = self._unpack_encoder_output(encoder_output) + + supervise_intermediate_depth = False + + if gaussians.means.size(0) != batch["target"]["extrinsics"].size(0): + supervise_intermediate_depth = True + assert gaussians.means.size(0) % batch["target"]["extrinsics"].size(0) == 0 + num_depths = gaussians.means.size(0) // batch["target"]["extrinsics"].size( + 0 + ) + # add loss to intermediate depth predictions + target_extrinsics = torch.cat( + [batch["target"]["extrinsics"]] * num_depths, dim=0 + ) + target_intrinsics = torch.cat( + [batch["target"]["intrinsics"]] * num_depths, dim=0 + ) + target_near = torch.cat([batch["target"]["near"]] * num_depths, dim=0) + target_far = torch.cat([batch["target"]["far"]] * num_depths, dim=0) + + output_all = self.decoder.forward( + gaussians, + target_extrinsics, + target_intrinsics, + target_near, + target_far, + (h, w), + depth_mode=self.train_cfg.depth_mode, + ) + # split + batch_size = batch["target"]["extrinsics"].size(0) + # order: intermediate depth, final depth + output_intermediate = DecoderOutput( + color=output_all.color[:-batch_size], + depth=( + output_all.depth[:-batch_size] + if output_all.depth is not None + else None + ), + ) + output = DecoderOutput( + color=output_all.color[-batch_size:], + depth=( + output_all.depth[-batch_size:] + if output_all.depth is not None + else None + ), + ) + + else: + output = self.decoder.forward( + gaussians, + batch["target"]["extrinsics"], + batch["target"]["intrinsics"], + batch["target"]["near"], + batch["target"]["far"], + (h, w), + depth_mode=self.train_cfg.depth_mode, + ) + + target_gt = batch["target"]["image"] + + # Compute metrics. + psnr = compute_psnr( + rearrange(target_gt, "b v c h w -> (b v) c h w"), + rearrange(output.color, "b v c h w -> (b v) c h w"), + ) + self.log("train/psnr", psnr.mean()) + + # Compute and log loss. + total_loss_per_sample = torch.zeros(batch_size, device=target_gt.device, dtype=target_gt.dtype) + + valid_depth_mask = None + + for loss_fn in self.losses: + if loss_fn.name == "mse": + loss_per_sample = loss_fn.forward( + output, + batch, + gaussians, + self.global_step, + l1_loss=self.train_cfg.l1_loss, + clamp_large_error=self.train_cfg.train_ignore_large_loss, + valid_depth_mask=valid_depth_mask, + reduction="none", + ) + else: + loss_per_sample = loss_fn.forward( + output, + batch, + gaussians, + self.global_step, + valid_depth_mask=valid_depth_mask, + reduction="none", + ) + self.log(f"loss/{loss_fn.name}", loss_per_sample.mean()) + total_loss_per_sample = total_loss_per_sample + loss_per_sample + + if ( + codec_depth_loss_pred is not None + and codec_depth_loss_target is not None + and codec_depth_loss_pred.shape == codec_depth_loss_target.shape + ): + codec_depth_loss = F.l1_loss( + codec_depth_loss_pred, + codec_depth_loss_target.detach(), + reduction="none", + ).mean(dim=(1, 2, 3)) + self.log("loss/codec_depth", codec_depth_loss.mean()) + if ( + pred_depths is not None + and codec_depth_targets is not None + and pred_depths.shape == codec_depth_targets.shape + ): + codec_depth_metric_l1 = F.l1_loss( + pred_depths, + codec_depth_targets.detach(), + reduction="none", + ).mean(dim=(1, 2, 3)) + self.log("loss/codec_depth_metric_l1", codec_depth_metric_l1.mean()) + total_loss_per_sample = ( + total_loss_per_sample + + self._codec_depth_loss_weight() * codec_depth_loss + ) + + # color loss on intermediate output + if supervise_intermediate_depth: + for loss_fn in self.losses: + if output_intermediate.color.size(0) != batch_size: + assert output_intermediate.color.size(0) % batch_size == 0 + num_intermediate = output_intermediate.color.size(0) // batch_size + intermediate_loss = torch.zeros( + batch_size, + device=target_gt.device, + dtype=target_gt.dtype, + ) + for i in range(num_intermediate): + curr_output = DecoderOutput( + color=output_intermediate.color[ + (batch_size * i) : (batch_size * (i + 1)) + ], + depth=( + output_intermediate.depth[ + (batch_size * i) : (batch_size * (i + 1)) + ] + if output_intermediate.depth is not None + else None + ), + ) + curr_loss_weight = self.train_cfg.intermediate_loss_weight ** ( + num_intermediate - i + ) + + if loss_fn.name == "mse": + loss = loss_fn.forward( + curr_output, + batch, + gaussians, + self.global_step, + l1_loss=self.train_cfg.l1_loss, + clamp_large_error=self.train_cfg.train_ignore_large_loss, + valid_depth_mask=valid_depth_mask, + reduction="none", + ) + else: + loss = loss_fn.forward( + curr_output, + batch, + gaussians, + self.global_step, + valid_depth_mask=valid_depth_mask, + reduction="none", + ) + + intermediate_loss = intermediate_loss + curr_loss_weight * loss + + self.log(f"loss/{loss_fn.name}_intermediate", intermediate_loss.mean()) + total_loss_per_sample = total_loss_per_sample + intermediate_loss + else: + if loss_fn.name == "mse": + loss = loss_fn.forward( + output_intermediate, + batch, + gaussians, + self.global_step, + l1_loss=self.train_cfg.l1_loss, + clamp_large_error=self.train_cfg.train_ignore_large_loss, + valid_depth_mask=valid_depth_mask, + reduction="none", + ) + else: + loss = loss_fn.forward( + output_intermediate, + batch, + gaussians, + self.global_step, + valid_depth_mask=valid_depth_mask, + reduction="none", + ) + weighted_loss = self.train_cfg.intermediate_loss_weight * loss + self.log(f"loss/{loss_fn.name}_intermediate", weighted_loss.mean()) + total_loss_per_sample = ( + total_loss_per_sample + weighted_loss + ) + + distortion_loss_per_sample = total_loss_per_sample + + if codec_loss is not None: + codec_loss_mean = codec_loss.mean() + self.log("loss/codec", codec_loss_mean) + num_context_views = batch["context"]["image"].shape[1] + self.log("loss/codec_per_view", (codec_loss / num_context_views).mean()) + if codec_lmb is not None: + self.log("loss/codec_lmb", codec_lmb.mean()) + codec_lmb_weight = self._codec_lmb_weight(codec_lmb) + self.log("loss/codec_lmb_weight", codec_lmb_weight.mean()) + distortion_loss_per_sample = distortion_loss_per_sample * codec_lmb_weight + total_loss_per_sample = distortion_loss_per_sample + self._codec_loss_weight() * codec_loss + else: + total_loss_per_sample = distortion_loss_per_sample + + self.log("loss/distortion", distortion_loss_per_sample.mean()) + total_loss = total_loss_per_sample.mean() + + self.log("loss/total", total_loss) + + if ( + self.global_rank == 0 + and self.global_step % self.train_cfg.print_log_every_n_steps == 0 + ): + print( + f"train step {self.global_step}; " + f"scene = {[x[:20] for x in batch['scene']]}; " + f"context = {batch['context']['index'].tolist()}; " + f"bound = [{batch['context']['near'].detach().cpu().numpy().mean()} " + f"{batch['context']['far'].detach().cpu().numpy().mean()}]; " + f"loss = {total_loss:.6f}" + ) + self.log("info/near", batch["context"]["near"].detach().cpu().numpy().mean()) + self.log("info/far", batch["context"]["far"].detach().cpu().numpy().mean()) + self.log("info/global_step", self.global_step) # hack for ckpt monitor + + # Tell the data loader processes about the current step. + if self.step_tracker is not None: + self.step_tracker.set_step(self.global_step) + + return total_loss + + def test_step(self, batch, batch_idx): + batch: BatchedExample = self.data_shim(batch) + b, v, _, h, w = batch["target"]["image"].shape + assert b == 1 + + # save input views for visualization + if self.test_cfg.save_input_images: + (scene,) = batch["scene"] + path = Path(get_cfg()["output_dir"]) + + input_images = batch["context"]["image"][0] # [V, 3, H, W] + index = batch["context"]["index"][0] + for idx, color in zip(index, input_images): + save_image(color, path / "images" / scene / f"color/input_{idx:0>6}.png") + + # save depth vis + if self.test_cfg.save_depth or self.test_cfg.save_gaussian: + visualization_dump = {} + else: + visualization_dump = None + + # Render Gaussians. + with self.benchmarker.time("encoder"): + encoder_output = self.encoder( + batch["context"], + self.global_step, + visualization_dump=visualization_dump, + return_bitstream=self.test_cfg.save_bitstream, + ) + ( + gaussians, + _, + _, + _, + _, + _, + _, + codec_num_bytes, + codec_bpp, + ) = self._unpack_encoder_output(encoder_output) + codec_extra_metrics = {} + if isinstance(encoder_output, dict): + for metric_name in ( + "codec_encode_context_seconds", + "codec_compress_seconds", + "codec_decompress_seconds", + "codec_camera_num_bytes", + "codec_feature_num_bytes", + "codec_estimated_feature_num_bytes", + "codec_latent_feature_num_bytes", + "codec_depth_num_bytes", + "codec_feature_raw_num_bytes", + "codec_depth_raw_num_bytes", + ): + metric_value = encoder_output.get(metric_name) + if metric_value is not None: + codec_extra_metrics[metric_name] = metric_value + + if self.test_cfg.save_bitstream: + if not isinstance(encoder_output, dict) or "bitstream" not in encoder_output: + raise RuntimeError( + "test.save_bitstream=true requires an enabled learned feature codec." + ) + scene = batch["scene"][0] + bitstream_path = ( + Path(get_cfg()["output_dir"]) / "bitstreams" / f"{scene}.bin" + ) + bitstream_path.parent.mkdir(parents=True, exist_ok=True) + bitstream_path.write_bytes(encoder_output["bitstream"]) + + # save gaussians + if self.test_cfg.save_gaussian: + scene = batch["scene"][0] + save_path = Path(get_cfg()['output_dir']) / 'gaussians' / (scene + '.ply') + save_gaussian_ply( + gaussians, + visualization_dump, + batch, + save_path, + save_all_gaussians=self.test_cfg.gaussian_export_all_points, + align_to_view=self.test_cfg.gaussian_export_align_to_view, + ) + + with self.benchmarker.time("decoder", num_calls=v): + camera_poses = batch["target"]["extrinsics"] + + if self.test_cfg.stabilize_camera: + stable_poses = render_stabilization_path( + camera_poses[0].detach().cpu().numpy(), + k_size=self.test_cfg.stab_camera_kernel, + ) + + stable_poses = list( + map( + lambda x: np.concatenate( + (x, np.array([[0.0, 0.0, 0.0, 1.0]])), axis=0 + ), + stable_poses, + ) + ) + stable_poses = torch.from_numpy(np.stack(stable_poses, axis=0)).to( + camera_poses + ) + camera_poses = stable_poses.unsqueeze(0) + + if self.test_cfg.render_chunk_size is not None: + chunk_size = self.test_cfg.render_chunk_size + num_chunks = math.ceil(camera_poses.shape[1] / chunk_size) + + output = None + for i in range(num_chunks): + start = chunk_size * i + end = chunk_size * (i + 1) + + render_intrinsics = batch["target"]["intrinsics"] + render_near = batch["target"]["near"] + render_far = batch["target"]["far"] + + curr_output = self.decoder.forward( + gaussians, + camera_poses[:, start:end], + render_intrinsics[:, start:end], + render_near[:, start:end], + render_far[:, start:end], + (h, w), + depth_mode=None, + ) + + if i == 0: + output = curr_output + else: + # ignore depth + output.color = torch.cat( + (output.color, curr_output.color), dim=1 + ) + + else: + output = self.decoder.forward( + gaussians, + camera_poses, + batch["target"]["intrinsics"], + batch["target"]["near"], + batch["target"]["far"], + (h, w), + depth_mode=None, + ) + + (scene,) = batch["scene"] + path = Path(get_cfg()["output_dir"]) + + # save depth + if self.test_cfg.save_depth: + depth = ( + visualization_dump["depth"][0, :, :, :, 0, 0].cpu().detach() + ) # [V, H, W] + + index = batch["context"]["index"][0] + + if self.test_cfg.save_depth_concat_img: + # concat (img0, img1, depth0, depth1) + image = batch['context']['image'][0] # [V, 3, H, W] in [0,1] + image = rearrange(image, "b c h w -> h (b w) c") # [H, VW, 3] + image_concat = (image.detach().cpu().numpy() * 255).astype(np.uint8) # [H, VW, 3] + + depth_concat = [] + + for idx, depth_i in zip(index, depth): + depth_viz = viz_depth_tensor( + 1.0 / depth_i, return_numpy=True + ) # [H, W, 3] + + if self.test_cfg.save_depth_concat_img: + depth_concat.append(depth_viz) + + save_path = path / "images" / scene / "depth" / f"{idx:0>6}.png" + save_dir = os.path.dirname(save_path) + os.makedirs(save_dir, exist_ok=True) + Image.fromarray(depth_viz).save(save_path) + + # save depth as npy + if self.test_cfg.save_depth_npy: + depth_npy = depth_i.detach().cpu().numpy() + save_path = path / "images" / scene / "depth" / f"{idx:0>6}.npy" + save_dir = os.path.dirname(save_path) + os.makedirs(save_dir, exist_ok=True) + np.save(save_path, depth_npy) + + if self.test_cfg.save_depth_concat_img: + depth_concat = np.concatenate(depth_concat, axis=1) # [H, VW, 3] + concat = np.concatenate((image_concat, depth_concat), axis=0) # [2H, VW, 3] + + save_path = path / "images" / scene / "depth" / f"img_depth_{scene}.png" + save_dir = os.path.dirname(save_path) + os.makedirs(save_dir, exist_ok=True) + Image.fromarray(concat).save(save_path) + + images_prob = output.color[0] + rgb_gt = batch["target"]["image"][0] + + # Save images. + if self.test_cfg.save_image: + if self.test_cfg.save_gt_image: + for index, color, gt in zip( + batch["target"]["index"][0], images_prob, rgb_gt + ): + save_image(color, path / "images" / scene / f"color/{index:0>6}.png") + save_image(gt, path / "images" / scene / f"color/{index:0>6}_gt.png") + else: + for index, color in zip(batch["target"]["index"][0], images_prob): + save_image(color, path / "images" / scene / f"color/{index:0>6}.png") + + # save video + if self.test_cfg.save_video: + frame_str = "_".join([str(x.item()) for x in batch["context"]["index"][0]]) + save_video( + [a for a in images_prob], + path / "videos" / f"{scene}_frame_{frame_str}.mp4", + ) + + # compute scores + if self.test_cfg.compute_scores: + if batch_idx < self.test_cfg.eval_time_skip_steps: + self.time_skip_steps_dict["encoder"] += 1 + self.time_skip_steps_dict["decoder"] += v + + rgb = images_prob + + if f"psnr" not in self.test_step_outputs: + self.test_step_outputs[f"psnr"] = [] + if f"ssim" not in self.test_step_outputs: + self.test_step_outputs[f"ssim"] = [] + if f"lpips" not in self.test_step_outputs: + self.test_step_outputs[f"lpips"] = [] + if codec_bpp is not None and f"codec_bpp" not in self.test_step_outputs: + self.test_step_outputs[f"codec_bpp"] = [] + if ( + codec_num_bytes is not None + and f"codec_num_bytes" not in self.test_step_outputs + ): + self.test_step_outputs[f"codec_num_bytes"] = [] + for metric_name in codec_extra_metrics: + if metric_name not in self.test_step_outputs: + self.test_step_outputs[metric_name] = [] + + self.test_step_outputs[f"psnr"].append( + compute_psnr(rgb_gt, rgb).mean().item() + ) + self.test_step_outputs[f"ssim"].append( + compute_ssim(rgb_gt, rgb).mean().item() + ) + self.test_step_outputs[f"lpips"].append( + self.compute_lpips(rgb_gt, rgb).mean().item() + ) + if codec_bpp is not None: + self.test_step_outputs[f"codec_bpp"].append(codec_bpp.item()) + if codec_num_bytes is not None: + self.test_step_outputs[f"codec_num_bytes"].append( + codec_num_bytes.item() + ) + for metric_name, metric_value in codec_extra_metrics.items(): + self.test_step_outputs[metric_name].append(metric_value.item()) + + def on_test_end(self) -> None: + out_dir = Path(get_cfg()["output_dir"]) / "metrics" + saved_scores = {} + if self.test_cfg.compute_scores: + self.benchmarker.dump_memory(out_dir / "peak_memory.json") + self.benchmarker.dump(out_dir / "benchmark.json") + + for metric_name, metric_scores in self.test_step_outputs.items(): + avg_scores = sum(metric_scores) / len(metric_scores) + saved_scores[metric_name] = avg_scores + print(metric_name, avg_scores) + with (out_dir / f"scores_{metric_name}_all.json").open("w") as f: + json.dump(metric_scores, f) + metric_scores.clear() + + for tag, times in self.benchmarker.execution_times.items(): + times = times[int(self.time_skip_steps_dict[tag]) :] + saved_scores[tag] = [len(times), np.mean(times)] + print( + f"{tag}: {len(times)} calls, avg. {np.mean(times)} seconds per call" + ) + self.time_skip_steps_dict[tag] = 0 + + with (out_dir / f"scores_all_avg.json").open("w") as f: + json.dump(saved_scores, f) + self.benchmarker.clear_history() + else: + self.benchmarker.dump(out_dir / "benchmark.json") + self.benchmarker.dump_memory(out_dir / "peak_memory.json") + self.benchmarker.summarize() + + @rank_zero_only + def validation_step(self, batch, batch_idx): + batch: BatchedExample = self.data_shim(batch) + + if self.global_rank == 0: + print( + f"validation step {self.global_step}; " + f"scene = {[a[:20] for a in batch['scene']]}; " + f"context = {batch['context']['index'].tolist()}" + ) + + # Render Gaussians. + b, _, _, h, w = batch["target"]["image"].shape + assert b == 1 + gaussians_softmax = self.encoder( + batch["context"], + self.global_step, + ) + ( + gaussians_softmax, + pred_depths, + _, + _, + _, + _, + _, + codec_num_bytes, + codec_bpp, + ) = self._unpack_encoder_output( + gaussians_softmax + ) + + output_softmax = self.decoder.forward( + gaussians_softmax, + batch["target"]["extrinsics"], + batch["target"]["intrinsics"], + batch["target"]["near"], + batch["target"]["far"], + (h, w), + ) + rgb_softmax = output_softmax.color[0] + + # Compute validation metrics. + rgb_gt = batch["target"]["image"][0] + for tag, rgb in zip(("val",), (rgb_softmax,)): + psnr = compute_psnr(rgb_gt, rgb).mean() + self.log(f"val/psnr_{tag}", psnr) + lpips = self.compute_lpips(rgb_gt, rgb).mean() + self.log(f"val/lpips_{tag}", lpips) + ssim = compute_ssim(rgb_gt, rgb).mean() + self.log(f"val/ssim_{tag}", ssim) + if codec_bpp is not None: + self.log("val/codec_bpp", codec_bpp) + if codec_num_bytes is not None: + self.log("val/codec_num_bytes", codec_num_bytes) + + # viz depth + if pred_depths is not None: + # only visualize predicted depth + pred_depths = pred_depths[0] # [V, H, W] + + # gaussian downsample + if pred_depths.shape[1:] != batch["context"]["image"].shape[-2:]: + pred_depths = F.interpolate( + pred_depths.unsqueeze(1), + size=batch["context"]["image"].shape[-2:], + mode="bilinear", + align_corners=True, + ).squeeze(1) + + inverse_depth_pred = 1.0 / pred_depths + + concat = [] + for i in range(inverse_depth_pred.size(0)): + concat.append(inverse_depth_pred[i]) + + concat = torch.cat(concat, dim=1) # [H, W*N] + + depth_viz = viz_depth_tensor(concat.cpu().detach()) # [3, H, W*N] + + # also concat images + input_images = batch["context"]["image"][0] # [N, 3, H, W] + concat_img = [img for img in input_images] + concat_img = torch.cat(concat_img, dim=-1) * 255 # [3, H, W*N] + + concat = torch.cat( + (concat_img.cpu().detach(), depth_viz), dim=1 + ) # [3, H*2, W*N] + + self.logger.log_image( + "depth", + [concat], + step=self.global_step, + caption=batch["scene"], + ) + + # Construct comparison image. + comparison = hcat( + add_label(vcat(*batch["context"]["image"][0]), "Context"), + add_label(vcat(*rgb_gt), "Target (Ground Truth)"), + add_label(vcat(*rgb_softmax), "Target (Prediction)"), + ) + self.logger.log_image( + "comparison", + [prep_image(add_border(comparison))], + step=self.global_step, + caption=batch["scene"], + ) + + # Run video validation step. + if not self.train_cfg.no_viz_video: + self.render_video_interpolation(batch) + self.render_video_wobble(batch) + if self.train_cfg.extended_visualization: + self.render_video_interpolation_exaggerated(batch) + + def on_validation_epoch_end(self) -> None: + """hack to run the full validation""" + if self.trainer.sanity_checking and self.global_rank == 0: + print(self.encoder) # log the model to wandb log files + + if (not self.trainer.sanity_checking) and (self.eval_data_cfg is not None): + self.eval_cnt = self.eval_cnt + 1 + if self.eval_cnt % self.train_cfg.eval_model_every_n_val == 0: + # backup current ckpt before running full test sets eval + if self.train_cfg.eval_save_model: + ckpt_saved_path = ( + self.trainer.checkpoint_callback.format_checkpoint_name( + dict( + epoch=self.trainer.current_epoch, + step=self.trainer.global_step, + ) + ) + ) + backup_dir = str( + Path(ckpt_saved_path).parent.parent / "checkpoints_backups" + ) + if self.global_rank == 0: + os.makedirs(backup_dir, exist_ok=True) + ckpt_saved_path = os.path.join( + backup_dir, os.path.basename(ckpt_saved_path) + ) + # call save_checkpoint on ALL process as suggested by pytorch_lightning + self.trainer.save_checkpoint( + ckpt_saved_path, + weights_only=True, + ) + if self.global_rank == 0: + print(f"backup model to {ckpt_saved_path}.") + + # run full test sets eval on rank=0 device + self.run_full_test_sets_eval() + + @rank_zero_only + def run_full_test_sets_eval(self) -> None: + start_t = time.time() + + pred_depths = None + depth_gt = None + + full_testsets = self.trainer.datamodule.test_dataloader( + dataset_cfg=self.eval_data_cfg + ) + scores_dict: dict[str, list[float]] = {} + + for score_tag in ("psnr", "ssim", "lpips"): + scores_dict[score_tag] = [] + + # evaluate depth + for score_tag in ("abs_rel", "rmse", "a1"): + scores_dict[score_tag] = [] + + self.benchmarker.clear_history() + time_skip_first_n_steps = min( + self.train_cfg.eval_time_skip_steps, len(full_testsets) + ) + time_skip_steps_dict = {"encoder": 0, "decoder": 0} + for batch_idx, batch in tqdm( + enumerate(full_testsets), + total=min(len(full_testsets), self.train_cfg.eval_data_length), + ): + if batch_idx >= self.train_cfg.eval_data_length: + break + + batch = self.data_shim(batch) + batch = self.transfer_batch_to_device(batch, "cuda", dataloader_idx=0) + + # Render Gaussians. + b, v, _, h, w = batch["target"]["image"].shape + assert b == 1 + if batch_idx < time_skip_first_n_steps: + time_skip_steps_dict["encoder"] += 1 + time_skip_steps_dict["decoder"] += v + + with self.benchmarker.time("encoder"): + encoder_output = self.encoder( + batch["context"], + self.global_step, + ) + ( + gaussians, + pred_depths, + _, + _, + _, + _, + _, + codec_num_bytes, + codec_bpp, + ) = self._unpack_encoder_output(encoder_output) + + with self.benchmarker.time("decoder", num_calls=v): + output = self.decoder.forward( + gaussians, + batch["target"]["extrinsics"], + batch["target"]["intrinsics"], + batch["target"]["near"], + batch["target"]["far"], + (h, w), + ) + + # Compute validation metrics. + rgb_gt = batch["target"]["image"][0] + rgb = output.color[0] + scores_dict["psnr"].append(compute_psnr(rgb_gt, rgb).mean().item()) + scores_dict["lpips"].append( + self.compute_lpips(rgb_gt, rgb).mean().item() + ) + scores_dict["ssim"].append(compute_ssim(rgb_gt, rgb).mean().item()) + + if codec_bpp is not None: + scores_dict.setdefault("codec_bpp", []).append(codec_bpp.item()) + if codec_num_bytes is not None: + scores_dict.setdefault("codec_num_bytes", []).append( + codec_num_bytes.item() + ) + + # summarise scores and log to logger + for score_tag, scores in scores_dict.items(): + if scores: + self.log(f"test/{score_tag}", sum(scores) / len(scores)) + # summarise run time + for tag, times in self.benchmarker.execution_times.items(): + times = times[int(time_skip_steps_dict[tag]) :] + print(f"{tag}: {len(times)} calls, avg. {np.mean(times)} seconds per call") + self.log(f"test/runtime_avg_{tag}", np.mean(times)) + self.benchmarker.clear_history() + + overall_eval_time = time.time() - start_t + print(f"Eval total time cost: {overall_eval_time:.3f}s") + self.log("test/runtime_all", overall_eval_time) + + @rank_zero_only + def render_video_wobble(self, batch: BatchedExample) -> None: + # Two views are needed to get the wobble radius. + _, v, _, _ = batch["context"]["extrinsics"].shape + if v != 2: + return + + def trajectory_fn(t): + origin_a = batch["context"]["extrinsics"][:, 0, :3, 3] + origin_b = batch["context"]["extrinsics"][:, 1, :3, 3] + delta = (origin_a - origin_b).norm(dim=-1) + extrinsics = generate_wobble( + batch["context"]["extrinsics"][:, 0], + delta * 0.25, + t, + ) + intrinsics = repeat( + batch["context"]["intrinsics"][:, 0], + "b i j -> b v i j", + v=t.shape[0], + ) + return extrinsics, intrinsics + + return self.render_video_generic(batch, trajectory_fn, "wobble", num_frames=60) + + @rank_zero_only + def render_video_interpolation(self, batch: BatchedExample) -> None: + _, v, _, _ = batch["context"]["extrinsics"].shape + + def trajectory_fn(t): + extrinsics = interpolate_extrinsics( + batch["context"]["extrinsics"][0, 0], + ( + batch["context"]["extrinsics"][0, 1] + if v == 2 + else batch["target"]["extrinsics"][0, 0] + ), + t, + ) + intrinsics = interpolate_intrinsics( + batch["context"]["intrinsics"][0, 0], + ( + batch["context"]["intrinsics"][0, 1] + if v == 2 + else batch["target"]["intrinsics"][0, 0] + ), + t, + ) + return extrinsics[None], intrinsics[None] + + return self.render_video_generic(batch, trajectory_fn, "rgb") + + @rank_zero_only + def render_video_interpolation_exaggerated(self, batch: BatchedExample) -> None: + # Two views are needed to get the wobble radius. + _, v, _, _ = batch["context"]["extrinsics"].shape + if v != 2: + return + + def trajectory_fn(t): + origin_a = batch["context"]["extrinsics"][:, 0, :3, 3] + origin_b = batch["context"]["extrinsics"][:, 1, :3, 3] + delta = (origin_a - origin_b).norm(dim=-1) + tf = generate_wobble_transformation( + delta * 0.5, + t, + 5, + scale_radius_with_t=False, + ) + extrinsics = interpolate_extrinsics( + batch["context"]["extrinsics"][0, 0], + ( + batch["context"]["extrinsics"][0, 1] + if v == 2 + else batch["target"]["extrinsics"][0, 0] + ), + t * 5 - 2, + ) + intrinsics = interpolate_intrinsics( + batch["context"]["intrinsics"][0, 0], + ( + batch["context"]["intrinsics"][0, 1] + if v == 2 + else batch["target"]["intrinsics"][0, 0] + ), + t * 5 - 2, + ) + return extrinsics @ tf, intrinsics[None] + + return self.render_video_generic( + batch, + trajectory_fn, + "interpolation_exaggerated", + num_frames=300, + smooth=False, + loop_reverse=False, + ) + + @rank_zero_only + def render_video_generic( + self, + batch: BatchedExample, + trajectory_fn: TrajectoryFn, + name: str, + num_frames: int = 30, + smooth: bool = True, + loop_reverse: bool = True, + ) -> None: + encoder_output = self.encoder(batch["context"], self.global_step) + gaussians, _, _, _, _, _, _, _, _ = self._unpack_encoder_output(encoder_output) + + t = torch.linspace(0, 1, num_frames, dtype=torch.float32, device=self.device) + if smooth: + t = (torch.cos(torch.pi * (t + 1)) + 1) / 2 + + extrinsics, intrinsics = trajectory_fn(t) + + _, _, _, h, w = batch["context"]["image"].shape + + # Color-map the result. + def depth_map(result): + near = result[result > 0][:16_000_000].quantile(0.01).log() + far = result.view(-1)[:16_000_000].quantile(0.99).log() + result = result.log() + result = 1 - (result - near) / (far - near) + return apply_color_map_to_image(result, "turbo") + + near = repeat(batch["context"]["near"][:, 0], "b -> b v", v=num_frames) + far = repeat(batch["context"]["far"][:, 0], "b -> b v", v=num_frames) + output = self.decoder.forward( + gaussians, extrinsics, intrinsics, near, far, (h, w), "depth" + ) + predictions = [ + vcat(rgb, depth) + for rgb, depth in zip(output.color[0], depth_map(output.depth[0])) + ] + + images = [ + add_border(hcat(add_label(prediction, "Prediction"))) + for prediction in predictions + ] + + video = torch.stack(images) + if loop_reverse: + video = torch.cat((video, video.flip(0)[1:-1]), dim=0) + + key = f"video/{name}" + if isinstance(self.logger, WandbLogger): + video_uint8 = ( + (video.clip(min=0, max=1) * 255).type(torch.uint8).cpu().numpy() + ) + wandb.log( + {key: wandb.Video(video_uint8[None], fps=30, format="mp4")} + ) + else: + assert isinstance(self.logger, LocalLogger) + save_video( + list(video), + self.logger.log_path / key / f"{self.global_step:0>6}.mp4", + fps=30, + ) + + def configure_optimizers(self): + if self.optimizer_cfg.lr_monodepth > 0: + pretrained_params = [] + new_params = [] + + for name, param in self.named_parameters(): + if not param.requires_grad: + continue + if "pretrained" in name: + pretrained_params.append(param) + else: + new_params.append(param) + + optimizer = torch.optim.AdamW( + [ + { + "params": pretrained_params, + "lr": self.optimizer_cfg.lr_monodepth, + }, + {"params": new_params, "lr": self.optimizer_cfg.lr}, + ], + weight_decay=self.optimizer_cfg.weight_decay, + ) + + scheduler = torch.optim.lr_scheduler.OneCycleLR( + optimizer, + [self.optimizer_cfg.lr_monodepth, self.optimizer_cfg.lr], + self.trainer.max_steps + 10, + pct_start=0.01, + div_factor=self.optimizer_cfg.div_factor, + final_div_factor=self.optimizer_cfg.final_div_factor, + cycle_momentum=False, + anneal_strategy="cos", + ) + + else: + optimizer = optim.AdamW( + [param for param in self.parameters() if param.requires_grad], + lr=self.optimizer_cfg.lr, + weight_decay=self.optimizer_cfg.weight_decay, + ) + + scheduler = torch.optim.lr_scheduler.OneCycleLR( + optimizer, + self.optimizer_cfg.lr, + self.trainer.max_steps + 10, + pct_start=0.01, + div_factor=self.optimizer_cfg.div_factor, + final_div_factor=self.optimizer_cfg.final_div_factor, + cycle_momentum=False, + anneal_strategy="cos", + ) + + return { + "optimizer": optimizer, + "lr_scheduler": { + "scheduler": scheduler, + "interval": "step", + "frequency": 1, + }, + } diff --git a/src/model/ply_export.py b/src/model/ply_export.py new file mode 100644 index 0000000000000000000000000000000000000000..e0a80f8b3f3dd2200561653aaa5470095b7bafa3 --- /dev/null +++ b/src/model/ply_export.py @@ -0,0 +1,259 @@ +from pathlib import Path + +import numpy as np +import torch +from einops import einsum, rearrange, repeat +from jaxtyping import Float +from plyfile import PlyData, PlyElement +from scipy.spatial.transform import Rotation as R +from torch import Tensor + +from ..misc.sh_rotation import rotate_sh + + +def construct_list_of_attributes(num_rest: int) -> list[str]: + attributes = ["x", "y", "z", "nx", "ny", "nz"] + for i in range(3): + attributes.append(f"f_dc_{i}") + for i in range(num_rest): + attributes.append(f"f_rest_{i}") + attributes.append("opacity") + for i in range(3): + attributes.append(f"scale_{i}") + for i in range(4): + attributes.append(f"rot_{i}") + return attributes + + +def export_ply( + extrinsics: Float[Tensor, "4 4"], + means: Float[Tensor, "gaussian 3"], + scales: Float[Tensor, "gaussian 3"], + rotations: Float[Tensor, "gaussian 4"], + harmonics: Float[Tensor, "gaussian 3 d_sh"], + opacities: Float[Tensor, " gaussian"], + path: Path, + align_to_view: bool = False, +): + path = Path(path) + with torch.autocast(device_type=means.device.type, enabled=False): + means = means.float() + scales = scales.float() + rotations = rotations.float() + harmonics = harmonics.float() + opacities = opacities.float() + extrinsics = extrinsics.to(means) + rotations = rotations / rotations.norm(dim=-1, keepdim=True).clamp_min(1e-8) + + if align_to_view: + view_rotation = extrinsics[:3, :3].transpose(-1, -2) + view_origin = extrinsics[:3, 3] + means = einsum( + view_rotation, + means - view_origin, + "i j, ... j -> ... i", + ) + harmonics = rotate_sh(harmonics, view_rotation) + rotations = R.from_quat(rotations.detach().cpu().numpy()).as_matrix() + rotations = view_rotation.detach().cpu().numpy() @ rotations + rotations = R.from_matrix(rotations).as_quat() + else: + rotations = rotations.detach().cpu().numpy() + + x, y, z, w = rearrange(rotations, "g xyzw -> xyzw g") + rotations = np.stack((w, x, y, z), axis=-1) + means = means.detach().cpu().numpy() + normals = np.zeros_like(means) + harmonics_dc = harmonics[..., 0].detach().cpu().contiguous().numpy() + harmonics_rest = ( + harmonics[..., 1:].flatten(start_dim=1).detach().cpu().contiguous().numpy() + ) + logit_opacities = ( + torch.logit(opacities[..., None].clamp(1e-6, 1 - 1e-6)) + .detach() + .cpu() + .numpy() + ) + log_scales = scales.clamp_min(1e-12).log().detach().cpu().numpy() + + num_rest = 3 * (harmonics.shape[-1] - 1) + dtype_full = [ + (attribute, "f4") for attribute in construct_list_of_attributes(num_rest) + ] + elements = np.empty(means.shape[0], dtype=dtype_full) + attributes = ( + means, + normals, + harmonics_dc, + harmonics_rest, + logit_opacities, + log_scales, + rotations, + ) + attributes = np.concatenate(attributes, axis=1) + elements[:] = list(map(tuple, attributes)) + path.parent.mkdir(exist_ok=True, parents=True) + PlyData([PlyElement.describe(elements, "vertex")]).write(path) + + +def infer_gaussian_grid( + num_gaussians: int, + num_views: int, + image_height: int, + image_width: int, +) -> tuple[int, int, int]: + if num_gaussians % num_views != 0: + raise ValueError( + f"Cannot split {num_gaussians} Gaussians over {num_views} views." + ) + + gaussians_per_view = num_gaussians // num_views + for downsample in (1, 2, 4, 8, 16, 32): + if image_height % downsample != 0 or image_width % downsample != 0: + continue + height = image_height // downsample + width = image_width // downsample + pixels = height * width + if pixels > 0 and gaussians_per_view % pixels == 0: + return height, width, gaussians_per_view // pixels + + raise ValueError( + "Could not infer Gaussian grid layout from " + f"num_gaussians={num_gaussians}, num_views={num_views}, " + f"image_shape=({image_height}, {image_width})." + ) + + +def select_gaussians( + tensor: Tensor, + mask: Tensor, + num_views: int, + height: int, + width: int, + points_per_pixel: int, +) -> Tensor: + tensor = rearrange( + tensor, + "() (v h w spp) ... -> h w spp v ...", + v=num_views, + h=height, + w=width, + spp=points_per_pixel, + ) + return tensor[mask] + + +def local_to_world_rotations( + rotations: Tensor, + context_extrinsics: Tensor, + num_views: int, + height: int, + width: int, + points_per_pixel: int, + mask: Tensor | None, +) -> Tensor: + if mask is None: + local_rotations = rotations[0] + c2w_rotations = repeat( + context_extrinsics[0, :, :3, :3], + "v i j -> (v h w spp) i j", + h=height, + w=width, + spp=points_per_pixel, + ) + else: + local_rotations = select_gaussians( + rotations, + mask, + num_views, + height, + width, + points_per_pixel, + ) + c2w_rotations = repeat( + context_extrinsics[0, :, :3, :3], + "v i j -> h w spp v i j", + h=height, + w=width, + spp=points_per_pixel, + )[mask] + + local_rotation_mats = R.from_quat( + local_rotations.detach().cpu().numpy() + ).as_matrix() + world_rotation_mats = c2w_rotations.detach().cpu().numpy() @ local_rotation_mats + world_rotations = R.from_matrix(world_rotation_mats).as_quat() + return torch.from_numpy(world_rotations).to( + device=rotations.device, + dtype=rotations.dtype, + ) + + +def save_gaussian_ply( + gaussians, + visualization_dump, + example, + save_path, + save_all_gaussians: bool = True, + align_to_view: bool = False, +): + del visualization_dump + + v, _, image_height, image_width = example["context"]["image"].shape[1:] + h, w, points_per_pixel = infer_gaussian_grid( + gaussians.means.shape[1], + v, + image_height, + image_width, + ) + + if gaussians.scales is None or gaussians.rotations is None: + raise ValueError("PLY export requires Gaussians.scales and Gaussians.rotations.") + + if save_all_gaussians: + mask = None + selected_means = gaussians.means[0] + selected_scales = gaussians.scales[0] + selected_harmonics = gaussians.harmonics[0] + selected_opacities = gaussians.opacities[0] + else: + mask = torch.zeros( + (h, w, points_per_pixel, v), + dtype=torch.bool, + device=gaussians.means.device, + ) + gaussian_trim = 8 + mask[gaussian_trim:-gaussian_trim, gaussian_trim:-gaussian_trim, :, :] = True + selected_means = select_gaussians( + gaussians.means, mask, v, h, w, points_per_pixel + ) + selected_scales = select_gaussians( + gaussians.scales, mask, v, h, w, points_per_pixel + ) + selected_harmonics = select_gaussians( + gaussians.harmonics, mask, v, h, w, points_per_pixel + ) + selected_opacities = select_gaussians( + gaussians.opacities, mask, v, h, w, points_per_pixel + ) + + world_rotations = local_to_world_rotations( + gaussians.rotations, + example["context"]["extrinsics"], + v, + h, + w, + points_per_pixel, + mask, + ) + + export_ply( + example["context"]["extrinsics"][0, v // 2], + selected_means, + selected_scales, + world_rotations, + selected_harmonics, + selected_opacities, + save_path, + align_to_view=align_to_view, + ) diff --git a/src/model/types.py b/src/model/types.py new file mode 100644 index 0000000000000000000000000000000000000000..45a0bb4219bb81510bbd1affb509ab130bfbcf7f --- /dev/null +++ b/src/model/types.py @@ -0,0 +1,15 @@ +from dataclasses import dataclass + +from jaxtyping import Float +from torch import Tensor + + +@dataclass +class Gaussians: + means: Float[Tensor, "batch gaussian dim"] + covariances: Float[Tensor, "batch gaussian dim dim"] + harmonics: Float[Tensor, "batch gaussian 3 d_sh"] + opacities: Float[Tensor, "batch gaussian"] + scales: Float[Tensor, "batch gaussian 3"] | None = None + rotations: Float[Tensor, "batch gaussian 4"] | None = None + rotations_unnorm: Float[Tensor, "batch gaussian 4"] | None = None diff --git a/src/visualization/annotation.py b/src/visualization/annotation.py new file mode 100644 index 0000000000000000000000000000000000000000..ae2a2ca8df66fa8aaf5cad40b50b94e151d226f4 --- /dev/null +++ b/src/visualization/annotation.py @@ -0,0 +1,49 @@ +from pathlib import Path +from string import ascii_letters, digits, punctuation + +import numpy as np +import torch +from einops import rearrange +from jaxtyping import Float +from PIL import Image, ImageDraw, ImageFont +from torch import Tensor + +from .layout import vcat + +EXPECTED_CHARACTERS = digits + punctuation + ascii_letters + + +def draw_label( + text: str, + font: Path, + font_size: int, + device: torch.device = torch.device("cpu"), +) -> Float[Tensor, "3 height width"]: + """Draw a black label on a white background with no border.""" + try: + font = ImageFont.truetype(str(font), font_size) + except OSError: + font = ImageFont.load_default() + left, _, right, _ = font.getbbox(text) + width = right - left + _, top, _, bottom = font.getbbox(EXPECTED_CHARACTERS) + height = bottom - top + image = Image.new("RGB", (width, height), color="white") + draw = ImageDraw.Draw(image) + draw.text((0, 0), text, font=font, fill="black") + image = torch.tensor(np.array(image) / 255, dtype=torch.float32, device=device) + return rearrange(image, "h w c -> c h w") + + +def add_label( + image: Float[Tensor, "3 height width"], + label: str, + font: Path = Path("assets/Inter-Regular.otf"), + font_size: int = 24, +) -> Float[Tensor, "3 height_with_label width_with_label"]: + return vcat( + draw_label(label, font, font_size, image.device), + image, + align="left", + gap=4, + ) diff --git a/src/visualization/camera_trajectory/interpolation.py b/src/visualization/camera_trajectory/interpolation.py new file mode 100644 index 0000000000000000000000000000000000000000..7c5d1781b18d5d31ec1862b27cb44d4c73045005 --- /dev/null +++ b/src/visualization/camera_trajectory/interpolation.py @@ -0,0 +1,255 @@ +import torch +from einops import einsum, rearrange, reduce +from jaxtyping import Float +from scipy.spatial.transform import Rotation as R +from torch import Tensor + + +def interpolate_intrinsics( + initial: Float[Tensor, "*#batch 3 3"], + final: Float[Tensor, "*#batch 3 3"], + t: Float[Tensor, " time_step"], +) -> Float[Tensor, "*batch time_step 3 3"]: + initial = rearrange(initial, "... i j -> ... () i j") + final = rearrange(final, "... i j -> ... () i j") + t = rearrange(t, "t -> t () ()") + return initial + (final - initial) * t + + +def intersect_rays( + a_origins: Float[Tensor, "*#batch dim"], + a_directions: Float[Tensor, "*#batch dim"], + b_origins: Float[Tensor, "*#batch dim"], + b_directions: Float[Tensor, "*#batch dim"], +) -> Float[Tensor, "*batch dim"]: + """Compute the least-squares intersection of rays. Uses the math from here: + https://math.stackexchange.com/a/1762491/286022 + """ + + # Broadcast and stack the tensors. + a_origins, a_directions, b_origins, b_directions = torch.broadcast_tensors( + a_origins, a_directions, b_origins, b_directions + ) + origins = torch.stack((a_origins, b_origins), dim=-2) + directions = torch.stack((a_directions, b_directions), dim=-2) + + # Compute n_i * n_i^T - eye(3) from the equation. + n = einsum(directions, directions, "... n i, ... n j -> ... n i j") + n = n - torch.eye(3, dtype=origins.dtype, device=origins.device) + + # Compute the left-hand side of the equation. + lhs = reduce(n, "... n i j -> ... i j", "sum") + + # Compute the right-hand side of the equation. + rhs = einsum(n, origins, "... n i j, ... n j -> ... n i") + rhs = reduce(rhs, "... n i -> ... i", "sum") + + # Left-matrix-multiply both sides by the inverse of lhs to find p. + return torch.linalg.lstsq(lhs, rhs).solution + + +def normalize(a: Float[Tensor, "*#batch dim"]) -> Float[Tensor, "*#batch dim"]: + return a / a.norm(dim=-1, keepdim=True) + + +def generate_coordinate_frame( + y: Float[Tensor, "*#batch 3"], + z: Float[Tensor, "*#batch 3"], +) -> Float[Tensor, "*batch 3 3"]: + """Generate a coordinate frame given perpendicular, unit-length Y and Z vectors.""" + y, z = torch.broadcast_tensors(y, z) + return torch.stack([y.cross(z), y, z], dim=-1) + + +def generate_rotation_coordinate_frame( + a: Float[Tensor, "*#batch 3"], + b: Float[Tensor, "*#batch 3"], + eps: float = 1e-4, +) -> Float[Tensor, "*batch 3 3"]: + """Generate a coordinate frame where the Y direction is normal to the plane defined + by unit vectors a and b. The other axes are arbitrary.""" + device = a.device + + # Replace every entry in b that's parallel to the corresponding entry in a with an + # arbitrary vector. + b = b.detach().clone() + parallel = (einsum(a, b, "... i, ... i -> ...").abs() - 1).abs() < eps + b[parallel] = torch.tensor([0, 0, 1], dtype=b.dtype, device=device) + parallel = (einsum(a, b, "... i, ... i -> ...").abs() - 1).abs() < eps + b[parallel] = torch.tensor([0, 1, 0], dtype=b.dtype, device=device) + + # Generate the coordinate frame. The initial cross product defines the plane. + return generate_coordinate_frame(normalize(a.cross(b)), a) + + +def matrix_to_euler( + rotations: Float[Tensor, "*batch 3 3"], + pattern: str, +) -> Float[Tensor, "*batch 3"]: + *batch, _, _ = rotations.shape + rotations = rotations.reshape(-1, 3, 3) + angles_np = R.from_matrix(rotations.detach().cpu().numpy()).as_euler(pattern) + rotations = torch.tensor(angles_np, dtype=rotations.dtype, device=rotations.device) + return rotations.reshape(*batch, 3) + + +def euler_to_matrix( + rotations: Float[Tensor, "*batch 3"], + pattern: str, +) -> Float[Tensor, "*batch 3 3"]: + *batch, _ = rotations.shape + rotations = rotations.reshape(-1, 3) + matrix_np = R.from_euler(pattern, rotations.detach().cpu().numpy()).as_matrix() + rotations = torch.tensor(matrix_np, dtype=rotations.dtype, device=rotations.device) + return rotations.reshape(*batch, 3, 3) + + +def extrinsics_to_pivot_parameters( + extrinsics: Float[Tensor, "*#batch 4 4"], + pivot_coordinate_frame: Float[Tensor, "*#batch 3 3"], + pivot_point: Float[Tensor, "*#batch 3"], +) -> Float[Tensor, "*batch 5"]: + """Convert the extrinsics to a representation with 5 degrees of freedom: + 1. Distance from pivot point in the "X" (look cross pivot axis) direction. + 2. Distance from pivot point in the "Y" (pivot axis) direction. + 3. Distance from pivot point in the Z (look) direction + 4. Angle in plane + 5. Twist (rotation not in plane) + """ + + # The pivot coordinate frame's Z axis is normal to the plane. + pivot_axis = pivot_coordinate_frame[..., :, 1] + + # Compute the translation elements of the pivot parametrization. + translation_frame = generate_coordinate_frame(pivot_axis, extrinsics[..., :3, 2]) + origin = extrinsics[..., :3, 3] + delta = pivot_point - origin + translation = einsum(translation_frame, delta, "... i j, ... i -> ... j") + + # Add the rotation elements of the pivot parametrization. + inverted = pivot_coordinate_frame.inverse() @ extrinsics[..., :3, :3] + y, _, z = matrix_to_euler(inverted, "YXZ").unbind(dim=-1) + + return torch.cat([translation, y[..., None], z[..., None]], dim=-1) + + +def pivot_parameters_to_extrinsics( + parameters: Float[Tensor, "*#batch 5"], + pivot_coordinate_frame: Float[Tensor, "*#batch 3 3"], + pivot_point: Float[Tensor, "*#batch 3"], +) -> Float[Tensor, "*batch 4 4"]: + translation, y, z = parameters.split((3, 1, 1), dim=-1) + + euler = torch.cat((y, torch.zeros_like(y), z), dim=-1) + rotation = pivot_coordinate_frame @ euler_to_matrix(euler, "YXZ") + + # The pivot coordinate frame's Z axis is normal to the plane. + pivot_axis = pivot_coordinate_frame[..., :, 1] + + translation_frame = generate_coordinate_frame(pivot_axis, rotation[..., :3, 2]) + delta = einsum(translation_frame, translation, "... i j, ... j -> ... i") + origin = pivot_point - delta + + *batch, _ = origin.shape + extrinsics = torch.eye(4, dtype=parameters.dtype, device=parameters.device) + extrinsics = extrinsics.broadcast_to((*batch, 4, 4)).clone() + extrinsics[..., 3, 3] = 1 + extrinsics[..., :3, :3] = rotation + extrinsics[..., :3, 3] = origin + return extrinsics + + +def interpolate_circular( + a: Float[Tensor, "*#batch"], + b: Float[Tensor, "*#batch"], + t: Float[Tensor, "*#batch"], +) -> Float[Tensor, " *batch"]: + a, b, t = torch.broadcast_tensors(a, b, t) + + tau = 2 * torch.pi + a = a % tau + b = b % tau + + # Consider piecewise edge cases. + d = (b - a).abs() + a_left = a - tau + d_left = (b - a_left).abs() + a_right = a + tau + d_right = (b - a_right).abs() + use_d = (d < d_left) & (d < d_right) + use_d_left = (d_left < d_right) & (~use_d) + use_d_right = (~use_d) & (~use_d_left) + + result = a + (b - a) * t + result[use_d_left] = (a_left + (b - a_left) * t)[use_d_left] + result[use_d_right] = (a_right + (b - a_right) * t)[use_d_right] + + return result + + +def interpolate_pivot_parameters( + initial: Float[Tensor, "*#batch 5"], + final: Float[Tensor, "*#batch 5"], + t: Float[Tensor, " time_step"], +) -> Float[Tensor, "*batch time_step 5"]: + initial = rearrange(initial, "... d -> ... () d") + final = rearrange(final, "... d -> ... () d") + t = rearrange(t, "t -> t ()") + ti, ri = initial.split((3, 2), dim=-1) + tf, rf = final.split((3, 2), dim=-1) + + t_lerp = ti + (tf - ti) * t + r_lerp = interpolate_circular(ri, rf, t) + + return torch.cat((t_lerp, r_lerp), dim=-1) + + +@torch.no_grad() +def interpolate_extrinsics( + initial: Float[Tensor, "*#batch 4 4"], + final: Float[Tensor, "*#batch 4 4"], + t: Float[Tensor, " time_step"], + eps: float = 1e-4, +) -> Float[Tensor, "*batch time_step 4 4"]: + """Interpolate extrinsics by rotating around their "focus point," which is the + least-squares intersection between the look vectors of the initial and final + extrinsics. + """ + + initial = initial.type(torch.float64) + final = final.type(torch.float64) + t = t.type(torch.float64) + + # Based on the dot product between the look vectors, pick from one of two cases: + # 1. Look vectors are parallel: interpolate about their origins' midpoint. + # 3. Look vectors aren't parallel: interpolate about their focus point. + initial_look = initial[..., :3, 2] + final_look = final[..., :3, 2] + dot_products = einsum(initial_look, final_look, "... i, ... i -> ...") + parallel_mask = (dot_products.abs() - 1).abs() < eps + + # Pick focus points. + initial_origin = initial[..., :3, 3] + final_origin = final[..., :3, 3] + pivot_point = 0.5 * (initial_origin + final_origin) + pivot_point[~parallel_mask] = intersect_rays( + initial_origin[~parallel_mask], + initial_look[~parallel_mask], + final_origin[~parallel_mask], + final_look[~parallel_mask], + ) + + # Convert to pivot parameters. + pivot_frame = generate_rotation_coordinate_frame(initial_look, final_look, eps=eps) + initial_params = extrinsics_to_pivot_parameters(initial, pivot_frame, pivot_point) + final_params = extrinsics_to_pivot_parameters(final, pivot_frame, pivot_point) + + # Interpolate the pivot parameters. + interpolated_params = interpolate_pivot_parameters(initial_params, final_params, t) + + # Convert back. + return pivot_parameters_to_extrinsics( + interpolated_params.type(torch.float32), + rearrange(pivot_frame, "... i j -> ... () i j").type(torch.float32), + rearrange(pivot_point, "... xyz -> ... () xyz").type(torch.float32), + ) diff --git a/src/visualization/camera_trajectory/spin.py b/src/visualization/camera_trajectory/spin.py new file mode 100644 index 0000000000000000000000000000000000000000..eadddcbc9b33075e902d610a3fc40ef0ba7bde2e --- /dev/null +++ b/src/visualization/camera_trajectory/spin.py @@ -0,0 +1,37 @@ +import numpy as np +import torch +from einops import repeat +from jaxtyping import Float +from scipy.spatial.transform import Rotation as R +from torch import Tensor + + +def generate_spin( + num_frames: int, + device: torch.device, + elevation: float, + radius: float, +) -> Float[Tensor, "frame 4 4"]: + # Translate back along the camera's look vector. + tf_translation = torch.eye(4, dtype=torch.float32, device=device) + tf_translation[:2] *= -1 + tf_translation[2, 3] = -radius + + # Generate the transformation for the azimuth. + phi = 2 * np.pi * (np.arange(num_frames) / num_frames) + rotation_vectors = np.stack([np.zeros_like(phi), phi, np.zeros_like(phi)], axis=-1) + + azimuth = R.from_rotvec(rotation_vectors).as_matrix() + azimuth = torch.tensor(azimuth, dtype=torch.float32, device=device) + tf_azimuth = torch.eye(4, dtype=torch.float32, device=device) + tf_azimuth = repeat(tf_azimuth, "i j -> b i j", b=num_frames).clone() + tf_azimuth[:, :3, :3] = azimuth + + # Generate the transformation for the elevation. + deg_elevation = np.deg2rad(elevation) + elevation = R.from_rotvec(np.array([deg_elevation, 0, 0], dtype=np.float32)) + elevation = torch.tensor(elevation.as_matrix()) + tf_elevation = torch.eye(4, dtype=torch.float32, device=device) + tf_elevation[:3, :3] = elevation + + return tf_azimuth @ tf_elevation @ tf_translation diff --git a/src/visualization/camera_trajectory/wobble.py b/src/visualization/camera_trajectory/wobble.py new file mode 100644 index 0000000000000000000000000000000000000000..5fd71c8f8e3d808561894f993f6cd4988e469b39 --- /dev/null +++ b/src/visualization/camera_trajectory/wobble.py @@ -0,0 +1,32 @@ +import torch +from einops import rearrange +from jaxtyping import Float +from torch import Tensor + + +@torch.no_grad() +def generate_wobble_transformation( + radius: Float[Tensor, "*#batch"], + t: Float[Tensor, " time_step"], + num_rotations: int = 1, + scale_radius_with_t: bool = True, +) -> Float[Tensor, "*batch time_step 4 4"]: + # Generate a translation in the image plane. + tf = torch.eye(4, dtype=torch.float32, device=t.device) + tf = tf.broadcast_to((*radius.shape, t.shape[0], 4, 4)).clone() + radius = radius[..., None] + if scale_radius_with_t: + radius = radius * t + tf[..., 0, 3] = torch.sin(2 * torch.pi * num_rotations * t) * radius + tf[..., 1, 3] = -torch.cos(2 * torch.pi * num_rotations * t) * radius + return tf + + +@torch.no_grad() +def generate_wobble( + extrinsics: Float[Tensor, "*#batch 4 4"], + radius: Float[Tensor, "*#batch"], + t: Float[Tensor, " time_step"], +) -> Float[Tensor, "*batch time_step 4 4"]: + tf = generate_wobble_transformation(radius, t) + return rearrange(extrinsics, "... i j -> ... () i j") @ tf diff --git a/src/visualization/color_map.py b/src/visualization/color_map.py new file mode 100644 index 0000000000000000000000000000000000000000..c130fa7509939520ac5060b7853ee163264acaac --- /dev/null +++ b/src/visualization/color_map.py @@ -0,0 +1,48 @@ +import torch +from colorspacious import cspace_convert +from einops import rearrange +from jaxtyping import Float +from matplotlib import cm +from torch import Tensor + + +def apply_color_map( + x: Float[Tensor, " *batch"], + color_map: str = "inferno", +) -> Float[Tensor, "*batch 3"]: + cmap = cm.get_cmap(color_map) + + # Convert to NumPy so that Matplotlib color maps can be used. + mapped = cmap(x.detach().clip(min=0, max=1).cpu().numpy())[..., :3] + + # Convert back to the original format. + return torch.tensor(mapped, device=x.device, dtype=torch.float32) + + +def apply_color_map_to_image( + image: Float[Tensor, "*batch height width"], + color_map: str = "inferno", +) -> Float[Tensor, "*batch 3 height width"]: + image = apply_color_map(image, color_map) + return rearrange(image, "... h w c -> ... c h w") + + +def apply_color_map_2d( + x: Float[Tensor, "*#batch"], + y: Float[Tensor, "*#batch"], +) -> Float[Tensor, "*batch 3"]: + red = cspace_convert((189, 0, 0), "sRGB255", "CIELab") + blue = cspace_convert((0, 45, 255), "sRGB255", "CIELab") + white = cspace_convert((255, 255, 255), "sRGB255", "CIELab") + x_np = x.detach().clip(min=0, max=1).cpu().numpy()[..., None] + y_np = y.detach().clip(min=0, max=1).cpu().numpy()[..., None] + + # Interpolate between red and blue on the x axis. + interpolated = x_np * red + (1 - x_np) * blue + + # Interpolate between color and white on the y axis. + interpolated = y_np * interpolated + (1 - y_np) * white + + # Convert to RGB. + rgb = cspace_convert(interpolated, "CIELab", "sRGB1") + return torch.tensor(rgb, device=x.device, dtype=torch.float32).clip(min=0, max=1) diff --git a/src/visualization/colors.py b/src/visualization/colors.py new file mode 100644 index 0000000000000000000000000000000000000000..ef82964e0bf71f147d41d95a6e389b1f22ba57b5 --- /dev/null +++ b/src/visualization/colors.py @@ -0,0 +1,32 @@ +from PIL import ImageColor + +# https://sashamaps.net/docs/resources/20-colors/ +DISTINCT_COLORS = [ + "#e6194b", + "#3cb44b", + "#ffe119", + "#4363d8", + "#f58231", + "#911eb4", + "#46f0f0", + "#f032e6", + "#bcf60c", + "#fabebe", + "#008080", + "#e6beff", + "#9a6324", + "#fffac8", + "#800000", + "#aaffc3", + "#808000", + "#ffd8b1", + "#000075", + "#808080", + "#ffffff", + "#000000", +] + + +def get_distinct_color(index: int) -> tuple[float, float, float]: + hex = DISTINCT_COLORS[index % len(DISTINCT_COLORS)] + return tuple(x / 255 for x in ImageColor.getcolor(hex, "RGB")) diff --git a/src/visualization/drawing/cameras.py b/src/visualization/drawing/cameras.py new file mode 100644 index 0000000000000000000000000000000000000000..6a60beb53a1a24931cd4b822e5ebe65b6425a6d6 --- /dev/null +++ b/src/visualization/drawing/cameras.py @@ -0,0 +1,195 @@ +from typing import Optional + +import torch +from einops import einsum, rearrange, repeat +from jaxtyping import Float +from torch import Tensor + +from ...geometry.projection import unproject +from ..annotation import add_label +from .lines import draw_lines +from .types import Scalar, sanitize_scalar + + +def draw_cameras( + resolution: int, + extrinsics: Float[Tensor, "batch 4 4"], + intrinsics: Float[Tensor, "batch 3 3"], + color: Float[Tensor, "batch 3"], + near: Optional[Scalar] = None, + far: Optional[Scalar] = None, + margin: float = 0.1, # relative to AABB + frustum_scale: float = 0.05, # relative to image resolution +) -> Float[Tensor, "3 3 height width"]: + device = extrinsics.device + + # Compute scene bounds. + minima, maxima = compute_aabb(extrinsics, intrinsics, near, far) + scene_minima, scene_maxima = compute_equal_aabb_with_margin( + minima, maxima, margin=margin + ) + span = (scene_maxima - scene_minima).max() + + # Compute frustum locations. + corner_depth = (span * frustum_scale)[None] + frustum_corners = unproject_frustum_corners(extrinsics, intrinsics, corner_depth) + if near is not None: + near_corners = unproject_frustum_corners(extrinsics, intrinsics, near) + if far is not None: + far_corners = unproject_frustum_corners(extrinsics, intrinsics, far) + + # Project the cameras onto each axis-aligned plane. + projections = [] + for projected_axis in range(3): + image = torch.zeros( + (3, resolution, resolution), + dtype=torch.float32, + device=device, + ) + image_x_axis = (projected_axis + 1) % 3 + image_y_axis = (projected_axis + 2) % 3 + + def project(points: Float[Tensor, "*batch 3"]) -> Float[Tensor, "*batch 2"]: + x = points[..., image_x_axis] + y = points[..., image_y_axis] + return torch.stack([x, y], dim=-1) + + x_range, y_range = torch.stack( + (project(scene_minima), project(scene_maxima)), dim=-1 + ) + + # Draw near and far planes. + if near is not None: + projected_near_corners = project(near_corners) + image = draw_lines( + image, + rearrange(projected_near_corners, "b p xy -> (b p) xy"), + rearrange(projected_near_corners.roll(1, 1), "b p xy -> (b p) xy"), + color=0.25, + width=2, + x_range=x_range, + y_range=y_range, + ) + if far is not None: + projected_far_corners = project(far_corners) + image = draw_lines( + image, + rearrange(projected_far_corners, "b p xy -> (b p) xy"), + rearrange(projected_far_corners.roll(1, 1), "b p xy -> (b p) xy"), + color=0.25, + width=2, + x_range=x_range, + y_range=y_range, + ) + if near is not None and far is not None: + image = draw_lines( + image, + rearrange(projected_near_corners, "b p xy -> (b p) xy"), + rearrange(projected_far_corners, "b p xy -> (b p) xy"), + color=0.25, + width=2, + x_range=x_range, + y_range=y_range, + ) + + # Draw the camera frustums themselves. + projected_origins = project(extrinsics[:, :3, 3]) + projected_frustum_corners = project(frustum_corners) + start = [ + repeat(projected_origins, "b xy -> (b p) xy", p=4), + rearrange(projected_frustum_corners.roll(1, 1), "b p xy -> (b p) xy"), + ] + start = rearrange(torch.cat(start, dim=0), "(r b p) xy -> (b r p) xy", r=2, p=4) + image = draw_lines( + image, + start, + repeat(projected_frustum_corners, "b p xy -> (b r p) xy", r=2), + color=repeat(color, "b c -> (b r p) c", r=2, p=4), + width=2, + x_range=x_range, + y_range=y_range, + ) + + x_name = "XYZ"[image_x_axis] + y_name = "XYZ"[image_y_axis] + image = add_label(image, f"{x_name}{y_name} Projection") + + # TODO: Draw axis indicators. + projections.append(image) + + return torch.stack(projections) + + +def compute_aabb( + extrinsics: Float[Tensor, "batch 4 4"], + intrinsics: Float[Tensor, "batch 3 3"], + near: Optional[Scalar] = None, + far: Optional[Scalar] = None, +) -> tuple[ + Float[Tensor, "3"], # minima of the scene + Float[Tensor, "3"], # maxima of the scene +]: + """Compute an axis-aligned bounding box for the camera frustums.""" + + device = extrinsics.device + + # These points are included in the AABB. + points = [extrinsics[:, :3, 3]] + + if near is not None: + near = sanitize_scalar(near, device) + corners = unproject_frustum_corners(extrinsics, intrinsics, near) + points.append(rearrange(corners, "b p xyz -> (b p) xyz")) + + if far is not None: + far = sanitize_scalar(far, device) + corners = unproject_frustum_corners(extrinsics, intrinsics, far) + points.append(rearrange(corners, "b p xyz -> (b p) xyz")) + + points = torch.cat(points, dim=0) + return points.min(dim=0).values, points.max(dim=0).values + + +def compute_equal_aabb_with_margin( + minima: Float[Tensor, "*#batch 3"], + maxima: Float[Tensor, "*#batch 3"], + margin: float = 0.1, +) -> tuple[ + Float[Tensor, "*batch 3"], # minima of the scene + Float[Tensor, "*batch 3"], # maxima of the scene +]: + midpoint = (maxima + minima) * 0.5 + span = (maxima - minima).max() * (1 + margin) + scene_minima = midpoint - 0.5 * span + scene_maxima = midpoint + 0.5 * span + return scene_minima, scene_maxima + + +def unproject_frustum_corners( + extrinsics: Float[Tensor, "batch 4 4"], + intrinsics: Float[Tensor, "batch 3 3"], + depth: Float[Tensor, "#batch"], +) -> Float[Tensor, "batch 4 3"]: + device = extrinsics.device + + # Get coordinates for the corners. Following them in a circle makes a rectangle. + xy = torch.linspace(0, 1, 2, device=device) + xy = torch.stack(torch.meshgrid(xy, xy, indexing="xy"), dim=-1) + xy = rearrange(xy, "i j xy -> (i j) xy") + xy = xy[torch.tensor([0, 1, 3, 2], device=device)] + + # Get ray directions in camera space. + directions = unproject( + xy, + torch.ones(1, dtype=torch.float32, device=device), + rearrange(intrinsics, "b i j -> b () i j"), + ) + + # Divide by the z coordinate so that multiplying by depth will produce orthographic + # depth (z depth) as opposed to Euclidean depth (distance from the camera). + directions = directions / directions[..., -1:] + directions = einsum(extrinsics[..., :3, :3], directions, "b i j, b r j -> b r i") + + origins = rearrange(extrinsics[:, :3, 3], "b xyz -> b () xyz") + depth = rearrange(depth, "b -> b () ()") + return origins + depth * directions diff --git a/src/visualization/drawing/coordinate_conversion.py b/src/visualization/drawing/coordinate_conversion.py new file mode 100644 index 0000000000000000000000000000000000000000..42ac5aca7cf903b61e6e9f9018c8aa1811300fa3 --- /dev/null +++ b/src/visualization/drawing/coordinate_conversion.py @@ -0,0 +1,44 @@ +from typing import Optional, Protocol, runtime_checkable + +import torch +from jaxtyping import Float +from torch import Tensor + +from .types import Pair, sanitize_pair + + +@runtime_checkable +class ConversionFunction(Protocol): + def __call__( + self, + xy: Float[Tensor, "*batch 2"], + ) -> Float[Tensor, "*batch 2"]: + pass + + +def generate_conversions( + shape: tuple[int, int], + device: torch.device, + x_range: Optional[Pair] = None, + y_range: Optional[Pair] = None, +) -> tuple[ + ConversionFunction, # conversion from world coordinates to pixel coordinates + ConversionFunction, # conversion from pixel coordinates to world coordinates +]: + h, w = shape + x_range = sanitize_pair((0, w) if x_range is None else x_range, device) + y_range = sanitize_pair((0, h) if y_range is None else y_range, device) + minima, maxima = torch.stack((x_range, y_range), dim=-1) + wh = torch.tensor((w, h), dtype=torch.float32, device=device) + + def convert_world_to_pixel( + xy: Float[Tensor, "*batch 2"], + ) -> Float[Tensor, "*batch 2"]: + return (xy - minima) / (maxima - minima) * wh + + def convert_pixel_to_world( + xy: Float[Tensor, "*batch 2"], + ) -> Float[Tensor, "*batch 2"]: + return xy / wh * (maxima - minima) + minima + + return convert_world_to_pixel, convert_pixel_to_world diff --git a/src/visualization/drawing/lines.py b/src/visualization/drawing/lines.py new file mode 100644 index 0000000000000000000000000000000000000000..85ce39825f456f5511049b2b1237836ded600416 --- /dev/null +++ b/src/visualization/drawing/lines.py @@ -0,0 +1,83 @@ +from typing import Literal, Optional + +import torch +from einops import einsum, repeat +from jaxtyping import Float +from torch import Tensor + +from .coordinate_conversion import generate_conversions +from .rendering import render_over_image +from .types import Pair, Scalar, Vector, sanitize_scalar, sanitize_vector + + +def draw_lines( + image: Float[Tensor, "3 height width"], + start: Vector, + end: Vector, + color: Vector, + width: Scalar, + cap: Literal["butt", "round", "square"] = "round", + num_msaa_passes: int = 1, + x_range: Optional[Pair] = None, + y_range: Optional[Pair] = None, +) -> Float[Tensor, "3 height width"]: + device = image.device + start = sanitize_vector(start, 2, device) + end = sanitize_vector(end, 2, device) + color = sanitize_vector(color, 3, device) + width = sanitize_scalar(width, device) + (num_lines,) = torch.broadcast_shapes( + start.shape[0], + end.shape[0], + color.shape[0], + width.shape, + ) + + # Convert world-space points to pixel space. + _, h, w = image.shape + world_to_pixel, _ = generate_conversions((h, w), device, x_range, y_range) + start = world_to_pixel(start) + end = world_to_pixel(end) + + def color_function( + xy: Float[Tensor, "point 2"], + ) -> Float[Tensor, "point 4"]: + # Define a vector between the start and end points. + delta = end - start + delta_norm = delta.norm(dim=-1, keepdim=True) + u_delta = delta / delta_norm + + # Define a vector between each sample and the start point. + indicator = xy - start[:, None] + + # Determine whether each sample is inside the line in the parallel direction. + extra = 0.5 * width[:, None] if cap == "square" else 0 + parallel = einsum(u_delta, indicator, "l xy, l s xy -> l s") + parallel_inside_line = (parallel <= delta_norm + extra) & (parallel > -extra) + + # Determine whether each sample is inside the line perpendicularly. + perpendicular = indicator - parallel[..., None] * u_delta[:, None] + perpendicular_inside_line = perpendicular.norm(dim=-1) < 0.5 * width[:, None] + + inside_line = parallel_inside_line & perpendicular_inside_line + + # Compute round caps. + if cap == "round": + near_start = indicator.norm(dim=-1) < 0.5 * width[:, None] + inside_line |= near_start + end_indicator = indicator = xy - end[:, None] + near_end = end_indicator.norm(dim=-1) < 0.5 * width[:, None] + inside_line |= near_end + + # Determine the sample's color. + selectable_color = color.broadcast_to((num_lines, 3)) + arrangement = inside_line * torch.arange(num_lines, device=device)[:, None] + top_color = selectable_color.gather( + dim=0, + index=repeat(arrangement.argmax(dim=0), "s -> s c", c=3), + ) + rgba = torch.cat((top_color, inside_line.any(dim=0).float()[:, None]), dim=-1) + + return rgba + + return render_over_image(image, color_function, device, num_passes=num_msaa_passes) diff --git a/src/visualization/drawing/points.py b/src/visualization/drawing/points.py new file mode 100644 index 0000000000000000000000000000000000000000..671db100d34cd9121cb2778dcdb7252ec915bb2d --- /dev/null +++ b/src/visualization/drawing/points.py @@ -0,0 +1,59 @@ +from typing import Optional + +import torch +from einops import repeat +from jaxtyping import Float +from torch import Tensor + +from .coordinate_conversion import generate_conversions +from .rendering import render_over_image +from .types import Pair, Scalar, Vector, sanitize_scalar, sanitize_vector + + +def draw_points( + image: Float[Tensor, "3 height width"], + points: Vector, + color: Vector = [1, 1, 1], + radius: Scalar = 1, + inner_radius: Scalar = 0, + num_msaa_passes: int = 1, + x_range: Optional[Pair] = None, + y_range: Optional[Pair] = None, +) -> Float[Tensor, "3 height width"]: + device = image.device + points = sanitize_vector(points, 2, device) + color = sanitize_vector(color, 3, device) + radius = sanitize_scalar(radius, device) + inner_radius = sanitize_scalar(inner_radius, device) + (num_points,) = torch.broadcast_shapes( + points.shape[0], + color.shape[0], + radius.shape, + inner_radius.shape, + ) + + # Convert world-space points to pixel space. + _, h, w = image.shape + world_to_pixel, _ = generate_conversions((h, w), device, x_range, y_range) + points = world_to_pixel(points) + + def color_function( + xy: Float[Tensor, "point 2"], + ) -> Float[Tensor, "point 4"]: + # Define a vector between the start and end points. + delta = xy[:, None] - points[None] + delta_norm = delta.norm(dim=-1) + mask = (delta_norm >= inner_radius[None]) & (delta_norm <= radius[None]) + + # Determine the sample's color. + selectable_color = color.broadcast_to((num_points, 3)) + arrangement = mask * torch.arange(num_points, device=device) + top_color = selectable_color.gather( + dim=0, + index=repeat(arrangement.argmax(dim=1), "s -> s c", c=3), + ) + rgba = torch.cat((top_color, mask.any(dim=1).float()[:, None]), dim=-1) + + return rgba + + return render_over_image(image, color_function, device, num_passes=num_msaa_passes) diff --git a/src/visualization/drawing/rendering.py b/src/visualization/drawing/rendering.py new file mode 100644 index 0000000000000000000000000000000000000000..65842a8ec5b9c6b109d19ef9509bc04ca0437ea7 --- /dev/null +++ b/src/visualization/drawing/rendering.py @@ -0,0 +1,152 @@ +from typing import Protocol, runtime_checkable + +import torch +from einops import rearrange, reduce +from jaxtyping import Bool, Float +from torch import Tensor + + +@runtime_checkable +class ColorFunction(Protocol): + def __call__( + self, + xy: Float[Tensor, "point 2"], + ) -> Float[Tensor, "point 4"]: # RGBA color + pass + + +def generate_sample_grid( + shape: tuple[int, int], + device: torch.device, +) -> Float[Tensor, "height width 2"]: + h, w = shape + x = torch.arange(w, device=device) + 0.5 + y = torch.arange(h, device=device) + 0.5 + x, y = torch.meshgrid(x, y, indexing="xy") + return torch.stack([x, y], dim=-1) + + +def detect_msaa_pixels( + image: Float[Tensor, "batch 4 height width"], +) -> Bool[Tensor, "batch height width"]: + b, _, h, w = image.shape + + mask = torch.zeros((b, h, w), dtype=torch.bool, device=image.device) + + # Detect horizontal differences. + horizontal = (image[:, :, :, 1:] != image[:, :, :, :-1]).any(dim=1) + mask[:, :, 1:] |= horizontal + mask[:, :, :-1] |= horizontal + + # Detect vertical differences. + vertical = (image[:, :, 1:, :] != image[:, :, :-1, :]).any(dim=1) + mask[:, 1:, :] |= vertical + mask[:, :-1, :] |= vertical + + # Detect diagonal (top left to bottom right) differences. + tlbr = (image[:, :, 1:, 1:] != image[:, :, :-1, :-1]).any(dim=1) + mask[:, 1:, 1:] |= tlbr + mask[:, :-1, :-1] |= tlbr + + # Detect diagonal (top right to bottom left) differences. + trbl = (image[:, :, :-1, 1:] != image[:, :, 1:, :-1]).any(dim=1) + mask[:, :-1, 1:] |= trbl + mask[:, 1:, :-1] |= trbl + + return mask + + +def reduce_straight_alpha( + rgba: Float[Tensor, "batch 4 height width"], +) -> Float[Tensor, "batch 4"]: + color, alpha = rgba.split((3, 1), dim=1) + + # Color becomes a weighted average of color (weighted by alpha). + weighted_color = reduce(color * alpha, "b c h w -> b c", "sum") + alpha_sum = reduce(alpha, "b c h w -> b c", "sum") + color = weighted_color / (alpha_sum + 1e-10) + + # Alpha becomes mean alpha. + alpha = reduce(alpha, "b c h w -> b c", "mean") + + return torch.cat((color, alpha), dim=-1) + + +@torch.no_grad() +def run_msaa_pass( + xy: Float[Tensor, "batch height width 2"], + color_function: ColorFunction, + scale: float, + subdivision: int, + remaining_passes: int, + device: torch.device, + batch_size: int = int(2**16), +) -> Float[Tensor, "batch 4 height width"]: # color (RGBA with straight alpha) + # Sample the color function. + b, h, w, _ = xy.shape + color = [ + color_function(batch) + for batch in rearrange(xy, "b h w xy -> (b h w) xy").split(batch_size) + ] + color = torch.cat(color, dim=0) + color = rearrange(color, "(b h w) c -> b c h w", b=b, h=h, w=w) + + # If any MSAA passes remain, subdivide. + if remaining_passes > 0: + mask = detect_msaa_pixels(color) + batch_index, row_index, col_index = torch.where(mask) + xy = xy[batch_index, row_index, col_index] + + offsets = generate_sample_grid((subdivision, subdivision), device) + offsets = (offsets / subdivision - 0.5) * scale + + color_fine = run_msaa_pass( + xy[:, None, None] + offsets, + color_function, + scale / subdivision, + subdivision, + remaining_passes - 1, + device, + batch_size=batch_size, + ) + color[batch_index, :, row_index, col_index] = reduce_straight_alpha(color_fine) + + return color + + +@torch.no_grad() +def render( + shape: tuple[int, int], + color_function: ColorFunction, + device: torch.device, + subdivision: int = 8, + num_passes: int = 2, +) -> Float[Tensor, "4 height width"]: # color (RGBA with straight alpha) + xy = generate_sample_grid(shape, device) + return run_msaa_pass( + xy[None], + color_function, + 1.0, + subdivision, + num_passes, + device, + )[0] + + +def render_over_image( + image: Float[Tensor, "3 height width"], + color_function: ColorFunction, + device: torch.device, + subdivision: int = 8, + num_passes: int = 1, +) -> Float[Tensor, "3 height width"]: + _, h, w = image.shape + overlay = render( + (h, w), + color_function, + device, + subdivision=subdivision, + num_passes=num_passes, + ) + color, alpha = overlay.split((3, 1), dim=0) + return image * (1 - alpha) + color * alpha diff --git a/src/visualization/drawing/types.py b/src/visualization/drawing/types.py new file mode 100644 index 0000000000000000000000000000000000000000..81a1899301d3b8fdbbcd266b8f543771d6a5066f --- /dev/null +++ b/src/visualization/drawing/types.py @@ -0,0 +1,69 @@ +from typing import Iterable, Union + +import torch +from einops import repeat +from jaxtyping import Float, Shaped +from torch import Tensor + +Real = Union[float, int] + +Vector = Union[ + Real, + Iterable[Real], + Shaped[Tensor, "2"], + Shaped[Tensor, "3"], + Shaped[Tensor, "batch 2"], + Shaped[Tensor, "batch 3"], +] + + +def sanitize_vector( + vector: Vector, + dim: int, + device: torch.device, +) -> Float[Tensor, "*#batch dim"]: + if isinstance(vector, Tensor): + vector = vector.type(torch.float32).to(device) + else: + vector = torch.tensor(vector, dtype=torch.float32, device=device) + while vector.ndim < 2: + vector = vector[None] + if vector.shape[-1] == 1: + vector = repeat(vector, "... () -> ... c", c=dim) + assert vector.shape[-1] == dim + assert vector.ndim == 2 + return vector + + +Scalar = Union[ + Real, + Iterable[Real], + Shaped[Tensor, ""], + Shaped[Tensor, " batch"], +] + + +def sanitize_scalar(scalar: Scalar, device: torch.device) -> Float[Tensor, "*#batch"]: + if isinstance(scalar, Tensor): + scalar = scalar.type(torch.float32).to(device) + else: + scalar = torch.tensor(scalar, dtype=torch.float32, device=device) + while scalar.ndim < 1: + scalar = scalar[None] + assert scalar.ndim == 1 + return scalar + + +Pair = Union[ + Iterable[Real], + Shaped[Tensor, "2"], +] + + +def sanitize_pair(pair: Pair, device: torch.device) -> Float[Tensor, "2"]: + if isinstance(pair, Tensor): + pair = pair.type(torch.float32).to(device) + else: + pair = torch.tensor(pair, dtype=torch.float32, device=device) + assert pair.shape == (2,) + return pair diff --git a/src/visualization/layout.py b/src/visualization/layout.py new file mode 100644 index 0000000000000000000000000000000000000000..7c2d461b3bcf2667a939e14d99aeb44d59da4e50 --- /dev/null +++ b/src/visualization/layout.py @@ -0,0 +1,228 @@ +"""This file contains useful layout utilities for images. They are: + +- add_border: Add a border to an image. +- cat/hcat/vcat: Join images by arranging them in a line. If the images have different + sizes, they are aligned as specified (start, end, center). Allows you to specify a gap + between images. + +Images are assumed to be float32 tensors with shape (channel, height, width). +""" + +from typing import Any, Generator, Iterable, Literal, Optional, Union + +import torch +import torch.nn.functional as F +from jaxtyping import Float +from torch import Tensor + +Alignment = Literal["start", "center", "end"] +Axis = Literal["horizontal", "vertical"] +Color = Union[ + int, + float, + Iterable[int], + Iterable[float], + Float[Tensor, "#channel"], + Float[Tensor, ""], +] + + +def _sanitize_color(color: Color) -> Float[Tensor, "#channel"]: + # Convert tensor to list (or individual item). + if isinstance(color, torch.Tensor): + color = color.tolist() + + # Turn iterators and individual items into lists. + if isinstance(color, Iterable): + color = list(color) + else: + color = [color] + + return torch.tensor(color, dtype=torch.float32) + + +def _intersperse(iterable: Iterable, delimiter: Any) -> Generator[Any, None, None]: + it = iter(iterable) + yield next(it) + for item in it: + yield delimiter + yield item + + +def _get_main_dim(main_axis: Axis) -> int: + return { + "horizontal": 2, + "vertical": 1, + }[main_axis] + + +def _get_cross_dim(main_axis: Axis) -> int: + return { + "horizontal": 1, + "vertical": 2, + }[main_axis] + + +def _compute_offset(base: int, overlay: int, align: Alignment) -> slice: + assert base >= overlay + offset = { + "start": 0, + "center": (base - overlay) // 2, + "end": base - overlay, + }[align] + return slice(offset, offset + overlay) + + +def overlay( + base: Float[Tensor, "channel base_height base_width"], + overlay: Float[Tensor, "channel overlay_height overlay_width"], + main_axis: Axis, + main_axis_alignment: Alignment, + cross_axis_alignment: Alignment, +) -> Float[Tensor, "channel base_height base_width"]: + # The overlay must be smaller than the base. + _, base_height, base_width = base.shape + _, overlay_height, overlay_width = overlay.shape + assert base_height >= overlay_height and base_width >= overlay_width + + # Compute spacing on the main dimension. + main_dim = _get_main_dim(main_axis) + main_slice = _compute_offset( + base.shape[main_dim], overlay.shape[main_dim], main_axis_alignment + ) + + # Compute spacing on the cross dimension. + cross_dim = _get_cross_dim(main_axis) + cross_slice = _compute_offset( + base.shape[cross_dim], overlay.shape[cross_dim], cross_axis_alignment + ) + + # Combine the slices and paste the overlay onto the base accordingly. + selector = [..., None, None] + selector[main_dim] = main_slice + selector[cross_dim] = cross_slice + result = base.clone() + result[tuple(selector)] = overlay + return result + + +def cat( + main_axis: Axis, + *images: Float[Tensor, "channel _ _"], + align: Alignment = "center", + gap: int = 8, + gap_color: Color = 1, +) -> Float[Tensor, "channel height width"]: + """Arrange images in a line. The interface resembles a CSS div with flexbox.""" + device = images[0].device + gap_color = _sanitize_color(gap_color).to(device) + + # Find the maximum image side length in the cross axis dimension. + cross_dim = _get_cross_dim(main_axis) + cross_axis_length = max(image.shape[cross_dim] for image in images) + + # Pad the images. + padded_images = [] + for image in images: + # Create an empty image with the correct size. + padded_shape = list(image.shape) + padded_shape[cross_dim] = cross_axis_length + base = torch.ones(padded_shape, dtype=torch.float32, device=device) + base = base * gap_color[:, None, None] + padded_images.append(overlay(base, image, main_axis, "start", align)) + + # Intersperse separators if necessary. + if gap > 0: + # Generate a separator. + c, _, _ = images[0].shape + separator_size = [gap, gap] + separator_size[cross_dim - 1] = cross_axis_length + separator = torch.ones((c, *separator_size), dtype=torch.float32, device=device) + separator = separator * gap_color[:, None, None] + + # Intersperse the separator between the images. + padded_images = list(_intersperse(padded_images, separator)) + + return torch.cat(padded_images, dim=_get_main_dim(main_axis)) + + +def hcat( + *images: Float[Tensor, "channel _ _"], + align: Literal["start", "center", "end", "top", "bottom"] = "start", + gap: int = 8, + gap_color: Color = 1, +): + """Shorthand for a horizontal linear concatenation.""" + return cat( + "horizontal", + *images, + align={ + "start": "start", + "center": "center", + "end": "end", + "top": "start", + "bottom": "end", + }[align], + gap=gap, + gap_color=gap_color, + ) + + +def vcat( + *images: Float[Tensor, "channel _ _"], + align: Literal["start", "center", "end", "left", "right"] = "start", + gap: int = 8, + gap_color: Color = 1, +): + """Shorthand for a horizontal linear concatenation.""" + return cat( + "vertical", + *images, + align={ + "start": "start", + "center": "center", + "end": "end", + "left": "start", + "right": "end", + }[align], + gap=gap, + gap_color=gap_color, + ) + + +def add_border( + image: Float[Tensor, "channel height width"], + border: int = 8, + color: Color = 1, +) -> Float[Tensor, "channel new_height new_width"]: + color = _sanitize_color(color).to(image) + c, h, w = image.shape + result = torch.empty( + (c, h + 2 * border, w + 2 * border), dtype=torch.float32, device=image.device + ) + result[:] = color[:, None, None] + result[:, border : h + border, border : w + border] = image + return result + + +def resize( + image: Float[Tensor, "channel height width"], + shape: Optional[tuple[int, int]] = None, + width: Optional[int] = None, + height: Optional[int] = None, +) -> Float[Tensor, "channel new_height new_width"]: + assert (shape is not None) + (width is not None) + (height is not None) == 1 + _, h, w = image.shape + + if width is not None: + shape = (int(h * width / w), width) + elif height is not None: + shape = (height, int(w * height / h)) + + return F.interpolate( + image[None], + shape, + mode="bilinear", + align_corners=False, + antialias="bilinear", + )[0] diff --git a/src/visualization/vis_depth.py b/src/visualization/vis_depth.py new file mode 100644 index 0000000000000000000000000000000000000000..45b28853cf3dce7f7c54a5f8dc38eec9fdc80b4c --- /dev/null +++ b/src/visualization/vis_depth.py @@ -0,0 +1,39 @@ +import torch +import torch.utils.data +import numpy as np +import torchvision.utils as vutils +import cv2 +from matplotlib.cm import get_cmap +import matplotlib as mpl +import matplotlib.cm as cm + + +# https://github.com/autonomousvision/unimatch/blob/master/utils/visualization.py + + +def vis_disparity(disp): + disp_vis = (disp - disp.min()) / (disp.max() - disp.min()) * 255.0 + disp_vis = disp_vis.astype("uint8") + disp_vis = cv2.applyColorMap(disp_vis, cv2.COLORMAP_INFERNO) + + return disp_vis + + +def viz_depth_tensor(disp, return_numpy=False, colormap="plasma"): + # visualize inverse depth + assert isinstance(disp, torch.Tensor) + + disp = disp.numpy() + vmax = np.percentile(disp, 95) + normalizer = mpl.colors.Normalize(vmin=disp.min(), vmax=vmax) + mapper = cm.ScalarMappable(norm=normalizer, cmap=colormap) + colormapped_im = (mapper.to_rgba(disp)[:, :, :3] * 255).astype( + np.uint8 + ) # [H, W, 3] + + if return_numpy: + return colormapped_im + + viz = torch.from_numpy(colormapped_im).permute(2, 0, 1) # [3, H, W] + + return viz