diff --git a/Dockerfile b/Dockerfile
index fc9d5710b3a7d19a43b0f8748c31d596b6a66e11..9c783f7801edd7423a91fb11eb1b3295dd3f7ee1 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,7 +1,7 @@
# Learn2Splat — interactive demo for a Hugging Face Space (Docker SDK, GPU).
#
# Installs the optgs package + prebuilt CUDA-extension wheels, then runs
-# demo.py's viser GUI: SfM-initialize a COLMAP scene and refine the Gaussians
+# demo.py's gradio GUI: SfM-initialize a COLMAP scene and refine the Gaussians
# with the learned optimizer live in the browser.
#
# The CUDA extensions are NOT compiled here — the HF Docker builder runs out of
@@ -53,15 +53,6 @@ RUN pip install torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 \
COPY --chown=user:user requirements.txt .
RUN pip install -r requirements.txt
-# HF serves the Space over HTTP/2, and the WebSocket subprotocol viser uses to
-# announce its client version doesn't survive the proxy — so viser's server
-# reads the client version as "unknown" and rejects the connection (the GUI
-# then hangs on "connecting"). Client and server are the same viser build
-# here, so treat an undeterminable client version as a match, not a reject.
-RUN VISER_INFRA="$(python -c 'import viser.infra._infra as m; print(m.__file__)')" \
- && sed -i 's/client_version_str = "unknown"/client_version_str = viser.__version__/' "$VISER_INFRA" \
- && grep -q 'client_version_str = viser.__version__' "$VISER_INFRA"
-
# Prebuilt CUDA-extension wheels — gsplat, nerfacc, pycolmap, fused-ssim,
# simple-knn, pointops, fused_knn_attn. Built on a matching machine (see
# DEPLOY.md) so the HF builder never compiles CUDA and never OOMs.
@@ -72,11 +63,10 @@ RUN pip install --no-deps ./wheels/*.whl
COPY --chown=user:user . .
RUN pip install --no-build-isolation --no-deps -e .
-# viser serves the GUI here — must equal app_port in README.md.
+# gradio serves the GUI here — must equal app_port in README.md.
EXPOSE 7860
-# server mode: the optgs decoder renders frames on the GPU and viser streams
-# them as images — no multi-MB splat-geometry transfer to the browser and no
-# client-side WebGL, which is far more robust behind HF's HTTP/2 proxy.
-# viser binds 0.0.0.0 by default.
-CMD ["python", "demo.py", "--with-gui", "server", "--gui-port", "7860"]
+# gradio mode: the optgs decoder renders the live optimization on the GPU and
+# gradio streams the frames as images (column 2); the finished splats load into
+# an interactive Model3D viewer (column 3). demo.py binds 0.0.0.0 on this port.
+CMD ["python", "demo.py", "--with-gui", "gradio", "--gui-port", "7860"]
diff --git a/README.md b/README.md
index 33ec5cfd686e01690bf8a78a40a1d3187ec90ada..2a44aa42e0ec6f4a86cae849336731941f60e894 100644
--- a/README.md
+++ b/README.md
@@ -14,11 +14,13 @@ short_description: Interactive demo of the Learn2Splat learned 3DGS optimizer
A learned optimizer for 3D Gaussian Splatting. This Space SfM-initializes a
COLMAP scene and refines the Gaussians live in your browser: pick the
Learn2Splat optimizer (dense or sparse checkpoint) or a 3DGS Adam baseline,
-press **Start**, and watch the splats converge.
+press **Start**, and watch the decoder render converge — then explore the
+finished splats in an interactive 3D viewer.
-Runs `demo.py --with-gui client` from the
+Runs `demo.py --with-gui gradio` from the
[Learn2Splat repository](https://github.com/autonomousvision/learn2splat);
-the splats are drawn by viser's in-browser WebGL renderer.
+the optimization is rendered on the GPU and streamed by gradio, and the
+result loads into a `Model3D` splat viewer.
> Requires GPU hardware. The demo holds two checkpoints in VRAM at once —
> an A10G (24 GB) is recommended.
diff --git a/demo.py b/demo.py
index 13325f206b2417b85d9211e37cd405ef80e908fb..50fb83e43d7491c50040cec69119275f3fdc5b0e 100644
--- a/demo.py
+++ b/demo.py
@@ -27,6 +27,7 @@ Usage (run from the repo root, with ``optgs`` importable):
python demo.py # headless: dense + sparse checkpoints + an Adam baseline
python demo.py --with-gui server # interactive viser GUI (frames rendered by the decoder)
python demo.py --with-gui client # interactive viser GUI (viser's WebGL splat renderer)
+ python demo.py --with-gui gradio # interactive gradio GUI (streamed renders + Model3D splats)
The demo scene and the checkpoints are fetched from the Hugging Face Hub on
first run (cached under ./data and ./checkpoints). A CUDA device is required.
@@ -118,11 +119,13 @@ class Config:
seed: int = 42
# --- Interactive GUI ---
- # Launch a viser GUI instead of the headless comparison. "server" renders
- # frames with the optgs decoder; "client" uses viser's built-in WebGL
- # Gaussian-splat renderer. Unset = headless run.
- with_gui: Optional[Literal["client", "server"]] = None
- # Port for the viser GUI web server (--with-gui only).
+ # Launch an interactive GUI instead of the headless comparison. viser:
+ # "server" renders frames with the optgs decoder, "client" uses viser's
+ # built-in WebGL Gaussian-splat renderer. "gradio" runs a browser GUI
+ # (decoder renders streamed live + an interactive Model3D splat viewer for
+ # the result). Unset = headless run.
+ with_gui: Optional[Literal["client", "server", "gradio"]] = None
+ # Port for the GUI web server (--with-gui only).
gui_port: int = 8080
# --- OptGS learned optimizer ---
@@ -600,6 +603,373 @@ def run_gui(
console.print("\n[yellow]GUI stopped.[/]")
+def run_gradio_gui(
+ instances: dict,
+ gaussians: Gaussians,
+ train_bv: dict,
+ cfg: Config,
+ device: torch.device,
+) -> None:
+ """Interactive gradio GUI — a browser port of :func:`run_gui` (viser).
+
+ gradio can't stream the camera back to Python, so there is no free-camera
+ server rendering. Instead the optimization is *watched* as a streamed
+ decoder render from a chosen training view (``gr.Image``, refreshed every
+ step), and the finished scene is handed to an interactive ``gr.Model3D``
+ splat viewer (orbit / zoom in the browser). The controls mirror the viser
+ GUI: pick the optimizer (Learn2Splat dense/sparse or a 3DGS Adam baseline),
+ set the step budget / view-minibatch size / sampling strategy, Start, Reset.
+
+ ``instances`` maps "dense"/"sparse" to their initialized ``OptGS``.
+ """
+ import gc
+
+ import gradio as gr
+
+ from optgs.experimental.api.integration.config_bridge import build_adam_baseline
+ from optgs.misc.image_io import prep_image
+ from optgs.model.ply_export import save_gaussian_ply
+
+ # Optimizer dropdown label -> (instances key, swap in a 3DGS Adam baseline);
+ # mirrors run_gui's OPTIONS.
+ OPTIONS: Dict[str, Tuple[str, bool]] = {
+ "Learn2Splat (dense)": ("dense", False),
+ "Learn2Splat (sparse)": ("sparse", False),
+ "Adam (3DGS)": ("dense", True),
+ }
+
+ n_train_views = int(train_bv["image"].shape[1])
+ h_full, w_full = train_bv["image"].shape[3], train_bv["image"].shape[4]
+ init_gaussians = gaussians.clone() # pristine copy; each Start re-inits from it
+
+ # Shared state (single-GPU, single-session demo, like run_gui's globals): the
+ # Gaussians currently shown, the OptGS rendering them, and a counter for
+ # unique PLY filenames (so Model3D reloads instead of serving a stale cache).
+ holder = {"current": init_gaussians, "active": instances["dense"], "ply": 0}
+
+ ply_dir = os.path.join(cfg.result_dir, "gradio")
+ os.makedirs(ply_dir, exist_ok=True)
+
+ @torch.no_grad()
+ def render_decoder(inst, g: Gaussians, view_idx: float, height: float) -> np.ndarray:
+ """Decoder-render Gaussians ``g`` from training view ``view_idx``.
+
+ Normalized intrinsics make the render resolution-independent; the width
+ is derived from ``height`` at the training views' aspect ratio.
+ """
+ h = int(height)
+ w = max(1, round(h * w_full / h_full))
+ sl = slice(int(view_idx), int(view_idx) + 1)
+ out = inst.decoder.forward(
+ g,
+ train_bv["extrinsics"][:, sl],
+ train_bv["intrinsics"][:, sl],
+ train_bv["near"][:, sl],
+ train_bv["far"][:, sl],
+ image_shape=(h, w),
+ )
+ return prep_image(out.color[0, 0]) # [H, W, 3] uint8
+
+ def reorient_for_viewer(g: Gaussians) -> Gaussians:
+ """Reorient a copy of ``g`` from the COLMAP world (this scene is Z-up)
+ into the Y-up frame the gradio Model3D shows upright.
+
+ gradio/Babylon flips the loaded splats' Y (``scaling.y *= -1``), so
+ exporting through the reflection E(p)=(x,-z,-y) makes the *displayed*
+ scene N(p)=(x,z,-y) — the world's Z-up mapped onto Babylon's Y-up. E's
+ point-reflection part leaves the covariance unchanged; its proper-rotation
+ part R_p=-E rotates the splat orientations to match.
+ """
+ from scipy.spatial.transform import Rotation as Rsp
+
+ g = g.clone()
+ m = g.means[0]
+ g.means = torch.stack([m[:, 0], -m[:, 2], -m[:, 1]], dim=1)[None] # E
+ q = F.normalize(g.rotations_unnorm[0], dim=-1).detach().cpu().numpy() # xyzw
+ R_p = np.array([[-1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 1.0, 0.0]])
+ q = (Rsp.from_matrix(R_p) * Rsp.from_quat(q)).as_quat() # R_p @ R, xyzw
+ g.rotations_unnorm = torch.from_numpy(q).to(g.means)[None]
+ return g
+
+ def export_ply(g: Gaussians) -> str:
+ """Write ``g`` (reoriented to Y-up) to a fresh PLY path for the viewer."""
+ from pathlib import Path
+
+ holder["ply"] += 1
+ path = os.path.join(ply_dir, f"result_{holder['ply']}.ply")
+ save_gaussian_ply(reorient_for_viewer(g), save_path=Path(path))
+ return path
+
+ def start(optimizer_label, max_steps, batch_size, strategy, view_idx, height):
+ """Generator: run the picked optimizer, streaming a decoder render each
+ step, then load the finished splats into the Model3D viewer."""
+ name, use_adam = OPTIONS[optimizer_label]
+ inst = instances[name]
+ holder["active"] = inst
+ # Re-init from the pristine copy so repeated Starts share one start point.
+ inst.initialize_from_tensors(init_gaussians.clone(), train_bv)
+ inst.num_refine = int(max_steps)
+ inst.opt_batch_size = min(int(batch_size), n_train_views)
+ inst.opt_batch_strategy = strategy
+ opt = build_adam_baseline(inst.num_refine).to(device) if use_adam else None
+
+ try:
+ # Disable Start, hide the viewer, keep the placeholder while running.
+ yield (
+ render_decoder(inst, init_gaussians, view_idx, height),
+ f"**optimizing** — step 0/{inst.num_refine} — "
+ f"{init_gaussians.means.shape[1]} Gaussians",
+ gr.update(visible=False, value=None), # hide the viewer
+ gr.update(interactive=False), # disable Start
+ gr.update(visible=True), # keep the placeholder
+ )
+
+ g = init_gaussians
+ for step, g in inst.optimize_iter(optimizer=opt):
+ holder["current"] = g
+ yield (
+ render_decoder(inst, g, view_idx, height),
+ f"**optimizing** — step {step + 1}/{inst.num_refine} — "
+ f"{g.means.shape[1]} Gaussians",
+ gr.update(), gr.update(), gr.update(), # no change mid-run
+ )
+
+ yield (
+ render_decoder(inst, g, view_idx, height),
+ f"**done** — {inst.num_refine} steps — {g.means.shape[1]} Gaussians",
+ gr.update(visible=True, value=export_ply(g)), # reveal the viewer
+ gr.update(interactive=True), # re-enable Start
+ gr.update(visible=False), # hide the placeholder
+ )
+ finally:
+ # Free the run's CUDA work (also runs if the user hits Stop mid-run),
+ # so GPU memory doesn't accumulate across repeated runs.
+ opt = None
+ gc.collect()
+ torch.cuda.empty_cache()
+
+ def reset(view_idx, height):
+ """Restore the initialization: re-render it, hide the viewer.
+
+ No CUDA cleanup here — ``start``'s ``finally`` already frees each run's
+ work; reset only re-renders the init.
+ """
+ holder["current"] = init_gaussians
+ holder["active"] = instances["dense"]
+ return (
+ render_decoder(instances["dense"], init_gaussians, view_idx, height),
+ "**Initialized.** Pick a method, then Start.",
+ gr.update(visible=False, value=None), # hide the viewer
+ gr.update(interactive=True), # enable Start
+ gr.update(visible=True), # show the placeholder
+ )
+
+ def rerender(view_idx, height):
+ """Re-render the current Gaussians (preview view / height changed)."""
+ return render_decoder(holder["active"], holder["current"], view_idx, height)
+
+ initial_img = render_decoder(instances["dense"], init_gaussians, 0, 540)
+
+ # Open the interactive viewer on the same vantage as the live render's
+ # default preview (view 0). The viewer shows splats in the reoriented frame
+ # N(p)=(x,z,-y) (see reorient_for_viewer); the orbit camera sits at the
+ # training view's position and looks at the scene centroid. babylon_camera
+ # maps the world camera into N and inverts Babylon's ArcRotateCamera position
+ # formula rel=(r·cosα·sinβ, r·cosβ, r·sinα·sinβ) into (alpha°, beta°, radius).
+ def babylon_camera(c2w: np.ndarray, centroid: np.ndarray) -> tuple:
+ p = c2w[:3, 3] - centroid
+ rel = np.array([p[0], p[2], -p[1]], dtype=np.float64) # N applied to (cam - centroid)
+ radius = float(np.linalg.norm(rel)) or 1e-3
+ beta = float(np.degrees(np.arccos(np.clip(rel[1] / radius, -1.0, 1.0))))
+ alpha = float(np.degrees(np.arctan2(rel[2], rel[0])))
+ return (alpha, beta, radius)
+
+ means0 = init_gaussians.means[0].detach().float().cpu().numpy()
+ centroid0 = (means0.min(0) + means0.max(0)) / 2.0
+ cam0_c2w = train_bv["extrinsics"][0, 0].detach().float().cpu().numpy()
+ init_camera = babylon_camera(cam0_c2w, centroid0)
+
+ # --- Visual style: lifted from the Learn2Splat project page
+ # (https://naamapearl.github.io/learn2splat/) — plum accent (#B04080) with an
+ # indigo hover, a light slate canvas with white cards, Source Serif 4
+ # headings / Inter body / JetBrains Mono labels. ---
+ plum = gr.themes.Color(
+ c50="#faf0f6", c100="#f5e6ef", c200="#eccadd", c300="#dda3c2",
+ c400="#c66ba0", c500="#b04080", c600="#9b3570", c700="#7f2a5b",
+ c800="#6a2550", c900="#581f43", c950="#350e26", name="plum",
+ )
+ slate = gr.themes.Color(
+ c50="#f7f8fc", c100="#eef0f8", c200="#e2e5ef", c300="#c8cde0",
+ c400="#8890b0", c500="#5a6080", c600="#454b6b", c700="#343a56",
+ c800="#262b42", c900="#1a1d2e", c950="#0f1120", name="slate",
+ )
+ theme = gr.themes.Soft(
+ primary_hue=plum,
+ neutral_hue=slate,
+ font=["Inter", "system-ui", "sans-serif"],
+ font_mono=["JetBrains Mono", "ui-monospace", "monospace"],
+ ).set(
+ body_background_fill="#f7f8fc",
+ body_text_color="#1a1d2e",
+ block_background_fill="#ffffff",
+ block_border_color="#e2e5ef",
+ block_radius="12px",
+ block_label_text_color="#5a6080",
+ block_title_text_color="#1a1d2e",
+ border_color_primary="#e2e5ef",
+ input_background_fill="#ffffff",
+ button_primary_background_fill="#b04080",
+ button_primary_background_fill_hover="#3d50c0",
+ button_primary_text_color="#ffffff",
+ button_secondary_background_fill="#ffffff",
+ button_secondary_border_color="#c8cde0",
+ button_secondary_text_color="#1a1d2e",
+ slider_color="#b04080",
+ link_text_color="#b04080",
+ )
+ # Source Serif 4 / Inter / JetBrains Mono, loaded into like the page.
+ fonts_head = (
+ '
"
+ "Extending the Horizon of Learned 3DGS Optimization "
+ "SfM-initialize a COLMAP scene, then refine the Gaussians with the "
+ "meta-learned optimizer — pick a method, press Start , and watch the "
+ "decoder render converge. The finished splats load in the interactive 3D "
+ "viewer. Project page ↗
"
+ )
+
+ with gr.Blocks(
+ title="Learn2Splat — Demo", theme=theme, css=css, head=fonts_head,
+ analytics_enabled=False,
+ ) as ui:
+ gr.HTML(hero_html, elem_id="l2s-hero")
+ with gr.Row(equal_height=False):
+ # Column 1 — controls.
+ with gr.Column(scale=3, min_width=300):
+ with gr.Group():
+ gr.HTML("Optimizer
")
+ optimizer_dd = gr.Dropdown(
+ list(OPTIONS), value=next(iter(OPTIONS)), label="Method"
+ )
+ with gr.Row():
+ max_steps_input = gr.Number(
+ value=cfg.max_steps, minimum=1, maximum=1000, step=1,
+ precision=0, label="Max steps",
+ )
+ batch_size_input = gr.Number(
+ value=min(cfg.opt_batch_size, n_train_views),
+ minimum=1, maximum=n_train_views, step=1, precision=0,
+ label="Opt batch size",
+ )
+ strategy_dd = gr.Dropdown(
+ ["random", "sequential", "fps"],
+ value=cfg.opt_batch_strategy, label="Batch strategy",
+ )
+ with gr.Group():
+ gr.HTML("Preview
")
+ view_slider = gr.Slider(
+ 0, n_train_views - 1, value=0, step=1, label="Preview view"
+ )
+ height_slider = gr.Slider(
+ 240, 1080, value=540, step=60, label="Render height"
+ )
+ with gr.Row():
+ start_btn = gr.Button(
+ "Start optimization", variant="primary",
+ elem_id="l2s-start", scale=2,
+ )
+ reset_btn = gr.Button(
+ "Reset", variant="secondary", elem_id="l2s-reset", scale=1
+ )
+ status_md = gr.Markdown(
+ "**Initialized.** Pick a method, then Start.", elem_id="l2s-status"
+ )
+ # Column 2 — live decoder render (streamed during optimization).
+ with gr.Column(scale=5, min_width=380):
+ image_out = gr.Image(
+ value=initial_img, label="Optimizer · live",
+ height=540, format="jpeg", interactive=False,
+ )
+ # Column 3 — interactive splats (hidden until a run finishes; a
+ # placeholder holds the column so the 3-up layout stays balanced).
+ with gr.Column(scale=5, min_width=380):
+ model3d_out = gr.Model3D(
+ label="Refined splats · interactive", height=540,
+ visible=False, camera_position=init_camera,
+ )
+ placeholder = gr.HTML(
+ "The interactive 3D splats
"
+ )
+
+ start_inputs = [
+ optimizer_dd, max_steps_input, batch_size_input, strategy_dd,
+ view_slider, height_slider,
+ ]
+ gui_outputs = [image_out, status_md, model3d_out, start_btn, placeholder]
+ # One shared GPU lane (concurrency_id) so Start / Reset / preview re-renders
+ # never run on the GPU at the same time — overlapping runs were the path to
+ # runaway VRAM growth.
+ start_btn.click(
+ start, inputs=start_inputs, outputs=gui_outputs, concurrency_id="gpu"
+ )
+ reset_btn.click(
+ reset, inputs=[view_slider, height_slider], outputs=gui_outputs,
+ concurrency_id="gpu",
+ )
+ # Slider release (not change) — re-render once when the user lets go.
+ view_slider.release(
+ rerender, [view_slider, height_slider], image_out, concurrency_id="gpu"
+ )
+ height_slider.release(
+ rerender, [view_slider, height_slider], image_out, concurrency_id="gpu"
+ )
+
+ console.print(
+ f"[green]✓[/] gradio GUI on port [cyan]{cfg.gui_port}[/]"
+ f" — forward the port over SSH and open the printed URL"
+ )
+ ui.queue(default_concurrency_limit=1).launch(
+ server_name="0.0.0.0", server_port=cfg.gui_port, share=False,
+ show_error=True,
+ )
+
+
def main(cfg: Config) -> None:
# Fetch the demo scene on first run, before anything else touches it.
ensure_data(cfg.data_dir)
@@ -663,7 +1033,10 @@ def main(cfg: Config) -> None:
for inst in instances.values():
inst.initialize_from_tensors(gaussians, train_bv)
- run_gui(instances, gaussians, train_bv, cfg, device, dtype)
+ if cfg.with_gui == "gradio":
+ run_gradio_gui(instances, gaussians, train_bv, cfg, device)
+ else:
+ run_gui(instances, gaussians, train_bv, cfg, device, dtype)
return
val_c2w, val_Ks, val_images = collect_cameras(dataset, val_idx)
diff --git a/optgs/config.py b/optgs/config.py
index 9ded2f8a1407b5ab718005ae1b94a9bff605284d..ce955211bc401235512c858ee56d77c8d4b9fab9 100644
--- a/optgs/config.py
+++ b/optgs/config.py
@@ -1,29 +1,34 @@
-import importlib
+"""Typed configuration schema and Hydra-to-dataclass loading.
+
+`RootCfg` is the top-level config dataclass (dataset, scene_trainer, meta_trainer, loss, checkpointing,
+...); `load_typed_root_config` converts the raw Hydra `DictConfig` into it via dacite, after running
+`config_migrate` so older checkpoints' configs load. `__post_init__` also resolves the output directory
+for `mode=test`. The per-component schemas live next to their components (e.g.
+`scene_trainer_cfg.py`, `meta_trainer_cfg.py`).
+"""
+
from copy import deepcopy
from dataclasses import dataclass
from pathlib import Path
-from typing import Literal, Optional, Type, TypeVar, Any, Callable
+from typing import Literal, Optional, Type, TypeVar
-import hydra
import torch
from dacite import Config, from_dict, UnionMatchError
-from hydra.core.global_hydra import GlobalHydra
from hydra.core.hydra_config import HydraConfig
-from hydra.types import RunMode
from omegaconf import DictConfig
from omegaconf import OmegaConf
from pytorch_lightning.strategies import DDPStrategy, FSDPStrategy
from .config_migrate import migrate, CURRENT_CFG_VERSION
from .dataset.data_module import DataLoaderCfg, DatasetCfg
-from .global_cfg import set_cfg
from .loss import LossCfgWrapper
from .misc.io import CustomPath
from .misc.io import cyan, read_omega_cfg
from .misc.checkpointing import find_latest_ckpt
-from .misc.hf_ckpt import maybe_resolve_hf_ref
-from .paths import CKPT_DIR, RESULTS_DIR
-from .scene_trainer.scene_trainer_cfg import SceneTrainerCfg, MetaOptimizerCfg, TestCfg, TrainCfg
+from .misc.hf_ckpt import hf_sibling_config, is_hf_ref, maybe_resolve_hf_ref
+from .paths import CKPT_DIR, RESULTS_DIR, DEBUG
+from .scene_trainer.scene_trainer_cfg import SceneTrainerCfg
+from .meta_trainer.meta_trainer_cfg import MetaOptimizerCfg, TestCfg, TrainCfg
# In order to extract filename or dirname from a path in the config
@@ -58,15 +63,28 @@ class CheckpointingCfg:
resume_update_module: str | None
load_existing_cfg: bool
+ # Checkpoints whose sibling config.yaml carries the model architecture that
+ # mode=test rebuilds from (the optimizer/initializer the Gaussians were
+ # trained with).
+ _ARCH_CHECKPOINTS = ("pretrained_model", "pretrained_optimizer", "pretrained_initializer")
+
def __post_init__(self):
# Resolve any Hugging Face Hub references (hf://org/repo/file[@rev]) to
# local cached paths so all downstream torch.load calls work unchanged.
+ # Rebuilding the model architecture needs both the .ckpt and its sibling
+ # config.yaml, but hf_hub_download pulls only the .ckpt, so fetch the
+ # sibling config.yaml too. It lands in the same checkpoints/ layout as
+ # the resolved .ckpt, so _find_config_for_checkpoint later discovers it
+ # on the local path.
for attr in ("pretrained_model", "pretrained_optimizer", "pretrained_initializer",
"pretrained_monodepth", "pretrained_mvdepth", "pretrained_depth",
"pretrained_scale_predictor", "pretrained_depth_teacher",
"resume_update_module"):
- resolved = maybe_resolve_hf_ref(getattr(self, attr))
- if resolved != getattr(self, attr):
+ ref = getattr(self, attr)
+ if attr in self._ARCH_CHECKPOINTS and is_hf_ref(ref):
+ hf_sibling_config(ref)
+ resolved = maybe_resolve_hf_ref(ref)
+ if resolved != ref:
setattr(self, attr, resolved)
for attr in ("pretrained_model", "pretrained_optimizer", "pretrained_initializer"):
@@ -90,6 +108,7 @@ class MetaTrainerCfg:
eval_index: str | None
limit_test_batches: int | float
limit_train_batches: int | float
+ meta_optimizer: MetaOptimizerCfg
test: TestCfg
train: TrainCfg
@@ -109,9 +128,9 @@ class MetaTrainerCfg:
return has_trainable
dist_strategy = FSDPStrategy(auto_wrap_policy=only_wrap_trainable)
- if self.train.use_replay_buffer:
- # When resampling from the replay buffer,
- # we don't project the condition_features to state, so the update_proj is not used
+ if self.train.use_ckpt_buffer:
+ # When resampling from the ckpt buffer,
+ # we don't project the condition_features to state, so the state_proj is not used
dist_strategy = "ddp_find_unused_parameters_true"
return dist_strategy
@@ -123,14 +142,13 @@ class RootCfg:
dataset: DatasetCfg
data_loader: DataLoaderCfg
scene_trainer: SceneTrainerCfg
- meta_optimizer: MetaOptimizerCfg ## TODO Naama: should we move under meta trainer config?
checkpointing: CheckpointingCfg
meta_trainer: MetaTrainerCfg
loss: list[LossCfgWrapper]
seed: int
use_plugins: bool
output_dir: str
- version: int | None
+ version: int | float | None # point versions like 2.1 are floats
debug_cfg: bool
def __post_init__(self):
@@ -202,12 +220,6 @@ TYPE_HOOKS = {
T = TypeVar("T")
-def get_class_by_path(path: str):
- module_path, class_name = path.rsplit('.', 1)
- module = importlib.import_module(module_path)
- return getattr(module, class_name)
-
-
def _diagnose_union_error(e: UnionMatchError, data: dict, dacite_config: Config) -> str:
"""Try each union member individually and report per-member errors."""
import dataclasses
@@ -273,32 +285,6 @@ def separate_loss_cfg_wrappers(joined: dict) -> list[LossCfgWrapper]:
]
-def universal_target_hook(cfg: dict, _: Type) -> Any:
- """Generic hook to construct config objects from `__target__`."""
- if not isinstance(cfg, dict):
- return None
- if "__target__" not in cfg:
- return None # Let decite handle it
-
- cfg_copy = deepcopy(cfg) # avoid mutating original
- target = cfg_copy.pop("__target__")
-
- if isinstance(target, str):
- target_type = get_class_by_path(target)
- else:
- target_type = target
-
- # Use recursive loading with known additional hooks
- return load_typed_config(
- DictConfig(cfg_copy),
- target_type,
- )
-
-
-def make_target_hook_for_type(t: Type) -> Callable:
- return lambda cfg: universal_target_hook(cfg, t)
-
-
def load_typed_root_config(cfg: DictConfig) -> RootCfg:
# scene_trainer/scene_optimizer=none loads a full dict from none.yaml;
# dacite can't match that dict to the None arm of SceneOptimizerCfg | None.
@@ -316,30 +302,17 @@ def load_typed_root_config(cfg: DictConfig) -> RootCfg:
)
-def should_run(cfg_dict):
- if cfg_dict.mode == "test":
- if cfg_dict.meta_trainer.test.skip_if_outputs_exist:
- output_dir = cfg_dict.output_dir
- if not output_dir.exists():
- return True
- metrics_path_pattern = output_dir / "metrics" / "target_*_psnr.json"
- metric_paths = list(metrics_path_pattern.parent.glob(metrics_path_pattern.name))
- if len(metric_paths) > 0:
- print(cyan(f"Test metrics already exist at {metric_paths}."))
- return False
- return True
-
-
def setup_cfg(cfg_dict):
# Get the original config from the output directory, when testing or resuming.
cfg_dict = merge_config_from_file(cfg_dict)
eval_cfg = get_eval_cfg(cfg_dict)
cfg = load_typed_root_config(cfg_dict)
- # Set global cfg object.
- set_cfg(cfg_dict)
# Set up the output directory.
setup_output_dir(cfg, cfg_dict)
- return cfg, cfg_dict, eval_cfg # TODO Naama: why do we need both cfg and cfg_dict?
+ # cfg is the typed config, used for attribute access throughout the code. cfg_dict is the
+ # raw OmegaConf tree: wandb logs it in full, the eval-config and migration helpers operate
+ # on it as nested dicts, and it holds fields outside the typed schema (e.g. profiling).
+ return cfg, cfg_dict, eval_cfg
def flatten_wandb(cfg):
@@ -375,6 +348,11 @@ def _apply_cli_overrides(merged_cfg: DictConfig, orig_cli_cfg: DictConfig, raw_o
print(cyan(f"Re-applying {len(raw_overrides)} CLI overrides onto merged config."))
OmegaConf.set_struct(merged_cfg, False)
+ # Per-key merge decisions are noisy on every run; show them only under a debugger.
+ def _trace(msg: str) -> None:
+ if DEBUG:
+ print(cyan(msg))
+
# Architecture subtrees: CLI group default fills in *new* fields only;
# checkpoint values win for fields that already exist.
ARCH_KEYS = {"scene_optimizer", "scene_initializer"}
@@ -390,7 +368,7 @@ def _apply_cli_overrides(merged_cfg: DictConfig, orig_cli_cfg: DictConfig, raw_o
if cli_val is None:
# No direct config path — e.g. +experiment=re10k is a defaults-list override
# whose effect is already baked into orig_cli_cfg; nothing to apply.
- print(cyan(f" Skipping '{key}' (no direct config path in cli)"))
+ _trace(f" Skipping '{key}' (no direct config path in cli)")
continue
# For architecture group overrides: fill in missing fields from CLI defaults
@@ -401,7 +379,7 @@ def _apply_cli_overrides(merged_cfg: DictConfig, orig_cli_cfg: DictConfig, raw_o
dotkey_parts = set(dotkey.split("."))
if dotkey_parts & CLI_WINS_SUBKEYS:
OmegaConf.update(merged_cfg, dotkey, cli_val, merge=False)
- print(cyan(f" '{dotkey}': replace from cli (CLI wins)"))
+ _trace(f" '{dotkey}': replace from cli (CLI wins)")
continue
existing_val = OmegaConf.select(merged_cfg, dotkey, default=None)
@@ -414,15 +392,15 @@ def _apply_cli_overrides(merged_cfg: DictConfig, orig_cli_cfg: DictConfig, raw_o
if cli_subval is not None:
OmegaConf.set_struct(new_val, False)
OmegaConf.update(new_val, subkey, cli_subval, merge=False)
- print(cyan(f" '{dotkey}.{subkey}': CLI override applied (CLI wins)"))
+ _trace(f" '{dotkey}.{subkey}': CLI override applied (CLI wins)")
OmegaConf.update(merged_cfg, dotkey, new_val, merge=False)
- print(cyan(f" '{dotkey}': fill-missing from cli (checkpoint values preserved)"))
+ _trace(f" '{dotkey}': fill-missing from cli (checkpoint values preserved)")
continue
# Group overrides and complex values replace the whole subtree;
# scalars are merged so sibling keys are preserved.
replace = is_group_override
- print(cyan(f" '{dotkey}': {'replace' if replace else 'update'} from cli"))
+ _trace(f" '{dotkey}': {'replace' if replace else 'update'} from cli")
OmegaConf.update(merged_cfg, dotkey, cli_val, merge=not replace)
OmegaConf.set_struct(merged_cfg, True)
@@ -548,12 +526,16 @@ def _merge_test_mode(
) -> tuple[DictConfig, DictConfig]:
"""
Test mode: CLI config is the base for all settings (dataset, test flags, etc.).
- Only optimizer and initializer *architecture* are patched in from checkpoint configs.
+ The optimizer and initializer *architecture* and the decoder are patched in from the
+ checkpoint config, so the model is rebuilt and rendered the way it was trained. An explicit
+ CLI override still wins for any of these via _apply_cli_overrides.
Initializer source priority:
1. separate initializer_config_path (pretrained_initializer ckpt with a config file)
- 2. main loaded_cfg (optimizer checkpoint's bundled initializer)
- 3. CLI config as-is (pretrained_initializer set but has no config file)
+ 2. full pretrained_model checkpoint (it bundles the initializer weights, so its
+ architecture must match the checkpoint config)
+ 3. CLI config as-is (pretrained_optimizer alone carries no initializer
+ weights; or pretrained_initializer has no config file)
Returns (merged_cfg, orig_cli_cfg); orig_cli_cfg is the snapshot taken before any
checkpoint patches so that _apply_cli_overrides can restore explicit CLI values.
@@ -572,17 +554,28 @@ def _merge_test_mode(
print(cyan("Test mode: patching scene_trainer.scene_optimizer from checkpoint config."))
OmegaConf.update(merged_cfg, "scene_trainer.scene_optimizer", optimizer_subcfg, merge=True)
+ # Patch decoder from checkpoint so rendering uses the rasterizer the Gaussians were trained
+ # for. Merge over the CLI decoder (checkpoint values win) so any field the checkpoint lacks
+ # is filled from the CLI default -- old checkpoints predate some DecoderCfg fields and there
+ # is no decoder migration. An explicit CLI decoder override still wins via _apply_cli_overrides.
+ decoder_subcfg = OmegaConf.select(loaded_cfg, "scene_trainer.decoder", default=None)
+ if decoder_subcfg is not None:
+ print(cyan("Test mode: patching scene_trainer.decoder from checkpoint config (CLI fills missing fields)."))
+ OmegaConf.update(merged_cfg, "scene_trainer.decoder", decoder_subcfg, merge=True)
+
# Patch initializer architecture (priority order above)
if initializer_config_path is not None and initializer_config_path.exists():
_patch_scene_initializer(merged_cfg, initializer_config_path, context="Test mode")
elif pretrained_initializer is None:
- pass
- # TODO Naama
- # No explicit initializer checkpoint — fall back to the optimizer checkpoint's initializer
- # initializer_subcfg = OmegaConf.select(loaded_cfg, "scene_trainer.scene_initializer", default=None)
- # if initializer_subcfg is not None:
- # print(cyan("Test mode: patching scene_trainer.scene_initializer from checkpoint config."))
- # OmegaConf.update(merged_cfg, "scene_trainer.scene_initializer", initializer_subcfg, merge=True)
+ # No separate initializer checkpoint. A full pretrained_model bundles the initializer
+ # weights (load_full_model), so its architecture must come from the checkpoint config.
+ # A pretrained_optimizer checkpoint carries no initializer weights, so the CLI
+ # scene_initializer config is kept as-is.
+ if cli_cfg.checkpointing.pretrained_model is not None:
+ initializer_subcfg = OmegaConf.select(loaded_cfg, "scene_trainer.scene_initializer", default=None)
+ if initializer_subcfg is not None:
+ print(cyan("Test mode: patching scene_trainer.scene_initializer from full-model checkpoint config."))
+ OmegaConf.update(merged_cfg, "scene_trainer.scene_initializer", initializer_subcfg, merge=True)
else:
print(cyan("pretrained_initializer set but has no config file; using CLI scene_initializer config."))
@@ -677,14 +670,8 @@ def setup_output_dir(cfg, cfg_dict):
output_dir = CustomPath(output_dir)
output_dir.mkdir(exist_ok=True, parents=True)
- if HydraConfig.get().mode == RunMode.MULTIRUN and output_dir == "placeholder":
- # Hack to overcome multirun issues
- # TODO Naama, need to move to post_init of cfg
- output_dir = CustomPath(hydra.core.hydra_config.HydraConfig.get()["run"]["dir"])
- print(cyan(f"Multirun detected, setting output_dir to {CustomPath(output_dir):link}"))
- # save checkoint path to a file for debugging
- ckpt_path = cfg.checkpointing.pretrained_model or cfg.checkpointing.pretrained_optimizer
- (output_dir / "ckpt_dir.txt").write_text(str(ckpt_path))
+ # Single point where output_dir is finalized. Downstream code reads cfg.output_dir; the
+ # value is also written into cfg_dict so the saved config.yaml records it. Keep both in sync.
cfg_dict.output_dir = output_dir
cfg.output_dir = output_dir
output_dir.mkdir(exist_ok=True, parents=True)
diff --git a/optgs/config/dataset/dl3dv.yaml b/optgs/config/dataset/dl3dv.yaml
index f709e9e48a9eda6ff5a4c80ebab1f91c23c62228..1f3a32624911c643447d54d8058bd55f5f0aa008 100644
--- a/optgs/config/dataset/dl3dv.yaml
+++ b/optgs/config/dataset/dl3dv.yaml
@@ -4,19 +4,16 @@ defaults:
name: dl3dv
roots: [datasets/dl3dv]
-make_baseline_1: false
augment: true
image_shape: [270, 480]
-baseline_epsilon: 1e-3
max_fov: 100.0
skip_bad_shape: true
near: -1.
far: -1.
-baseline_scale_bounds: false
shuffle_val: true
test_len: -1
test_chunk_interval: 1
@@ -27,35 +24,11 @@ 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
-mix_tartanair: false
no_mix_test_set: true
-load_depth: false
-center_pose: false
-
-pose_align_first_view: false
-
-scale_extrinsics: 1.
-metric_scale_align_dl3dv: false
# view filtering
min_views: 0
max_views: 0
-highres: false
-
-# mix re10k & dl3dv
-mix_re10k: false
-re10k_min_view_dist: 40
-re10k_max_view_dist: 300
-
-# load remaining context views
-load_remain_context: false
-num_remain_context: 8
-
-# random crop in training
-random_crop: false
-min_size: null
-max_size: null
-index_name: index.json
\ No newline at end of file
+index_name: index.json
diff --git a/optgs/config/dataset/re10k.yaml b/optgs/config/dataset/re10k.yaml
index be4e8f29ba74f7e8fd6f946c575183d441cb865d..3558ab33712bb7387bbe3ff52c392b9f7d013de7 100644
--- a/optgs/config/dataset/re10k.yaml
+++ b/optgs/config/dataset/re10k.yaml
@@ -4,24 +4,16 @@ defaults:
name: re10k
roots: [datasets/re10k]
-make_baseline_1: false
augment: true
image_shape: [180, 320]
highres: false
-baseline_epsilon: 1e-3
max_fov: 100.0
skip_bad_shape: true
near: -1.
far: -1.
-baseline_scale_bounds: true
shuffle_val: true
test_len: -1
-test_chunk_interval: 1
-
-use_index_to_load_chunk: false
-
-average_pose: false
-center_pose: false
\ No newline at end of file
+test_chunk_interval: 1
\ No newline at end of file
diff --git a/optgs/config/experiment/re10k_unified.yaml b/optgs/config/experiment/re10k_unified.yaml
index 1d9b1a659a9e2b8934566b1aa7338ba23365fdcc..c9cedcb251178049c0d338516e3ea5667c41a490 100644
--- a/optgs/config/experiment/re10k_unified.yaml
+++ b/optgs/config/experiment/re10k_unified.yaml
@@ -63,16 +63,11 @@ dataset:
roots: [datasets/re10k]
near: 0.01
far: 100.
- baseline_scale_bounds: false
- make_baseline_1: false
train_times_per_scene: 1
highres: false
scannet: false
tartanair: false
load_depth: false
- pose_align_first_view: false
- scale_extrinsics: 1.
- load_remain_context: false
pose_align_middle_view: false
overfit_to_scene: null
opencv_pose_format: false
diff --git a/optgs/config/experiment/train_dl3dv.yaml b/optgs/config/experiment/train_dl3dv.yaml
index 160f746688a3b752b777d155a54fb4666da47dbb..7aa917eef336e80734207e99e3ef6b18e88516eb 100644
--- a/optgs/config/experiment/train_dl3dv.yaml
+++ b/optgs/config/experiment/train_dl3dv.yaml
@@ -21,7 +21,6 @@ meta_trainer:
max_steps: 50_000
val_check_interval: 0.25
train:
- l1_loss: true
depth_smooth_loss_weight: 0.0
test:
eval_time_skip_steps: 0
@@ -31,6 +30,8 @@ meta_trainer:
# lpips loss
loss:
+ mse:
+ l1_loss: true
lpips:
apply_after_step: 0
weight: 0.5
@@ -40,8 +41,6 @@ dataset:
roots: [ datasets/dl3dv-480p-chunks ]
near: 0.01
far: 200.
- min_size: [ 384,512 ]
- max_size: [ 512,960 ]
image_shape: [ 256, 448 ]
view_sampler:
num_context_views: 8
diff --git a/optgs/config/experiment/train_l2s_sparse_dl3dv.yaml b/optgs/config/experiment/train_l2s_sparse_dl3dv.yaml
index 232e95237564afac18aae815bf1da5e44a518500..6a364c9bac7256fc21643e010c904c7d63b74504 100644
--- a/optgs/config/experiment/train_l2s_sparse_dl3dv.yaml
+++ b/optgs/config/experiment/train_l2s_sparse_dl3dv.yaml
@@ -2,8 +2,8 @@
defaults:
- train_dl3dv
- - override /meta_trainer/train/replay_buffer_cfg: default
- - override /loss: [ mse, lpips, deltas ]
+ - override /meta_trainer/train/ckpt_buffer_cfg: default
+ - override /loss: [ mse, lpips, deltas, stability ]
loss:
mse:
@@ -23,19 +23,24 @@ meta_trainer:
train:
loss_on_input_views: true
loss_on_input_views_num: 4
- use_replay_buffer: true
+ use_ckpt_buffer: true
+ meta_optimizer:
+ lr: 1e-4
+ lr_monodepth: 0.0
scene_trainer:
train_scene_opt: true
- num_update_steps: 4
+ num_update_steps: 2000
+ opt_batch_size: 8
train_max_refine: 6
train_min_refine: 1
-
-meta_optimizer:
- lr: 1e-4
- lr_monodepth: 0.0
-
+ # Decoder settings for the learn2splat paper weights (the generic gsplat default is
+ # antialiased / eps2d=0.3; these weights are trained with classic rasterization).
+ decoder:
+ rasterize_mode: classic
+ eps2d: 1e-4
checkpointing:
- pretrained_initializer: checkpoints/optgs/unified-dl3dv-8views/init/checkpoints/epoch_20-step_100000.ckpt # resplat inititalizer
+ # ReSplat initializer
+ pretrained_initializer: hf://autonomousvision/learn2splat/resplat_init/checkpoints/epoch_20-step_100000.ckpt
no_strict_load: false
\ No newline at end of file
diff --git a/optgs/config/experiment/train_l2s_sparse_dl3dv_no_delta.yaml b/optgs/config/experiment/train_l2s_sparse_dl3dv_no_delta.yaml
index bc160feac7dd19ee8f524fe16e5e2a1dd440ebc4..f6e5f14954c22cd6b49f5af097113696c2680907 100644
--- a/optgs/config/experiment/train_l2s_sparse_dl3dv_no_delta.yaml
+++ b/optgs/config/experiment/train_l2s_sparse_dl3dv_no_delta.yaml
@@ -2,7 +2,7 @@
defaults:
- train_dl3dv
- - override /meta_trainer/train/replay_buffer_cfg: default
+ - override /meta_trainer/train/ckpt_buffer_cfg: default
- override /loss: [ mse, lpips ]
loss:
@@ -17,7 +17,10 @@ meta_trainer:
train:
loss_on_input_views: true
loss_on_input_views_num: 4
- use_replay_buffer: true
+ use_ckpt_buffer: true
+ meta_optimizer:
+ lr: 1e-4
+ lr_monodepth: 0.0
scene_trainer:
train_scene_opt: true
@@ -25,11 +28,6 @@ scene_trainer:
train_max_refine: 6
train_min_refine: 1
-meta_optimizer:
- lr: 1e-4
- lr_monodepth: 0.0
-
-
checkpointing:
pretrained_initializer: checkpoints/optgs/unified-dl3dv-8views/init/checkpoints/epoch_20-step_100000.ckpt # resplat inititalizer
no_strict_load: false
\ No newline at end of file
diff --git a/optgs/config/experiment/train_l2s_sparse_dl3dv_no_loss.yaml b/optgs/config/experiment/train_l2s_sparse_dl3dv_no_loss.yaml
index bc160feac7dd19ee8f524fe16e5e2a1dd440ebc4..f6e5f14954c22cd6b49f5af097113696c2680907 100644
--- a/optgs/config/experiment/train_l2s_sparse_dl3dv_no_loss.yaml
+++ b/optgs/config/experiment/train_l2s_sparse_dl3dv_no_loss.yaml
@@ -2,7 +2,7 @@
defaults:
- train_dl3dv
- - override /meta_trainer/train/replay_buffer_cfg: default
+ - override /meta_trainer/train/ckpt_buffer_cfg: default
- override /loss: [ mse, lpips ]
loss:
@@ -17,7 +17,10 @@ meta_trainer:
train:
loss_on_input_views: true
loss_on_input_views_num: 4
- use_replay_buffer: true
+ use_ckpt_buffer: true
+ meta_optimizer:
+ lr: 1e-4
+ lr_monodepth: 0.0
scene_trainer:
train_scene_opt: true
@@ -25,11 +28,6 @@ scene_trainer:
train_max_refine: 6
train_min_refine: 1
-meta_optimizer:
- lr: 1e-4
- lr_monodepth: 0.0
-
-
checkpointing:
pretrained_initializer: checkpoints/optgs/unified-dl3dv-8views/init/checkpoints/epoch_20-step_100000.ckpt # resplat inititalizer
no_strict_load: false
\ No newline at end of file
diff --git a/optgs/config/loss/lpips.yaml b/optgs/config/loss/lpips.yaml
index cb3aafcc76f7a2219e6c2ab178bd37a2af9f6472..bf1343a8d82eb3858011ed60dac7be9a75384cf9 100644
--- a/optgs/config/loss/lpips.yaml
+++ b/optgs/config/loss/lpips.yaml
@@ -2,3 +2,4 @@ lpips:
weight: 0.05
apply_after_step: 150_000
perceptual_loss: false
+ half_res: false
diff --git a/optgs/config/loss/mse.yaml b/optgs/config/loss/mse.yaml
index 80cc0be6dc7950661998336bd2bb5cb4ff06ba07..53f0313259737568ca909196e4566e9aa53c3212 100644
--- a/optgs/config/loss/mse.yaml
+++ b/optgs/config/loss/mse.yaml
@@ -1,2 +1,4 @@
mse:
weight: 1.0
+ l1_loss: false
+ clamp_large_error: 0.0
\ No newline at end of file
diff --git a/optgs/config/loss/sgd.yaml b/optgs/config/loss/sgd.yaml
index cb59020610feeeaabd07e7c6ac05b11ad8094712..83aa71f5b22e0a9b4cad3a10373e0f7a94edc336 100644
--- a/optgs/config/loss/sgd.yaml
+++ b/optgs/config/loss/sgd.yaml
@@ -1,2 +1,4 @@
sgd:
- weight: 1.0
\ No newline at end of file
+ weight: 1.0
+ l1_loss: false
+ clamp_large_error: 0.0
\ No newline at end of file
diff --git a/optgs/config/loss/stability.yaml b/optgs/config/loss/stability.yaml
index a4a9abc4af85dec6c037aeb1b029e297ca401bbe..db11aaaaa49c286487c96e5531f598958f638b15 100644
--- a/optgs/config/loss/stability.yaml
+++ b/optgs/config/loss/stability.yaml
@@ -1,2 +1,4 @@
stability:
weight: 1.0
+ # Also penalize inputs optimized on per-step view subsets (not fully tested in training).
+ subset_aware: false
diff --git a/optgs/config/main.yaml b/optgs/config/main.yaml
index bbd128507fc8b4c60e520875c503df4445e17ce9..67d3b1637c036bb1399af0969ad7bdef44169deb 100644
--- a/optgs/config/main.yaml
+++ b/optgs/config/main.yaml
@@ -5,7 +5,7 @@ defaults:
- scene_trainer/scene_optimizer: null
- scene_trainer/decoder: gsplat
- meta_trainer/test/postprocessing: none
- - meta_trainer/train/replay_buffer_cfg: none
+ - meta_trainer/train/ckpt_buffer_cfg: none
wandb:
project: placeholder
@@ -34,15 +34,6 @@ data_loader:
batch_size: 1
seed: 3456
-meta_optimizer:
- lr: 2.e-4
- lr_monodepth: 2.e-6
- lr_depth: 0.
- warm_up_steps: 2000
- weight_decay: 0.01
- warm_up_ratio: 0.01
- adamw_8bit: false
-
checkpointing:
load: null
every_n_train_steps: 1000
@@ -75,25 +66,30 @@ meta_trainer:
limit_test_batches: 1.0
limit_train_batches: 1.0
num_nodes: 1
+ meta_optimizer:
+ lr: 2.e-4
+ lr_monodepth: 2.e-6
+ lr_depth: 0.
+ warm_up_steps: 2000
+ weight_decay: 0.01
+ warm_up_ratio: 0.01
+ adamw_8bit: 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_deterministic: false
eval_time_skip_steps: 3
eval_save_model: true
- l1_loss: false
intermediate_loss_weight: 0.9
no_viz_video: false
eval_depth: false
- train_ignore_large_loss: 0.
no_log_projections: true
no_log_video: true
depth_loss_weight: 0.
log_depth_loss: true
- depth_smooth_loss_weight: 0.01
+ depth_smooth_loss_weight: 0. # only meaningful when train_scene_init=True
depth_smooth_loss_nonorm: false
depth_smooth_loss_weight_nvs: 0. # for novel views
monodepth_loss_weight: 0. # for monocular depth loss
@@ -103,21 +99,17 @@ meta_trainer:
render_depth_loss_weight: 0.
viz_render_depth: false
use_gt_depth_range: false
- depth_range_from_disparity: false
- max_disparity: 128.
- min_disparity: 4.
loss_on_input_views: false
loss_on_target_views: true
loss_on_input_views_num: 1
loss_on_target_views_num: -1
train_window_size: null
- half_res_lpips_loss: false
viz_depth_separate: false
# L2 weight decay on Gaussian properties (meta-loss)
scale_l2_loss_weight: 0.
sh_l2_loss_weight: 0.
opacity_l2_loss_weight: 0.
- use_replay_buffer: false
+ use_ckpt_buffer: false
test:
output_path: null
compute_scores: true
@@ -130,22 +122,22 @@ meta_trainer:
save_gt_image: false
save_render_depth: false
save_gt_depth: false
+ save_init_pred_depth: false
save_error_image: false
save_video: false
- save_video_fixed_view: false
- save_video_fixed_view_index: 0
- save_video_fixed_view_duplicate: 0
- save_video_fixed_iteration: false
- save_video_fixed_iteration_indices: null
- save_video_fixed_iteration_render_fixed_view: false
- save_video_combined: false
- save_video_combined_iterations: null
- save_video_combined_fixed_iteration_length: 50
+ save_video_optim: false
+ save_video_view_index: 0
+ save_video_frame_repeat: 1
+ save_video_orbit: false
+ save_video_orbit_steps: null
+ save_video_orbit_with_optim: false
+ save_video_optim_orbit: false
+ save_video_optim_orbit_steps: null
+ save_video_orbit_span: 50
save_gaussian: false
save_poses: false
save_cameras_json: true
save_cameras_npz: true
- save_point_cloud: false
render_chunk_size: null
dec_chunk_size: null
stablize_camera: false
diff --git a/optgs/config/meta_trainer/train/replay_buffer_cfg/default.yaml b/optgs/config/meta_trainer/train/ckpt_buffer_cfg/default.yaml
similarity index 60%
rename from optgs/config/meta_trainer/train/replay_buffer_cfg/default.yaml
rename to optgs/config/meta_trainer/train/ckpt_buffer_cfg/default.yaml
index df3d239d2ea4207f77ce8208ce39c58a9d7fa87a..91bd3944d4626740e56cd3a7bce60bb36428ff4f 100644
--- a/optgs/config/meta_trainer/train/replay_buffer_cfg/default.yaml
+++ b/optgs/config/meta_trainer/train/ckpt_buffer_cfg/default.yaml
@@ -3,10 +3,10 @@ sample_batch_size: 1
sample_prob: 0.7
insert_prob: 0.7
return_prob: 0.99
-simulate_ahead: true
-simulate_ahead_min_steps: 1
-simulate_ahead_max_steps: 50
-simulate_ahead_grow: 10000
+rollout: true
+rollout_min_steps: 1
+rollout_max_steps: 50
+rollout_grow: 10000
max_t: null
push_only_if_not_full: false
remove_strategy_when_full: oldest
\ No newline at end of file
diff --git a/optgs/config/meta_trainer/train/replay_buffer_cfg/none.yaml b/optgs/config/meta_trainer/train/ckpt_buffer_cfg/none.yaml
similarity index 61%
rename from optgs/config/meta_trainer/train/replay_buffer_cfg/none.yaml
rename to optgs/config/meta_trainer/train/ckpt_buffer_cfg/none.yaml
index 06e1f12d212f6187e107b6175a85a0dd0ae4f724..c12e50a48be8a38dac6e457a5e209c5d4dece7d8 100644
--- a/optgs/config/meta_trainer/train/replay_buffer_cfg/none.yaml
+++ b/optgs/config/meta_trainer/train/ckpt_buffer_cfg/none.yaml
@@ -3,10 +3,10 @@ sample_batch_size: 1
sample_prob: 0.0
insert_prob: 0.0
return_prob: 0.0
-simulate_ahead: false
-simulate_ahead_min_steps: 0
-simulate_ahead_max_steps: 0
-simulate_ahead_grow: 0
+rollout: false
+rollout_min_steps: 0
+rollout_max_steps: 0
+rollout_grow: 0
max_t: null
push_only_if_not_full: false
remove_strategy_when_full: oldest
\ No newline at end of file
diff --git a/optgs/config/scene_trainer/scene_initializer/edgs.yaml b/optgs/config/scene_trainer/scene_initializer/edgs.yaml
index 265bae8d718fae849f69ec40ae72f0a087f77d7e..8897d7c7d02aa4ee19d8b7a77534711bfe3ad03b 100644
--- a/optgs/config/scene_trainer/scene_initializer/edgs.yaml
+++ b/optgs/config/scene_trainer/scene_initializer/edgs.yaml
@@ -7,4 +7,5 @@ sh_degree: 3
init_opacity: 0.5
scaling_factor: 0.5
roma_model_type: outdoors
-sample_init_gaussians: -1
\ No newline at end of file
+sample_init_gaussians: -1
+cache_dir: cache/edgs
\ No newline at end of file
diff --git a/optgs/config/scene_trainer/scene_initializer/resplat_v1.yaml b/optgs/config/scene_trainer/scene_initializer/resplat_v1.yaml
index ce18ce1a2f8a9c0e340c799a4aa4e8f698eac0af..60295c4a7043330db7960b74154b48c1a194b2da 100644
--- a/optgs/config/scene_trainer/scene_initializer/resplat_v1.yaml
+++ b/optgs/config/scene_trainer/scene_initializer/resplat_v1.yaml
@@ -71,19 +71,13 @@ init_pt_with_mv_attn_lowres: false
pt_head_channels: null
pt_head_concat_img: false
pt_head_conv: false
-multi_scale_pt: false
attn_proj_channels: 64
-fps_num_samples: null
knn_samples: 16
post_norm: false
no_rpe: true
no_knn_attn: false
num_blocks: 4
-pt_downsample: 0
-fps_agg_func: attn
-subsample_method: fps
add_pt_residual: true
-pt_pred_residual_position: false
# freeze depth
freeze_depth: false
@@ -99,10 +93,6 @@ bilinear_upsample_depth: false
no_upsample_depth: false
return_lowres_depth: false
-# lvsm gaussian regressor
-lvsm_gaussian_regressor: false
-lvsm_layers: 6
-
# latent gaussian instead of pixel aligned gaussians
latent_gs: true
latent_downsample: 4
diff --git a/optgs/config/scene_trainer/scene_initializer/resplat_v2.yaml b/optgs/config/scene_trainer/scene_initializer/resplat_v2.yaml
index 7fe1337fc7f108d18dcbef59cd09977c307cf263..a4888b8f446fa364ddd80261565f7f0cee5f9c0c 100644
--- a/optgs/config/scene_trainer/scene_initializer/resplat_v2.yaml
+++ b/optgs/config/scene_trainer/scene_initializer/resplat_v2.yaml
@@ -77,19 +77,13 @@ init_pt_with_mv_attn_lowres: true
pt_head_channels: null
pt_head_concat_img: false
pt_head_conv: false
-multi_scale_pt: false
attn_proj_channels: 64
-fps_num_samples: null
knn_samples: 16
post_norm: false
no_rpe: true
no_knn_attn: false
num_blocks: 6
-pt_downsample: 0
-fps_agg_func: attn
-subsample_method: fps
add_pt_residual: true
-pt_pred_residual_position: false
# freeze depth
freeze_depth: false
@@ -105,10 +99,6 @@ bilinear_upsample_depth: false
no_upsample_depth: false
return_lowres_depth: false
-# lvsm gaussian regressor
-lvsm_gaussian_regressor: false
-lvsm_layers: 6
-
# latent gaussian instead of pixel aligned gaussians
latent_gs: true
latent_downsample: 4
diff --git a/optgs/config/scene_trainer/scene_optimizer/3dgs.yaml b/optgs/config/scene_trainer/scene_optimizer/3dgs.yaml
index 46f55dac1b966cb9a3491ec4ef8f3357fcccefbd..88f627b6aedcc848565282fe1de055dc9c9f5418 100644
--- a/optgs/config/scene_trainer/scene_optimizer/3dgs.yaml
+++ b/optgs/config/scene_trainer/scene_optimizer/3dgs.yaml
@@ -1,22 +1,11 @@
defaults:
- base
- override refiner: default
+ - override lr_scheduler: expon
name: adam
# Adam optimizer
betas: [0.9, 0.999]
eps: 1e-15
-weight_decay: 0.0
-
-# learning rates (gsplat)
-base_lr: 1
-means_lr_init: 1.6e-4
-means_lr_final: 1.6e-6
-means_lr_delay_mult: 1.0
-means_lr_max_steps: 30000 # should be equal to total optimization steps
-scales_lr: 5e-3
-rotations_lr: 1e-3
-opacities_lr: 5e-2
-sh0s_lr: 2.5e-3
-shNs_lr: 1.25e-4
+weight_decay: 0.0
\ No newline at end of file
diff --git a/optgs/config/scene_trainer/scene_optimizer/3dgs_star.yaml b/optgs/config/scene_trainer/scene_optimizer/3dgs_star.yaml
index 44d4c33fd69fadee8ea7af12f94b0867606b19f0..d4cbff5c5feca8a4a4e71715e91d6b5533c9476a 100644
--- a/optgs/config/scene_trainer/scene_optimizer/3dgs_star.yaml
+++ b/optgs/config/scene_trainer/scene_optimizer/3dgs_star.yaml
@@ -1,6 +1,7 @@
defaults:
- base
- override refiner: default
+ - override lr_scheduler: expon
name: adam
@@ -9,14 +10,8 @@ betas: [0.99, 0.999]
eps: 1e-15
weight_decay: 0.0
-base_lr: 5
-# 3dgs defaults
-means_lr_init: 1.6e-4 # Setting same as final to have constant LR
-means_lr_final: 1.6e-4
-means_lr_delay_mult: 0.0
-means_lr_max_steps: 30000 # should be equal to total optimization steps
-scales_lr: 5e-3
-rotations_lr: 1e-3
-opacities_lr: 5e-2
-sh0s_lr: 2.5e-3
-shNs_lr: 1.25e-4
\ No newline at end of file
+# 5x base LR; means held constant (lr_final == _means, so means init 5*1.6e-4 == final 5*1.6e-4).
+lr_scheduler:
+ lr_data:
+ _base: 5
+ lr_final: 1.6e-4
\ No newline at end of file
diff --git a/optgs/config/scene_trainer/scene_optimizer/base.yaml b/optgs/config/scene_trainer/scene_optimizer/base.yaml
index 813184bc9dfd4a981b75d8f2767e20850c5b25fc..306a541504dabcc9881a37cc3a725fb5097e634c 100644
--- a/optgs/config/scene_trainer/scene_optimizer/base.yaml
+++ b/optgs/config/scene_trainer/scene_optimizer/base.yaml
@@ -6,12 +6,12 @@ defaults:
input_gradients_chunk_size: -1
# iterative refine
-no_refine_mean: false
-no_refine_scale: false
-no_refine_rotation: false
-no_refine_opacity: false
-no_refine_sh0: false
-no_refine_shN: false
+freeze_mean: false
+freeze_scale: false
+freeze_rotation: false
+freeze_opacity: false
+freeze_sh0: false
+freeze_shN: false
zero_state_on_densify: false
diff --git a/optgs/config/scene_trainer/scene_optimizer/knn_based.yaml b/optgs/config/scene_trainer/scene_optimizer/knn_based.yaml
index 4ff20fd09c73e600eb5f0f3e32e89d46c39c7a16..80288e1416565a4f8f8ef6e83a14c3cab13b0c7d 100644
--- a/optgs/config/scene_trainer/scene_optimizer/knn_based.yaml
+++ b/optgs/config/scene_trainer/scene_optimizer/knn_based.yaml
@@ -5,19 +5,18 @@ name: knn_based
# iterative refine
no_render_error: false
-refine_sh_only: false
-num_basic_refine_blocks: 4
-num_refine_blocks: 1
+num_basic_blocks: 4
+num_blocks: 1
concat_init_state: false
replace_init_state: false
state_channels: 256
-refine_block_rmsnorm: false
-refine_block_layernorm: false
+block_rmsnorm: false
+block_layernorm: false
pt_qk_norm: false
norm_pt_block: false
-refine_gaussian_multiple: 1
-refine_residual_init_state: false
-clamp_refine_max_scale: 3.0
+delta_gaussian_multiple: 1
+residual_init_state: false
+clamp_max_scale: 3.0
clamp_min_scale: 1e-6
clamp_min_raw_scales: -1e10
clamp_max_raw_scales: 1e10
@@ -30,15 +29,14 @@ clamp_min_shs: -1e10
clamp_max_shs: 1e10
clamp_shs_soft: false
-update_attn_proj_channels: 64
-update_no_knn_attn: false
-update_no_tran_block_norm: false
-update_tran_block_act: gelu
+attn_proj_channels: 64
+no_knn_attn: false
+no_tran_block_norm: false
+tran_block_act: gelu
multi_gaussian_scale_smaller: false
init_gaussian_multiple: 1
-refine_condition_pt_feature: true
-reinit_gaussian_when_refine_multiple: false
-refine_same_num_points: false
+condition_pt_feature: true
+same_num_points: false
input_error_rgb_no_shuffle: false
input_error_cache_resnet_feature: false
@@ -50,20 +48,19 @@ init_state_scale: 0
pt_heads: 1
# refine with mv attention
-refine_with_mv_attn: false
-refine_with_mv_attn_lowres: false
-refine_no_mv_attn: false
+with_mv_attn: false
+with_mv_attn_lowres: false
+no_mv_attn: false
mv_attn_conv_with_norm: false
-refine_mv_shuffle_attn: false
-refine_mv_attn_with_pos_enc: false
-refine_shuffle_attn_no_norm: false
-refine_mv_unimatch_attn: false
-refine_knn_samples: 16
-refine_multi_scale_pt: false
+mv_shuffle_attn: false
+mv_attn_with_pos_enc: false
+shuffle_attn_no_norm: false
+mv_unimatch_attn: false
+knn_samples: 16
+max_active_gaussians: 100_000
# KNN
use_fused_attn: true # fused KNN gather + attention CUDA kernel (faster, less memory)
-prune_invisible_gaussians: false
knn_idx_update_every: 1
# inputs
@@ -103,25 +100,17 @@ input_gradient: false
input_gradient_log: false
input_gradient_log_clip_deltas: 0.001
input_gradient_scale: 1.
-gradient_update_scale: 1.
+residual_grad_scale: 1.
input_gradient_with_ssim_loss: false
input_gradient_same_loss: false
input_gradient_loss_reduction: mean
scale_residual_grads: false
-window_local_refine: false
-window_global_refine: false
-window_local_global_refine: false
-
-# sliding window update to save training memory
-update_window_size: 0
-local_gaussian_render: false
-
-train_global_update_only: false
+train_global_only: false
# random size refine
# update more for low resolution, less for high
-random_update_with_size: false
+random_step_with_size: false
use_amp: true
@@ -131,9 +120,6 @@ pt_update_amp: true
use_checkpointing: false
recurrent_use_checkpointing: false
-# Debugging
-debug_refine_update_module: true
-
# Normalizing input
input_gradient_normalize: false
input_gradient_normalize_type: layer
@@ -146,21 +132,20 @@ predict_state_scale: false
predict_state_scale_norm: false
# Update head
-update_head_concat_img: false
-update_head_layer_num: 2
-update_head_act: gelu
-update_head_final_act: identity
-update_head_hidden_dim_matches: "input" # rebuttal version. switch to "output" for submission version
+delta_head_concat_img: false
+delta_head_layer_num: 2
+delta_head_act: gelu
+delta_head_final_act: identity
+delta_head_hidden_dim_matches: "input" # rebuttal version. switch to "output" for submission version
-update_head_scale_mag: false
-update_head_scalar_scale: false
-update_head_scalar_scale_act: relu
+delta_head_scalar_scale: false
+delta_head_scalar_scale_act: relu
# Per-parameter-group heads (Feature A): separate heads per param group, each with own normalize+scale
-update_head_per_param_heads: false
-update_head_per_param_hidden_dim: 48 # tuned so total params ≈ baseline head (~81K)
-# Per-parameter scalar scales (Feature B): per-group scalar scales (requires update_head_scalar_scale=true)
-update_head_per_param_scales: false
+delta_head_per_param_heads: false
+delta_head_per_param_hidden_dim: 48 # tuned so total params ≈ baseline head (~81K)
+# Per-parameter scalar scales (Feature B): per-group scalar scales (requires delta_head_scalar_scale=true)
+delta_head_per_param_scales: false
opt_scales_before_act: false
diff --git a/optgs/config/scene_trainer/scene_optimizer/learn2splat.yaml b/optgs/config/scene_trainer/scene_optimizer/learn2splat.yaml
index e730c789a0af3b7ae89747a91d54eca1273c68f5..b22c2016c115684b66bf6385bc2a54e41ba05a30 100644
--- a/optgs/config/scene_trainer/scene_optimizer/learn2splat.yaml
+++ b/optgs/config/scene_trainer/scene_optimizer/learn2splat.yaml
@@ -5,13 +5,25 @@ name: l2s
# General optimization settings
opt_scales_before_act: true
+# Scales are refined in log space, so clamp in log space too.
+clamp_min_raw_scales: -8.0
+clamp_max_raw_scales: 1.0986 # ln(3.0)
sh_d: 16 # should be the default, but just in case
+# SH coefficients are clamped to a fixed range during refinement.
+clamp_min_shs: -2.0
+clamp_max_shs: 2.0
+
+# KNN neighborhood size for the point-transformer attention.
+knn_samples: 4
+
# Input gradient settings
input_gradient: true
input_gradient_normalize: true
input_gradient_normalize_type: adam
input_gradient_with_ssim_loss: true
+# Average the input gradient over pixels, then sum over the supervised views.
+input_gradient_loss_reduction: mean_pixels_sum_views
# Freeze zero-grad gaussians ❄️($G_{∇=0})
update_only_nonzero_grad: true
@@ -20,8 +32,8 @@ update_only_nonzero_grad: true
predict_state_scale: true
# Delta scale
-update_head_scalar_scale: true
-update_head_scalar_scale_act: relu # should be the default, but just in case
+delta_head_scalar_scale: true
+delta_head_scalar_scale_act: relu # should be the default, but just in case
diff --git a/optgs/config/scene_trainer/scene_optimizer/learn2splat_dense.yaml b/optgs/config/scene_trainer/scene_optimizer/learn2splat_dense.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..06e0d9312c46f134d0024b23aa89ebe42ff118f4
--- /dev/null
+++ b/optgs/config/scene_trainer/scene_optimizer/learn2splat_dense.yaml
@@ -0,0 +1,8 @@
+defaults:
+ - learn2splat
+
+# Dense variant (64-view SfM/COLMAP init): the optimizer state is initialized as random
+# features rather than the sparse model's constant zero state.
+init_state_wo_features: true
+init_state_type: random
+init_state_scale: 1.0
\ No newline at end of file
diff --git a/optgs/config/scene_trainer/scene_optimizer/lr_scheduler/expon.yaml b/optgs/config/scene_trainer/scene_optimizer/lr_scheduler/expon.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..54fcaf3f4b39c6fea91a81c4a071902c1420a198
--- /dev/null
+++ b/optgs/config/scene_trainer/scene_optimizer/lr_scheduler/expon.yaml
@@ -0,0 +1,28 @@
+defaults:
+ - base
+
+name: expon
+
+# 3DGS learning rates (gsplat); means decays, the rest are constant
+lr_data:
+ _base: 1
+ _means: 1.6e-4
+ _scales: 5e-3
+ _quats: 1e-3
+ _opacities: 5e-2
+ _sh0: 2.5e-3
+ _shN: 1.25e-4
+apply_scheduler:
+ _base: true # gates all params (Bool3DGSCfg.,
+ (3) the averaged initializer/decoder/optimization runtime in ms (first
+ eval_time_skip_steps batches are excluded as warmup), and
+ (4) a per-inner-step PSNR summary line + total wall-clock.
+ """
print(
- f"Validation step at global step {self.global_step}. Running evaluation on {self.train_cfg.eval_data_length} test sets...")
+ f"Validation step at global step {self.global_step}. "
+ f"Running evaluation on {self.train_cfg.eval_data_length} test sets..."
+ )
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 = defaultdict(list)
+ self.benchmarker.clear_history()
- full_testsets = self.trainer.datamodule.test_dataloader(
- dataset_cfg=self.eval_data_cfg
- )
- scores_dict = defaultdict(lambda: defaultdict(list))
+ time_skip_first_n_steps = min(self.train_cfg.eval_time_skip_steps, len(full_testsets))
+ time_skip_steps_dict = {"initializer": 0, "decoder": 0, "optimization": 0}
- 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
+ self._eval_one_test_batch(
+ batch, batch_idx, scores_dict, time_skip_steps_dict, time_skip_first_n_steps,
+ )
- batch = self.transfer_batch_to_device(batch, self.device, dataloader_idx=0)
- batch = self.on_after_batch_transfer(batch, dataloader_idx=batch_idx)
- batch = self.initializer_data_shim(batch)
+ self._log_inner_iteration_table(scores_dict)
+ self._summarize_test_scores(scores_dict, time_skip_steps_dict, start_t)
- # use gt depth range instead of a fixed one
- self.scene_initializer.preprocessing(batch, self.train_cfg)
+ def _eval_one_test_batch(
+ self,
+ batch,
+ batch_idx: int,
+ scores_dict: dict,
+ time_skip_steps_dict: dict,
+ time_skip_first_n_steps: int,
+ ) -> None:
+ """Evaluate one test-set batch and append metrics into `scores_dict` in-place.
- # 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
+ Counts initializer/decoder calls into `time_skip_steps_dict` only for the warmup window
+ (`batch_idx < time_skip_first_n_steps`) so the final timing summary can trim them.
+ Skips the batch on SkipBatchException from initialization or optimization.
+ """
+ batch = self.transfer_batch_to_device(batch, self.device, dataloader_idx=0)
+ batch = self.on_after_batch_transfer(batch, dataloader_idx=batch_idx)
- with self.benchmarker.time("encoder"):
- init_output = self.get_init_gaussians(batch, is_training=False)
+ v = batch["target"]["image"].shape[1]
+ if batch_idx < time_skip_first_n_steps:
+ time_skip_steps_dict["initializer"] += 1
+ time_skip_steps_dict["decoder"] += v
+ time_skip_steps_dict["optimization"] += 1
- with self.benchmarker.time("decoder", num_calls=v):
- output_probabilistic = self.scene_decoder.forward_target(
- init_output.gaussians, batch, (h, w),
- depth_mode='depth' if self.train_cfg.eval_render_depth or self.train_cfg.viz_render_depth else None,
- )
+ try:
+ batch, init_output, final_render, optimizer_output = self._run_eval_pipeline(
+ batch, benchmark=True)
+ except SkipBatchException as e:
+ warn(f'Skipping batch due to SkipBatch: {e}')
+ return
- init_rgb = output_probabilistic.color[0]
+ # RGBs
+ # target_render_list holds the init render first (position 0), then one render per
+ # optimizer step. When no optimizer runs (num_update_steps == 0) there is no list, so
+ # score the init render alone at step 0.
+ if optimizer_output is not None:
+ rgbs = [render.color[0] for render in optimizer_output.target_render_list]
+ steps = self.scene_optimizer.save_every.get_iterations(len(rgbs))
+ else:
+ rgbs = [init_output.target_render.color[0]]
+ steps = [0]
+
+ rgb_gt = batch["target"]["image"][0]
+ self._accumulate_rgb_metrics(rgbs, rgb_gt, steps, scores_dict)
+
+ # Depths
+ # Two distinct depth metrics: init-depth scores the initializer's context-view
+ # prediction (pred_depths); render-depth scores the final target render's depth
+ # (init's target render, or the optimizer's last render).
+ pred_depths = init_output.depths
+ depth_gt = batch["context"].get("depth")
+ if pred_depths is not None and depth_gt is not None and depth_gt.max() > 0:
+ self._accumulate_init_depth_metrics(pred_depths, depth_gt, batch, scores_dict)
+ target_depth_gt = batch["target"].get("depth")
+ if self.train_cfg.eval_render_depth and target_depth_gt is not None and target_depth_gt.max() > 0:
+ self._accumulate_render_depth_metrics(final_render, batch, scores_dict)
- # refine
- if self.scene_optimizer is not None:
- try:
- optimizer_output = self.get_optimized_gaussians(batch, init_output)
- except SkipBatchException as e:
- warn(f'Skipping batch due to SkipBatch during optimization: {e}')
- continue
- render_output = optimizer_output.target_render_list
- output_probabilistic = render_output[-1]
-
- rgbs = [init_rgb]
- if self.scene_trainer_cfg.num_update_steps > 0:
- rgbs += [render.color[0] for render in render_output]
- tags = ["probabilistic"] * len(rgbs)
-
- if self.train_cfg.eval_deterministic:
- gaussians_deterministic = self.encoder(
- batch["context"],
- self.global_step,
- deterministic=True,
- )
- output_deterministic = self.scene_decoder.forward(
- gaussians_deterministic,
- batch["target"]["extrinsics"],
- batch["target"]["intrinsics"],
- batch["target"]["near"],
- batch["target"]["far"],
- (h, w),
- )
- rgbs.append(output_deterministic.color[0])
- tags.append("deterministic")
+ @staticmethod
+ def _score_renders(
+ renders: list[Tensor],
+ rgb_gt: Tensor,
+ metrics: list[str],
+ iter_batch_size: int = -1,
+ ) -> list[dict[str, float]]:
+ """Compute RGB metrics for each render in `renders` against `rgb_gt`.
+
+ Returns a list parallel to `renders`; each entry maps metric name to a Python float.
+ `lpips` is unpacked into separate `alex_lpips`/`vgg_lpips` entries so callers can
+ treat all metrics uniformly.
+ """
+ per_render: list[dict[str, float]] = []
+ for rgb in renders:
+ raw = compute_rgb_metrics(rgb, rgb_gt, metrics=metrics, iter_batch_size=iter_batch_size)
+ flat: dict[str, float] = {}
+ for name, score in raw.items():
+ if name == "lpips":
+ alex, vgg = score
+ flat["alex_lpips"] = alex.item()
+ flat["vgg_lpips"] = vgg.item()
+ else:
+ flat[name] = score.item()
+ per_render.append(flat)
+ return per_render
+
+ def _accumulate_rgb_metrics(self, rgbs: list, rgb_gt, steps: list[int], scores_dict: dict) -> None:
+ """For each refinement checkpoint k, compute psnr/ssim/lpips(alex,vgg) and append into
+ scores_dict[f"{metric}_{steps[k]}"]. Also mirror the last step's value under the
+ unsuffixed `metric` key for between-run comparison."""
+ per_render = self._score_renders(rgbs, rgb_gt, metrics=["psnr", "ssim", "lpips"])
+ for i, scores in enumerate(per_render):
+ is_last = (i == len(per_render) - 1)
+ for name, val in scores.items():
+ scores_dict[f"{name}_{steps[i]}"].append(val)
+ if is_last:
+ scores_dict[name].append(val)
- # Compute validation metrics.
- rgb_gt = batch["target"]["image"][0]
- if self.scene_optimizer is not None:
- steps = self.scene_optimizer.save_every.get_iterations(len(rgbs))
- else:
- steps = [0]
- for i, (tag, rgb) in enumerate(zip(tags, rgbs)):
- # Move prediction back to device
- rgb = rgb.to(batch["target"]["image"].device)
- metric_scores: dict = compute_rgb_metrics(
- rgb, rgb_gt,
- metrics=["psnr", "ssim", "lpips"],
- iter_batch_size=-1,
- )
- for name, score in metric_scores.items():
- if name == "lpips":
- # tuple of (alex, vgg)
- scores_dict[f"alex_lpips_{steps[i]}"][tag].append(score[0].item())
- scores_dict[f"vgg_lpips_{steps[i]}"][tag].append(score[1].item())
- else:
- scores_dict[f"{name}_{steps[i]}"][tag].append(score.item())
- # log the last step metrics to compare between runs
- if i == len(rgbs) - 1:
- if name == "lpips":
- # tuple of (alex, vgg)
- scores_dict[f"alex_lpips"][tag].append(score[0].item())
- scores_dict[f"vgg_lpips"][tag].append(score[1].item())
- else:
- scores_dict[f"{name}"][tag].append(score.item())
-
- # compute depth metrics
- if pred_depths is not None and depth_gt is not None and depth_gt.max() > 0:
- assert pred_depths is not None and depth_gt is not None
-
- 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)
-
- depth_gt = depth_gt[0] # [V, H, W]
-
- near = batch["context"]["near"][...,
- None, None][0] # [V, 1, 1]
- far = batch["context"]["far"][..., None, None][0] # [V, 1, 1]
-
- valid = (depth_gt >= near) & (depth_gt <= far)
-
- all_metrics = compute_depth_errors(depth_gt[valid].detach().cpu().numpy(),
- pred_depths[valid].detach().cpu().numpy())
- scores_dict["abs_rel"]["probabilistic"].append(all_metrics[0])
- scores_dict["rmse"]["probabilistic"].append(all_metrics[2])
- scores_dict["a1"]["probabilistic"].append(all_metrics[4])
-
- # compute rendered depth metrics
- if self.train_cfg.eval_render_depth:
- render_depth = output_probabilistic.depth
- target_depth_gt = batch["target"]["depth"]
-
- pred_depths = render_depth[0] # [V, H, W]
- depth_gt = target_depth_gt[0] # [V, H, W]
-
- near = batch["target"]["near"][..., None, None][0] # [V, 1, 1]
- far = batch["target"]["far"][..., None, None][0] # [V, 1, 1]
-
- valid = (depth_gt >= near) & (depth_gt <= far)
-
- all_metrics = compute_depth_errors(depth_gt[valid].detach().cpu().numpy(),
- pred_depths[valid].detach().cpu().numpy())
-
- scores_dict["render_abs_rel"]["probabilistic"].append(all_metrics[0])
- scores_dict["render_rmse"]["probabilistic"].append(all_metrics[2])
- scores_dict["render_a1"]["probabilistic"].append(all_metrics[4])
-
- # summarise scores and log to logger
- # Create wandb table for inner iteration visualization
- # For now, log only psnr
- if hasattr(self.logger, 'experiment') and self.logger.experiment is not None:
- # Extract metrics that have step numbers (e.g., "psnr_0", "psnr_1", etc.)
- inner_iteration_data = []
- for score_tag, methods in scores_dict.items():
- # Check if this is a step-specific metric (e.g., "psnr_0", "psnr_1", etc.)
- if '_' in score_tag and score_tag.split('_')[-1].isdigit():
- metric_name, step_str = score_tag.rsplit('_', 1)
- inner_step = int(step_str)
-
- if metric_name not in ["psnr"]:
- continue
+ @staticmethod
+ def _compute_init_depth_metrics(pred_depths, depth_gt, batch) -> dict[str, float]:
+ """abs_rel / rmse / a1 between the initializer's context-view depth prediction and GT.
+ Depth is upsampled to context-image resolution to match GT if the initializer downsampled."""
+ pred_depths = pred_depths[0] # [V, H, W]
+ 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)
+ depth_gt = depth_gt[0] # [V, H, W]
+ near = batch["context"]["near"][..., None, None][0] # [V, 1, 1]
+ far = batch["context"]["far"][..., None, None][0]
+ valid = (depth_gt >= near) & (depth_gt <= far)
+
+ all_metrics = compute_depth_errors(
+ depth_gt[valid].detach().cpu().numpy(),
+ pred_depths[valid].detach().cpu().numpy(),
+ )
+ return {"abs_rel": all_metrics[0], "rmse": all_metrics[2], "a1": all_metrics[4]}
+
+ def _accumulate_init_depth_metrics(self, pred_depths, depth_gt, batch, scores_dict: dict) -> None:
+ """Append init-depth metrics (abs_rel/rmse/a1) for one scene into scores_dict."""
+ for name, val in self._compute_init_depth_metrics(pred_depths, depth_gt, batch).items():
+ scores_dict[name].append(val)
+
+ @staticmethod
+ def _accumulate_render_depth_metrics(target_render, batch, scores_dict: dict) -> None:
+ """render_abs_rel / render_rmse / render_a1 between the final target-render depth and GT.
+
+ The render depth comes from whichever module produced the final target render: the
+ initializer's target render, or the optimizer's last step if an optimizer ran.
+
+ (`pred_depths` is the initializer's context-view prediction)
+ """
+ render_depth = target_render.depth[0] # [V, H, W]
+ depth_gt = batch["target"]["depth"][0] # [V, H, W]
+ near = batch["target"]["near"][..., None, None][0] # [V, 1, 1]
+ far = batch["target"]["far"][..., None, None][0]
+ valid = (depth_gt >= near) & (depth_gt <= far)
+
+ all_metrics = compute_depth_errors(
+ depth_gt[valid].detach().cpu().numpy(),
+ render_depth[valid].detach().cpu().numpy(),
+ )
+ scores_dict["render_abs_rel"].append(all_metrics[0])
+ scores_dict["render_rmse"].append(all_metrics[2])
+ scores_dict["render_a1"].append(all_metrics[4])
+
+ def _log_inner_iteration_table(self, scores_dict: dict) -> None:
+ """Log a per-meta-step PSNR-vs-inner-iteration series under test/psnr/meta_,
+ so wandb can render the optimization trajectory for the current meta step.
+
+ Picks the step-suffixed psnr_ entries out of scores_dict and logs (inner_step, mean psnr).
+ """
+ if not (hasattr(self.logger, 'experiment') and self.logger.experiment is not None):
+ return
+
+ # Step-suffixed metrics look like "psnr_0", "psnr_1", ...; collect (inner_step, mean psnr).
+ series = []
+ for score_tag, scores in scores_dict.items():
+ if '_' not in score_tag or not score_tag.split('_')[-1].isdigit():
+ continue
+ metric_name, step_str = score_tag.rsplit('_', 1)
+ if metric_name != "psnr":
+ continue
+ if scores:
+ series.append((int(step_str), sum(scores) / len(scores)))
+
+ if len(series) <= 1:
+ return
+
+ try:
+ run = self.logger.experiment
+ run.define_metric("inner_iteration")
+ run.define_metric(f"test/psnr/meta_{self.global_step}", step_metric="inner_iteration")
+ for inner_step, value in series:
+ run.log({
+ "inner_iteration": inner_step,
+ f"test/psnr/meta_{self.global_step}": value,
+ })
+ except Exception as e:
+ warn(f"Could not create automatic charts: {e}")
+
+ def _summarize_test_scores(self, scores_dict: dict, time_skip_steps_dict: dict, start_t: float) -> None:
+ """Average each metric across batches, log under test/, log per-tag avg runtimes
+ (trimming the first time_skip_steps_dict[tag] entries as warmup), and print the per-step
+ PSNR summary + total wall-clock."""
+ for score_tag, cur_scores in scores_dict.items():
+ if len(cur_scores) > 0:
+ self.log(f"test/{score_tag}", sum(cur_scores) / len(cur_scores))
- for method_tag, cur_scores in methods.items():
- if len(cur_scores) > 0:
- cur_mean = sum(cur_scores) / len(cur_scores)
- inner_iteration_data.append({
- 'meta_iteration': self.global_step,
- 'inner_iteration': inner_step,
- 'metric_name': metric_name,
- 'method': method_tag,
- 'value': cur_mean
- })
-
- # Log the table if we have inner iteration data
- if inner_iteration_data:
- try:
- # Rewrite the chart (wandb cannot append to the current figure (?))
- df = pd.DataFrame(inner_iteration_data)
- df["meta_iteration_str"] = df["meta_iteration"].astype(str)
- metric_to_plot = "psnr"
- df_metric = df[df["metric_name"] == metric_to_plot]
-
- table = wandb.Table(dataframe=df_metric)
-
- # self.logger.experiment.log({f"{metric_to_plot}_line": wandb.plot.line(
- # table,
- # x="inner_iteration",
- # y="value",
- # title=f"{metric_to_plot} per inner iteration",
- # stroke="meta_iteration_str",
- # )})
-
- # Plot psnr for current meta iteration in a separate chart
- # run = self.logger.experiment
- current_meta = self.global_step
-
- df_current = df_metric[df_metric["meta_iteration"] == current_meta]
-
- if len(df_current) > 1:
- run = self.logger.experiment
- run.define_metric("inner_iteration")
- run.define_metric(f"test/psnr/meta_{current_meta}", step_metric="inner_iteration")
- for _, row in df_current.iterrows():
- run.log({
- "inner_iteration": row["inner_iteration"],
- f"test/psnr/meta_{current_meta}": row["value"],
- })
-
- except Exception as e:
- warn(f"Could not create automatic charts: {e}")
- # Fallback: just log the table
- pass
-
- # Keep the original logging
- for score_tag, methods in scores_dict.items():
- for method_tag, cur_scores in methods.items():
- if len(cur_scores) > 0:
- cur_mean = sum(cur_scores) / len(cur_scores)
- self.log(f"test/{score_tag}", cur_mean)
- # 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")
+ times = times[int(time_skip_steps_dict.get(tag, 0)):] # drop the warmup calls
+ if len(times) == 0:
+ continue
+ print(f"{tag}: {len(times)} calls, avg. {np.mean(times):.1f} ms per call")
self.log(f"test/runtime_avg_{tag}", np.mean(times))
self.benchmarker.clear_history()
overall_eval_time = time.time() - start_t
- psnr_list = [scores_dict[f"psnr_{i}"]["probabilistic"] for i in
- range(self.scene_trainer_cfg.num_update_steps + 1)]
- psnr_list = [sum(pnsr) / len(pnsr) for pnsr in psnr_list if len(pnsr) > 0]
- psnr_str = ", ".join(f"psnr_{i}: {np.mean(pnsr):.3f}" for i, pnsr in enumerate(psnr_list))
- example_num = len(scores_dict['psnr_0']['probabilistic'])
+ # Keep the step index i tied to its psnr_ label; skip steps that recorded no scores.
+ psnr_str = ", ".join(
+ f"psnr_{i}: {sum(scores) / len(scores):.3f}"
+ for i in range(self.scene_trainer_cfg.num_update_steps + 1)
+ if (scores := scores_dict[f"psnr_{i}"])
+ )
+ example_num = len(scores_dict['psnr_0'])
print(f"Eval total time cost: {overall_eval_time:.3f}s, {psnr_str}, example_num: {example_num} ")
self.log("test/runtime_all", overall_eval_time)
@@ -1738,8 +1786,9 @@ class MetaTrainer(LightningModule):
elif isinstance(module, (Optimizer, PostProcessing3DGS)):
nr_gaussians_log = module.nr_gaussians_log
nr_nonzero_grads_log = module.nr_nonzero_grad_log
+ # The one-time GPU warm-up is already zeroed at its source (the first optimized
+ # scene's first iteration, in _run_optimizer), so the timer is read as-is here.
iter_time_log = module.iter_time_log
- iter_time_log[0] = 0.0
else:
raise ValueError(f"Unknown module type: {type(module)}")
@@ -1747,42 +1796,35 @@ class MetaTrainer(LightningModule):
output_dict = self.test_step_outputs_context if input_str == "context" else self.test_step_outputs_target
out_dir.mkdir(parents=True, exist_ok=True)
- for i, step in tqdm(enumerate(iterations), desc=f"Evaluating {input_str}", total=len(iterations)):
+ renders = [renders_list[i].color[0] for i in range(len(iterations))] # each (V, 3, H, W)
+ per_render = self._score_renders(
+ renders, rgb_gt,
+ metrics=self.test_cfg.compute_scores_metrics,
+ iter_batch_size=self.test_cfg.metrics_batch_size,
+ )
+
+ for i, nr_iter in tqdm(enumerate(iterations), desc=f"Evaluating {input_str}", total=len(iterations)):
is_last = (i == len(iterations) - 1)
- nr_iter = iterations[i]
# j: index into per-step logs; last step uses nr_iter-1 because logs are 0-indexed up to nr_iter
j = nr_iter - 1 if is_last else nr_iter
- iter_rgb = renders_list[i].color[0] # (V, 3, H, W)
+ scores = dict(per_render[i])
- scores: dict = compute_rgb_metrics(
- iter_rgb,
- rgb_gt,
- metrics=self.test_cfg.compute_scores_metrics,
- iter_batch_size=self.test_cfg.metrics_batch_size,
- )
-
- if nr_gaussians_log is not None:
+ if nr_gaussians_log:
assert j <= len(nr_gaussians_log), f"{j}, {len(nr_gaussians_log)}"
- scores["gaussians"] = torch.tensor(nr_gaussians_log[j])
+ scores["gaussians"] = nr_gaussians_log[j]
- if nr_nonzero_grads_log is not None and nr_nonzero_grads_log:
+ if nr_nonzero_grads_log:
assert j <= len(nr_nonzero_grads_log), f"{j}, {len(nr_nonzero_grads_log)}"
- scores["nonzero_grads"] = torch.tensor(nr_nonzero_grads_log[j])
+ scores["nonzero_grads"] = nr_nonzero_grads_log[j]
- if iter_time_log is not None:
+ if iter_time_log:
assert j <= len(iter_time_log), f"{j}, {len(iter_time_log)}"
- scores["time"] = torch.tensor(sum(iter_time_log[:j + 1]))
+ scores["time"] = sum(iter_time_log[:j + 1])
- for name, score in scores.items():
- if name == "lpips":
- output_dict[f"{module_name}_alex_lpips"][-1].append(score[0].item())
- output_dict[f"{module_name}_vgg_lpips"][-1].append(score[1].item())
- else:
- output_dict[f"{module_name}_{name}"][-1].append(score.item())
+ for name, val in scores.items():
+ output_dict[f"{module_name}_{name}"][-1].append(val)
output_dict[f"{module_name}_iterations"][-1].append(nr_iter)
- del iter_rgb
-
# Save per-scene metrics to JSON
last_scene_metrics = {key: vals[-1] for key, vals in output_dict.items()}
if extra_scene_metrics:
@@ -1801,8 +1843,8 @@ class MetaTrainer(LightningModule):
output: OptimizerOutput | InitializerOutput,
output_path: Path,
extra_scene_metrics: dict | None = None,
- ) -> dict:
- """Evaluate and save results. Returns collected metrics dict (keyed by module_name_metric)."""
+ ) -> None:
+ """Evaluate and save renders/depths/scores for one module's output (init / optimizer / postprocess)."""
module_name = module.__class__.__name__.lower()
output_path = CustomPath(output_path / module_name)
@@ -1843,6 +1885,16 @@ class MetaTrainer(LightningModule):
else:
raise ValueError(f"Unknown output type: {type(output)}")
+ # Score the initializer's predicted context-view depth against the GT depth (initialization
+ # quality, one value per scene).
+ if isinstance(module, Initializer) and self.test_cfg.compute_scores:
+ pred_depths = output.depths
+ depth_gt = batch["context"].get("depth")
+ if pred_depths is not None and depth_gt is not None and depth_gt.max() > 0:
+ for name, val in self._compute_init_depth_metrics(pred_depths, depth_gt, batch).items():
+ # one row per scene (length-1, like the init's other per-scene metrics)
+ self.test_step_outputs_context[f"{module_name}_{name}"].append([float(val)])
+
input_strs = ["target"]
if self.test_cfg.eval_context_views:
input_strs.insert(0, "context")
@@ -1850,8 +1902,19 @@ class MetaTrainer(LightningModule):
depth_vmin, depth_vmax = self._compute_depth_range(output, input_strs, module)
error_vmax = self._compute_error_vmax(output, input_strs, module, batch)
+ # Save the initializer's predicted input-view (context) depth — distinct from the rendered
+ # depth saved in the loop below. Only the initializer carries a predicted depth.
+ if (self.test_cfg.save_init_pred_depth and isinstance(output, InitializerOutput)
+ and output.depths is not None):
+ ctx_indices = batch["context"]["index"][0]
+ self.test_save_pred_depth(output.depths[0], ctx_indices, output_path, scene_name, "context",
+ vmin=depth_vmin, vmax=depth_vmax)
+
for input_str in input_strs:
- indices = batch[input_str]["index"][0] # (V,) # TODO Naama: bug when using opt bs > 0
+ # Full-V frame indices for filename labelling. Safe to use even when
+ # opt_batch_size < V because Optimizer._save_post_update_renders always
+ # renders the full V views at test time.
+ indices = batch[input_str]["index"][0] # (V,)
renders_list, iterations = self._get_renders_list(output, input_str, module)
if renders_list is None:
continue
@@ -1890,36 +1953,6 @@ class MetaTrainer(LightningModule):
if self.test_cfg.save_gt_depth and depth_gt is not None:
self.test_save_gt_depth(depth_gt, indices, output_path, scene_name, input_str,
vmin=depth_vmin, vmax=depth_vmax)
-
- # save video
- # TODO Naama: reorganize video rendering
- # Note: when video mode is enabled this returns early, skipping score computation.
- if module is not None and self.test_cfg.save_video and isinstance(output, OptimizerOutput):
- # Generate only for the first view in the batch
- # Generate a video with optimization trajectory for the first view (using ffmpeg)
- if input_str == "target":
- if self.test_cfg.save_video_fixed_view:
- self.render_supp_videos(batch, h, input_str, iterations, output.gaussian_list, output_path,
- scene_name, v, w, fixed_view_video=True, video_type="fixed_view")
- if self.test_cfg.save_video_fixed_iteration:
- for t in self.test_cfg.save_video_fixed_iteration_indices:
- self.render_supp_videos(batch, h, input_str, iterations, output.gaussian_list, output_path,
- scene_name, v, w,
- fixed_view_video=self.test_cfg.save_video_fixed_iteration_render_fixed_view,
- # render a fixed view until the required iteration
- fixed_iteration_video=True,
- fixed_iteration_indices=[t],
- video_type="fixed_iteration")
- if self.test_cfg.save_video_combined:
- self.render_supp_videos(batch, h, input_str, iterations, output.gaussian_list, output_path,
- scene_name, v, w,
- fixed_view_video=True,
- fixed_iteration_video=True,
- fixed_iteration_indices=self.test_cfg.save_video_combined_iterations,
- fixed_iteration_length=self.test_cfg.save_video_combined_fixed_iteration_length,
- video_type="combined")
- return
-
# Compute scores
if self.test_cfg.compute_scores:
print("\nComputing scores...")
@@ -1929,137 +1962,58 @@ class MetaTrainer(LightningModule):
extra_scene_metrics=extra_scene_metrics,
)
- # Merge metrics from target (and context if evaluated) for combined plotting
- all_metrics = {}
- if self.test_cfg.compute_scores:
- for key, vals in self.test_step_outputs_target.items():
- if vals:
- all_metrics[key] = vals[-1]
- for key, vals in self.test_step_outputs_context.items():
- if vals:
- all_metrics[key] = vals[-1]
- return all_metrics
+ # Save the optimization videos for the target view (one mp4 per enabled mode).
+ if module is not None and self.test_cfg.save_video and isinstance(output, OptimizerOutput):
+ self._save_optimization_videos(batch, output, module, h, v, w, scene_name, output_path)
+ # region ==================== Save Results Methods =======================
@staticmethod
- def _plot_combined_metrics(
- output_path: Path,
- scene_name: str,
- phases: list[tuple[str, dict]],
- ):
- """Create a combined metrics plot for optimizer + postprocessing per scene.
-
- Args:
- output_path: Root output directory.
- scene_name: Name of the current scene.
- phases: List of (label, metrics_dict) tuples in order. Each metrics_dict
- has keys like "{label}_psnr", "{label}_iterations", etc.
- """
- try:
- MetaTrainer._plot_combined_metrics_impl(output_path, scene_name, phases)
- except Exception as e:
- warn(f"[plot] failed to create combined metrics plot: {e}")
- import traceback
- traceback.print_exc()
+ def test_save_cameras_json(batch: BatchedExample, output_path: Path) -> None:
+ """Save raw extrinsics/intrinsics (cam-to-world, normalized) for context + target as JSON."""
+ scene_name = batch["scene"][0]
+ cameras_dir = output_path / "cameras"
+ cameras_dir.mkdir(parents=True, exist_ok=True)
+ relevant_keys = ["extrinsics", "intrinsics"]
+ cameras_data = {
+ "scene": scene_name,
+ "context": {key: batch["context"][key][0].cpu().tolist() for key in relevant_keys},
+ "target": {key: batch["target"][key][0].cpu().tolist() for key in relevant_keys},
+ "resolution": list(batch["context"]["image"].shape[-2:]),
+ }
+ cameras_path = cameras_dir / f"{scene_name}_cameras.json"
+ with open(cameras_path, "w") as f:
+ json.dump(cameras_data, f, indent=4)
+ print(f"Saved cameras JSON to {cameras_path}")
@staticmethod
- def _plot_combined_metrics_impl(
- output_path: Path,
- scene_name: str,
- phases: list[tuple[str, dict]],
- ):
- from matplotlib import pyplot as plt
-
- # Filter out phases with no data
- phases = [(label, data) for label, data in phases if data]
- if not phases:
- print("[plot] No metrics data available, skipping combined plot.")
- return
-
- plot_metrics = ["psnr", "ssim", "alex_lpips", "vgg_lpips", "gaussians"]
-
- # Build combined series for each metric
- plots = [] # (title, x_values_list, y_values_list, labels_list, divider_x)
- for metric in plot_metrics:
- combined_x = []
- combined_y = []
- combined_labels = []
- x_offset = 0
- divider_x = None
-
- for phase_idx, (label, data) in enumerate(phases):
- iter_key = f"{label}_iterations"
- metric_key = f"{label}_{metric}"
-
- iterations = data.get(iter_key, [])
- values = data.get(metric_key, [])
-
- if not iterations or not values:
- continue
-
- n = min(len(iterations), len(values))
- xs = [x_offset + iterations[j] for j in range(n)]
- ys = values[:n]
-
- combined_x.append(xs)
- combined_y.append(ys)
- combined_labels.append(label)
-
- if phase_idx < len(phases) - 1 and xs:
- divider_x = xs[-1]
- x_offset = divider_x
-
- if combined_x:
- plots.append((metric, combined_x, combined_y, combined_labels, divider_x))
-
- if not plots:
- print("[plot] No plottable metrics found, skipping combined plot.")
- return
-
- fig, axes = plt.subplots(len(plots), 1, figsize=(10, 3.5 * len(plots)), squeeze=False)
- axes = axes[:, 0]
-
- # Metrics where lower is better
- lower_is_better = {"alex_lpips", "vgg_lpips"}
-
- for ax, (metric_name, x_lists, y_lists, labels, divider_x) in zip(axes, plots):
- for xs, ys, label in zip(x_lists, y_lists, labels):
- ax.plot(xs, ys, marker=".", markersize=3, label=label)
- if divider_x is not None:
- ax.axvline(x=divider_x, color="gray", linestyle="--", linewidth=1, alpha=0.7)
-
- # Find and annotate the best value across all phases
- all_ys = [v for ys in y_lists for v in ys]
- if all_ys:
- if metric_name in lower_is_better:
- best_val = min(all_ys)
- else:
- best_val = max(all_ys)
- ax.axhline(y=best_val, color="red", linestyle=":", linewidth=1, alpha=0.6)
- ax.text(
- 1.0, best_val, f" best={best_val:.4f}",
- transform=ax.get_yaxis_transform(),
- va="bottom", ha="right", fontsize=7, color="red",
- )
-
- ax.set_title(metric_name)
- ax.set_xlabel("iteration")
- ax.set_ylabel(metric_name)
- ax.legend(fontsize=8)
- ax.grid(True, alpha=0.3)
-
- fig.suptitle(f"Scene: {scene_name}", fontsize=12, y=1.0)
- fig.tight_layout()
-
- plot_dir = output_path / "plots" / scene_name
- plot_dir.mkdir(parents=True, exist_ok=True)
- save_path = plot_dir / "combined_metrics.png"
- plt.savefig(save_path, dpi=150, bbox_inches="tight")
- plt.close(fig)
- print(f"Saved combined metrics plot to {save_path}")
+ def test_save_cameras_npz(batch: BatchedExample, output_path: Path) -> None:
+ """Save cameras in renderer-ready form: viewmats=inverse(extrinsics) (world-to-cam),
+ Ks=intrinsics * diag(W, H, 1) (pixel-space). Mirrors
+ GSplatDecoderSplattingCUDA.forward (gsplat_decoder_splatting_cuda.py:137-140).
+ """
+ scene_name = batch["scene"][0]
+ cameras_dir = output_path / "cameras"
+ cameras_dir.mkdir(parents=True, exist_ok=True)
+ npz_data = {"scene": scene_name}
+ for input_str in ("context", "target"):
+ view = batch[input_str]
+ extrinsics = view["extrinsics"][0] # [V, 4, 4] cam-to-world
+ intrinsics = view["intrinsics"][0] # [V, 3, 3] normalized
+ h, w = view["image"].shape[-2:]
+ viewmats = extrinsics.inverse() # [V, 4, 4] world-to-cam
+ scale = intrinsics.new_tensor([[w], [h], [1]])
+ Ks = intrinsics * scale # [V, 3, 3] pixel-space
+ npz_data[f"{input_str}_viewmats"] = viewmats.cpu().numpy()
+ npz_data[f"{input_str}_Ks"] = Ks.cpu().numpy()
+ npz_data[f"{input_str}_image_shape"] = np.array([h, w], dtype=np.int64)
+ cameras_npz_path = cameras_dir / f"{scene_name}_cameras.npz"
+ np.savez(cameras_npz_path, **npz_data)
+ print(f"Saved renderer-ready cameras NPZ to {cameras_npz_path}")
- # region ==================== Save Results Methods =======================
@staticmethod
def test_save_rendered_images(renders_list: list, indices, output_path, scene_name, input_str):
+ """Save the per-view optimization trajectory as /images//color_/.png,
+ with iterations concatenated along width. Also dumps the last iteration separately."""
out_dir = output_path / "images" / scene_name / f"color_{input_str}"
for i, index in tqdm(enumerate(indices), desc=f"Saving {input_str} images"):
color = []
@@ -2075,6 +2029,7 @@ class MetaTrainer(LightningModule):
@staticmethod
def test_save_last_rendered_images(renders_list: list, indices, output_path, scene_name, input_str):
+ """Save only the final-iteration render per view at /images//last/color_/.png."""
out_dir = output_path / "images" / scene_name / "last" / f"color_{input_str}"
out_dir.mkdir(parents=True, exist_ok=True)
for i, index in tqdm(enumerate(indices), desc=f"Saving {input_str} last images"):
@@ -2087,7 +2042,9 @@ class MetaTrainer(LightningModule):
@staticmethod
def test_save_gt_images(rgb_gt, indices, output_path, scene_name, input_str):
- out_dir = output_path / "images" / scene_name / f"color_{input_str}"
+ """Save GT images alongside the last render at /images//last/color_/_gt.png."""
+ out_dir = output_path / "images" / scene_name / "last" / f"color_{input_str}"
+ out_dir.mkdir(parents=True, exist_ok=True)
for index, gt in tqdm(zip(indices, rgb_gt), desc=f"Saving {input_str} GT images"):
save_image(gt, out_dir / f"{index:06d}_gt.png")
@@ -2095,32 +2052,55 @@ class MetaTrainer(LightningModule):
def test_save_rendered_depth(renders_list: list, indices, output_path, scene_name, input_str,
vmin: float = 0.0,
vmax: float = 1.0):
- out_dir = output_path / "images" / scene_name / f"depth_{input_str}"
- for i, index in tqdm(enumerate(indices), desc=f"Saving {input_str} depths"):
+ """Save the rendered-depth trajectory (depth rasterized from the Gaussians) as colormapped PNGs at
+ /images//rendered_depth_/.png, iterations concatenated along width.
+ This is the rendered depth, not the initializer's predicted depth. vmin/vmax control the per-scene
+ depth normalization."""
+ out_dir = output_path / "images" / scene_name / f"rendered_depth_{input_str}"
+ for i, index in tqdm(enumerate(indices), desc=f"Saving {input_str} rendered depths"):
depth = []
for iter_renders in renders_list:
- iter_depths = iter_renders.depth # (1, V, 3, H, W)
+ iter_depths = iter_renders.depth # (1, V, H, W)
assert iter_depths is not None, "Depths not found in renders."
- iter_depths = iter_depths[0] # (V, 3, H, W)
- depth.append(iter_depths[i])
+ iter_depths = iter_depths[0] # (V, H, W)
+ depth.append(iter_depths[i]) # (H, W)
depth = torch.cat(depth, dim=-1) # concat along width
color = viz_depth_tensor(depth, return_numpy=False, as_uint8=False, vmin=vmin, vmax=vmax)
save_image(color, out_dir / f"{index:06d}.png")
del iter_depths
del color
+ @staticmethod
+ def test_save_pred_depth(pred_depths, indices, output_path, scene_name, input_str,
+ vmin: float = 0.0, vmax: float = 1.0):
+ """Save the initializer's predicted input-view depth (the encoder/monodepth estimate used to place
+ Gaussians) as colormapped PNGs at /images//pred_depth_/.png.
+ Distinct from the rendered depth above; uses the same vmin/vmax for comparability."""
+ out_dir = output_path / "images" / scene_name / f"pred_depth_{input_str}"
+ out_dir.mkdir(parents=True, exist_ok=True)
+ for i, index in tqdm(enumerate(indices), desc=f"Saving {input_str} pred depths"):
+ color = viz_depth_tensor(pred_depths[i], return_numpy=False, as_uint8=False, vmin=vmin, vmax=vmax)
+ save_image(color, out_dir / f"{index:06d}.png")
+ del color
+
@staticmethod
def test_save_gt_depth(depth_gt, indices, output_path, scene_name, input_str, vmin: float = 0.0,
vmax: float = 1.0):
- out_dir = output_path / "images" / scene_name / f"depth_{input_str}"
+ """Save GT depth colormapped under /images//rendered_depth_/_gt.png
+ using the same vmin/vmax as the rendered depths."""
+ out_dir = output_path / "images" / scene_name / f"rendered_depth_{input_str}"
+ depth_gt = depth_gt[0] # (B, V, H, W) -> (V, H, W); iterate views, not the batch dim
for index, gt in tqdm(zip(indices, depth_gt), desc=f"Saving {input_str} GT depths"):
- color = viz_depth_tensor(gt, return_numpy=False, as_uint8=False, vmin=vmin, vmax=vmax)
+ color = viz_depth_tensor(gt, return_numpy=False, as_uint8=False, vmin=vmin, vmax=vmax) # gt is (H, W)
save_image(color, out_dir / f"{index:06d}_gt.png")
@staticmethod
def test_save_rendered_errors(renders_list: list, rgb_gt, indices, output_path, scene_name, input_str,
vmin: float = 0.0,
vmax: float = 1.0):
+ """Save per-view RGB-error magma maps at /images//error_/.png,
+ iterations concatenated along width. vmin/vmax come from the per-scene 99th-percentile error
+ computed in `_compute_error_vmax` for consistent visualization."""
# rgb_gt is (V, 3, H, W). Looping per view (outer) and per iteration (inner)
# keeps only one view's iterations in memory at a time.
out_dir = output_path / "images" / scene_name / f"error_{input_str}"
@@ -2135,64 +2115,11 @@ class MetaTrainer(LightningModule):
save_image(color, out_dir / f"{index:06d}.png")
del iter_rgbs, error_maps, error_map
- def save_colmap_test_train_views(self, batch, h, w):
- # load the distortion parameters from the original colmap data
- assert self.test_cfg.ori_colmap_data_path is not None
- (scene_name,) = batch["scene"]
- output_path = self.test_cfg.output_path
- # training views
- input_images = batch["context"]["image"][0] # [V, 3, H, W]
- index = batch["context"]["index"][0]
- for idx, color in zip(index, input_images):
- # NOTE: the original image id starts from 1
- save_image(color, output_path / scene_name / "images_train" / f"frame_{idx + 1:05d}.png")
- # testing views
- target_images = batch["target"]["image"][0] # [V, 3, H, W]
- index = batch["target"]["index"][0]
- for idx, color in zip(index, target_images):
- # NOTE: the original image id starts from 1
- save_image(color, output_path / scene_name / "images_test" / f"frame_{idx + 1:05d}.png")
- # save the camera intrinsics
- intrinsics = batch["context"]["intrinsics"][0][0].clone() # [3, 3]
- # need to rescale to the image size
- intrinsics[0, :] *= w
- intrinsics[1, :] *= h
- # distortion parameters
- json_path = os.path.join(self.test_cfg.ori_colmap_data_path, scene_name, "nerfstudio", "transforms.json")
- assert os.path.exists(json_path), f"Cannot find {json_path}"
- sparse_save_dir = output_path / scene_name / "sparse" / "0"
- sparse_save_dir_train = output_path / scene_name / "sparse_train" / "0"
- # save to cameras.bin
- save_opencv_camera(intrinsics.cpu().numpy(), json_path, sparse_save_dir, image_size=(w, h))
- save_opencv_camera(intrinsics.cpu().numpy(), json_path, sparse_save_dir_train, image_size=(w, h))
- # extract extrinsics from the dense view images.bin
- dense_view_extrinsics = os.path.join(self.test_cfg.ori_colmap_data_path, scene_name,
- "nerfstudio/colmap/sparse/0")
- selected_train_ids = [idx + 1 for idx in batch["context"]["index"][0].tolist()]
- selected_test_ids = [idx + 1 for idx in batch["target"]["index"][0].tolist()]
- selected_ids = selected_train_ids + selected_test_ids
- # also save the sparse features and points3D
- extract_sparse_images_bin(dense_view_extrinsics, sparse_save_dir, selected_ids, keep_features=False)
- # only for training views: reconstruct the sparse point cloud only from training views
- extract_sparse_images_bin(dense_view_extrinsics, sparse_save_dir_train, selected_train_ids)
- return
-
- def compute_depth_scores(self, batch, depth_gt, init_pred_depths):
- pred_depths = init_pred_depths[0] # [V, H, W]
- depth_gt = depth_gt[0] # [V, H, W]
- near = batch["context"]["near"][...,
- None, None][0] # [V, 1, 1]
- far = batch["context"]["far"][..., None, None][0] # [V, 1, 1]
- valid = (depth_gt >= near) & (depth_gt <= far)
- all_metrics = compute_depth_errors(depth_gt[valid].detach().cpu().numpy(),
- pred_depths[valid].detach().cpu().numpy())
- print(all_metrics)
- self.test_step_outputs_target[f"abs_rel"].append(
- float(all_metrics[0]))
- self.test_step_outputs_target[f"rmse"].append(float(all_metrics[2]))
- self.test_step_outputs_target[f"a1"].append(float(all_metrics[4]))
-
def init_output_dict_for_new_scene(self, input_str, tag=None):
+ """Append an empty per-iteration sublist to test_step_outputs_ for every metric we
+ will collect (psnr, ssim, lpips×2, iterations, time, gaussians, nonzero_grads). Called once at
+ the start of each scene's evaluation; `tag` (e.g. module class name) prefixes each key so
+ init / optimizer / postprocess buckets don't collide."""
tag = "" if tag is None else f"{tag}_"
if input_str == "target":
@@ -2217,67 +2144,148 @@ class MetaTrainer(LightningModule):
# endregion
# region ==================== Video Rendering Methods ====================
- def render_supp_videos(self, batch, h, input_str, all_iterations, gaussian_list, output_path, scene_name,
- v, w, fixed_view_video=False, fixed_iteration_video=False,
- fixed_iteration_indices=None,
- fixed_iteration_length=-1,
- video_type=None):
- out_dir = output_path / "supp_videos" / scene_name
- out_dir.mkdir(parents=True, exist_ok=True)
- combined_iterations = []
+ def _save_optimization_videos(
+ self, batch, output, module, h, v, w, scene_name, output_path,
+ ) -> None:
+ """Render and save the enabled optimization videos for the target view.
+
+ Each enabled mode writes one mp4 (plus a matching per-frame iteration-label JSON) to
+ videos/:
+ - save_video_optim: the chosen view (save_video_view_index) rendered at every save checkpoint.
+ - save_video_orbit: at each step in save_video_orbit_steps, a full camera orbit around the
+ scene, optionally preceded by the per-checkpoint frames (save_video_orbit_with_optim).
+ - save_video_optim_orbit: the per-checkpoint frames with a short orbit spliced in at
+ save_video_optim_orbit_steps (spanning view -> view + save_video_orbit_span).
+ """
+ _, iterations = self._get_renders_list(output, "target", module)
+ gaussian_list = output.gaussian_list
+
+ if self.test_cfg.save_video_optim:
+ self.render_optim_video(
+ batch, h, v, w, "target", iterations, gaussian_list, scene_name, output_path,
+ )
+ if self.test_cfg.save_video_orbit:
+ for t in self.test_cfg.save_video_orbit_steps:
+ self.render_orbit_video(
+ batch, h, v, w, "target", iterations, gaussian_list, scene_name, output_path,
+ orbit_step=t,
+ with_optim=self.test_cfg.save_video_orbit_with_optim,
+ )
+ if self.test_cfg.save_video_optim_orbit:
+ self.render_optim_orbit_video(
+ batch, h, v, w, "target", iterations, gaussian_list, scene_name, output_path,
+ orbit_steps=self.test_cfg.save_video_optim_orbit_steps,
+ orbit_span=self.test_cfg.save_video_orbit_span,
+ )
- view = self.test_cfg.save_video_fixed_view_index
+ def render_optim_video(
+ self, batch, h, v, w, input_str, all_iterations, gaussian_list, scene_name, output_path,
+ ) -> None:
+ """For each save-checkpoint, render only `save_video_view_index` and append it
+ (duplicated `save_video_frame_repeat` times so the optimization trajectory at
+ that single view stands out)."""
+ view = self.test_cfg.save_video_view_index
+ duplicate = self.test_cfg.save_video_frame_repeat
- all_frames = []
+ all_frames, combined_iterations = [], []
+ for i, t in enumerate(all_iterations):
+ decoder_output = self.test_render_videos_views(
+ batch, gaussian_list[i], h, v, w, input_str, start=view, end=view + 1,
+ )
+ frame = decoder_output.color[0].detach().cpu() # (1, 3, H, W)
+ assert frame.shape[0] == 1, f"{frame.shape}"
+ all_frames += [frame[0]] * duplicate # (3, H, W)
+ combined_iterations.extend([t] * duplicate)
- duplicate = self.test_cfg.save_video_fixed_view_duplicate # to focus on the optimization steps
+ self._save_video_file(input_str, "optim", all_frames, combined_iterations,
+ scene_name, output_path)
- if fixed_iteration_length > 0:
- start = view
- end = view + fixed_iteration_length
- else:
- start = None
- end = None
+ def render_orbit_video(
+ self, batch, h, v, w, input_str, all_iterations, gaussian_list, scene_name, output_path,
+ orbit_step: int, with_optim: bool,
+ ) -> None:
+ """At one chosen iteration, orbit the camera around the scene; optionally also append
+ the per-iteration fixed-view frames leading up to it."""
+ view = self.test_cfg.save_video_view_index
+ duplicate = self.test_cfg.save_video_frame_repeat
+ all_frames, combined_iterations = [], []
for i, t in enumerate(all_iterations):
- # Render only the view
- if fixed_view_video:
- decoder_output = self.test_render_videos_views(batch, gaussian_list[i], h, v, w, input_str,
- start=view, end=view + 1)
- frames_t = decoder_output.color[0].detach().cpu() # (1, 3, H, W)
- assert frames_t.shape[0] == 1, f"{frames_t.shape}"
- all_frames += [frames_t[0]] * duplicate # (3, H, W)
+ if with_optim:
+ decoder_output = self.test_render_videos_views(
+ batch, gaussian_list[i], h, v, w, input_str, start=view, end=view + 1,
+ )
+ frame = decoder_output.color[0].detach().cpu()
+ all_frames += [frame[0]] * duplicate
combined_iterations.extend([t] * duplicate)
- if fixed_iteration_video:
- if t in fixed_iteration_indices:
- # Render a trajectory around the scene
- decoder_output = self.test_render_videos_views(batch, gaussian_list[i], h, v, w, input_str,
- start=start, end=end)
- frames_t = decoder_output.color[0].detach().cpu() # (num_frames, 3, H, W)
- for i in range(3): # forward and backward
- for frame in frames_t:
- all_frames += [frame] * 3
- for frame in frames_t.flip(0):
- all_frames += [frame] * 3
- combined_iterations.extend(['orbit'] * frames_t.shape[0] * 2)
- if fixed_iteration_video and not fixed_view_video:
- break # no need to continue
-
- if video_type == "combined":
- save_str = f"_combined_{view}"
- elif video_type == "fixed_view":
- save_str = f"_fixed_view_{view}"
- elif video_type == "fixed_iteration":
- assert len(fixed_iteration_indices) == 1, f"{fixed_iteration_indices}"
- save_str = f"_fixed_iteration_{fixed_iteration_indices[0]}"
- else:
- raise ValueError
+ if t == orbit_step:
+ decoder_output = self.test_render_videos_views(
+ batch, gaussian_list[i], h, v, w, input_str, start=None, end=None,
+ )
+ self._collect_orbit_frames(decoder_output, all_frames, combined_iterations)
+ if not with_optim:
+ break # nothing more to render
+
+ self._save_video_file(input_str, f"orbit_{orbit_step}",
+ all_frames, combined_iterations, scene_name, output_path)
+
+ def render_optim_orbit_video(
+ self, batch, h, v, w, input_str, all_iterations, gaussian_list, scene_name, output_path,
+ orbit_steps, orbit_span: int,
+ ) -> None:
+ """Render the per-iteration fixed-view trajectory; at the given iterations also splice
+ in a short orbit (rendered from view → view+orbit_span)."""
+ view = self.test_cfg.save_video_view_index
+ duplicate = self.test_cfg.save_video_frame_repeat
+ orbit_start = view if orbit_span > 0 else None
+ orbit_end = view + orbit_span if orbit_span > 0 else None
+
+ all_frames, combined_iterations = [], []
+ for i, t in enumerate(all_iterations):
+ decoder_output = self.test_render_videos_views(
+ batch, gaussian_list[i], h, v, w, input_str, start=view, end=view + 1,
+ )
+ frame = decoder_output.color[0].detach().cpu()
+ all_frames += [frame[0]] * duplicate
+ combined_iterations.extend([t] * duplicate)
+
+ if t in orbit_steps:
+ decoder_output = self.test_render_videos_views(
+ batch, gaussian_list[i], h, v, w, input_str, start=orbit_start, end=orbit_end,
+ )
+ self._collect_orbit_frames(decoder_output, all_frames, combined_iterations)
+
+ self._save_video_file(input_str, "optim_orbit", all_frames, combined_iterations,
+ scene_name, output_path)
+
+ @staticmethod
+ def _collect_orbit_frames(decoder_output, all_frames: list, combined_iterations: list) -> None:
+ """Append a forward+backward camera-orbit sweep (3 repeats, 3x temporal duplication per frame)."""
+ frames_t = decoder_output.color[0].detach().cpu() # (num_frames, 3, H, W)
+ for _ in range(3): # repeat forward + backward 3 times
+ for frame in frames_t:
+ all_frames += [frame] * 3
+ for frame in frames_t.flip(0):
+ all_frames += [frame] * 3
+ combined_iterations.extend(['orbit'] * frames_t.shape[0] * 2)
+
+ @staticmethod
+ def _save_video_file(input_str: str, save_str: str, all_frames: list, combined_iterations: list,
+ scene_name: str, output_path: Path) -> None:
+ """Write the mp4 and the matching per-frame iteration-label JSON to videos/."""
+ out_dir = output_path / "videos" / scene_name
+ out_dir.mkdir(parents=True, exist_ok=True)
save_video(all_frames, out_dir / f"{input_str}_{save_str}.mp4")
- with (open(out_dir / f"{input_str}_{save_str}_iterations.json", 'w')) as f:
+ with open(out_dir / f"{input_str}_{save_str}_iterations.json", 'w') as f:
json.dump(combined_iterations, f, indent=4)
def test_render_videos_views(self, batch, gaussians, h, v, w, input_str="target", poses=None, start=None, end=None):
+ """Render the Gaussians along a camera path and return the decoder output.
+
+ The path is the input_str views' poses (optionally camera-stabilized) or the explicit `poses`;
+ rendering is split into chunks of test_cfg.render_chunk_size when that is set.
+ """
gaussians = gaussians.to(batch["target"]["image"].device)
with self.benchmarker.time("decoder", num_calls=v):
if poses is None:
@@ -2333,33 +2341,9 @@ class MetaTrainer(LightningModule):
end=end)
return output
- @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:
+ """Log a video interpolating the camera between the two context views (or context->target when there is one context view)."""
_, v, _, _ = batch["context"]["extrinsics"].shape
def trajectory_fn(t):
@@ -2387,6 +2371,7 @@ class MetaTrainer(LightningModule):
@rank_zero_only
def render_video_interpolation_exaggerated(self, batch: BatchedExample) -> None:
+ """Log a context-view interpolation video with an added exaggerated wobble transform along the path (needs 2 context views)."""
# Two views are needed to get the wobble radius.
_, v, _, _ = batch["context"]["extrinsics"].shape
if v != 2:
@@ -2441,14 +2426,11 @@ class MetaTrainer(LightningModule):
smooth: bool = True,
loop_reverse: bool = True,
) -> None:
+ """Render num_frames of the init Gaussians along the trajectory_fn camera path and log an
+ RGB+depth video to wandb under video/{name}."""
if self.train_cfg.no_log_video:
return
- # Render probabilistic estimate of scene.
- gaussians_prob = self.encoder(batch["context"], self.global_step, False)
- # gaussians_det = self.encoder(batch["context"], self.global_step, True)
-
- if isinstance(gaussians_prob, dict):
- gaussians_prob = gaussians_prob["gaussians"]
+ gaussians = self.get_init_gaussians(batch, is_training=False).gaussians
t = torch.linspace(0, 1, num_frames, dtype=torch.float32, device=self.device)
if smooth:
@@ -2468,12 +2450,12 @@ class MetaTrainer(LightningModule):
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_prob = self.scene_decoder.forward(
- gaussians_prob, extrinsics, intrinsics, near, far, (h, w), "depth"
+ output = self.scene_decoder.forward(
+ gaussians, extrinsics, intrinsics, near, far, (h, w), "depth"
)
rgb_pred = [
vcat(rgb, depth)
- for rgb, depth in zip(output_prob.color[0], depth_map(output_prob.depth[0]))
+ for rgb, depth in zip(output.color[0], depth_map(output.depth[0]))
]
images = [
diff --git a/optgs/meta_trainer/meta_trainer_cfg.py b/optgs/meta_trainer/meta_trainer_cfg.py
new file mode 100644
index 0000000000000000000000000000000000000000..1dac9b168b04496e76b16b0f25166ead776477de
--- /dev/null
+++ b/optgs/meta_trainer/meta_trainer_cfg.py
@@ -0,0 +1,116 @@
+from dataclasses import dataclass
+from pathlib import Path
+
+from optgs.meta_trainer.ckpt_buffer import CkptBufferCfg
+from optgs.scene_trainer.postprocessing import PostProcessCfg
+from optgs.model.decoder.decoder import DepthRenderingMode
+
+
+@dataclass
+class MetaOptimizerCfg:
+ lr: float
+ warm_up_steps: int
+ lr_monodepth: float
+ lr_depth: float
+ weight_decay: float
+ warm_up_ratio: float
+ adamw_8bit: bool
+
+
+@dataclass
+class TestCfg:
+ output_path: Path | None
+ compute_scores: bool
+ compute_scores_metrics: list[str] | None
+ metrics_batch_size: int
+ eval_initialization: bool
+ save_render_image: bool
+ save_render_image_last_only: bool
+ save_gt_image: bool
+ save_render_depth: bool
+ save_gt_depth: bool
+ save_init_pred_depth: bool # save the initializer's predicted input-view depth (not the rendered depth)
+ save_error_image: bool
+ save_video: bool
+ save_video_optim: bool
+ save_video_view_index: int
+ save_video_frame_repeat: int
+ save_video_orbit: bool
+ save_video_orbit_steps: list | None
+ save_video_orbit_with_optim: bool
+ save_video_optim_orbit: bool
+ save_video_optim_orbit_steps: list | None
+ save_video_orbit_span: int
+ eval_time_skip_steps: int
+ save_gaussian: bool
+ save_poses: bool
+ save_cameras_json: bool
+ save_cameras_npz: bool
+ render_chunk_size: int | None
+ stablize_camera: bool
+ stab_camera_kernel: int
+ eval_context_views: bool
+ inference_window_size: int | None
+ profile_model: bool
+ save_colmap_train_test_views: bool
+ ori_colmap_data_path: str | None
+ postprocessing: PostProcessCfg | None
+ save_at_iters: list[int] | None
+ save_every_freq: list[int] | None
+ save_every_steps: list[int] | None
+ skip_if_outputs_exist: bool
+ scenes_filter: list[str] | None
+
+ experimental_add_noise_to_images: bool
+ experimental_add_noise_to_images_std: float | int | None
+
+
+@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
+ intermediate_loss_weight: float
+ no_viz_video: bool
+ eval_depth: bool
+ no_log_projections: bool
+
+ depth_loss_weight: float
+ log_depth_loss: bool
+ depth_smooth_loss_weight: float
+ depth_teacher_loss_weight: float
+ viz_depth_teacher: bool
+ depth_smooth_loss_nonorm: bool
+ depth_smooth_loss_weight_nvs: float # for novel views
+ monodepth_loss_weight: float # for monocular depth loss
+
+ eval_render_depth: bool
+ render_depth_loss_weight: float
+ viz_render_depth: bool
+ viz_depth_separate: bool
+
+ use_gt_depth_range: bool
+
+ no_log_video: bool
+
+ # when doing refinement, supervise input view or not since we also render input views
+ loss_on_target_views: bool
+ loss_on_target_views_num: int
+ loss_on_input_views: bool
+ loss_on_input_views_num: int
+
+ # local window training
+ train_window_size: int | None
+
+ # Ckpt buffer
+ use_ckpt_buffer: bool
+ ckpt_buffer_cfg: CkptBufferCfg | None
+
+ # L2 weight decay regularization on Gaussian properties (meta-loss)
+ scale_l2_loss_weight: float
+ sh_l2_loss_weight: float
+ opacity_l2_loss_weight: float
\ No newline at end of file
diff --git a/optgs/misc/batchify.py b/optgs/misc/batchify.py
index f1c26dc23cddb4b79283244273afb6900485bfa5..c4f3f87edf465a000441f8ea92293da936c32f12 100644
--- a/optgs/misc/batchify.py
+++ b/optgs/misc/batchify.py
@@ -1,23 +1,6 @@
import torch
-def split_to_minibatch(batch_split, iter_context_idxs):
- minibatch = {
- "image": batch_split["image"][0][iter_context_idxs].unsqueeze(
- 0
- ), # [1, Vc', 3, Hc, Wc]
- "extrinsics": batch_split["extrinsics"][0][iter_context_idxs].unsqueeze(
- 0
- ), # [1, Vc', 4, 4]
- "intrinsics": batch_split["intrinsics"][0][iter_context_idxs].unsqueeze(
- 0
- ), # [1, Vc', 4, 4]
- "near": batch_split["near"][0][iter_context_idxs].unsqueeze(0), # [1, Vc']
- "far": batch_split["far"][0][iter_context_idxs].unsqueeze(0), # [1, Vc']
- }
- return minibatch
-
-
def batched_select(data, indices):
"""
Select data[i, indices[i]] for each batch element i.
diff --git a/optgs/misc/benchmarker.py b/optgs/misc/benchmarker.py
index 644347a28b809c770b6677523191767ca0c24405..289ec0cb03682685070db4b9c3135055730b386e 100644
--- a/optgs/misc/benchmarker.py
+++ b/optgs/misc/benchmarker.py
@@ -2,34 +2,77 @@ import json
from collections import defaultdict
from contextlib import contextmanager
from pathlib import Path
-from time import time
+from time import perf_counter
import numpy as np
import torch
class Benchmarker:
+ """Collects named timing measurements (in milliseconds) and peak memory.
+
+ `time(...)` measures a code block. By default it uses CUDA events, which give real GPU time:
+ recording an event is asynchronous, so timing many blocks does not stall the GPU — the single
+ `torch.cuda.synchronize()` needed to read the elapsed times is deferred until the results are
+ read (`dump`/`summarize`/`execution_times`). It falls back to wall-clock when CUDA is
+ unavailable. `record(...)` stores a pre-measured millisecond value (e.g. the optimizer's own
+ per-iteration CUDA-event timings). All stored values are milliseconds.
+ """
+
def __init__(self):
- self.execution_times = defaultdict(list)
+ self._execution_times = defaultdict(list)
+ # CUDA-event measurements awaiting a synchronize(): (tag, start_event, end_event, num_calls).
+ self._pending = []
+
+ @property
+ def execution_times(self):
+ """The recorded times per tag (ms). Resolves any pending CUDA-event measurements first."""
+ self._flush()
+ return self._execution_times
+
+ def _flush(self) -> None:
+ """Read all pending CUDA-event timings with a single synchronize()."""
+ if not self._pending:
+ return
+ torch.cuda.synchronize()
+ for tag, start_event, end_event, num_calls in self._pending:
+ ms = start_event.elapsed_time(end_event) # CUDA events measure milliseconds
+ for _ in range(num_calls):
+ self._execution_times[tag].append(ms / num_calls)
+ self._pending = []
@contextmanager
- def time(self, tag: str, num_calls: int = 1):
- try:
- start_time = time()
+ def time(self, tag: str, num_calls: int = 1, disable: bool = False, cuda: bool = True):
+ if disable:
yield
- finally:
- end_time = time()
- for _ in range(num_calls):
- self.execution_times[tag].append((end_time - start_time) / num_calls)
+ return
+ if cuda and torch.cuda.is_available():
+ start_event = torch.cuda.Event(enable_timing=True)
+ end_event = torch.cuda.Event(enable_timing=True)
+ start_event.record()
+ try:
+ yield
+ finally:
+ # record() is async; the elapsed time is read later in _flush() after one sync.
+ end_event.record()
+ self._pending.append((tag, start_event, end_event, num_calls))
+ else:
+ start = perf_counter()
+ try:
+ yield
+ finally:
+ ms = (perf_counter() - start) * 1000.0
+ for _ in range(num_calls):
+ self._execution_times[tag].append(ms / num_calls)
- def record(self, tag: str, elapsed_ms: float) -> None:
- """Record a pre-measured elapsed time (in milliseconds) under the given tag."""
- self.execution_times[tag].append(elapsed_ms)
+ def record(self, tag: str, value: float | int) -> None:
+ """Record a pre-measured value under the given tag: time, memory, per-iteration count (e.g. number of gaussians)."""
+ self._execution_times[tag].append(value)
def merge(self, other: "Benchmarker") -> None:
"""Merge another benchmarker's recorded times into this one."""
- for tag, times in other.execution_times.items():
- self.execution_times[tag].extend(times)
+ for tag, times in other.execution_times.items(): # property flushes `other`'s pending events
+ self._execution_times[tag].extend(times)
def dump(self, path: Path) -> None:
path.parent.mkdir(exist_ok=True, parents=True)
@@ -46,4 +89,5 @@ class Benchmarker:
print(f"{tag}: {len(times)} calls, avg {np.mean(times):.1f} ms/call, total {sum(times)/1000:.1f} s")
def clear_history(self) -> None:
- self.execution_times = defaultdict(list)
+ self._execution_times = defaultdict(list)
+ self._pending = []
diff --git a/optgs/misc/checkpointing.py b/optgs/misc/checkpointing.py
index 8f2dc9ea7ff80b16be57acce5c28ea2f2855d17b..e9149c1cd8c4a0e45cfb8c879cf270792e794706 100644
--- a/optgs/misc/checkpointing.py
+++ b/optgs/misc/checkpointing.py
@@ -7,6 +7,35 @@ import torch
from optgs.misc.io import cyan
+# TODO (release): remove ckpt migration
+# Optimizer module attributes renamed in the update_/refine_ cleanup pass.
+# Maps every old spelling (including the resplat-era "render_error_mv_attn") to the
+# current name. Applied per dot-delimited key segment so it works whether checkpoint
+# keys are relative to the optimizer ("update_head.0.weight") or prefixed
+# ("optimizer.update_head.0.weight"). Idempotent: current names never match an old
+# segment, so re-running on an already-migrated checkpoint is a no-op.
+_OPTIMIZER_ATTR_RENAMES = {
+ "update_proj": "state_proj",
+ "update_feature": "error_feature_extractor",
+ "update_rgb_error_proj": "rgb_error_proj",
+ "update_input_norm": "input_norm",
+ "update_module": "point_transformer",
+ "update_head_list": "extra_gaussian_heads",
+ "update_head": "delta_head",
+ "update_error_attn": "error_mv_attn",
+ "render_error_mv_attn": "error_mv_attn", # resplat-era name
+}
+
+
+def _rename_optimizer_attrs(state_dict):
+ """Rewrite checkpoint keys for renamed optimizer module attributes (see map above)."""
+ renamed = {}
+ for k, v in state_dict.items():
+ parts = [_OPTIMIZER_ATTR_RENAMES.get(p, p) for p in k.split(".")]
+ renamed[".".join(parts)] = v
+ return renamed
+
+
# Function to extract the step number from the filename
def extract_step(file_name):
step_str = file_name.split("-")[1].split("_")[1].replace(".ckpt", "")
@@ -41,9 +70,6 @@ def load_partial_state_dict(model, pretrained_state_dict):
k: v for k, v in pretrained_state_dict.items()
if k in model_state_dict and v.shape == model_state_dict[k].shape
}
- # for key in model_state_dict:
- # if key not in filtered_state_dict:
- # print(key)
model_state_dict.update(filtered_state_dict)
model.load_state_dict(model_state_dict)
@@ -71,24 +97,13 @@ def load_optimizer(cfg, scene_trainer, strict_load):
# Resplat repo format: keys are encoder.* (before init/opt split).
# Strip encoder. prefix; init-related keys will be ignored via strict=False.
optimizer_state_dict = {k[len("encoder."):]: v for k, v in pretrained_model.items() if k.startswith("encoder.")}
- # Rename module attributes that changed when the encoder was split.
- _ORIG_OPTIMIZER_ATTR_RENAMES = {
- "render_error_mv_attn": "update_error_attn",
- }
- renamed = {}
- for k, v in optimizer_state_dict.items():
- for old, new in _ORIG_OPTIMIZER_ATTR_RENAMES.items():
- if k == old or k.startswith(old + "."):
- k = new + k[len(old):]
- break
- renamed[k] = v
- optimizer_state_dict = renamed
- # If init_state_wo_features is True, remove all feature-related parameters from the optimizer state dict
- print(cfg.scene_trainer.scene_optimizer.init_state_wo_features)
+ # Rename module attributes that were renamed in the update_/refine_ cleanup pass
+ optimizer_state_dict = _rename_optimizer_attrs(optimizer_state_dict)
+ # If init_state_wo_features is True, remove all feature-related parameters from the optimizer state dict
if getattr(cfg.scene_trainer.scene_optimizer, "init_state_wo_features", False):
- optimizer_state_dict = {k: v for k, v in optimizer_state_dict.items() if "update_proj" not in k}
+ optimizer_state_dict = {k: v for k, v in optimizer_state_dict.items() if "state_proj" not in k}
scene_trainer.optimizer.load_state_dict(optimizer_state_dict, strict=strict_load)
print(cyan(f"Loaded pretrained optimizer: {cfg.checkpointing.pretrained_optimizer}"))
@@ -116,6 +131,8 @@ def load_full_model(cfg, scene_trainer, strict_load):
pretrained_model = torch.load(cfg.checkpointing.pretrained_model, map_location='cpu')
if 'state_dict' in pretrained_model:
pretrained_model = pretrained_model['state_dict']
+ # Rename optimizer module attributes renamed in the update_/refine_ cleanup pass.
+ pretrained_model = _rename_optimizer_attrs(pretrained_model)
if cfg.checkpointing.partial_load:
print('partial load')
load_partial_state_dict(scene_trainer, pretrained_model)
@@ -231,7 +248,7 @@ def apply_freezes(cfg, scene_trainer):
print('train refine only, freezing scene optimizer')
for name, params in scene_trainer.optimizer.named_parameters():
params.requires_grad = False
- if cfg.scene_trainer.scene_optimizer.train_global_update_only:
+ if cfg.scene_trainer.scene_optimizer.train_global_only:
print('train global update only')
for name, params in scene_trainer.optimizer.named_parameters():
if 'global_update' not in name:
diff --git a/optgs/misc/detaching_cpu_list.py b/optgs/misc/detaching_cpu_list.py
index cd6e721e9e85d2c916d5fd96980d921b346fc2cc..53a4b9d3e99eed096bf7f4b86edb3f2f8c75e160 100644
--- a/optgs/misc/detaching_cpu_list.py
+++ b/optgs/misc/detaching_cpu_list.py
@@ -1,294 +1,10 @@
-# import os
-# import torch
-# from dataclasses import dataclass, is_dataclass, replace, field
-# from typing import TypeVar, Generic, Optional, Any, Callable, Iterable
-# from pathlib import Path
-# import pickle
-# import tempfile
-# import shutil
-# import atexit
-# import signal
-# from typing import TypeVar, Generic
-#
-# T = TypeVar('T')
-#
-#
-# @dataclass
-# class DetachingCPUList(list[T]):
-# cache_dir: Optional[Path] = field(default=None)
-# save_serializer: Optional[Callable[[Any, Path], None]] = field(default=None)
-# load_serializer: Optional[Callable[[Path], Any]] = field(default=None)
-# detach_func: Optional[Callable[[Any], Any]] = field(default=None)
-# remove_files_on_delete: bool = field(default=True)
-# verbose: bool = field(default=True)
-#
-# _cache: dict = field(init=False, repr=False, default_factory=dict)
-# _fd_map: dict = field(init=False, repr=False, default_factory=dict)
-# _no_cache_set: set = field(init=False, repr=False, default_factory=set) # Track paths marked as no_cache
-# _cache_dir: Path = field(init=False, repr=False)
-# _tmp_dir_created: bool = field(init=False, repr=False, default=False)
-#
-# def __post_init__(self):
-#
-# # Set default serializers if not provided
-# if self.save_serializer is None:
-# def _save_pickle(obj, path: Path):
-# path.parent.mkdir(parents=True, exist_ok=True)
-# with open(path, "wb") as f:
-# pickle.dump(obj, f, protocol=pickle.HIGHEST_PROTOCOL)
-# self.save_serializer = _save_pickle
-#
-# # Set default deserializers if not provided
-# if self.load_serializer is None:
-# def _load_pickle(source):
-# if isinstance(source, (str, Path)):
-# with open(str(source), "rb") as f:
-# return pickle.load(f)
-# else:
-# return pickle.load(source)
-# self.load_serializer = _load_pickle
-#
-# # Set default detach_func if not provided
-# if self.detach_func is None:
-# self.detach_func = self._detach_recursive
-#
-# if self.cache_dir is None:
-# tmp = tempfile.mkdtemp(prefix="detaching_cpu_list_")
-# self._cache_dir = Path(tmp)
-# self._tmp_dir_created = True
-# else:
-# self._cache_dir = Path(self.cache_dir)
-# self._cache_dir.mkdir(parents=True, exist_ok=True)
-#
-# atexit.register(self._cleanup)
-# for sig in (signal.SIGINT, signal.SIGTERM):
-# try:
-# old = signal.getsignal(sig)
-# def _handler(signum, frame, _old=old):
-# self._cleanup()
-# if callable(_old) and _old not in (signal.SIG_DFL, signal.SIG_IGN):
-# _old(signum, frame)
-# signal.signal(sig, _handler)
-# except Exception:
-# pass
-#
-# # --------------------------
-# # Core cleanup
-# # --------------------------
-# def _cleanup(self):
-# # Close all file descriptors
-# for fd in list(self._fd_map.values()):
-# try:
-# os.close(fd)
-# except Exception:
-# pass
-# self._fd_map.clear()
-# # Remove on-disk files (non-guaranteed mode)
-# if self.remove_files_on_delete and self._tmp_dir_created and self._cache_dir.exists():
-# shutil.rmtree(self._cache_dir, ignore_errors=True)
-#
-# # --------------------------
-# # Save / load helpers
-# # --------------------------
-# def _save_unlinked(self, item: Any):
-# """Save item to disk with guaranteed deletion (fd-based approach)."""
-# tmp = tempfile.NamedTemporaryFile(delete=False, dir=str(self._cache_dir))
-# temp_path = Path(tmp.name)
-# tmp.close()
-#
-# assert self.save_serializer is not None, "save_serializer must be defined"
-# self.save_serializer(item, temp_path)
-# fd = os.open(str(temp_path), os.O_RDONLY)
-# os.unlink(str(temp_path)) # unlink immediately - kernel guarantees cleanup on fd close
-# pseudo = Path(f"/proc/self/fd/{fd}")
-# token = pseudo.as_posix()
-# self._fd_map[token] = fd
-# if self.verbose:
-# print(f"Saved item to fd {fd} with path {pseudo}")
-# return pseudo
-#
-# def _load_from_fd(self, token: str):
-# """Load from file descriptor path."""
-# fd = self._fd_map.get(token)
-# if fd is None:
-# raise RuntimeError(f"FD {token} not available.")
-# dupfd = os.dup(fd)
-# assert self.load_serializer is not None, "load_serializer must be defined"
-# with os.fdopen(dupfd, "rb") as f:
-# f.seek(0)
-# obj = self.load_serializer(f)
-# if self.verbose:
-# print(f"Loaded item from fd {fd} (token {token})")
-# return obj
-#
-# def _load_from_disk(self, path: Path):
-# """Load from regular file path."""
-# assert self.load_serializer is not None, "load_serializer must be defined"
-# if self.verbose:
-# print(f"Loading from disk: {path}")
-# return self.load_serializer(path)
-#
-# # --------------------------
-# # Public interface
-# # --------------------------
-# def append(self, item, detach_and_cpu: bool = False, save_to_disk: bool = False, no_cache: bool = False):
-# """
-# Append an item to the list.
-#
-# Args:
-# item: The item to append
-# detach_and_cpu: If True, apply detach_func to move tensors to CPU
-# save_to_disk: If True, save to disk using fd-based guaranteed deletion
-# no_cache: If True, never cache this item in memory when accessed (always reload from disk)
-# """
-# # Validate save_to_disk requires detach capability
-# if save_to_disk and not detach_and_cpu:
-# raise ValueError("Cannot save to disk without detach_and_cpu=True")
-#
-# if not save_to_disk and no_cache:
-# print("Warning: no_cache=True has no effect when save_to_disk=False")
-#
-# if detach_and_cpu and self.detach_func:
-# item = self.detach_func(item)
-#
-# if save_to_disk:
-# # Always use fd-based guaranteed deletion
-# p = self._save_unlinked(item)
-# super().append(p)
-#
-# # Mark as no_cache if requested
-# if no_cache:
-# self._no_cache_set.add(p.as_posix())
-# else:
-# super().append(item)
-#
-# def insert(self, index: int, item, detach_and_cpu: bool = False, save_to_disk: bool = False, no_cache: bool = False):
-# """
-# Insert an item at a specific index in the list.
-#
-# Args:
-# index: The index to insert the item at
-# item: The item to insert
-# detach_and_cpu: If True, apply detach_func to move tensors to CPU
-# save_to_disk: If True, save to disk using fd-based guaranteed deletion
-# no_cache: If True, never cache this item in memory when accessed (always reload from disk)
-# """
-# # Validate save_to_disk requires detach capability
-# if save_to_disk and not detach_and_cpu:
-# raise ValueError("Cannot save to disk without detach_and_cpu=True")
-#
-# if detach_and_cpu and self.detach_func:
-# item = self.detach_func(item)
-#
-# if save_to_disk:
-# # Always use fd-based guaranteed deletion
-# p = self._save_unlinked(item)
-# super().insert(index, p)
-#
-# # Mark as no_cache if requested
-# if no_cache:
-# self._no_cache_set.add(p.as_posix())
-# else:
-# super().insert(index, item)
-#
-# def extend(self, items: Iterable[Any], **kwargs):
-# for it in items:
-# self.append(it, **kwargs)
-#
-# def __getitem__(self, index):
-# """
-# Return the item at `index`.
-#
-# - If the underlying stored value is a Path, load from disk/fd.
-# - If marked as no_cache, always reload from disk (never cache).
-# - Otherwise, cache after first load and return cached object on subsequent accesses.
-# - If not a Path, return the in-memory value directly.
-# """
-# raw = super().__getitem__(index)
-#
-# # If it's not a Path, it's an in-memory object: return directly
-# if not isinstance(raw, Path):
-# return raw
-#
-# # it's a Path -> use its string as cache key (works for both normal paths and /proc/self/fd/)
-# key = raw.as_posix()
-#
-# # If in cache, return cached object
-# if key in self._cache:
-# assert key not in self._no_cache_set, "Inconsistent state: item both cached and marked no_cache"
-# return self._cache[key]
-#
-# # Always reload from disk/fd, never cache
-# if str(raw).startswith("/proc/self/fd/"):
-# obj = self._load_from_fd(key)
-# else:
-# raise NotImplementedError("Loading from disk is not implemented")
-# # return self._load_from_disk(raw)
-#
-# if key in self._no_cache_set:
-# pass # never cache
-# else:
-# # Cache it permanently (unless marked as no_cache)
-# self._cache[key] = obj
-#
-# return obj
-#
-# def pop(self, index: int = -1):
-# raw = super().pop(index)
-# # If it was a Path, return the loaded object (and keep the fd open / mapping intact)
-# if isinstance(raw, Path):
-# key = raw.as_posix()
-# # return cached value if present; else load now and cache it
-# if key in self._cache:
-# return self._cache[key]
-# if str(raw).startswith("/proc/self/fd/"):
-# obj = self._load_from_fd(key)
-# else:
-# obj = self._load_from_disk(raw)
-# self._cache[key] = obj
-# return obj
-# return raw
-#
-# def clear(self):
-# # do not close fds; keep them open until process exit as requested
-# # keep cache intact if you want (or clear it if you prefer)
-# # here we remove list entries but keep any cached objects and open fds
-# super().clear()
-#
-# def __iter__(self):
-# """Iterate over items, loading from disk as needed."""
-# for i in range(len(self)):
-# yield self[i]
-#
-# def __del__(self):
-# self._cleanup()
-#
-# def _detach_recursive(self, obj):
-# if isinstance(obj, torch.Tensor):
-# return obj.detach().cpu()
-# elif isinstance(obj, dict):
-# return {k: self._detach_recursive(v) for k, v in obj.items()}
-# elif isinstance(obj, (list, tuple)):
-# t = type(obj)
-# return t(self._detach_recursive(x) for x in obj)
-# elif is_dataclass(obj):
-# # Replace fields recursively (returns a new instance)
-# return replace(obj, **{
-# field.name: self._detach_recursive(getattr(obj, field.name))
-# for field in obj.__dataclass_fields__.values()
-# })
-# else:
-# return obj
-
-
from dataclasses import is_dataclass, replace, dataclass
import torch
@dataclass
class DetachingCPUList(list):
- # TODO Naama: Add back disk saving
- def append(self, item, detach_and_cpu=False, save_to_disk=False, no_cache=False):
+ def append(self, item, detach_and_cpu=False):
if detach_and_cpu:
item = self._detach_recursive(item)
super().append(item)
@@ -298,7 +14,7 @@ class DetachingCPUList(list):
iterable = (self._detach_recursive(x) for x in iterable)
super().extend(iterable)
- def insert(self, index, item, detach_and_cpu=False, save_to_disk=False, no_cache=False):
+ def insert(self, index, item, detach_and_cpu=False):
if detach_and_cpu:
item = self._detach_recursive(item)
super().insert(index, item)
@@ -318,4 +34,4 @@ class DetachingCPUList(list):
for field in obj.__dataclass_fields__.values()
})
else:
- return obj
+ return obj
\ No newline at end of file
diff --git a/optgs/misc/discrete_probability_distribution.py b/optgs/misc/discrete_probability_distribution.py
deleted file mode 100644
index 1b9fcb228098da9a12ce7f7b59fe43c3562083f5..0000000000000000000000000000000000000000
--- a/optgs/misc/discrete_probability_distribution.py
+++ /dev/null
@@ -1,33 +0,0 @@
-import torch
-from einops import reduce
-from jaxtyping import Float, Int64
-from torch import Tensor
-
-
-def sample_discrete_distribution(
- pdf: Float[Tensor, "*batch bucket"],
- num_samples: int,
- eps: float = torch.finfo(torch.float32).eps,
-) -> tuple[
- Int64[Tensor, "*batch sample"], # index
- Float[Tensor, "*batch sample"], # probability density
-]:
- *batch, bucket = pdf.shape
- normalized_pdf = pdf / (eps + reduce(pdf, "... bucket -> ... ()", "sum"))
- cdf = normalized_pdf.cumsum(dim=-1)
- samples = torch.rand((*batch, num_samples), device=pdf.device)
- index = torch.searchsorted(cdf, samples, right=True).clip(max=bucket - 1)
- return index, normalized_pdf.gather(dim=-1, index=index)
-
-
-def gather_discrete_topk(
- pdf: Float[Tensor, "*batch bucket"],
- num_samples: int,
- eps: float = torch.finfo(torch.float32).eps,
-) -> tuple[
- Int64[Tensor, "*batch sample"], # index
- Float[Tensor, "*batch sample"], # probability density
-]:
- normalized_pdf = pdf / (eps + reduce(pdf, "... bucket -> ... ()", "sum"))
- index = pdf.topk(k=num_samples, dim=-1).indices
- return index, normalized_pdf.gather(dim=-1, index=index)
diff --git a/optgs/misc/general_utils.py b/optgs/misc/general_utils.py
index fa5944a05cea3fd1f4fda9c99591bb1bce3c310d..015bd7cbc16a3549801e43842f6b29b5e143cc18 100644
--- a/optgs/misc/general_utils.py
+++ b/optgs/misc/general_utils.py
@@ -1,7 +1,6 @@
import torch
import numpy as np
import torch.nn.functional as F
-from torch.optim.lr_scheduler import LambdaLR
def get_expon_lr_func(
diff --git a/optgs/misc/image_io.py b/optgs/misc/image_io.py
index 7fcdb2a133f7191fb62e435482f1807898af4164..f9846127fefab3e6e550b8031a3a2400a9b1d3cc 100644
--- a/optgs/misc/image_io.py
+++ b/optgs/misc/image_io.py
@@ -124,7 +124,7 @@ def save_video(images, path, fps=None, iterations=None):
cv2.putText(frame, f"Iter {iter_num}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
frames[i] = frame
- # TODO Naama: videos cannot be saved with odd dimensions
+ # NOTE: the encoder cannot save videos with odd frame dimensions
with imageio.get_writer(str(path), fps=fps) as writer:
for frame in frames:
writer.append_data(frame)
\ No newline at end of file
diff --git a/optgs/misc/io.py b/optgs/misc/io.py
index 15361bae4ae86039956f66d60fd6641a83a8dae5..37e661efed7f2f1cccaa092ce0be1b48bc484361 100644
--- a/optgs/misc/io.py
+++ b/optgs/misc/io.py
@@ -1,7 +1,5 @@
import pathlib
from copy import copy
-import numpy as np
-import torch
import yaml
from colorama import Fore
from omegaconf import OmegaConf
@@ -70,15 +68,6 @@ class CustomPath(pathlib.Path):
def __add__(self, other: str):
return CustomPath(str(self) + other)
- def is_json(self):
- return self.suffix == '.json'
-
- def is_yaml(self):
- return self.suffix == '.yaml'
-
- def json_encoder(self):
- return str(self)
-
def __sub__(self, other):
return CustomPath(self.resolve().relative_to(other.resolve()))
@@ -148,9 +137,6 @@ class FrequencyScheduler:
self.is_disabled = False
- def set_verbose(self, verbose: bool):
- self.verbose = verbose
-
def set_all_tags(self, enabled: bool):
for key in self.enabled_tags:
self.enabled_tags[key] = enabled
@@ -189,10 +175,13 @@ class FrequencyScheduler:
return iterations
def get_iterations(self, length_of_event: int) -> list[int]:
- """Returns a list of all iterations where an event occurs up to the given length."""
+ """Iteration indices for the first `length_of_event` saved events.
+
+ Callers pass len(render_list)/len(gaussian_list), whose position 0 is the init (iteration 0),
+ so the labels are the first `length_of_event` scheduled iterations -- including for a partial
+ scene that saved fewer than the full schedule.
+ """
if self.iterations is not None and len(self.iterations) >= length_of_event:
- if length_of_event == 1:
- return [self.iterations[-1]]
return self.iterations[:length_of_event]
else:
raise ValueError(
@@ -211,13 +200,6 @@ class FrequencyScheduler:
return f"FrequencyScheduler({self.iterations})"
-def log_mem(tag=""):
- torch.cuda.synchronize()
- print(f"{tag}: allocated={torch.cuda.memory_allocated() / 1e6:.1f}MB, "
- f"reserved={torch.cuda.memory_reserved() / 1e6:.1f}MB, "
- f"max_allocated={torch.cuda.max_memory_allocated() / 1e6:.1f}MB")
-
-
def read_omega_cfg(path: pathlib.Path) -> OmegaConf:
"""Reads an OmegaConf YAML file, handling custom tags safely."""
try:
@@ -235,7 +217,6 @@ def read_omega_cfg(path: pathlib.Path) -> OmegaConf:
path = CustomPath()
for part in seq:
path = path / str(part)
- print(path)
return path
else:
raise TypeError(f"Unsupported YAML node type for CustomPath: {type(node)}")
diff --git a/optgs/misc/render_utils.py b/optgs/misc/render_utils.py
deleted file mode 100644
index eb8025c10db948ff8eaad3dd76a4651bc0fed881..0000000000000000000000000000000000000000
--- a/optgs/misc/render_utils.py
+++ /dev/null
@@ -1,330 +0,0 @@
-# Copyright 2022 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
-# https://github.com/hbb1/2d-gaussian-splatting/blob/main/utils/render_utils.py
-
-import numpy as np
-import os
-import enum
-import types
-from typing import List, Mapping, Optional, Text, Tuple, Union
-import copy
-from PIL import Image
-# import mediapy as media
-from matplotlib import cm
-from tqdm import tqdm
-
-import torch
-
-
-def normalize(x: np.ndarray) -> np.ndarray:
- """Normalization helper function."""
- return x / np.linalg.norm(x)
-
-
-def pad_poses(p: np.ndarray) -> np.ndarray:
- """Pad [..., 3, 4] pose matrices with a homogeneous bottom row [0,0,0,1]."""
- bottom = np.broadcast_to([0, 0, 0, 1.], p[..., :1, :4].shape)
- return np.concatenate([p[..., :3, :4], bottom], axis=-2)
-
-
-def unpad_poses(p: np.ndarray) -> np.ndarray:
- """Remove the homogeneous bottom row from [..., 4, 4] pose matrices."""
- return p[..., :3, :4]
-
-
-def recenter_poses(poses: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
- """Recenter poses around the origin."""
- cam2world = average_pose(poses)
- transform = np.linalg.inv(pad_poses(cam2world))
- poses = transform @ pad_poses(poses)
- return unpad_poses(poses), transform
-
-
-def average_pose(poses: np.ndarray) -> np.ndarray:
- """New pose using average position, z-axis, and up vector of input poses."""
- position = poses[:, :3, 3].mean(0)
- z_axis = poses[:, :3, 2].mean(0)
- up = poses[:, :3, 1].mean(0)
- cam2world = viewmatrix(z_axis, up, position)
- return cam2world
-
-
-def viewmatrix(lookdir: np.ndarray, up: np.ndarray,
- position: np.ndarray) -> np.ndarray:
- """Construct lookat view matrix."""
- vec2 = normalize(lookdir)
- vec0 = normalize(np.cross(up, vec2))
- vec1 = normalize(np.cross(vec2, vec0))
- m = np.stack([vec0, vec1, vec2, position], axis=1)
- return m
-
-
-def focus_point_fn(poses: np.ndarray) -> np.ndarray:
- """Calculate nearest point to all focal axes in poses."""
- directions, origins = poses[:, :3, 2:3], poses[:, :3, 3:4]
- m = np.eye(3) - directions * np.transpose(directions, [0, 2, 1])
- mt_m = np.transpose(m, [0, 2, 1]) @ m
- focus_pt = np.linalg.inv(mt_m.mean(0)) @ (mt_m @ origins).mean(0)[:, 0]
- return focus_pt
-
-
-def transform_poses_pca(poses: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
- """Transforms poses so principal components lie on XYZ axes.
-
- Args:
- poses: a (N, 3, 4) array containing the cameras' camera to world transforms.
-
- Returns:
- A tuple (poses, transform), with the transformed poses and the applied
- camera_to_world transforms.
- """
- t = poses[:, :3, 3]
- t_mean = t.mean(axis=0)
- t = t - t_mean
-
- eigval, eigvec = np.linalg.eig(t.T @ t)
- # Sort eigenvectors in order of largest to smallest eigenvalue.
- inds = np.argsort(eigval)[::-1]
- eigvec = eigvec[:, inds]
- rot = eigvec.T
- if np.linalg.det(rot) < 0:
- rot = np.diag(np.array([1, 1, -1])) @ rot
-
- transform = np.concatenate([rot, rot @ -t_mean[:, None]], -1)
- poses_recentered = unpad_poses(transform @ pad_poses(poses))
- transform = np.concatenate([transform, np.eye(4)[3:]], axis=0)
-
- # Flip coordinate system if z component of y-axis is negative
- if poses_recentered.mean(axis=0)[2, 1] < 0:
- poses_recentered = np.diag(np.array([1, -1, -1])) @ poses_recentered
- transform = np.diag(np.array([1, -1, -1, 1])) @ transform
-
- return poses_recentered, transform
- # points = np.random.rand(3,100)
- # points_h = np.concatenate((points,np.ones_like(points[:1])), axis=0)
- # (poses_recentered @ points_h)[0]
- # (transform @ pad_poses(poses) @ points_h)[0,:3]
- # import pdb; pdb.set_trace()
-
- # # Just make sure it's it in the [-1, 1]^3 cube
- # scale_factor = 1. / np.max(np.abs(poses_recentered[:, :3, 3]))
- # poses_recentered[:, :3, 3] *= scale_factor
- # transform = np.diag(np.array([scale_factor] * 3 + [1])) @ transform
-
- # return poses_recentered, transform
-
-
-def generate_ellipse_path(poses: np.ndarray,
- n_frames: int = 120,
- const_speed: bool = True,
- z_variation: float = 0.,
- z_phase: float = 0.) -> np.ndarray:
- """Generate an elliptical render path based on the given poses."""
- # Calculate the focal point for the path (cameras point toward this).
- center = focus_point_fn(poses)
- # Path height sits at z=0 (in middle of zero-mean capture pattern).
- offset = np.array([center[0], center[1], 0])
-
- # Calculate scaling for ellipse axes based on input camera positions.
- sc = np.percentile(np.abs(poses[:, :3, 3] - offset), 90, axis=0)
- # Use ellipse that is symmetric about the focal point in xy.
- low = -sc + offset
- high = sc + offset
- # Optional height variation need not be symmetric
- z_low = np.percentile((poses[:, :3, 3]), 10, axis=0)
- z_high = np.percentile((poses[:, :3, 3]), 90, axis=0)
-
- def get_positions(theta):
- # Interpolate between bounds with trig functions to get ellipse in x-y.
- # Optionally also interpolate in z to change camera height along path.
- return np.stack([
- low[0] + (high - low)[0] * (np.cos(theta) * .5 + .5),
- low[1] + (high - low)[1] * (np.sin(theta) * .5 + .5),
- z_variation * (z_low[2] + (z_high - z_low)[2] *
- (np.cos(theta + 2 * np.pi * z_phase) * .5 + .5)),
- ], -1)
-
- theta = np.linspace(0, 2. * np.pi, n_frames + 1, endpoint=True)
- positions = get_positions(theta)
-
- # if const_speed:
-
- # # Resample theta angles so that the velocity is closer to constant.
- # lengths = np.linalg.norm(positions[1:] - positions[:-1], axis=-1)
- # theta = stepfun.sample(None, theta, np.log(lengths), n_frames + 1)
- # positions = get_positions(theta)
-
- # Throw away duplicated last position.
- positions = positions[:-1]
-
- # Set path's up vector to axis closest to average of input pose up vectors.
- avg_up = poses[:, :3, 1].mean(0)
- avg_up = avg_up / np.linalg.norm(avg_up)
- ind_up = np.argmax(np.abs(avg_up))
- up = np.eye(3)[ind_up] * np.sign(avg_up[ind_up])
-
- return np.stack([viewmatrix(p - center, up, p) for p in positions])
-
-
-def generate_path(viewpoint_cameras, n_frames=480):
- c2ws = np.array([np.linalg.inv(np.asarray(
- (cam.world_view_transform.T).cpu().numpy())) for cam in viewpoint_cameras])
- pose = c2ws[:, :3, :] @ np.diag([1, -1, -1, 1])
- pose_recenter, colmap_to_world_transform = transform_poses_pca(pose)
-
- # generate new poses
- new_poses = generate_ellipse_path(poses=pose_recenter, n_frames=n_frames)
- # warp back to orignal scale
- new_poses = np.linalg.inv(colmap_to_world_transform) @ pad_poses(new_poses)
-
- traj = []
- for c2w in new_poses:
- c2w = c2w @ np.diag([1, -1, -1, 1])
- cam = copy.deepcopy(viewpoint_cameras[0])
- cam.image_height = int(cam.image_height / 2) * 2
- cam.image_width = int(cam.image_width / 2) * 2
- cam.world_view_transform = torch.from_numpy(
- np.linalg.inv(c2w).T).float().cuda()
- cam.full_proj_transform = (cam.world_view_transform.unsqueeze(
- 0).bmm(cam.projection_matrix.unsqueeze(0))).squeeze(0)
- cam.camera_center = cam.world_view_transform.inverse()[3, :3]
- traj.append(cam)
-
- return traj
-
-
-def generate_video_render_path(c2ws, n_frames=480):
- # c2ws = np.array([np.linalg.inv(np.asarray(
- # (cam.world_view_transform.T).cpu().numpy())) for cam in viewpoint_cameras])
- # c2ws: [V, 4, 4]
- pose = c2ws[:, :3, :] @ np.diag([1, -1, -1, 1])
- pose_recenter, colmap_to_world_transform = transform_poses_pca(pose)
-
- # generate new poses
- new_poses = generate_ellipse_path(poses=pose_recenter, n_frames=n_frames)
- # warp back to orignal scale
- new_poses = np.linalg.inv(colmap_to_world_transform) @ pad_poses(new_poses)
-
- traj = []
- for c2w in new_poses:
- c2w = c2w @ np.diag([1, -1, -1, 1])
- # cam = copy.deepcopy(viewpoint_cameras[0])
- # cam.image_height = int(cam.image_height / 2) * 2
- # cam.image_width = int(cam.image_width / 2) * 2
- # cam.world_view_transform = torch.from_numpy(
- # np.linalg.inv(c2w).T).float().cuda()
- # cam.full_proj_transform = (cam.world_view_transform.unsqueeze(
- # 0).bmm(cam.projection_matrix.unsqueeze(0))).squeeze(0)
- # cam.camera_center = cam.world_view_transform.inverse()[3, :3]
- # traj.append(cam)
- traj.append(c2w)
-
- return traj
-
-
-def load_img(pth: str) -> np.ndarray:
- """Load an image and cast to float32."""
- with open(pth, 'rb') as f:
- image = np.array(Image.open(f), dtype=np.float32)
- return image
-
-
-def create_videos(base_dir, input_dir, out_name, num_frames=480):
- """Creates videos out of the images saved to disk."""
- # Last two parts of checkpoint path are experiment name and scene name.
- video_prefix = f'{out_name}'
-
- zpad = max(5, len(str(num_frames - 1)))
- def idx_to_str(idx): return str(idx).zfill(zpad)
-
- os.makedirs(base_dir, exist_ok=True)
- render_dist_curve_fn = np.log
-
- # Load one example frame to get image shape and depth range.
- depth_file = os.path.join(input_dir, 'vis', f'depth_{idx_to_str(0)}.tiff')
- depth_frame = load_img(depth_file)
- shape = depth_frame.shape
- p = 3
- distance_limits = np.percentile(depth_frame.flatten(), [p, 100 - p])
- lo, hi = [render_dist_curve_fn(x) for x in distance_limits]
- print(f'Video shape is {shape[:2]}')
-
- video_kwargs = {
- 'shape': shape[:2],
- 'codec': 'h264',
- 'fps': 60,
- 'crf': 18,
- }
-
- for k in ['depth', 'normal', 'color']:
- video_file = os.path.join(base_dir, f'{video_prefix}_{k}.mp4')
- input_format = 'gray' if k == 'alpha' else 'rgb'
-
- file_ext = 'png' if k in ['color', 'normal'] else 'tiff'
- idx = 0
-
- if k == 'color':
- file0 = os.path.join(input_dir, 'renders',
- f'{idx_to_str(0)}.{file_ext}')
- else:
- file0 = os.path.join(
- input_dir, 'vis', f'{k}_{idx_to_str(0)}.{file_ext}')
-
- if not os.path.exists(file0):
- print(f'Images missing for tag {k}')
- continue
- print(f'Making video {video_file}...')
- with media.VideoWriter(
- video_file, **video_kwargs, input_format=input_format) as writer:
- for idx in tqdm(range(num_frames)):
- # img_file = os.path.join(input_dir, f'{k}_{idx_to_str(idx)}.{file_ext}')
- if k == 'color':
- img_file = os.path.join(
- input_dir, 'renders', f'{idx_to_str(idx)}.{file_ext}')
- else:
- img_file = os.path.join(
- input_dir, 'vis', f'{k}_{idx_to_str(idx)}.{file_ext}')
-
- if not os.path.exists(img_file):
- ValueError(f'Image file {img_file} does not exist.')
- img = load_img(img_file)
- if k in ['color', 'normal']:
- img = img / 255.
- elif k.startswith('depth'):
- img = render_dist_curve_fn(img)
- img = np.clip((img - np.minimum(lo, hi)) /
- np.abs(hi - lo), 0, 1)
- img = cm.get_cmap('turbo')(img)[..., :3]
-
- frame = (np.clip(np.nan_to_num(img), 0., 1.)
- * 255.).astype(np.uint8)
- writer.add_image(frame)
- idx += 1
-
-
-def save_img_u8(img, pth):
- """Save an image (probably RGB) in [0, 1] to disk as a uint8 PNG."""
- with open(pth, 'wb') as f:
- Image.fromarray(
- (np.clip(np.nan_to_num(img), 0., 1.) * 255.).astype(np.uint8)).save(
- f, 'PNG')
-
-
-def save_img_f32(depthmap, pth):
- """Save an image (probably a depthmap) to disk as a float32 TIFF."""
- with open(pth, 'wb') as f:
- Image.fromarray(np.nan_to_num(depthmap).astype(
- np.float32)).save(f, 'TIFF')
diff --git a/optgs/misc/wandb_tools.py b/optgs/misc/wandb_tools.py
index fec2cd7eaf176f6038681897400117e3f080aa69..de7f620dcd7d174b45d1b81efd0648e63ea59acb 100644
--- a/optgs/misc/wandb_tools.py
+++ b/optgs/misc/wandb_tools.py
@@ -56,7 +56,7 @@ def setup_wandb_logger(cfg, cfg_dict) -> WandbLogger | LocalLogger:
if cfg_dict.wandb.id is None:
print(f"Resuming wandb run without id, using latest run in output directory.")
# Find the latest wandb run id in the output directory
- wandb_dir = cfg_dict.output_dir / "wandb" / "latest-run"
+ wandb_dir = cfg.output_dir / "wandb" / "latest-run"
# look for a file name in the format "run-######.wandb" file and extract the id
wandb_files = list(wandb_dir.glob("run-*.wandb"))
assert len(wandb_files) <= 1, "Multiple wandb files found in the latest run directory."
@@ -69,7 +69,7 @@ def setup_wandb_logger(cfg, cfg_dict) -> WandbLogger | LocalLogger:
print(f"Setting wandb run with id from cfg {cfg_dict.wandb.id}.")
wandb_extra_kwargs.update({'id': cfg_dict.wandb.id, 'resume': "must"})
- run_name = os.path.basename(cfg_dict.output_dir)
+ run_name = os.path.basename(cfg.output_dir)
if cfg_dict.log_slurm_id:
hostname = os.uname().nodename
@@ -90,7 +90,7 @@ def setup_wandb_logger(cfg, cfg_dict) -> WandbLogger | LocalLogger:
name=run_name,
tags=cfg_dict.wandb.get("tags", None),
log_model=False,
- save_dir=cfg_dict.output_dir,
+ save_dir=cfg.output_dir,
config=OmegaConf.to_container(cfg_dict),
**wandb_extra_kwargs,
)
@@ -101,12 +101,6 @@ def setup_wandb_logger(cfg, cfg_dict) -> WandbLogger | LocalLogger:
# Log notes
if cfg_dict.wandb.notes is not None:
logger.experiment.notes = cfg_dict.wandb.notes
- # Write wandb run ID to file for SLURM requeue resume
- wandb_id_file = os.environ.get("WANDB_ID_FILE")
- if wandb_id_file:
- with open(wandb_id_file, "w") as f:
- f.write(logger.experiment.id)
- print(f"Wrote wandb run ID {logger.experiment.id} to {wandb_id_file}")
return logger
diff --git a/optgs/model/colmap_utils/__init__.py b/optgs/model/backbones/__init__.py
similarity index 100%
rename from optgs/model/colmap_utils/__init__.py
rename to optgs/model/backbones/__init__.py
diff --git a/optgs/model/encoder/.deprecated/common/__init__.py b/optgs/model/backbones/depth_anything_v2/__init__.py
similarity index 100%
rename from optgs/model/encoder/.deprecated/common/__init__.py
rename to optgs/model/backbones/depth_anything_v2/__init__.py
diff --git a/optgs/model/encoder/depth_anything_v2/dinov2.py b/optgs/model/backbones/depth_anything_v2/dinov2.py
similarity index 100%
rename from optgs/model/encoder/depth_anything_v2/dinov2.py
rename to optgs/model/backbones/depth_anything_v2/dinov2.py
diff --git a/optgs/model/encoder/depth_anything_v2/dinov2_layers/__init__.py b/optgs/model/backbones/depth_anything_v2/dinov2_layers/__init__.py
similarity index 100%
rename from optgs/model/encoder/depth_anything_v2/dinov2_layers/__init__.py
rename to optgs/model/backbones/depth_anything_v2/dinov2_layers/__init__.py
diff --git a/optgs/model/encoder/depth_anything_v2/dinov2_layers/attention.py b/optgs/model/backbones/depth_anything_v2/dinov2_layers/attention.py
similarity index 100%
rename from optgs/model/encoder/depth_anything_v2/dinov2_layers/attention.py
rename to optgs/model/backbones/depth_anything_v2/dinov2_layers/attention.py
diff --git a/optgs/model/encoder/depth_anything_v2/dinov2_layers/block.py b/optgs/model/backbones/depth_anything_v2/dinov2_layers/block.py
similarity index 100%
rename from optgs/model/encoder/depth_anything_v2/dinov2_layers/block.py
rename to optgs/model/backbones/depth_anything_v2/dinov2_layers/block.py
diff --git a/optgs/model/encoder/depth_anything_v2/dinov2_layers/drop_path.py b/optgs/model/backbones/depth_anything_v2/dinov2_layers/drop_path.py
similarity index 100%
rename from optgs/model/encoder/depth_anything_v2/dinov2_layers/drop_path.py
rename to optgs/model/backbones/depth_anything_v2/dinov2_layers/drop_path.py
diff --git a/optgs/model/encoder/depth_anything_v2/dinov2_layers/layer_scale.py b/optgs/model/backbones/depth_anything_v2/dinov2_layers/layer_scale.py
similarity index 100%
rename from optgs/model/encoder/depth_anything_v2/dinov2_layers/layer_scale.py
rename to optgs/model/backbones/depth_anything_v2/dinov2_layers/layer_scale.py
diff --git a/optgs/model/encoder/depth_anything_v2/dinov2_layers/mlp.py b/optgs/model/backbones/depth_anything_v2/dinov2_layers/mlp.py
similarity index 100%
rename from optgs/model/encoder/depth_anything_v2/dinov2_layers/mlp.py
rename to optgs/model/backbones/depth_anything_v2/dinov2_layers/mlp.py
diff --git a/optgs/model/encoder/depth_anything_v2/dinov2_layers/patch_embed.py b/optgs/model/backbones/depth_anything_v2/dinov2_layers/patch_embed.py
similarity index 100%
rename from optgs/model/encoder/depth_anything_v2/dinov2_layers/patch_embed.py
rename to optgs/model/backbones/depth_anything_v2/dinov2_layers/patch_embed.py
diff --git a/optgs/model/encoder/depth_anything_v2/dinov2_layers/swiglu_ffn.py b/optgs/model/backbones/depth_anything_v2/dinov2_layers/swiglu_ffn.py
similarity index 100%
rename from optgs/model/encoder/depth_anything_v2/dinov2_layers/swiglu_ffn.py
rename to optgs/model/backbones/depth_anything_v2/dinov2_layers/swiglu_ffn.py
diff --git a/optgs/model/encoder/depth_anything_v2/dpt.py b/optgs/model/backbones/depth_anything_v2/dpt.py
similarity index 100%
rename from optgs/model/encoder/depth_anything_v2/dpt.py
rename to optgs/model/backbones/depth_anything_v2/dpt.py
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/__init__.py b/optgs/model/backbones/depth_anything_v2/util/__init__.py
similarity index 100%
rename from optgs/model/encoder/.deprecated/foundationstereo/__init__.py
rename to optgs/model/backbones/depth_anything_v2/util/__init__.py
diff --git a/optgs/model/encoder/depth_anything_v2/util/blocks.py b/optgs/model/backbones/depth_anything_v2/util/blocks.py
similarity index 100%
rename from optgs/model/encoder/depth_anything_v2/util/blocks.py
rename to optgs/model/backbones/depth_anything_v2/util/blocks.py
diff --git a/optgs/model/encoder/depth_anything_v2/util/transform.py b/optgs/model/backbones/depth_anything_v2/util/transform.py
similarity index 100%
rename from optgs/model/encoder/depth_anything_v2/util/transform.py
rename to optgs/model/backbones/depth_anything_v2/util/transform.py
diff --git a/optgs/model/encoder/layer.py b/optgs/model/backbones/layer.py
similarity index 100%
rename from optgs/model/encoder/layer.py
rename to optgs/model/backbones/layer.py
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/core/__init__.py b/optgs/model/backbones/point_transformer/__init__.py
similarity index 100%
rename from optgs/model/encoder/.deprecated/foundationstereo/core/__init__.py
rename to optgs/model/backbones/point_transformer/__init__.py
diff --git a/optgs/model/backbones/point_transformer/layer.py b/optgs/model/backbones/point_transformer/layer.py
new file mode 100644
index 0000000000000000000000000000000000000000..d826ad106a72847e30eec754c888469f50ae8b83
--- /dev/null
+++ b/optgs/model/backbones/point_transformer/layer.py
@@ -0,0 +1,541 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+
+import pointops
+from einops import rearrange
+
+from ..unimatch.dinov2.layers.block import Block as MultiViewBlock
+from ..unimatch.utils import mv_feature_add_position
+
+USE_PYTORCH_ATTN = False
+USE_FLASH_ATTN3 = False
+
+
+class KNNAttention(nn.Module):
+ def __init__(self, channels, knn_samples=16, no_rpe=True,
+ qk_norm=False,
+ num_heads=1,
+ proj_channels=None,
+ use_fused=False,
+ ):
+ super().__init__()
+ self.proj_channels = proj_channels
+
+ self.knn_samples = knn_samples
+ self.no_rpe = no_rpe
+ self.num_heads = num_heads
+ assert self.num_heads == 1
+
+ self.use_fused = use_fused
+ if use_fused:
+ try:
+ import sys
+ from optgs.paths import PROJECT_DIR
+ sys.path.append(str(PROJECT_DIR / "submodules"))
+ from fused_knn_attn import fused_knn_attention, FUSED_KNN_ATTN_CUDA_AVAILABLE
+ self._fused_knn_attention = fused_knn_attention
+ if not FUSED_KNN_ATTN_CUDA_AVAILABLE:
+ import warnings
+ warnings.warn(
+ "Fused KNN attention CUDA extension not available, "
+ "using PyTorch fallback (still avoids [N,K,C] intermediates)"
+ )
+ except ImportError:
+ import warnings
+ warnings.warn(
+ "fused_knn_attn package not found, falling back to unfused attention"
+ )
+ self.use_fused = False
+
+ self.qk_norm = qk_norm
+ if qk_norm:
+ self.q_norm = nn.RMSNorm(channels)
+ self.k_norm = nn.RMSNorm(channels)
+
+ if self.proj_channels is not None:
+ self.qkv = nn.Linear(channels, self.proj_channels * 3, bias=False)
+ self.proj = nn.Linear(self.proj_channels, channels)
+ else:
+ self.qkv = nn.Linear(channels, channels * 3, bias=False)
+ self.proj = nn.Linear(channels, channels)
+
+ if not self.no_rpe:
+ self.rpe = nn.Sequential(
+ nn.Linear(3, 32),
+ nn.GELU(),
+ nn.Linear(32, 1)
+ )
+
+ def forward(self, pxo, knn_idx=None):
+ # [N, 3], [N, C], [B]
+ p, x, o = pxo
+ c = x.size(1)
+
+ if self.proj_channels is not None:
+ c = self.proj_channels
+
+ assert c % self.num_heads == 0
+ head_dim = c // self.num_heads
+ scale_factor = head_dim ** -0.5
+
+ qkv = self.qkv(x) # [N, 3*C]
+ x_q, x_k, x_v = torch.chunk(qkv, chunks=3, dim=-1) # each [N, C]
+
+ # ---- Fused path: gather + attention in one kernel ----
+ if self.use_fused and self.no_rpe:
+ # Ensure we have KNN indices
+ if knn_idx is None:
+ knn_idx, _ = pointops.knn_query(
+ self.knn_samples, p, o, p, o
+ )
+
+ # qk_norm: RMSNorm normalizes each C-dim vector independently,
+ # so applying before gather is equivalent to applying after gather.
+ if self.qk_norm:
+ x_q = self.q_norm(x_q)
+ x_k = self.k_norm(x_k)
+
+ out = self._fused_knn_attention(
+ x_q.contiguous(), x_k.contiguous(), x_v.contiguous(),
+ knn_idx.contiguous(), scale_factor
+ )
+ out = self.proj(out)
+ return out
+
+ # ---- Original unfused path ----
+ # # [N, K, C], [N, K]
+ # x_k, idx = pointops.knn_query_and_group(
+ # x_k.contiguous(), p, o, new_xyz=p, new_offset=o,
+ # idx=knn_idx,
+ # nsample=self.knn_samples, with_xyz=False
+ # ) # [N, K, C]
+ #
+ # # [N, K, C]
+ # x_v, _ = pointops.knn_query_and_group(
+ # x_v.contiguous(),
+ # p,
+ # o,
+ # new_xyz=p,
+ # new_offset=o,
+ # idx=idx,
+ # nsample=self.knn_samples,
+ # with_xyz=False,
+ # )
+
+ # ---- Initial improved version ----
+ x_kv = torch.cat([x_k, x_v], dim=-1) # [N, 2C/3]
+ x_kv_query, _ = pointops.knn_query_and_group(
+ x_kv.contiguous(), p, o, new_xyz=p, new_offset=o,
+ idx=knn_idx, nsample=self.knn_samples, with_xyz=False
+ ) # [N, K, 2C/3]
+ x_k, x_v = torch.chunk(x_kv_query, chunks=2, dim=-1)
+
+ # [N, K, 3], [N, K, C]
+ # NOTE: without xyz in knn
+ # p_r, x_k = x_k[:, :, :3], x_k[:, :, 3:]
+
+ # [N, 1, K]
+ assert self.no_rpe
+ if not self.no_rpe:
+ rpe = self.rpe(p_r).permute(0, 2, 1)
+ else:
+ rpe = 0
+
+ if self.qk_norm:
+ x_q = self.q_norm(x_q)
+ x_k = self.k_norm(x_k)
+
+ n, k, c = x_k.shape
+
+ # attention
+ if USE_PYTORCH_ATTN:
+ out = F.scaled_dot_product_attention(
+ x_q.view(n, 1, c),
+ x_k.view(n, k, c),
+ x_v.view(n, k, c),
+ ).reshape(n, c) # [N, C]
+
+ elif (USE_FLASH_ATTN3 and FA3_AVAILABLE and self.no_rpe):
+ # no relative pos enc
+ out = flash_attn_func(
+ x_q.view(n, 1, self.num_heads, head_dim).to(torch.bfloat16),
+ x_k.view(n, k, self.num_heads, head_dim).to(torch.bfloat16),
+ x_v.view(n, k, self.num_heads, head_dim).to(torch.bfloat16),
+ )[0].reshape(n, c).float() # [N, C]
+ else:
+ # [N, 1, K]
+ scores = torch.matmul(x_q.unsqueeze(1), x_k.permute(0, 2, 1)) * scale_factor + rpe
+ # [N, C]
+ out = torch.matmul(torch.softmax(scores, dim=2), x_v).squeeze(1)
+
+ out = self.proj(out)
+
+ return out
+
+
+class MLP(nn.Module):
+ def __init__(
+ self,
+ channels,
+ act="gelu",
+ ):
+ super().__init__()
+
+ expansion = 4
+
+ self.fc1 = nn.Linear(channels, channels * expansion)
+ if act is None or act in ['none', 'identity']:
+ self.act = nn.Identity()
+ elif act == 'gelu':
+ self.act = nn.GELU()
+ elif act == 'tanh':
+ self.act = nn.Tanh()
+ else:
+ raise ValueError(f"unsupported activation {act}")
+ self.fc2 = nn.Linear(channels * expansion, channels)
+
+ def forward(self, x):
+ x = self.fc1(x)
+ x = self.act(x)
+ x = self.fc2(x)
+ return x
+
+
+class TransformerBlock(nn.Module):
+ def __init__(self, channels, knn_samples=16, post_norm=False,
+ no_rpe=False,
+ no_attn=False,
+ no_norm=False,
+ act="gelu",
+ qk_norm=False,
+ norm_pt_block=False,
+ num_heads=1,
+ attn_proj_channels=None,
+ use_fused_attn=False,
+ ):
+ super().__init__()
+ self.post_norm = post_norm
+ self.no_attn = no_attn
+ self.norm_pt_block = norm_pt_block
+
+ if no_norm:
+ self.norm1 = nn.Identity()
+ self.norm2 = nn.Identity()
+ else:
+ self.norm1 = nn.LayerNorm(channels)
+ self.norm2 = nn.LayerNorm(channels)
+
+ if self.no_attn:
+ self.linear = nn.Linear(channels, channels)
+ else:
+ self.attn = KNNAttention(channels, knn_samples=knn_samples, no_rpe=no_rpe,
+ qk_norm=qk_norm,
+ num_heads=num_heads,
+ proj_channels=attn_proj_channels,
+ use_fused=use_fused_attn,
+ )
+ self.mlp = MLP(channels, act=act)
+
+ if self.norm_pt_block:
+ self.norm3 = nn.LayerNorm(channels)
+
+ def forward(self, pxo, knn_idx=None):
+ p, x, o = pxo
+
+ if self.post_norm:
+ if self.no_attn:
+ x = x + self.norm1(self.linear(x))
+ else:
+ x = x + self.norm1(self.attn((p, x, o), knn_idx=knn_idx))
+ x = x + self.norm2(self.mlp(x))
+ else:
+ if self.no_attn:
+ x = x + self.linear(self.norm1(x))
+ else:
+ x = x + self.attn((p, self.norm1(x), o), knn_idx=knn_idx)
+ x = x + self.mlp(self.norm2(x))
+
+ if self.norm_pt_block:
+ x = self.norm3(x)
+
+ return x
+
+
+class PlainPointTransformer(nn.Module):
+ def __init__(self, channels, knn_samples=16, num_blocks=4, post_norm=False,
+ no_rpe=False,
+ no_attn=False,
+ no_norm=False,
+ act="gelu",
+ qk_norm=False,
+ norm_pt_block=False,
+ num_heads=1,
+ attn_proj_channels=None,
+ cache_knn_idx=None,
+ knn_idx_update_every=1,
+ with_mv_attn=False,
+ with_mv_attn_lowres=False,
+ mv_attn_first=False,
+ no_mv_attn=False,
+ conv_with_norm=False,
+ mv_shuffle_attn=False,
+ with_pos_enc=False,
+ shuffle_attn_no_norm=False,
+ mv_unimatch_attn=False,
+ use_checkpointing=False,
+ init_use_checkpointing=False,
+ use_fused_attn=False,
+ ):
+ super().__init__()
+
+ self.cache_knn_idx = cache_knn_idx
+ self.knn_idx_update_every = knn_idx_update_every
+ self.knn_samples = knn_samples
+ self.use_checkpointing = use_checkpointing
+ self.init_use_checkpointing = init_use_checkpointing
+
+ self.with_mv_attn = with_mv_attn
+ self.with_mv_attn_lowres = with_mv_attn_lowres
+ if with_pos_enc:
+ assert mv_shuffle_attn
+
+ self.blocks = nn.ModuleList()
+ for _ in range(num_blocks):
+ self.blocks.append(TransformerBlock(channels, knn_samples=knn_samples,
+ post_norm=post_norm,
+ no_rpe=no_rpe,
+ no_attn=no_attn,
+ no_norm=no_norm,
+ act=act,
+ qk_norm=qk_norm,
+ norm_pt_block=norm_pt_block,
+ num_heads=num_heads,
+ attn_proj_channels=attn_proj_channels,
+ use_fused_attn=use_fused_attn,
+ ))
+
+ # multi-view attention
+ if self.with_mv_attn:
+ self.mv_blocks = nn.ModuleList()
+ for _ in range(num_blocks):
+ # if mv_shuffle_attn:
+ if self.with_mv_attn_lowres:
+ self.mv_blocks.append(
+ MultViewLowresAttn(
+ channels,
+ )
+ )
+ else:
+ self.mv_blocks.append(
+ MultiViewBlock(
+ channels,
+ num_heads=4,
+ )
+ )
+
+ def forward(self, pxo, iter=0, b=None, v=None, h=None, w=None):
+ p, x, o = pxo
+ # compute knn idx here only once and pass it to the model
+ # the positions are not changed inside the blocks
+ if self.cache_knn_idx is None or (iter % self.knn_idx_update_every) == 0:
+ knn_idx, _ = pointops.knn_query(self.knn_samples, p, o, p, o)
+ self.cache_knn_idx = knn_idx
+ # print(knn_idx.float().mean().item())
+ else:
+ knn_idx = self.cache_knn_idx
+
+ if self.with_mv_attn:
+ assert b is not None and v is not None and h is not None and w is not None
+ if self.use_checkpointing:
+ raise NotImplementedError
+
+ for i in range(len(self.blocks)):
+ # knn attention
+ x = self.blocks[i]([p, x, o], knn_idx=knn_idx)
+ # global multi-view attention
+ x = rearrange(x, "(b v h w) c -> b (v h w) c", b=b, v=v, h=h, w=w)
+ if self.with_mv_attn_lowres:
+ x = self.mv_blocks[i](x, v=v, h=h, w=w)
+ else:
+ x = self.mv_blocks[i](x)
+ x = rearrange(x, "b (v h w) c -> (b v h w) c",
+ b=b, v=v, h=h, w=w)
+ else:
+ for blk in self.blocks:
+ if self.init_use_checkpointing:
+ # checkpointing the inital reconstruction model
+ # NOTE: cannot cache knn_idx here, otherwise index out error
+ def custom_forward(p, x, o):
+ return blk((p, x, o), knn_idx=None) # knn_idx is closed over
+
+ x = torch.utils.checkpoint.checkpoint(custom_forward, p, x, o)
+ else:
+ x = blk((p, x, o), knn_idx=knn_idx)
+
+ return x
+
+
+class MultViewLowresAttn(nn.Module):
+ def __init__(self, channels, no_mv_attn=False, with_pos_enc=False, shuffle_attn_no_norm=False,
+ down_factor=4,
+ attn_proj_channels=None,
+ ):
+ super().__init__()
+
+ self.down_factor = down_factor
+ self.with_pos_enc = with_pos_enc
+
+ self.attn_proj_channels = attn_proj_channels
+
+ if attn_proj_channels:
+ ori_channels = channels
+ self.proj0 = nn.Linear(channels, attn_proj_channels)
+ channels = attn_proj_channels
+
+ if self.down_factor == 8:
+ down_factor = 4
+ else:
+ down_factor = self.down_factor
+
+ self.proj1 = nn.Linear(channels * down_factor ** 2, channels)
+ if shuffle_attn_no_norm:
+ self.norm1 = nn.Identity()
+ else:
+ self.norm1 = nn.LayerNorm(channels)
+
+ self.proj2 = nn.Linear(channels, channels * down_factor ** 2)
+
+ if shuffle_attn_no_norm:
+ self.norm2 = nn.Identity()
+ else:
+ self.norm2 = nn.LayerNorm(channels * down_factor ** 2)
+
+ self.conv = nn.Conv2d(channels, channels, 3, 1, 1)
+
+ if attn_proj_channels:
+ self.proj3 = nn.Linear(channels, ori_channels)
+
+ if no_mv_attn:
+ self.attn = nn.Identity()
+ else:
+ num_heads = 1 if self.attn_proj_channels else 4
+ self.attn = MultiViewBlock(channels, num_heads, no_attn=no_mv_attn)
+
+ def forward(self, x, v=None, h=None, w=None, y=None):
+ if y is not None:
+ return self.forward_cross_attn(x, y, v, h, w)
+ residual = x
+ if self.attn_proj_channels:
+ x = self.proj0(x)
+
+ x = rearrange(x, "b (v h w) c -> (b v) c h w", v=v, h=h, w=w)
+
+ if self.with_pos_enc:
+ x = mv_feature_add_position(x, attn_splits=1, feature_channels=x.size(1))
+ # print(x.shape)
+
+ if self.down_factor == 8:
+ # bilinear to half first to save channels
+ x = F.interpolate(x, scale_factor=0.5, mode='bilinear', align_corners=True)
+ down_factor = 4
+ else:
+ down_factor = self.down_factor
+
+ x = F.pixel_unshuffle(x, down_factor)
+
+ x = rearrange(x, "(b v) c h w -> b (v h w) c", v=v)
+ x = self.proj1(x)
+ x = self.norm1(x)
+
+ x = self.attn(x)
+
+ x = self.proj2(x)
+ x = self.norm2(x)
+
+ x = rearrange(x, "b (v h w) c -> (b v) c h w", v=v, h=h // self.down_factor, w=w // self.down_factor)
+ x = F.pixel_shuffle(x, down_factor)
+ x = self.conv(x)
+ if self.down_factor == 8:
+ # bilinear to full
+ x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True)
+ x = rearrange(x, "(b v) c h w -> b (v h w) c", v=v)
+ if self.attn_proj_channels:
+ x = self.proj3(x)
+ x = x + residual
+
+ return x
+
+ def forward_cross_attn(self, x, y, v=None, h=None, w=None):
+ residual = x
+ if self.attn_proj_channels:
+ x = self.proj0(x)
+
+ assert y is not None
+ y = rearrange(y, "b (v h w) c -> (b v) c h w", h=h, w=w) # different v with x
+ num_cross_view = y.shape[0] // x.shape[0]
+
+ x = rearrange(x, "b (v h w) c -> (b v) c h w", v=v, h=h, w=w)
+
+ if self.with_pos_enc:
+ x = mv_feature_add_position(x, attn_splits=1, feature_channels=x.size(1))
+ # print(x.shape)
+
+ if self.down_factor == 8:
+ # bilinear to half first to save channels
+ x = F.interpolate(x, scale_factor=0.5, mode='bilinear', align_corners=True)
+ y = F.interpolate(y, scale_factor=0.5, mode='bilinear', align_corners=True)
+ down_factor = 4
+ else:
+ down_factor = self.down_factor
+
+ x = F.pixel_unshuffle(x, down_factor)
+ y = F.pixel_unshuffle(y, down_factor)
+
+ x = rearrange(x, "(b v) c h w -> b (v h w) c", v=v)
+ y = rearrange(y, "(b v) c h w -> b (v h w) c", v=num_cross_view)
+ x = self.proj1(x)
+ x = self.norm1(x)
+
+ y = self.proj1(y)
+ y = self.norm1(y)
+
+ # x_tmp = self.attn(x)
+
+ # print((x - y).abs().max().item())
+
+ x = self.attn(x, y)
+
+ # there will be slight diff for self and cross attn caused by flash3
+ # print((x_tmp - x).abs().max().item())
+
+ x = self.proj2(x)
+ x = self.norm2(x)
+
+ x = rearrange(x, "b (v h w) c -> (b v) c h w", v=v, h=h // self.down_factor, w=w // self.down_factor)
+ x = F.pixel_shuffle(x, down_factor)
+ x = self.conv(x)
+ if self.down_factor == 8:
+ # bilinear to full
+ x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True)
+ x = rearrange(x, "(b v) c h w -> b (v h w) c", v=v)
+ if self.attn_proj_channels:
+ x = self.proj3(x)
+ x = x + residual
+
+ return x
+
+
+class PointLinearWrapper(nn.Module):
+ def __init__(self, in_channels, out_channels):
+ super().__init__()
+
+ self.linear = nn.Linear(in_channels, out_channels)
+
+ def forward(self, pxo, b=None, v=None, h=None, w=None):
+ p, x, o = pxo
+ x = self.linear(x)
+
+ return [p, x, o]
+
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/core/utils/__init__.py b/optgs/model/backbones/unimatch/__init__.py
similarity index 100%
rename from optgs/model/encoder/.deprecated/foundationstereo/core/utils/__init__.py
rename to optgs/model/backbones/unimatch/__init__.py
diff --git a/optgs/model/encoder/unimatch/backbone.py b/optgs/model/backbones/unimatch/backbone.py
old mode 100644
new mode 100755
similarity index 100%
rename from optgs/model/encoder/unimatch/backbone.py
rename to optgs/model/backbones/unimatch/backbone.py
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/depth_anything/__init__.py b/optgs/model/backbones/unimatch/dinov2/__init__.py
similarity index 100%
rename from optgs/model/encoder/.deprecated/foundationstereo/depth_anything/__init__.py
rename to optgs/model/backbones/unimatch/dinov2/__init__.py
diff --git a/optgs/model/encoder/unimatch/dinov2/dinov2.py b/optgs/model/backbones/unimatch/dinov2/dinov2.py
similarity index 100%
rename from optgs/model/encoder/unimatch/dinov2/dinov2.py
rename to optgs/model/backbones/unimatch/dinov2/dinov2.py
diff --git a/optgs/model/encoder/unimatch/dinov2/layers/__init__.py b/optgs/model/backbones/unimatch/dinov2/layers/__init__.py
similarity index 100%
rename from optgs/model/encoder/unimatch/dinov2/layers/__init__.py
rename to optgs/model/backbones/unimatch/dinov2/layers/__init__.py
diff --git a/optgs/model/encoder/unimatch/dinov2/layers/attention.py b/optgs/model/backbones/unimatch/dinov2/layers/attention.py
similarity index 100%
rename from optgs/model/encoder/unimatch/dinov2/layers/attention.py
rename to optgs/model/backbones/unimatch/dinov2/layers/attention.py
diff --git a/optgs/model/encoder/unimatch/dinov2/layers/block.py b/optgs/model/backbones/unimatch/dinov2/layers/block.py
similarity index 100%
rename from optgs/model/encoder/unimatch/dinov2/layers/block.py
rename to optgs/model/backbones/unimatch/dinov2/layers/block.py
diff --git a/optgs/model/encoder/unimatch/dinov2/layers/dino_head.py b/optgs/model/backbones/unimatch/dinov2/layers/dino_head.py
similarity index 100%
rename from optgs/model/encoder/unimatch/dinov2/layers/dino_head.py
rename to optgs/model/backbones/unimatch/dinov2/layers/dino_head.py
diff --git a/optgs/model/encoder/unimatch/dinov2/layers/drop_path.py b/optgs/model/backbones/unimatch/dinov2/layers/drop_path.py
similarity index 100%
rename from optgs/model/encoder/unimatch/dinov2/layers/drop_path.py
rename to optgs/model/backbones/unimatch/dinov2/layers/drop_path.py
diff --git a/optgs/model/encoder/unimatch/dinov2/layers/layer_scale.py b/optgs/model/backbones/unimatch/dinov2/layers/layer_scale.py
similarity index 100%
rename from optgs/model/encoder/unimatch/dinov2/layers/layer_scale.py
rename to optgs/model/backbones/unimatch/dinov2/layers/layer_scale.py
diff --git a/optgs/model/encoder/unimatch/dinov2/layers/mlp.py b/optgs/model/backbones/unimatch/dinov2/layers/mlp.py
similarity index 100%
rename from optgs/model/encoder/unimatch/dinov2/layers/mlp.py
rename to optgs/model/backbones/unimatch/dinov2/layers/mlp.py
diff --git a/optgs/model/encoder/unimatch/dinov2/layers/patch_embed.py b/optgs/model/backbones/unimatch/dinov2/layers/patch_embed.py
similarity index 100%
rename from optgs/model/encoder/unimatch/dinov2/layers/patch_embed.py
rename to optgs/model/backbones/unimatch/dinov2/layers/patch_embed.py
diff --git a/optgs/model/encoder/unimatch/dinov2/layers/swiglu_ffn.py b/optgs/model/backbones/unimatch/dinov2/layers/swiglu_ffn.py
similarity index 100%
rename from optgs/model/encoder/unimatch/dinov2/layers/swiglu_ffn.py
rename to optgs/model/backbones/unimatch/dinov2/layers/swiglu_ffn.py
diff --git a/optgs/model/encoder/unimatch/dpt_head.py b/optgs/model/backbones/unimatch/dpt_head.py
similarity index 100%
rename from optgs/model/encoder/unimatch/dpt_head.py
rename to optgs/model/backbones/unimatch/dpt_head.py
diff --git a/optgs/model/encoder/unimatch/feature_upsampler.py b/optgs/model/backbones/unimatch/feature_upsampler.py
similarity index 100%
rename from optgs/model/encoder/unimatch/feature_upsampler.py
rename to optgs/model/backbones/unimatch/feature_upsampler.py
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/depth_anything/util/__init__.py b/optgs/model/backbones/unimatch/ldm_unet/__init__.py
similarity index 100%
rename from optgs/model/encoder/.deprecated/foundationstereo/depth_anything/util/__init__.py
rename to optgs/model/backbones/unimatch/ldm_unet/__init__.py
diff --git a/optgs/model/encoder/unimatch/ldm_unet/attention.py b/optgs/model/backbones/unimatch/ldm_unet/attention.py
similarity index 100%
rename from optgs/model/encoder/unimatch/ldm_unet/attention.py
rename to optgs/model/backbones/unimatch/ldm_unet/attention.py
diff --git a/optgs/model/encoder/unimatch/ldm_unet/cross_attention.py b/optgs/model/backbones/unimatch/ldm_unet/cross_attention.py
old mode 100644
new mode 100755
similarity index 100%
rename from optgs/model/encoder/unimatch/ldm_unet/cross_attention.py
rename to optgs/model/backbones/unimatch/ldm_unet/cross_attention.py
diff --git a/optgs/model/encoder/unimatch/ldm_unet/unet.py b/optgs/model/backbones/unimatch/ldm_unet/unet.py
similarity index 80%
rename from optgs/model/encoder/unimatch/ldm_unet/unet.py
rename to optgs/model/backbones/unimatch/ldm_unet/unet.py
index 7fe6483313e07de89e2ea8cd7e840f5027606d1e..339201c86f190976854645e26a5a1e19ab22a736 100644
--- a/optgs/model/encoder/unimatch/ldm_unet/unet.py
+++ b/optgs/model/backbones/unimatch/ldm_unet/unet.py
@@ -1,33 +1,29 @@
-from abc import abstractmethod
-from functools import partial
import math
-from typing import Iterable
-from einops import rearrange
+from abc import abstractmethod
import numpy as np
+import torch
import torch as th
import torch.nn as nn
import torch.nn.functional as F
-import torch
+from einops import rearrange
+from .attention import SpatialTransformer
+from .cross_attention import UNetCrossAttentionBlock
from .util import (
checkpoint,
conv_nd,
- linear,
avg_pool_nd,
zero_module,
normalization,
- timestep_embedding,
)
-from .attention import SpatialTransformer
-
-from .cross_attention import UNetCrossAttentionBlock
# dummy replace
def convert_module_to_f16(x):
pass
+
def convert_module_to_f32(x):
pass
@@ -39,11 +35,11 @@ class AttentionPool2d(nn.Module):
"""
def __init__(
- self,
- spacial_dim: int,
- embed_dim: int,
- num_heads_channels: int,
- output_dim: int = None,
+ self,
+ spacial_dim: int,
+ embed_dim: int,
+ num_heads_channels: int,
+ output_dim: int = None,
):
super().__init__()
self.positional_embedding = nn.Parameter(th.randn(embed_dim, spacial_dim ** 2 + 1) / embed_dim ** 0.5)
@@ -102,8 +98,8 @@ class Upsample(nn.Module):
"""
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
- ):
+ downsample_3ddim=False, # downsample all 3d dims instead of only spatial dims
+ ):
super().__init__()
self.channels = channels
self.out_channels = out_channels or channels
@@ -126,16 +122,18 @@ class Upsample(nn.Module):
x = self.conv(x)
return x
+
class TransposedUpsample(nn.Module):
'Learned 2x upsampling without padding'
+
def __init__(self, channels, out_channels=None, ks=5):
super().__init__()
self.channels = channels
self.out_channels = out_channels or channels
- self.up = nn.ConvTranspose2d(self.channels,self.out_channels,kernel_size=ks,stride=2)
+ self.up = nn.ConvTranspose2d(self.channels, self.out_channels, kernel_size=ks, stride=2)
- def forward(self,x):
+ def forward(self, x):
return self.up(x)
@@ -148,9 +146,9 @@ class Downsample(nn.Module):
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
- ):
+ 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
@@ -192,20 +190,20 @@ class ResBlock(TimestepBlock):
"""
def __init__(
- self,
- channels,
- emb_channels,
- dropout,
- out_channels=None,
- use_conv=False,
- use_scale_shift_norm=False,
- dims=2,
- use_checkpoint=False,
- up=False,
- down=False,
- postnorm=False,
- channels_per_group=None,
- kernel_size=3,
+ self,
+ channels,
+ emb_channels,
+ dropout,
+ out_channels=None,
+ use_conv=False,
+ use_scale_shift_norm=False,
+ dims=2,
+ use_checkpoint=False,
+ up=False,
+ down=False,
+ postnorm=False,
+ channels_per_group=None,
+ kernel_size=3,
):
super().__init__()
self.channels = channels
@@ -255,7 +253,7 @@ class ResBlock(TimestepBlock):
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),
@@ -317,16 +315,16 @@ class AttentionBlock(nn.Module):
"""
def __init__(
- self,
- channels,
- num_heads=1,
- num_head_channels=-1,
- use_checkpoint=False,
- use_new_attention_order=False,
- postnorm=False,
- channels_per_group=None,
- num_frames=2,
- use_cross_view_self_attn=False,
+ self,
+ channels,
+ num_heads=1,
+ num_head_channels=-1,
+ use_checkpoint=False,
+ use_new_attention_order=False,
+ postnorm=False,
+ channels_per_group=None,
+ num_frames=2,
+ use_cross_view_self_attn=False,
):
super().__init__()
@@ -336,7 +334,7 @@ class AttentionBlock(nn.Module):
self.num_heads = num_heads
else:
assert (
- channels % num_head_channels == 0
+ 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.use_checkpoint = use_checkpoint
@@ -362,8 +360,7 @@ class AttentionBlock(nn.Module):
self.postnorm = postnorm
def forward(self, x):
- # return checkpoint(self._forward, (x,), self.parameters(), True) # TODO: check checkpoint usage, is True # TODO: fix the .half call!!!
- #return pt_checkpoint(self._forward, x) # pytorch
+ # return pt_checkpoint(self._forward, x) # pytorch
return self._forward(x)
def _forward(self, x):
@@ -392,19 +389,19 @@ class CrossAttentionBlock(nn.Module):
"""
def __init__(
- self,
- channels,
- condition_channels,
- num_heads=8,
- proj_channels=512,
- num_views=3,
- num_head_channels=-1,
- use_checkpoint=False,
- use_new_attention_order=False,
- channels_per_group=None,
- with_norm=False,
- tanh_gating=False, # following Flamingo
- ffn_after_cross_attn=False, # following Flamingo
+ self,
+ channels,
+ condition_channels,
+ num_heads=8,
+ proj_channels=512,
+ num_views=3,
+ num_head_channels=-1,
+ use_checkpoint=False,
+ use_new_attention_order=False,
+ channels_per_group=None,
+ with_norm=False,
+ tanh_gating=False, # following Flamingo
+ ffn_after_cross_attn=False, # following Flamingo
):
super().__init__()
@@ -420,7 +417,6 @@ class CrossAttentionBlock(nn.Module):
self.k_proj = nn.Linear(condition_channels, proj_channels)
self.v_proj = nn.Linear(condition_channels, proj_channels)
- # TODO: whether need norm layer
# self.norm = normalization(proj_channels, channels_per_group=channels_per_group)
if self.tanh_gating:
@@ -436,14 +432,13 @@ class CrossAttentionBlock(nn.Module):
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)
- )
+ 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 checkpoint(self._forward, (x,), self.parameters(), True) # TODO: check checkpoint usage, is True # TODO: fix the .half call!!!
- #return pt_checkpoint(self._forward, x) # pytorch
+ # return pt_checkpoint(self._forward, x) # pytorch
return self._forward(x, y)
def _forward(self, x, y=None):
@@ -462,7 +457,7 @@ class CrossAttentionBlock(nn.Module):
x = x.permute(0, 2, 3, 4, 1).reshape(b * d, h * w, c1) # [B*D, H*W, C1]
y = y.permute(0, 2, 1, 3).reshape(b * d, ly, c2) # [B*D, H*W, C2]
-
+
c = self.proj_channels
q = self.q_proj(x) # [B*D, H*W, C]
@@ -475,7 +470,8 @@ class CrossAttentionBlock(nn.Module):
k = k.view(b * d, ly, self.num_head, c // self.num_head) # [B*D, H*W, N, C]
v = v.view(b * d, ly, self.num_head, c // self.num_head) # [B*D, H*W, N, C]
- scores = torch.matmul(q.permute(0, 2, 1, 3), k.permute(0, 2, 3, 1)) / ((c // self.num_head) ** 0.5) # [B*D, N, H*W, H*W]
+ scores = torch.matmul(q.permute(0, 2, 1, 3), k.permute(0, 2, 3, 1)) / (
+ (c // self.num_head) ** 0.5) # [B*D, N, H*W, H*W]
prob = torch.softmax(scores, dim=-1)
out = torch.matmul(prob, v.permute(0, 2, 1, 3)) # [B*D, H*W, N, C]
out = out.view(b * d, lx, -1) # [B*D, H*W, C]
@@ -639,51 +635,51 @@ class UNetModel(nn.Module):
"""
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,
- use_checkpoint=False,
- use_fp16=False,
- 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,
+ 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,
+ use_checkpoint=False,
+ use_fp16=False,
+ 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:
@@ -722,7 +718,7 @@ class UNetModel(nn.Module):
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
@@ -775,7 +771,7 @@ class UNetModel(nn.Module):
num_heads = ch // num_head_channels
dim_head = num_head_channels
if legacy:
- #num_heads = 1
+ # num_heads = 1
dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
if not no_self_attn: # only cross attn, without self attn
@@ -798,12 +794,12 @@ class UNetModel(nn.Module):
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,
- )
+ 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))
@@ -844,7 +840,7 @@ class UNetModel(nn.Module):
num_heads = ch // num_head_channels
dim_head = num_head_channels
if legacy:
- #num_heads = 1
+ # num_heads = 1
dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
if self.middle_block_attn:
@@ -872,16 +868,16 @@ class UNetModel(nn.Module):
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
- ),
+ 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(),
+ 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,
@@ -897,28 +893,28 @@ class UNetModel(nn.Module):
else:
if self.middle_block_no_identity:
self.middle_block = nn.Sequential(
- ResBlock(
- ch,
- time_embed_dim,
- dropout,
- dims=dims,
- use_checkpoint=use_checkpoint,
- 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_checkpoint=use_checkpoint,
- 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_checkpoint=use_checkpoint,
+ 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_checkpoint=use_checkpoint,
+ 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(
@@ -934,12 +930,12 @@ class UNetModel(nn.Module):
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(),
+ 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,
@@ -980,7 +976,7 @@ class UNetModel(nn.Module):
num_heads = ch // num_head_channels
dim_head = num_head_channels
if legacy:
- #num_heads = 1
+ # num_heads = 1
dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
if not no_self_attn: # only cross attn, without self attn
@@ -1003,14 +999,14 @@ class UNetModel(nn.Module):
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,
- )
+ 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(
@@ -1029,8 +1025,8 @@ class UNetModel(nn.Module):
)
if resblock_updown
else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch,
- downsample_3ddim=downsample_3ddim,
- )
+ downsample_3ddim=downsample_3ddim,
+ )
)
ds //= 2
self.output_blocks.append(nn.Sequential(*layers))
@@ -1047,7 +1043,9 @@ class UNetModel(nn.Module):
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)),
+ 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:
@@ -1067,10 +1065,10 @@ class UNetModel(nn.Module):
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),
- #nn.LogSoftmax(dim=1) # change to cross_entropy and produce non-normalized logits
- )
+ normalization(ch, channels_per_group=channels_per_group),
+ conv_nd(dims, model_channels, n_embed, 1),
+ # nn.LogSoftmax(dim=1) # change to cross_entropy and produce non-normalized logits
+ )
def convert_to_fp16(self):
"""
@@ -1088,7 +1086,7 @@ class UNetModel(nn.Module):
self.middle_block.apply(convert_module_to_f32)
self.output_blocks.apply(convert_module_to_f32)
- def forward(self, x, num_views=None, timesteps=None, context=None, y=None,**kwargs):
+ 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.
@@ -1098,7 +1096,7 @@ class UNetModel(nn.Module):
:return: an [N x C x ...] Tensor of outputs.
"""
assert (y is not None) == (
- self.num_classes is not None
+ self.num_classes is not None
), "must specify y if and only if the model is class-conditional"
hs = []
# t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
@@ -1174,15 +1172,15 @@ class StackUNet(nn.Module):
zero_final_layer=False,
resblock_updown=False,
channels_per_group=None,
- cross_attn_condition=False,
+ cross_attn_condition=False,
cross_attn_with_norm=False,
- condition_channels=128,
- tanh_gating=False,
- ffn_after_cross_attn=False,
- condition_num_views=3,
- no_self_attn=False,
- middle_block_no_identity=False,
- conv_kernel_size=3,
+ condition_channels=128,
+ tanh_gating=False,
+ ffn_after_cross_attn=False,
+ condition_num_views=3,
+ no_self_attn=False,
+ middle_block_no_identity=False,
+ conv_kernel_size=3,
):
super().__init__()
@@ -1195,31 +1193,31 @@ class StackUNet(nn.Module):
for i in range(num_stacks):
self.stacks.append(UNetModel(image_size=None,
- in_channels=in_channels,
- model_channels=model_channels,
- out_channels=out_channels,
- num_res_blocks=num_res_blocks,
- attention_resolutions=attention_resolutions,
- channel_mult=channel_mult,
- num_head_channels=num_head_channels,
- dims=dims,
- middle_block_attn=middle_block_attn,
- middle_block_no_identity=middle_block_no_identity,
- postnorm=postnorm,
- attn_prenorm=attn_prenorm,
- zero_final_layer=zero_final_layer and i == 0,
- resblock_updown=resblock_updown,
- channels_per_group=channels_per_group,
- cross_attn_condition=cross_attn_condition,
- tanh_gating=tanh_gating,
- ffn_after_cross_attn=ffn_after_cross_attn,
- cross_attn_with_norm=cross_attn_with_norm,
- condition_channels=condition_channels,
- condition_num_views=condition_num_views,
- no_self_attn=no_self_attn,
- conv_kernel_size=conv_kernel_size,
- )
- )
+ in_channels=in_channels,
+ model_channels=model_channels,
+ out_channels=out_channels,
+ num_res_blocks=num_res_blocks,
+ attention_resolutions=attention_resolutions,
+ channel_mult=channel_mult,
+ num_head_channels=num_head_channels,
+ dims=dims,
+ middle_block_attn=middle_block_attn,
+ middle_block_no_identity=middle_block_no_identity,
+ postnorm=postnorm,
+ attn_prenorm=attn_prenorm,
+ zero_final_layer=zero_final_layer and i == 0,
+ resblock_updown=resblock_updown,
+ channels_per_group=channels_per_group,
+ cross_attn_condition=cross_attn_condition,
+ tanh_gating=tanh_gating,
+ ffn_after_cross_attn=ffn_after_cross_attn,
+ cross_attn_with_norm=cross_attn_with_norm,
+ condition_channels=condition_channels,
+ condition_num_views=condition_num_views,
+ no_self_attn=no_self_attn,
+ conv_kernel_size=conv_kernel_size,
+ )
+ )
in_channels = out_channels
@@ -1230,7 +1228,6 @@ class StackUNet(nn.Module):
dims, out_channels, in_channels, 3, padding=1
)))
-
def forward(self, x, context=None):
x = self.stacks[0](x, context=context)
for i in range(self.num_stacks - 1):
@@ -1238,6 +1235,3 @@ class StackUNet(nn.Module):
x = x + residual
return x
-
-
-
diff --git a/optgs/model/encoder/unimatch/ldm_unet/util.py b/optgs/model/backbones/unimatch/ldm_unet/util.py
similarity index 100%
rename from optgs/model/encoder/unimatch/ldm_unet/util.py
rename to optgs/model/backbones/unimatch/ldm_unet/util.py
diff --git a/optgs/model/encoder/unimatch/matching.py b/optgs/model/backbones/unimatch/matching.py
old mode 100644
new mode 100755
similarity index 100%
rename from optgs/model/encoder/unimatch/matching.py
rename to optgs/model/backbones/unimatch/matching.py
diff --git a/optgs/model/encoder/unimatch/mv_transformer.py b/optgs/model/backbones/unimatch/mv_transformer.py
similarity index 100%
rename from optgs/model/encoder/unimatch/mv_transformer.py
rename to optgs/model/backbones/unimatch/mv_transformer.py
diff --git a/optgs/model/encoder/unimatch/mv_unimatch.py b/optgs/model/backbones/unimatch/mv_unimatch.py
old mode 100644
new mode 100755
similarity index 99%
rename from optgs/model/encoder/unimatch/mv_unimatch.py
rename to optgs/model/backbones/unimatch/mv_unimatch.py
index faf4a551df696b0af5395fae1440a30ed15b8d34..611a6b893ffebdb4bbfffa4fdfd5bfb516beebd8
--- a/optgs/model/encoder/unimatch/mv_unimatch.py
+++ b/optgs/model/backbones/unimatch/mv_unimatch.py
@@ -346,7 +346,6 @@ class MultiViewUniMatch(nn.Module):
# mono feature
ori_h, ori_w = images.shape[-2:]
- # TODO: support portrait images later
assert ori_h <= ori_w
if ori_w > self.max_mono_vit_input_size:
resize_w = self.max_mono_vit_input_size // 14 * 14
@@ -401,7 +400,6 @@ class MultiViewUniMatch(nn.Module):
results_dict.update({"raw_mono_features": raw_mono_features})
# last mono feature
- # TODO: use all the intermediate features for depth estimation?
mono_features = mono_intermediate_features[-1]
if self.lowest_feature_resolution == 4:
diff --git a/optgs/model/encoder/unimatch/position.py b/optgs/model/backbones/unimatch/position.py
old mode 100644
new mode 100755
similarity index 100%
rename from optgs/model/encoder/unimatch/position.py
rename to optgs/model/backbones/unimatch/position.py
diff --git a/optgs/model/encoder/unimatch/utils.py b/optgs/model/backbones/unimatch/utils.py
old mode 100644
new mode 100755
similarity index 100%
rename from optgs/model/encoder/unimatch/utils.py
rename to optgs/model/backbones/unimatch/utils.py
diff --git a/optgs/model/encoder/unimatch/vit_fpn.py b/optgs/model/backbones/unimatch/vit_fpn.py
similarity index 100%
rename from optgs/model/encoder/unimatch/vit_fpn.py
rename to optgs/model/backbones/unimatch/vit_fpn.py
diff --git a/optgs/model/colmap_utils/convert_to_colmap.py b/optgs/model/colmap_utils/convert_to_colmap.py
deleted file mode 100644
index 7a6ba58fef9070bb96da63f680f92421c434279c..0000000000000000000000000000000000000000
--- a/optgs/model/colmap_utils/convert_to_colmap.py
+++ /dev/null
@@ -1,57 +0,0 @@
-import json
-import numpy as np
-from pathlib import Path
-from .read_write_model import Camera, write_cameras_binary
-
-
-def save_opencv_camera(K: np.ndarray, json_path: str, output_path, image_size: tuple, camera_id=1):
- assert K.shape == (3, 3), "Intrinsic matrix K must be 3x3"
-
- # Extract intrinsics from K
- fx = K[0, 0]
- fy = K[1, 1]
- cx = K[0, 2]
- cy = K[1, 2]
-
- # Load distortion coefficients from JSON
- with open(json_path, 'r') as f:
- data = json.load(f)
-
- k1 = data.get("k1", 0.0)
- k2 = data.get("k2", 0.0)
- p1 = data.get("p1", 0.0)
- p2 = data.get("p2", 0.0)
-
- # Image resolution
- w, h = image_size
-
- # COLMAP OPENCV model requires 8 parameters
- params = (fx, fy, cx, cy, k1, k2, p1, p2)
-
- # Create and write Camera object
- camera = Camera(
- id=camera_id,
- model="OPENCV",
- width=w,
- height=h,
- params=params
- )
-
- output_path = Path(output_path)
- output_path.mkdir(parents=True, exist_ok=True)
- write_cameras_binary({camera_id: camera}, output_path / "cameras.bin")
-
-# Example usage
-if __name__ == "__main__":
- K = np.array([
- [629.23, 0, 480.0],
- [0, 625.73, 270.0],
- [0, 0, 1.0]
- ])
- save_opencv_camera(
- K=K,
- json_path="transforms.json",
- output_path="output_model",
- image_size=(960, 540), # e.g., 4x downsampled from 3840x2160
- camera_id=1
- )
diff --git a/optgs/model/colmap_utils/extract_sparse_view_extrinsics.py b/optgs/model/colmap_utils/extract_sparse_view_extrinsics.py
deleted file mode 100644
index 9db9c22273c2403adc1ce8afd1e890336a5007f5..0000000000000000000000000000000000000000
--- a/optgs/model/colmap_utils/extract_sparse_view_extrinsics.py
+++ /dev/null
@@ -1,54 +0,0 @@
-import numpy as np
-from pathlib import Path
-from .read_write_model import (
- read_images_binary,
- write_images_binary,
- write_points3D_binary,
- Image as ColmapImage
-)
-
-def extract_sparse_images_bin(input_model_dir, output_model_dir, selected_image_ids, keep_features=False):
- input_model_dir = Path(input_model_dir)
- output_model_dir = Path(output_model_dir)
- output_model_dir.mkdir(parents=True, exist_ok=True)
-
- # Load full images.bin
- images = read_images_binary(input_model_dir / "images.bin")
-
- # Select and blank the images
- sparse_images = {}
- for image_id in selected_image_ids:
- image = images[image_id]
-
- if keep_features:
- xys = image.xys
- point3D_ids = image.point3D_ids
- else:
- xys = np.empty((0, 2))
- point3D_ids = np.empty((0,), dtype=int)
-
- blank_image = ColmapImage(
- id=image.id,
- qvec=image.qvec,
- tvec=image.tvec,
- camera_id=image.camera_id,
- name=image.name,
- xys=xys,
- point3D_ids=point3D_ids
- )
- sparse_images[image_id] = blank_image
-
- # Save sparse images.bin
- write_images_binary(sparse_images, output_model_dir / "images.bin")
-
- # Save empty points3D.bin
- write_points3D_binary({}, output_model_dir / "points3D.bin")
-
-# Example usage
-if __name__ == "__main__":
- selected_ids = [1, 4, 10, 20] # Replace with your sparse frame IDs
- extract_sparse_images_bin(
- input_model_dir="dense_model",
- output_model_dir="sparse_model",
- selected_image_ids=selected_ids
- )
diff --git a/optgs/model/colmap_utils/read_write_model.py b/optgs/model/colmap_utils/read_write_model.py
deleted file mode 100644
index 1510f12ce5d623a703f9cf812b34b697d3514d5c..0000000000000000000000000000000000000000
--- a/optgs/model/colmap_utils/read_write_model.py
+++ /dev/null
@@ -1,605 +0,0 @@
-# Copyright (c), ETH Zurich and UNC Chapel Hill.
-# All rights reserved.
-#
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-#
-# * Redistributions of source code must retain the above copyright
-# notice, this list of conditions and the following disclaimer.
-#
-# * Redistributions in binary form must reproduce the above copyright
-# notice, this list of conditions and the following disclaimer in the
-# documentation and/or other materials provided with the distribution.
-#
-# * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of
-# its contributors may be used to endorse or promote products derived
-# from this software without specific prior written permission.
-#
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
-
-
-import argparse
-import collections
-import os
-import struct
-
-import numpy as np
-
-CameraModel = collections.namedtuple(
- "CameraModel", ["model_id", "model_name", "num_params"]
-)
-Camera = collections.namedtuple(
- "Camera", ["id", "model", "width", "height", "params"]
-)
-BaseImage = collections.namedtuple(
- "Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"]
-)
-Point3D = collections.namedtuple(
- "Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"]
-)
-
-
-class Image(BaseImage):
- def qvec2rotmat(self):
- return qvec2rotmat(self.qvec)
-
-
-CAMERA_MODELS = {
- CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3),
- CameraModel(model_id=1, model_name="PINHOLE", num_params=4),
- CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4),
- CameraModel(model_id=3, model_name="RADIAL", num_params=5),
- CameraModel(model_id=4, model_name="OPENCV", num_params=8),
- CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8),
- CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12),
- CameraModel(model_id=7, model_name="FOV", num_params=5),
- CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4),
- CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5),
- CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12),
-}
-CAMERA_MODEL_IDS = dict(
- [(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS]
-)
-CAMERA_MODEL_NAMES = dict(
- [(camera_model.model_name, camera_model) for camera_model in CAMERA_MODELS]
-)
-
-
-def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
- """Read and unpack the next bytes from a binary file.
- :param fid:
- :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc.
- :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.
- :param endian_character: Any of {@, =, <, >, !}
- :return: Tuple of read and unpacked values.
- """
- data = fid.read(num_bytes)
- return struct.unpack(endian_character + format_char_sequence, data)
-
-
-def write_next_bytes(fid, data, format_char_sequence, endian_character="<"):
- """pack and write to a binary file.
- :param fid:
- :param data: data to send, if multiple elements are sent at the same time,
- they should be encapsuled either in a list or a tuple
- :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.
- should be the same length as the data list or tuple
- :param endian_character: Any of {@, =, <, >, !}
- """
- if isinstance(data, (list, tuple)):
- bytes = struct.pack(endian_character + format_char_sequence, *data)
- else:
- bytes = struct.pack(endian_character + format_char_sequence, data)
- fid.write(bytes)
-
-
-def read_cameras_text(path):
- """
- see: src/colmap/scene/reconstruction.cc
- void Reconstruction::WriteCamerasText(const std::string& path)
- void Reconstruction::ReadCamerasText(const std::string& path)
- """
- cameras = {}
- with open(path, "r") as fid:
- while True:
- line = fid.readline()
- if not line:
- break
- line = line.strip()
- if len(line) > 0 and line[0] != "#":
- elems = line.split()
- camera_id = int(elems[0])
- model = elems[1]
- width = int(elems[2])
- height = int(elems[3])
- params = np.array(tuple(map(float, elems[4:])))
- cameras[camera_id] = Camera(
- id=camera_id,
- model=model,
- width=width,
- height=height,
- params=params,
- )
- return cameras
-
-
-def read_cameras_binary(path_to_model_file):
- """
- see: src/colmap/scene/reconstruction.cc
- void Reconstruction::WriteCamerasBinary(const std::string& path)
- void Reconstruction::ReadCamerasBinary(const std::string& path)
- """
- cameras = {}
- with open(path_to_model_file, "rb") as fid:
- num_cameras = read_next_bytes(fid, 8, "Q")[0]
- for _ in range(num_cameras):
- camera_properties = read_next_bytes(
- fid, num_bytes=24, format_char_sequence="iiQQ"
- )
- camera_id = camera_properties[0]
- model_id = camera_properties[1]
- model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name
- width = camera_properties[2]
- height = camera_properties[3]
- num_params = CAMERA_MODEL_IDS[model_id].num_params
- params = read_next_bytes(
- fid,
- num_bytes=8 * num_params,
- format_char_sequence="d" * num_params,
- )
- cameras[camera_id] = Camera(
- id=camera_id,
- model=model_name,
- width=width,
- height=height,
- params=np.array(params),
- )
- assert len(cameras) == num_cameras
- return cameras
-
-
-def write_cameras_text(cameras, path):
- """
- see: src/colmap/scene/reconstruction.cc
- void Reconstruction::WriteCamerasText(const std::string& path)
- void Reconstruction::ReadCamerasText(const std::string& path)
- """
- HEADER = (
- "# Camera list with one line of data per camera:\n"
- + "# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n"
- + "# Number of cameras: {}\n".format(len(cameras))
- )
- with open(path, "w") as fid:
- fid.write(HEADER)
- for _, cam in cameras.items():
- to_write = [cam.id, cam.model, cam.width, cam.height, *cam.params]
- line = " ".join([str(elem) for elem in to_write])
- fid.write(line + "\n")
-
-
-def write_cameras_binary(cameras, path_to_model_file):
- """
- see: src/colmap/scene/reconstruction.cc
- void Reconstruction::WriteCamerasBinary(const std::string& path)
- void Reconstruction::ReadCamerasBinary(const std::string& path)
- """
- with open(path_to_model_file, "wb") as fid:
- write_next_bytes(fid, len(cameras), "Q")
- for _, cam in cameras.items():
- model_id = CAMERA_MODEL_NAMES[cam.model].model_id
- camera_properties = [cam.id, model_id, cam.width, cam.height]
- write_next_bytes(fid, camera_properties, "iiQQ")
- for p in cam.params:
- write_next_bytes(fid, float(p), "d")
- return cameras
-
-
-def read_images_text(path):
- """
- see: src/colmap/scene/reconstruction.cc
- void Reconstruction::ReadImagesText(const std::string& path)
- void Reconstruction::WriteImagesText(const std::string& path)
- """
- images = {}
- with open(path, "r") as fid:
- while True:
- line = fid.readline()
- if not line:
- break
- line = line.strip()
- if len(line) > 0 and line[0] != "#":
- elems = line.split()
- image_id = int(elems[0])
- qvec = np.array(tuple(map(float, elems[1:5])))
- tvec = np.array(tuple(map(float, elems[5:8])))
- camera_id = int(elems[8])
- image_name = elems[9]
- elems = fid.readline().split()
- xys = np.column_stack(
- [
- tuple(map(float, elems[0::3])),
- tuple(map(float, elems[1::3])),
- ]
- )
- point3D_ids = np.array(tuple(map(int, elems[2::3])))
- images[image_id] = Image(
- id=image_id,
- qvec=qvec,
- tvec=tvec,
- camera_id=camera_id,
- name=image_name,
- xys=xys,
- point3D_ids=point3D_ids,
- )
- return images
-
-
-def read_images_binary(path_to_model_file):
- """
- see: src/colmap/scene/reconstruction.cc
- void Reconstruction::ReadImagesBinary(const std::string& path)
- void Reconstruction::WriteImagesBinary(const std::string& path)
- """
- images = {}
- with open(path_to_model_file, "rb") as fid:
- num_reg_images = read_next_bytes(fid, 8, "Q")[0]
- for _ in range(num_reg_images):
- binary_image_properties = read_next_bytes(
- fid, num_bytes=64, format_char_sequence="idddddddi"
- )
- image_id = binary_image_properties[0]
- qvec = np.array(binary_image_properties[1:5])
- tvec = np.array(binary_image_properties[5:8])
- camera_id = binary_image_properties[8]
- binary_image_name = b""
- current_char = read_next_bytes(fid, 1, "c")[0]
- while current_char != b"\x00": # look for the ASCII 0 entry
- binary_image_name += current_char
- current_char = read_next_bytes(fid, 1, "c")[0]
- image_name = binary_image_name.decode("utf-8")
- num_points2D = read_next_bytes(
- fid, num_bytes=8, format_char_sequence="Q"
- )[0]
- x_y_id_s = read_next_bytes(
- fid,
- num_bytes=24 * num_points2D,
- format_char_sequence="ddq" * num_points2D,
- )
- xys = np.column_stack(
- [
- tuple(map(float, x_y_id_s[0::3])),
- tuple(map(float, x_y_id_s[1::3])),
- ]
- )
- point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3])))
- images[image_id] = Image(
- id=image_id,
- qvec=qvec,
- tvec=tvec,
- camera_id=camera_id,
- name=image_name,
- xys=xys,
- point3D_ids=point3D_ids,
- )
- return images
-
-
-def write_images_text(images, path):
- """
- see: src/colmap/scene/reconstruction.cc
- void Reconstruction::ReadImagesText(const std::string& path)
- void Reconstruction::WriteImagesText(const std::string& path)
- """
- if len(images) == 0:
- mean_observations = 0
- else:
- mean_observations = sum(
- (len(img.point3D_ids) for _, img in images.items())
- ) / len(images)
- HEADER = (
- "# Image list with two lines of data per image:\n"
- + "# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n"
- + "# POINTS2D[] as (X, Y, POINT3D_ID)\n"
- + "# Number of images: {}, mean observations per image: {}\n".format(
- len(images), mean_observations
- )
- )
-
- with open(path, "w") as fid:
- fid.write(HEADER)
- for _, img in images.items():
- image_header = [
- img.id,
- *img.qvec,
- *img.tvec,
- img.camera_id,
- img.name,
- ]
- first_line = " ".join(map(str, image_header))
- fid.write(first_line + "\n")
-
- points_strings = []
- for xy, point3D_id in zip(img.xys, img.point3D_ids):
- points_strings.append(" ".join(map(str, [*xy, point3D_id])))
- fid.write(" ".join(points_strings) + "\n")
-
-
-def write_images_binary(images, path_to_model_file):
- """
- see: src/colmap/scene/reconstruction.cc
- void Reconstruction::ReadImagesBinary(const std::string& path)
- void Reconstruction::WriteImagesBinary(const std::string& path)
- """
- with open(path_to_model_file, "wb") as fid:
- write_next_bytes(fid, len(images), "Q")
- for _, img in images.items():
- write_next_bytes(fid, img.id, "i")
- write_next_bytes(fid, img.qvec.tolist(), "dddd")
- write_next_bytes(fid, img.tvec.tolist(), "ddd")
- write_next_bytes(fid, img.camera_id, "i")
- for char in img.name:
- write_next_bytes(fid, char.encode("utf-8"), "c")
- write_next_bytes(fid, b"\x00", "c")
- write_next_bytes(fid, len(img.point3D_ids), "Q")
- for xy, p3d_id in zip(img.xys, img.point3D_ids):
- write_next_bytes(fid, [*xy, p3d_id], "ddq")
-
-
-def read_points3D_text(path):
- """
- see: src/colmap/scene/reconstruction.cc
- void Reconstruction::ReadPoints3DText(const std::string& path)
- void Reconstruction::WritePoints3DText(const std::string& path)
- """
- points3D = {}
- with open(path, "r") as fid:
- while True:
- line = fid.readline()
- if not line:
- break
- line = line.strip()
- if len(line) > 0 and line[0] != "#":
- elems = line.split()
- point3D_id = int(elems[0])
- xyz = np.array(tuple(map(float, elems[1:4])))
- rgb = np.array(tuple(map(int, elems[4:7])))
- error = float(elems[7])
- image_ids = np.array(tuple(map(int, elems[8::2])))
- point2D_idxs = np.array(tuple(map(int, elems[9::2])))
- points3D[point3D_id] = Point3D(
- id=point3D_id,
- xyz=xyz,
- rgb=rgb,
- error=error,
- image_ids=image_ids,
- point2D_idxs=point2D_idxs,
- )
- return points3D
-
-
-def read_points3D_binary(path_to_model_file):
- """
- see: src/colmap/scene/reconstruction.cc
- void Reconstruction::ReadPoints3DBinary(const std::string& path)
- void Reconstruction::WritePoints3DBinary(const std::string& path)
- """
- points3D = {}
- with open(path_to_model_file, "rb") as fid:
- num_points = read_next_bytes(fid, 8, "Q")[0]
- for _ in range(num_points):
- binary_point_line_properties = read_next_bytes(
- fid, num_bytes=43, format_char_sequence="QdddBBBd"
- )
- point3D_id = binary_point_line_properties[0]
- xyz = np.array(binary_point_line_properties[1:4])
- rgb = np.array(binary_point_line_properties[4:7])
- error = np.array(binary_point_line_properties[7])
- track_length = read_next_bytes(
- fid, num_bytes=8, format_char_sequence="Q"
- )[0]
- track_elems = read_next_bytes(
- fid,
- num_bytes=8 * track_length,
- format_char_sequence="ii" * track_length,
- )
- image_ids = np.array(tuple(map(int, track_elems[0::2])))
- point2D_idxs = np.array(tuple(map(int, track_elems[1::2])))
- points3D[point3D_id] = Point3D(
- id=point3D_id,
- xyz=xyz,
- rgb=rgb,
- error=error,
- image_ids=image_ids,
- point2D_idxs=point2D_idxs,
- )
- return points3D
-
-
-def write_points3D_text(points3D, path):
- """
- see: src/colmap/scene/reconstruction.cc
- void Reconstruction::ReadPoints3DText(const std::string& path)
- void Reconstruction::WritePoints3DText(const std::string& path)
- """
- if len(points3D) == 0:
- mean_track_length = 0
- else:
- mean_track_length = sum(
- (len(pt.image_ids) for _, pt in points3D.items())
- ) / len(points3D)
- HEADER = (
- "# 3D point list with one line of data per point:\n"
- + "# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n"
- + "# Number of points: {}, mean track length: {}\n".format(
- len(points3D), mean_track_length
- )
- )
-
- with open(path, "w") as fid:
- fid.write(HEADER)
- for _, pt in points3D.items():
- point_header = [pt.id, *pt.xyz, *pt.rgb, pt.error]
- fid.write(" ".join(map(str, point_header)) + " ")
- track_strings = []
- for image_id, point2D in zip(pt.image_ids, pt.point2D_idxs):
- track_strings.append(" ".join(map(str, [image_id, point2D])))
- fid.write(" ".join(track_strings) + "\n")
-
-
-def write_points3D_binary(points3D, path_to_model_file):
- """
- see: src/colmap/scene/reconstruction.cc
- void Reconstruction::ReadPoints3DBinary(const std::string& path)
- void Reconstruction::WritePoints3DBinary(const std::string& path)
- """
- with open(path_to_model_file, "wb") as fid:
- write_next_bytes(fid, len(points3D), "Q")
- for _, pt in points3D.items():
- write_next_bytes(fid, pt.id, "Q")
- write_next_bytes(fid, pt.xyz.tolist(), "ddd")
- write_next_bytes(fid, pt.rgb.tolist(), "BBB")
- write_next_bytes(fid, pt.error, "d")
- track_length = pt.image_ids.shape[0]
- write_next_bytes(fid, track_length, "Q")
- for image_id, point2D_id in zip(pt.image_ids, pt.point2D_idxs):
- write_next_bytes(fid, [image_id, point2D_id], "ii")
-
-
-def detect_model_format(path, ext):
- if (
- os.path.isfile(os.path.join(path, "cameras" + ext))
- and os.path.isfile(os.path.join(path, "images" + ext))
- and os.path.isfile(os.path.join(path, "points3D" + ext))
- ):
- print("Detected model format: '" + ext + "'")
- return True
-
- return False
-
-
-def read_model(path, ext=""):
- # try to detect the extension automatically
- if ext == "":
- if detect_model_format(path, ".bin"):
- ext = ".bin"
- elif detect_model_format(path, ".txt"):
- ext = ".txt"
- else:
- print("Provide model format: '.bin' or '.txt'")
- return
-
- if ext == ".txt":
- cameras = read_cameras_text(os.path.join(path, "cameras" + ext))
- images = read_images_text(os.path.join(path, "images" + ext))
- points3D = read_points3D_text(os.path.join(path, "points3D") + ext)
- else:
- cameras = read_cameras_binary(os.path.join(path, "cameras" + ext))
- images = read_images_binary(os.path.join(path, "images" + ext))
- points3D = read_points3D_binary(os.path.join(path, "points3D") + ext)
- return cameras, images, points3D
-
-
-def write_model(cameras, images, points3D, path, ext=".bin"):
- if ext == ".txt":
- write_cameras_text(cameras, os.path.join(path, "cameras" + ext))
- write_images_text(images, os.path.join(path, "images" + ext))
- write_points3D_text(points3D, os.path.join(path, "points3D") + ext)
- else:
- write_cameras_binary(cameras, os.path.join(path, "cameras" + ext))
- write_images_binary(images, os.path.join(path, "images" + ext))
- write_points3D_binary(points3D, os.path.join(path, "points3D") + ext)
- return cameras, images, points3D
-
-
-def qvec2rotmat(qvec):
- return np.array(
- [
- [
- 1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2,
- 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
- 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2],
- ],
- [
- 2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
- 1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2,
- 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1],
- ],
- [
- 2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],
- 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
- 1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2,
- ],
- ]
- )
-
-
-def rotmat2qvec(R):
- Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat
- K = (
- np.array(
- [
- [Rxx - Ryy - Rzz, 0, 0, 0],
- [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0],
- [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0],
- [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz],
- ]
- )
- / 3.0
- )
- eigvals, eigvecs = np.linalg.eigh(K)
- qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)]
- if qvec[0] < 0:
- qvec *= -1
- return qvec
-
-
-def main():
- parser = argparse.ArgumentParser(
- description="Read and write COLMAP binary and text models"
- )
- parser.add_argument("--input_model", help="path to input model folder")
- parser.add_argument(
- "--input_format",
- choices=[".bin", ".txt"],
- help="input model format",
- default="",
- )
- parser.add_argument("--output_model", help="path to output model folder")
- parser.add_argument(
- "--output_format",
- choices=[".bin", ".txt"],
- help="output model format",
- default=".txt",
- )
- args = parser.parse_args()
-
- cameras, images, points3D = read_model(
- path=args.input_model, ext=args.input_format
- )
-
- print("num_cameras:", len(cameras))
- print("num_images:", len(images))
- print("num_points3D:", len(points3D))
-
- if args.output_model is not None:
- write_model(
- cameras,
- images,
- points3D,
- path=args.output_model,
- ext=args.output_format,
- )
-
-
-if __name__ == "__main__":
- main()
\ No newline at end of file
diff --git a/optgs/model/decoder/diffgs_cuda_splatting.py b/optgs/model/decoder/diffgs_cuda_splatting.py
deleted file mode 100644
index 8663d799d4d5db66e8da248a6d7763d7a6c04633..0000000000000000000000000000000000000000
--- a/optgs/model/decoder/diffgs_cuda_splatting.py
+++ /dev/null
@@ -1,136 +0,0 @@
-from math import isqrt
-from typing import Literal
-
-import torch
-from diff_gs 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
-
-
-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 = False,
- use_sh: bool = True,
-):
- assert use_sh or gaussian_sh_coefficients.shape[-1] == 1
-
- assert scale_invariant is False
-
- # 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 = []
- all_depths = []
- 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,
- debug=False,
- antialiasing=False,
- )
- rasterizer = GaussianRasterizer(settings)
-
- row, col = torch.triu_indices(3, 3)
-
- image, radii, depth = 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)
- all_depths.append(depth)
- return {
- 'image': torch.stack(all_images),
- 'depth': torch.stack(all_depths),
- 'radii': torch.stack(all_radii),
- }
-
diff --git a/optgs/model/decoder/diffgs_decoder_splatting_cuda.py b/optgs/model/decoder/diffgs_decoder_splatting_cuda.py
deleted file mode 100644
index f03f2588c0d76aa02a3554e0484c7053f0866d37..0000000000000000000000000000000000000000
--- a/optgs/model/decoder/diffgs_decoder_splatting_cuda.py
+++ /dev/null
@@ -1,68 +0,0 @@
-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 .decoder import DepthRenderingMode
-from .diffgs_cuda_splatting import render_cuda
-from .decoder import Decoder, DecoderOutput
-
-
-@dataclass
-class DiffgsDecoderSplattingCUDACfg:
- name: Literal["diffgs"]
- scale_invariant: bool
-
-
-class DiffgsDecoderSplattingCUDA(Decoder[DiffgsDecoderSplattingCUDACfg]):
- background_color: Float[Tensor, "3"]
-
- def __init__(
- self,
- cfg: DiffgsDecoderSplattingCUDACfg,
- 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
- out = 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),
- scale_invariant=self.cfg.scale_invariant,
- )
- color = rearrange(out['image'], "(b v) c h w -> b v c h w", b=b, v=v)
- # the output is inverse depth (c == 1)
- depth = 1. / rearrange(out['depth'], "(b v) c h w -> b (v c) h w", b=b, v=v).clamp(min=1e-6)
-
- return DecoderOutput(
- color,
- depth
- )
diff --git a/optgs/model/decoder/gsplat_decoder_splatting_cuda.py b/optgs/model/decoder/gsplat_decoder_splatting_cuda.py
index 123c72c6cc626b1107f437f4d63dceeb7516630f..9324501ffabb8c1a04fdb4244114d63dfe400fa5 100644
--- a/optgs/model/decoder/gsplat_decoder_splatting_cuda.py
+++ b/optgs/model/decoder/gsplat_decoder_splatting_cuda.py
@@ -4,7 +4,6 @@ from typing import Literal
import math
import torch
# import torch.nn.functional as F
-from einops import repeat
from gsplat.rendering import rasterization
from jaxtyping import Float
from torch import Tensor
diff --git a/optgs/model/encoder/.deprecated/common/sampler.py b/optgs/model/encoder/.deprecated/common/sampler.py
deleted file mode 100644
index 1bf5b61674093d00ef1d1705835f43648e4bc679..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/common/sampler.py
+++ /dev/null
@@ -1,42 +0,0 @@
-from jaxtyping import Float, Int64, Shaped
-from torch import Tensor, nn
-
-from ....misc.discrete_probability_distribution import (
- gather_discrete_topk,
- sample_discrete_distribution,
-)
-
-
-class Sampler(nn.Module):
- def forward(
- self,
- probabilities: Float[Tensor, "*batch bucket"],
- num_samples: int,
- deterministic: bool,
- ) -> tuple[
- Int64[Tensor, "*batch 1"], # index
- Float[Tensor, "*batch 1"], # probability density
- ]:
- return (
- gather_discrete_topk(probabilities, num_samples)
- if deterministic
- else sample_discrete_distribution(probabilities, num_samples)
- )
-
- def gather(
- self,
- index: Int64[Tensor, "*batch sample"],
- target: Shaped[Tensor, "..."], # *batch bucket *shape
- ) -> Shaped[Tensor, "..."]: # *batch sample *shape
- """Gather from the target according to the specified index. Handle the
- broadcasting needed for the gather to work. See the comments for the actual
- expected input/output shapes since jaxtyping doesn't support multiple variadic
- lengths in annotations.
- """
- bucket_dim = index.ndim - 1
- while len(index.shape) < len(target.shape):
- index = index[..., None]
- broadcasted_index_shape = list(target.shape)
- broadcasted_index_shape[bucket_dim] = index.shape[bucket_dim]
- index = index.broadcast_to(broadcasted_index_shape)
- return target.gather(dim=bucket_dim, index=index)
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/Utils.py b/optgs/model/encoder/.deprecated/foundationstereo/Utils.py
deleted file mode 100644
index a76573aa69b26a0a785f969ecbeabe67975221eb..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/foundationstereo/Utils.py
+++ /dev/null
@@ -1,134 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
-#
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto. Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
-
-
-import os, sys, time,torch,torchvision,pickle,itertools,datetime,imageio,logging,importlib,argparse
-import torch.nn.functional as F
-import torch.nn as nn
-from functools import partial
-# import open3d as o3d
-import cv2
-import numpy as np
-# from transformations import *
-code_dir = os.path.dirname(os.path.realpath(__file__))
-sys.path.append(code_dir)
-
-
-
-def set_logging_format(level=logging.INFO):
- importlib.reload(logging)
- FORMAT = '%(message)s'
- logging.basicConfig(level=level, format=FORMAT, datefmt='%m-%d|%H:%M:%S')
-
-set_logging_format()
-
-
-
-def set_seed(random_seed):
- import torch,random
- np.random.seed(random_seed)
- random.seed(random_seed)
- torch.manual_seed(random_seed)
- torch.cuda.manual_seed_all(random_seed)
- torch.backends.cudnn.deterministic = True
- torch.backends.cudnn.benchmark = False
-
-
-def toOpen3dCloud(points,colors=None,normals=None):
- cloud = o3d.geometry.PointCloud()
- cloud.points = o3d.utility.Vector3dVector(points.astype(np.float64))
- if colors is not None:
- if colors.max()>1:
- colors = colors/255.0
- cloud.colors = o3d.utility.Vector3dVector(colors.astype(np.float64))
- if normals is not None:
- cloud.normals = o3d.utility.Vector3dVector(normals.astype(np.float64))
- return cloud
-
-
-
-def depth2xyzmap(depth:np.ndarray, K, uvs:np.ndarray=None, zmin=0.1):
- invalid_mask = (depthmax_H or W_resize>max_W:
- if H_resize>W_resize:
- W_resize = round_by_divider(W_resize*max_H/H_resize)
- H_resize = max_H
- else:
- H_resize = round_by_divider(H_resize*max_W/W_resize)
- W_resize = max_W
- return int(H_resize), int(W_resize)
-
-
-def vis_disparity(disp, min_val=None, max_val=None, invalid_thres=np.inf, color_map=cv2.COLORMAP_TURBO, cmap=None, other_output={}):
- """
- @disp: np array (H,W)
- @invalid_thres: > thres is invalid
- """
- disp = disp.copy()
- H,W = disp.shape[:2]
- invalid_mask = disp>=invalid_thres
- if (invalid_mask==0).sum()==0:
- other_output['min_val'] = None
- other_output['max_val'] = None
- return np.zeros((H,W,3))
- if min_val is None:
- min_val = disp[invalid_mask==0].min()
- if max_val is None:
- max_val = disp[invalid_mask==0].max()
- other_output['min_val'] = min_val
- other_output['max_val'] = max_val
- vis = ((disp-min_val)/(max_val-min_val)).clip(0,1) * 255
- if cmap is None:
- vis = cv2.applyColorMap(vis.clip(0, 255).astype(np.uint8), color_map)[...,::-1]
- else:
- vis = cmap(vis.astype(np.uint8))[...,:3]*255
- if invalid_mask.any():
- vis[invalid_mask] = 0
- return vis.astype(np.uint8)
-
-
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/config/large.yaml b/optgs/model/encoder/.deprecated/foundationstereo/config/large.yaml
deleted file mode 100644
index 39b74b93de8f590517fb709ee6a814037feba4c9..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/foundationstereo/config/large.yaml
+++ /dev/null
@@ -1,33 +0,0 @@
-corr_implementation: reg
-corr_levels: 2
-corr_radius: 4
-finetune_ckpt_name: model_best_bp2.pth
-finetune_from: null
-hidden_dims:
-- 128
-- 128
-- 128
-img_gamma: null
-inference_tile: 0
-low_memory: 0
-max_disp: 416
-max_val_sample: null
-mixed_precision: true
-n_downsample: 2
-n_gru_layers: 3
-disp_head_dim: 1
-notes: ''
-num_steps: 200000
-num_worker: 8
-slow_fast_gru: false
-tags_more: []
-tile_min_overlap:
-- 16
-- 16
-tile_wtype: gaussian
-time_limit: 14400
-train_iters: 22
-val_interval: 1
-valid_iters: 32
-wdecay: 0
-world_size: 32
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/config/small.yaml b/optgs/model/encoder/.deprecated/foundationstereo/config/small.yaml
deleted file mode 100644
index 4358c390672e3432f5a34476eaa74900460bc312..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/foundationstereo/config/small.yaml
+++ /dev/null
@@ -1,34 +0,0 @@
-corr_implementation: reg
-corr_levels: 2
-corr_radius: 4
-finetune_ckpt_name: model_best_bp2.pth
-finetune_from: null
-hidden_dims:
-- 128
-- 128
-- 128
-img_gamma: null
-inference_tile: 0
-low_memory: 0
-max_disp: 416
-max_val_sample: null
-mixed_precision: true
-n_downsample: 2
-n_gru_layers: 3
-disp_head_dim: 1
-notes: ''
-num_steps: 200000
-num_worker: 8
-slow_fast_gru: false
-tags_more: []
-tile_min_overlap:
-- 16
-- 16
-tile_wtype: gaussian
-time_limit: 14400
-train_iters: 22
-val_interval: 1
-valid_iters: 32
-wdecay: 0
-world_size: 32
-vit_type: vits
\ No newline at end of file
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/core/extractor.py b/optgs/model/encoder/.deprecated/foundationstereo/core/extractor.py
deleted file mode 100644
index bdee3e6b444a2bc797f41df76e6e46f91d340745..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/foundationstereo/core/extractor.py
+++ /dev/null
@@ -1,382 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
-#
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto. Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
-
-
-import torch,logging,os,sys,urllib,warnings
-import torch.nn as nn
-import torch.nn.functional as F
-code_dir = os.path.dirname(os.path.realpath(__file__))
-sys.path.append(f'{code_dir}/../')
-from core.submodule import *
-from Utils import *
-import timm
-
-
-class ResidualBlock(nn.Module):
- def __init__(self, in_planes, planes, norm_fn='group', stride=1):
- super(ResidualBlock, self).__init__()
-
- self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride)
- self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)
- self.relu = nn.ReLU(inplace=True)
-
- num_groups = planes // 8
-
- if norm_fn == 'group':
- self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
- self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
- if not (stride == 1 and in_planes == planes):
- self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
-
- elif norm_fn == 'batch':
- self.norm1 = nn.BatchNorm2d(planes)
- self.norm2 = nn.BatchNorm2d(planes)
- if not (stride == 1 and in_planes == planes):
- self.norm3 = nn.BatchNorm2d(planes)
-
- elif norm_fn == 'instance':
- self.norm1 = nn.InstanceNorm2d(planes)
- self.norm2 = nn.InstanceNorm2d(planes)
- if not (stride == 1 and in_planes == planes):
- self.norm3 = nn.InstanceNorm2d(planes)
-
- elif norm_fn=='layer':
- self.norm1 = LayerNorm2d(planes)
- self.norm2 = LayerNorm2d(planes)
- if not (stride == 1 and in_planes == planes):
- self.norm3 = LayerNorm2d(planes)
-
- elif norm_fn == 'none':
- self.norm1 = nn.Sequential()
- self.norm2 = nn.Sequential()
- if not (stride == 1 and in_planes == planes):
- self.norm3 = nn.Sequential()
-
- 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.conv1(y)
- y = self.norm1(y)
- y = self.relu(y)
- y = self.conv2(y)
- y = self.norm2(y)
- y = self.relu(y)
-
- if self.downsample is not None:
- x = self.downsample(x)
-
- return self.relu(x+y)
-
-
-
-class MultiBasicEncoder(nn.Module):
- def __init__(self, output_dim=[128], norm_fn='batch', dropout=0.0, downsample=3):
- super(MultiBasicEncoder, self).__init__()
- self.norm_fn = norm_fn
- self.downsample = downsample
-
- if self.norm_fn == 'group':
- self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64)
-
- elif self.norm_fn == 'batch':
- self.norm1 = nn.BatchNorm2d(64)
-
- elif self.norm_fn == 'instance':
- self.norm1 = nn.InstanceNorm2d(64)
-
- elif self.norm_fn=='layer':
- self.norm1 = LayerNorm2d(64)
-
- elif self.norm_fn == 'none':
- self.norm1 = nn.Sequential()
-
- self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=1 + (downsample > 2), padding=3)
- self.relu1 = nn.ReLU(inplace=True)
-
- self.in_planes = 64
- self.layer1 = self._make_layer(64, stride=1)
- self.layer2 = self._make_layer(96, stride=1 + (downsample > 1))
- self.layer3 = self._make_layer(128, stride=1 + (downsample > 0))
- self.layer4 = self._make_layer(128, stride=2)
- self.layer5 = self._make_layer(128, stride=2)
-
- output_list = []
-
- for dim in output_dim:
- conv_out = nn.Sequential(
- ResidualBlock(128, 128, self.norm_fn, stride=1),
- nn.Conv2d(128, dim[2], 3, padding=1))
- output_list.append(conv_out)
-
- self.outputs04 = nn.ModuleList(output_list)
-
- output_list = []
- for dim in output_dim:
- conv_out = nn.Sequential(
- ResidualBlock(128, 128, self.norm_fn, stride=1),
- nn.Conv2d(128, dim[1], 3, padding=1))
- output_list.append(conv_out)
-
- self.outputs08 = nn.ModuleList(output_list)
-
- output_list = []
- for dim in output_dim:
- conv_out = nn.Conv2d(128, dim[0], 3, padding=1)
- output_list.append(conv_out)
-
- self.outputs16 = nn.ModuleList(output_list)
-
- if dropout > 0:
- self.dropout = nn.Dropout2d(p=dropout)
- else:
- self.dropout = None
-
- 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):
- layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
- layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
- layers = (layer1, layer2)
-
- self.in_planes = dim
- return nn.Sequential(*layers)
-
- def forward(self, x, dual_inp=False, num_layers=3):
-
- x = self.conv1(x)
- x = self.norm1(x)
- x = self.relu1(x)
- x = self.layer1(x)
- x = self.layer2(x)
- x = self.layer3(x)
- if dual_inp:
- v = x
- x = x[:(x.shape[0]//2)]
-
- outputs04 = [f(x) for f in self.outputs04]
- if num_layers == 1:
- return (outputs04, v) if dual_inp else (outputs04,)
-
- y = self.layer4(x)
- outputs08 = [f(y) for f in self.outputs08]
-
- if num_layers == 2:
- return (outputs04, outputs08, v) if dual_inp else (outputs04, outputs08)
-
- z = self.layer5(y)
- outputs16 = [f(z) for f in self.outputs16]
-
- return (outputs04, outputs08, outputs16, v) if dual_inp else (outputs04, outputs08, outputs16)
-
-
-
-class ContextNetDino(MultiBasicEncoder):
- def __init__(self, output_dim=[128], norm_fn='batch', downsample=3):
- nn.Module.__init__(self)
- self.patch_size = 14
- self.image_size = 518
- self.vit_feat_dim = 384
- code_dir = os.path.dirname(os.path.realpath(__file__))
-
- self.out_dims = output_dim
-
- self.norm_fn = norm_fn
-
- if self.norm_fn == 'group':
- self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64)
-
- elif self.norm_fn == 'batch':
- self.norm1 = nn.BatchNorm2d(64)
-
- elif self.norm_fn == 'instance':
- self.norm1 = nn.InstanceNorm2d(64)
-
- elif self.norm_fn=='layer':
- self.norm1 = LayerNorm2d(64)
-
- elif self.norm_fn == 'none':
- self.norm1 = nn.Sequential()
-
- self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=1 + (downsample > 2), padding=3)
- self.relu1 = nn.ReLU(inplace=True)
-
- self.in_planes = 64
- self.layer1 = self._make_layer(64, stride=1)
- self.layer2 = self._make_layer(96, stride=1 + (downsample > 1))
- self.layer3 = self._make_layer(128, stride=1 + (downsample > 0))
- self.layer4 = self._make_layer(128, stride=2)
- self.layer5 = self._make_layer(128, stride=2)
- # unused parameters
- # self.down = nn.Sequential(
- # nn.Conv2d(128, 128, kernel_size=4, stride=4, padding=0),
- # nn.BatchNorm2d(128),
- # )
- self.conv2 = BasicConv(128+128, 128, kernel_size=3, padding=1)
- self.norm = nn.BatchNorm2d(256)
-
- output_list = []
- for dim in output_dim:
- conv_out = nn.Sequential(
- ResidualBlock(128, 128, self.norm_fn, stride=1),
- nn.Conv2d(128, dim[2], 3, padding=1))
- output_list.append(conv_out)
-
- self.outputs04 = nn.ModuleList(output_list)
-
- output_list = []
- for dim in output_dim:
- conv_out = nn.Sequential(
- ResidualBlock(128, 128, self.norm_fn, stride=1),
- nn.Conv2d(128, dim[1], 3, padding=1))
- output_list.append(conv_out)
-
- self.outputs08 = nn.ModuleList(output_list)
-
- output_list = []
- for dim in output_dim:
- conv_out = nn.Conv2d(128, dim[0], 3, padding=1)
- output_list.append(conv_out)
-
- self.outputs16 = nn.ModuleList(output_list)
-
- def forward(self, x_in, vit_feat, dual_inp=False, num_layers=3):
- B,C,H,W = x_in.shape
- x = self.conv1(x_in)
- x = self.norm1(x)
- x = self.relu1(x)
- x = self.layer1(x)
- x = self.layer2(x)
- x = self.layer3(x)
-
- divider = np.lcm(self.patch_size, 16)
- H_resize, W_resize = get_resize_keep_aspect_ratio(H,W, divider=divider, max_H=1344, max_W=1344)
- x = torch.cat([x, vit_feat], dim=1)
- x = self.conv2(x)
- outputs04 = [f(x) for f in self.outputs04]
-
- y = self.layer4(x)
- outputs08 = [f(y) for f in self.outputs08]
-
- z = self.layer5(y)
- outputs16 = [f(z) for f in self.outputs16]
-
- return (outputs04, outputs08, outputs16)
-
-
-class DepthAnythingFeature(nn.Module):
- def __init__(self, encoder='vits'):
- super().__init__()
- model_configs = {
- 'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
- 'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
- 'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]}
- }
- from depth_anything.dpt import DepthAnything
- self.encoder = encoder
- depth_anything = DepthAnything(model_configs[encoder])
- self.depth_anything = depth_anything
-
- self.intermediate_layer_idx = { #!NOTE For V2
- 'vits': [2, 5, 8, 11],
- 'vitb': [2, 5, 8, 11],
- 'vitl': [4, 11, 17, 23],
- 'vitg': [9, 19, 29, 39]
- }
-
-
- def forward(self, x):
- """
- @x: (B,C,H,W)
- """
- h, w = x.shape[-2:]
- features = self.depth_anything.pretrained.get_intermediate_layers(x, self.intermediate_layer_idx[self.encoder], return_class_token=True)
-
-
- patch_size = self.depth_anything.pretrained.patch_size
- patch_h, patch_w = h // patch_size, w // patch_size
- out, path_1, path_2, path_3, path_4, disp = self.depth_anything.depth_head.forward(features, patch_h, patch_w, return_intermediate=True)
-
- return {'out':out, 'path_1':path_1, 'path_2':path_2, 'path_3':path_3, 'path_4':path_4, 'features':features, 'disp':disp} # path_1 is 1/2; path_2 is 1/4
-
-
-class Feature(nn.Module):
- def __init__(self, vit_type='vitl', no_freeze_mono=False):
- super(Feature, self).__init__()
- model = timm.create_model('edgenext_small', pretrained=True, features_only=False)
- self.stem = model.stem
- self.stages = model.stages
- chans = [48, 96, 160, 304]
- self.chans = chans
- self.no_freeze_mono = no_freeze_mono
-
- self.deconv32_16 = Conv2x_IN(chans[3], chans[2], deconv=True, concat=True)
- self.deconv16_8 = Conv2x_IN(chans[2]*2, chans[1], deconv=True, concat=True)
- self.deconv8_4 = Conv2x_IN(chans[1]*2, chans[0], deconv=True, concat=True)
- self.conv4 = nn.Sequential(
- BasicConv(chans[0]*2+128, chans[0]*2+128, kernel_size=3, stride=1, padding=1, norm='instance'),
- ResidualBlock(chans[0]*2+128, chans[0]*2+128, norm_fn='instance'),
- ResidualBlock(chans[0]*2+128, chans[0]*2+128, norm_fn='instance'),
- )
-
- self.dino = DepthAnythingFeature(encoder=vit_type)
- if not no_freeze_mono:
- self.dino = freeze_model(self.dino)
- self.patch_size = 14
- self.d_out = [chans[0]*2+128, chans[1]*2, chans[2]*2, chans[3]]
-
- self.vit_type = vit_type
- if vit_type == 'vits':
- self.vit_proj = nn.Conv2d(32, 128, 1)
- elif vit_type == 'vitb':
- self.vit_proj = nn.Conv2d(64, 128, 1)
-
- def forward(self, x):
- B,C,H,W = x.shape
- divider = np.lcm(self.patch_size, 16)
- H_resize, W_resize = get_resize_keep_aspect_ratio(H,W, divider=divider, max_H=1344, max_W=1344)
- x_in_ = F.interpolate(x, size=(H_resize, W_resize), mode='bicubic', align_corners=False)
- if self.no_freeze_mono:
- output = self.dino(x_in_)
- else:
- self.dino = self.dino.eval()
- with torch.no_grad():
- output = self.dino(x_in_)
- vit_feat = output['out']
- if self.vit_type in ['vits', 'vitb']:
- # TODO: or just copy the smaller features to 128
- vit_feat = self.vit_proj(vit_feat)
-
- vit_feat = F.interpolate(vit_feat, size=(H//4,W//4), mode='bilinear', align_corners=True)
- x = self.stem(x)
- x4 = self.stages[0](x)
- x8 = self.stages[1](x4)
- x16 = self.stages[2](x8)
- x32 = self.stages[3](x16)
-
- x16 = self.deconv32_16(x32, x16)
- x8 = self.deconv16_8(x16, x8)
- x4 = self.deconv8_4(x8, x4)
- x4 = torch.cat([x4, vit_feat], dim=1)
- x4 = self.conv4(x4)
- return [x4, x8, x16, x32], vit_feat
-
-
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/core/foundation_stereo.py b/optgs/model/encoder/.deprecated/foundationstereo/core/foundation_stereo.py
deleted file mode 100644
index ea45d81ad96da8041036ef652e27997f3a3980a6..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/foundationstereo/core/foundation_stereo.py
+++ /dev/null
@@ -1,621 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
-#
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto. Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
-
-
-import torch,pdb,logging
-import torch.nn as nn
-import torch.nn.functional as F
-import sys,os
-code_dir = os.path.dirname(os.path.realpath(__file__))
-sys.path.append(f'{code_dir}/../')
-from core.update import *
-from core.extractor import *
-from core.geometry import Combined_Geo_Encoding_Volume
-from core.submodule import *
-from core.utils.utils import *
-from Utils import *
-import time
-import cv2
-
-from .unimatch_matching import group_correlation_softmax_depth, CorrBlock, coords_grid, warp_with_pose_depth_candidates
-from .unimatch_geometry import compute_flow_with_depth_pose
-
-
-count = 0
-
-
-try:
- # autocast = torch.cuda.amp.autocast
- autocast = torch.amp.autocast
-except:
- class autocast:
- def __init__(self, enabled):
- pass
- def __enter__(self):
- pass
- def __exit__(self, *args):
- pass
-
-
-def normalize_image(img):
- '''
- @img: (B,C,H,W) in range 0-255, RGB order
- '''
- tf = torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], inplace=False)
- return tf(img/255.0).contiguous()
-
-
-class hourglass(nn.Module):
- def __init__(self, cfg, in_channels, feat_dims=None):
- super().__init__()
- self.cfg = cfg
- self.conv1 = nn.Sequential(BasicConv(in_channels, in_channels*2, is_3d=True, bn=True, relu=True, kernel_size=3,
- padding=1, stride=2, dilation=1),
- Conv3dNormActReduced(in_channels*2, in_channels*2, kernel_size=3, kernel_disp=17))
-
- self.conv2 = nn.Sequential(BasicConv(in_channels*2, in_channels*4, is_3d=True, bn=True, relu=True, kernel_size=3,
- padding=1, stride=2, dilation=1),
- Conv3dNormActReduced(in_channels*4, in_channels*4, kernel_size=3, kernel_disp=17))
-
- self.conv3 = nn.Sequential(BasicConv(in_channels*4, in_channels*6, is_3d=True, bn=True, relu=True, kernel_size=3,
- padding=1, stride=2, dilation=1),
- Conv3dNormActReduced(in_channels*6, in_channels*6, kernel_size=3, kernel_disp=17))
-
-
- self.conv3_up = BasicConv(in_channels*6, in_channels*4, deconv=True, is_3d=True, bn=True,
- relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2))
-
- self.conv2_up = BasicConv(in_channels*4, in_channels*2, deconv=True, is_3d=True, bn=True,
- relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2))
-
- self.conv1_up = BasicConv(in_channels*2, in_channels, deconv=True, is_3d=True, bn=True,
- relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2))
- self.conv_out = nn.Sequential(
- Conv3dNormActReduced(in_channels, in_channels, kernel_size=3, kernel_disp=17),
- Conv3dNormActReduced(in_channels, in_channels, kernel_size=3, kernel_disp=17),
- )
-
- self.agg_0 = nn.Sequential(BasicConv(in_channels*8, in_channels*4, is_3d=True, kernel_size=1, padding=0, stride=1),
- Conv3dNormActReduced(in_channels*4, in_channels*4, kernel_size=3, kernel_disp=17),
- Conv3dNormActReduced(in_channels*4, in_channels*4, kernel_size=3, kernel_disp=17),)
-
- self.agg_1 = nn.Sequential(BasicConv(in_channels*4, in_channels*2, is_3d=True, kernel_size=1, padding=0, stride=1),
- Conv3dNormActReduced(in_channels*2, in_channels*2, kernel_size=3, kernel_disp=17),
- Conv3dNormActReduced(in_channels*2, in_channels*2, kernel_size=3, kernel_disp=17))
- self.atts = nn.ModuleDict({
- "4": CostVolumeDisparityAttention(d_model=in_channels, nhead=4, dim_feedforward=in_channels, norm_first=False, num_transformer=4, max_len=self.cfg['max_disp']//16),
- })
- self.conv_patch = nn.Sequential(
- nn.Conv3d(in_channels, in_channels, kernel_size=4, stride=4, padding=0, groups=in_channels),
- nn.BatchNorm3d(in_channels),
- )
-
- self.feature_att_8 = FeatureAtt(in_channels*2, feat_dims[1])
- self.feature_att_16 = FeatureAtt(in_channels*4, feat_dims[2])
- self.feature_att_32 = FeatureAtt(in_channels*6, feat_dims[3])
- self.feature_att_up_16 = FeatureAtt(in_channels*4, feat_dims[2])
- self.feature_att_up_8 = FeatureAtt(in_channels*2, feat_dims[1])
-
- def forward(self, x, features):
- conv1 = self.conv1(x)
- conv1 = self.feature_att_8(conv1, features[1])
-
- conv2 = self.conv2(conv1)
- conv2 = self.feature_att_16(conv2, features[2])
-
- conv3 = self.conv3(conv2)
- conv3 = self.feature_att_32(conv3, features[3])
-
- conv3_up = self.conv3_up(conv3)
- conv2 = torch.cat((conv3_up, conv2), dim=1)
- conv2 = self.agg_0(conv2)
- conv2 = self.feature_att_up_16(conv2, features[2])
-
- conv2_up = self.conv2_up(conv2)
- conv1 = torch.cat((conv2_up, conv1), dim=1)
- conv1 = self.agg_1(conv1)
- conv1 = self.feature_att_up_8(conv1, features[1])
-
- conv = self.conv1_up(conv1)
- x = self.conv_patch(x)
- x = self.atts["4"](x)
- x = F.interpolate(x, scale_factor=4, mode='trilinear', align_corners=False)
- conv = conv + x
- conv = self.conv_out(conv)
-
- return conv
-
-
-
-class FoundationStereo(nn.Module):
- def __init__(self, args,
- bilinear_init_depth=False,
- flow_corr=False,
- flow_corr_levels=4,
- amp_bf16=True,
- vit_type='vitl',
- no_geo_volume=False,
- concat_geo_volume=False,
- depth_sample_geo_volume=False,
- no_freeze_mono=False,
- bilinear_up_depth=False,
- local_match_radius=0,
- supervise_init_depth=False,
- sample_log_depth=False,
- ):
- super().__init__()
- self.args = args
-
- self.bilinear_init_depth = bilinear_init_depth
- self.flow_corr = flow_corr
- self.flow_corr_levels = flow_corr_levels
- self.amp_bf16 = amp_bf16
- self.no_geo_volume = no_geo_volume
- self.concat_geo_volume = concat_geo_volume
- self.depth_sample_geo_volume = depth_sample_geo_volume
- self.bilinear_up_depth = bilinear_up_depth
- self.local_match_radius = local_match_radius
- self.supervise_init_depth = supervise_init_depth
- self.sample_log_depth = sample_log_depth
-
- if local_match_radius > 0:
- self.flow_corr = flow_corr = False
-
- if concat_geo_volume:
- assert flow_corr or local_match_radius > 0
-
- if depth_sample_geo_volume:
- assert concat_geo_volume
-
- context_dims = args.hidden_dims
- self.cv_group = 8
- volume_dim = 28
-
- self.cnet = ContextNetDino(output_dim=[args.hidden_dims, context_dims], downsample=args.n_downsample)
- self.update_block = BasicSelectiveMultiUpdateBlock(self.args, self.args.hidden_dims[0], volume_dim=volume_dim,
- depth_head_dim=2 * local_match_radius + 1,
- )
- self.sam = SpatialAttentionExtractor()
- self.cam = ChannelAttentionEnhancement(self.args.hidden_dims[0])
-
- # unused parameters
- # self.context_zqr_convs = nn.ModuleList([nn.Conv2d(context_dims[i], args.hidden_dims[i]*3, kernel_size=3, padding=3//2) for i in range(self.args.n_gru_layers)])
-
- self.feature = Feature(vit_type=vit_type, no_freeze_mono=no_freeze_mono)
- self.proj_cmb = nn.Conv2d(self.feature.d_out[0], 12, kernel_size=1, padding=0)
-
- self.stem_2 = nn.Sequential(
- BasicConv_IN(3, 32, kernel_size=3, stride=2, padding=1),
- nn.Conv2d(32, 32, 3, 1, 1, bias=False),
- nn.InstanceNorm2d(32), nn.ReLU()
- )
- # self.stem_4 = nn.Sequential(
- # BasicConv_IN(32, 48, kernel_size=3, stride=2, padding=1),
- # nn.Conv2d(48, 48, 3, 1, 1, bias=False),
- # nn.InstanceNorm2d(48), nn.ReLU()
- # )
-
- self.spx_2_gru = Conv2x(32, 32, True, bn=False)
- self.spx_gru = nn.Sequential(
- nn.ConvTranspose2d(2*32, 9, kernel_size=4, stride=2, padding=1),
- )
-
-
- self.corr_stem = nn.Sequential(
- nn.Conv3d(32, volume_dim, kernel_size=1),
- BasicConv(volume_dim, volume_dim, kernel_size=3, padding=1, is_3d=True),
- ResnetBasicBlock3D(volume_dim, volume_dim, kernel_size=3, stride=1, padding=1),
- ResnetBasicBlock3D(volume_dim, volume_dim, kernel_size=3, stride=1, padding=1),
- )
- self.corr_feature_att = FeatureAtt(volume_dim, self.feature.d_out[0])
- self.cost_agg = hourglass(cfg=self.args, in_channels=volume_dim, feat_dims=self.feature.d_out)
- self.classifier = nn.Sequential(
- BasicConv(volume_dim, volume_dim//2, kernel_size=3, padding=1, is_3d=True),
- ResnetBasicBlock3D(volume_dim//2, volume_dim//2, kernel_size=3, stride=1, padding=1),
- nn.Conv3d(volume_dim//2, 1, kernel_size=7, padding=3),
- )
-
- r = self.args.corr_radius
- dx = torch.linspace(-r, r, 2*r+1, requires_grad=False).reshape(1, 1, 2*r+1, 1)
- self.dx = dx
-
- if self.flow_corr:
- flow_corr_channels = self.flow_corr_levels * (2 * 4 + 1) ** 2
- if self.concat_geo_volume:
- self.corr_proj = nn.Conv2d(flow_corr_channels + 504, 522, 1)
- else:
- self.corr_proj = nn.Conv2d(flow_corr_channels, 522, 1)
-
- if self.local_match_radius > 0:
- # TODO: multi-scale matching, maybe also combine geometry volume sampled from geo_volume
- corr_channels = 2 * self.local_match_radius + 1
- self.correlation_proj = nn.Conv2d(corr_channels, 522, 1)
-
- if self.no_geo_volume:
- self.corr_proj = nn.Conv2d(18, 522, 1)
-
- # unused parameters
- # del self.cnet.down[0].weight
- # del self.cnet.down[0].bias
- # del self.stem_4[0].conv.weight
- # del self.stem_4[1].weight
-
-
- def upsample_disp(self, disp, mask_feat_4, stem_2x, task='stereo'):
- assert task in ['stereo', 'depth']
-
- dtype = torch.bfloat16 if self.amp_bf16 else torch.float16
-
- with autocast('cuda', enabled=self.args.mixed_precision, dtype=dtype):
- xspx = self.spx_2_gru(mask_feat_4, stem_2x) # 1/2 resolution
- spx_pred = self.spx_gru(xspx)
- spx_pred = F.softmax(spx_pred, 1)
- if task == 'depth':
- # no 4x disp since we predict inverse depth
- up_disp = context_upsample(disp, spx_pred, task='depth').unsqueeze(1)
- else:
- up_disp = context_upsample(disp*4., spx_pred).unsqueeze(1)
-
- return up_disp.float()
-
-
- def forward(self, image1, image2, iters=12, flow_init=None, test_mode=False, low_memory=False, init_disp=None,
- no_norm_img=False,
- task='stereo',
- intrinsics=None,
- pose=None, # relative pose transform
- min_depth=1. / 0.5, # inverse depth range
- max_depth=1. / 10,
- num_depth_candidates=64,
- pred_bidir_depth=False,
- return_features=False,
- rectified_stereo=False,
- ):
- """ Estimate disparity between pair of frames """
- assert task in ['stereo', 'depth']
-
- if self.sample_log_depth:
- min_depth, max_depth = np.log(1. / max_depth), np.log(1. / min_depth)
- # print(min_depth, max_depth)
-
- if rectified_stereo:
- from PIL import Image
- ori_img1 = image1[0].permute(1, 2, 0).cpu().numpy()
- ori_img2 = image2[0].permute(1, 2, 0).cpu().numpy()
- ori_concat = np.concatenate((ori_img1, ori_img2), axis=1)
-
- save_dir = 'tmp_rectified_stereo'
- os.makedirs(save_dir, exist_ok=True)
-
- # Image.fromarray(ori_concat.astype(np.uint8)).save(save_dir + '/ori.png')
-
- img_left = ori_img1
- img_right = ori_img2
-
- ori_h, ori_w = image1.shape[-2:]
- tmp_intrinsics = intrinsics.clone()
- tmp_intrinsics[:, 0] *= ori_w
- tmp_intrinsics[:, 1] *= ori_h
-
- K1 = K2 = tmp_intrinsics[0].cpu().numpy()
- D1 = np.zeros(5) # Assuming no distortion
- D2 = np.zeros(5)
-
- R = pose[0].cpu().numpy()[:3, :3]
- T = pose[0].cpu().numpy()[:3, 3:]
-
- # TODO: choose left or right
- # if T[0] < 0:
- # img_left, img_right = ori_img1, ori_img2
- # else:
- # # swap
- # img_left, img_right = ori_img2, ori_img1
- # # inverse
- # R, T = R.T, -R.T @ T
-
- # [w, h]
- img_size = [ori_img1.shape[1], ori_img1.shape[0]]
-
- # === 3. Stereo rectification ===
- R1, R2, P1, P2, Q, _, _ = cv2.stereoRectify(K1, D1, K2, D2, img_size, R, T)
-
- # === 4. Create rectification maps ===
- map1x, map1y = cv2.initUndistortRectifyMap(K1, D1, R1, P1, img_size, cv2.CV_32FC1)
- map2x, map2y = cv2.initUndistortRectifyMap(K2, D2, R2, P2, img_size, cv2.CV_32FC1)
-
- # === 5. Remap (rectify) images ===
- rect_left = cv2.remap(img_left, map1x, map1y, cv2.INTER_LINEAR)
- rect_right = cv2.remap(img_right, map2x, map2y, cv2.INTER_LINEAR)
-
- # print(rect_left.shape, rect_right.shape)
-
- rect_concat = np.concatenate((rect_left, rect_right), axis=1)
-
- concat = np.concatenate((ori_concat, rect_concat), axis=0)
- global count
- Image.fromarray(concat.astype(np.uint8)).save(save_dir + f'/rect_{count}.png')
- count += 1
-
- if count > 20:
- assert False
-
- B = len(image1)
- low_memory = low_memory or (self.args.get('low_memory', False))
- if not no_norm_img:
- image1 = normalize_image(image1)
- image2 = normalize_image(image2)
-
- dtype = torch.bfloat16 if self.amp_bf16 else torch.float16
- with autocast('cuda', enabled=self.args.mixed_precision, dtype=dtype):
- out, vit_feat = self.feature(torch.cat([image1, image2], dim=0))
- if not pred_bidir_depth:
- vit_feat = vit_feat[:B]
- features_left = [o[:B] for o in out]
- features_right = [o[B:] for o in out]
- if pred_bidir_depth:
- stem_2x = self.stem_2(torch.cat([image1, image2], dim=0))
- for i in range(len(features_left)):
- features_left[i], features_right[i] = torch.cat([features_left[i], features_right[i]], dim=0), torch.cat([features_right[i], features_left[i]], dim=0)
- else:
- stem_2x = self.stem_2(image1)
-
- if task == 'depth':
- assert intrinsics is not None and pose is not None
- # NOTE: in this codebase, intrinsics are normalized by image width and height
- # in unimatch's codebase, no normalization
- ori_h, ori_w = image1.shape[-2:]
- intrinsics = intrinsics.clone()
- intrinsics[:, 0] *= ori_w
- intrinsics[:, 1] *= ori_h
-
- # scale intrinsics
- intrinsics_curr = intrinsics.clone()
- intrinsics_curr[:, :2] = intrinsics_curr[:, :2] / 4
-
- if pred_bidir_depth:
- intrinsics_curr = intrinsics_curr.repeat(2, 1, 1)
- pose = torch.cat((pose, torch.inverse(pose)), dim=0)
-
- b, _, h, w = features_left[0].shape
-
- depth_candidates = torch.linspace(min_depth, max_depth, num_depth_candidates).type_as(image1)
- depth_candidates = depth_candidates.view(1, num_depth_candidates, 1, 1).repeat(b, 1, h,
- w) # [B, D, H, W]
-
- # gwc_volume: [B, G, D, H, W]
- gwc_volume, warped_feature1 = group_correlation_softmax_depth(features_left[0], features_right[0],
- intrinsics_curr,
- pose,
- depth_candidates=depth_candidates,
- num_groups=self.cv_group,
- sample_log_depth=self.sample_log_depth,
- )
-
- left_tmp = self.proj_cmb(features_left[0]).unsqueeze(2).repeat(1, 1, num_depth_candidates, 1, 1) # [B, C, D, H, W]
- right_tmp = self.proj_cmb(warped_feature1.reshape(b, features_left[0].size(1), -1, w)).reshape(b, -1, num_depth_candidates, h, w)
- concat_volume = torch.cat((left_tmp, right_tmp), dim=1) # [B, 2C, D, H, W]
- del left_tmp, right_tmp
-
- else:
- gwc_volume = build_gwc_volume(features_left[0], features_right[0], self.args.max_disp//4, self.cv_group) # Group-wise correlation volume (B, N_group, max_disp, H, W)
- left_tmp = self.proj_cmb(features_left[0])
- right_tmp = self.proj_cmb(features_right[0])
- concat_volume = build_concat_volume(left_tmp, right_tmp, maxdisp=self.args.max_disp//4)
- del left_tmp, right_tmp
-
- comb_volume = torch.cat([gwc_volume, concat_volume], dim=1)
- comb_volume = self.corr_stem(comb_volume)
- comb_volume = self.corr_feature_att(comb_volume, features_left[0])
- comb_volume = self.cost_agg(comb_volume, features_left)
-
- # Init disp from geometry encoding volume
- prob = F.softmax(self.classifier(comb_volume).squeeze(1), dim=1) #(B, max_disp, H, W)
- if init_disp is None:
- if task == 'depth':
- init_disp = (prob * depth_candidates).sum(dim=1, keepdim=True) # [B, 1, H, W]
- else:
- init_disp = disparity_regression(prob, self.args.max_disp//4) # Weighted sum of disparity
-
- if pred_bidir_depth:
- cnet_list = self.cnet(torch.cat((image1, image2), dim=0), vit_feat=vit_feat, num_layers=self.args.n_gru_layers) #(1/4, 1/8, 1/16)
- else:
- cnet_list = self.cnet(image1, vit_feat=vit_feat, num_layers=self.args.n_gru_layers) #(1/4, 1/8, 1/16)
- cnet_list = list(cnet_list)
- net_list = [torch.tanh(x[0]) for x in cnet_list] # Hidden information
- inp_list = [torch.relu(x[1]) for x in cnet_list] # Context information list of pyramid levels
- inp_list = [self.cam(x) * x for x in inp_list]
- att = [self.sam(x) for x in inp_list]
-
- if self.flow_corr:
- geo_fn = CorrBlock(features_left[0].float(), features_right[0].float(), num_levels=self.flow_corr_levels)
- if self.concat_geo_volume:
- geo_volume_fn = Combined_Geo_Encoding_Volume(features_left[0].float(), features_right[0].float(), comb_volume.float(),
- num_levels=self.args.corr_levels, dx=self.dx, no_corr=True)
-
- else:
- geo_fn = Combined_Geo_Encoding_Volume(features_left[0].float(), features_right[0].float(), comb_volume.float(), num_levels=self.args.corr_levels, dx=self.dx)
-
- b, c, h, w = features_left[0].shape
- coords = torch.arange(w, dtype=torch.float, device=init_disp.device).reshape(1,1,w,1).repeat(b, h, 1, 1) # (B,H,W,1) Horizontal only
- disp = init_disp.float()
- disp_preds = []
-
- # GRUs iterations to update disparity (1/4 resolution)
- for itr in range(iters):
- disp = disp.detach()
- if self.flow_corr:
- proj_coords_from_depth = compute_flow_with_depth_pose(
- torch.exp(disp.squeeze(1)) if self.sample_log_depth else 1. / disp.squeeze(1),
- intrinsics_curr,
- extrinsics_rel=pose,
- return_coords=True,
- )
- geo_feat = geo_fn(proj_coords_from_depth)
-
- if self.concat_geo_volume:
- if self.depth_sample_geo_volume:
- indices = torch.sum(disp >= depth_candidates, dim=1, keepdim=True) - 1
- # Clamp to ensure indices are within [0, D-1]
- indices = indices.clamp(min=0, max=num_depth_candidates-1).float()
- tmp = geo_volume_fn(indices, coords)
- else:
- tmp = geo_volume_fn(disp, coords)
- geo_feat = torch.cat((geo_feat, tmp), dim=1)
-
- # use the pre-trained weights
- geo_feat = self.corr_proj(geo_feat)
-
- elif self.local_match_radius > 0:
- # 2x smaller interval for each iteration
- disp_interval = (max_depth - min_depth) / num_depth_candidates / (2 ** itr)
- disp_range_min = (disp - disp_interval * self.local_match_radius).clamp(min=min_depth) # [B, 1, H, W]
- disp_range_max = (disp + disp_interval * self.local_match_radius).clamp(max=max_depth)
- linear_space = torch.linspace(0, 1, 2 * self.local_match_radius + 1
- ).type_as(disp).view(1, -1, 1, 1) # [1, K, 1, 1]
- disp_candidates = disp_range_min + linear_space * (disp_range_max - disp_range_min) # [B, K, H, W]
-
- warped_feature1 = warp_with_pose_depth_candidates(features_right[0].float(),
- intrinsics_curr,
- pose,
- torch.exp(disp_candidates) if self.sample_log_depth else (1. / disp_candidates),
- ) # [B, C, K, H, W]
- corr = (F.normalize(features_left[0].float().unsqueeze(2), dim=1) * F.normalize(warped_feature1, dim=1)).sum(1) # [B, K, H, W]
- geo_feat = self.correlation_proj(corr)
-
- else:
- geo_feat = geo_fn(disp, coords, low_memory=low_memory, no_geo_volume=self.no_geo_volume)
-
- if self.no_geo_volume:
- geo_feat = self.corr_proj(geo_feat)
-
- dtype = torch.bfloat16 if self.amp_bf16 else torch.float16
- with autocast('cuda', enabled=self.args.mixed_precision, dtype=dtype):
- net_list, mask_feat_4, delta_disp = self.update_block(net_list, inp_list, geo_feat, disp, att)
-
- if self.local_match_radius > 0:
- match_prob = F.softmax(delta_disp.float(), dim=1)
- disp = (match_prob * disp_candidates).sum(1, keepdim=True)
- else:
- disp = disp + delta_disp.float()
-
- if task == 'depth':
- disp = disp.clamp(min=min_depth, max=max_depth)
-
- if test_mode and itr < iters-1:
- continue
-
- # upsample predictions
- if self.bilinear_up_depth:
- disp_up = F.interpolate(disp, scale_factor=4, mode='bilinear', align_corners=True)
- else:
- disp_up = self.upsample_disp(disp.float(), mask_feat_4.float(), stem_2x.float(), task=task)
-
- disp_preds.append(disp_up)
-
- if iters == 0 and task == 'depth':
- disp = F.interpolate(disp, scale_factor=4, mode='bilinear', align_corners=True).squeeze(1)
- disp_up = torch.exp(disp) if self.sample_log_depth else (1. / disp)
- # else:
- # # no refine, check the base model
- # disp = disp.detach()
- # geo_feat = geo_fn(disp, coords, low_memory=low_memory)
-
- # dtype = torch.bfloat16 if self.amp_bf16 else torch.float16
- # with autocast('cuda', enabled=self.args.mixed_precision, dtype=dtype):
- # net_list, mask_feat_4, delta_disp = self.update_block(net_list, inp_list, geo_feat, disp, att)
-
- # # upsample predictions
- # disp_up = 1. / self.upsample_disp(disp.float(), mask_feat_4.float(), stem_2x.float(), task=task).squeeze(1)
-
- if test_mode:
- if pred_bidir_depth:
- half = disp_up.size(0) // 2
-
- if return_features:
- return disp_up[:half], disp_up[half:], torch.cat((vit_feat, features_left[0]), dim=1), prob
-
- return disp_up
- else:
- depth_preds = [disp_up]
-
- if pred_bidir_depth:
- half = depth_preds[0].size(0) // 2
- fwd_depth_preds = [pred[:half] for pred in depth_preds]
- bwd_depth_preds = [pred[half:] for pred in depth_preds]
-
- if return_features:
- return fwd_depth_preds, bwd_depth_preds, torch.cat((vit_feat, features_left[0]), dim=1), prob
-
- return fwd_depth_preds, bwd_depth_preds
-
- return depth_preds
-
- if task == 'depth':
- # convert inverse depth to depth
- disp_up = torch.exp(disp_up.squeeze(1)) if self.sample_log_depth else (1. / disp_up.squeeze(1))
- init_disp = torch.exp(init_disp) if self.sample_log_depth else (1. / init_disp)
- for i in range(len(disp_preds)):
- disp_preds[i] = torch.exp(disp_preds[i].squeeze(1)) if self.sample_log_depth else (1. / disp_preds[i].squeeze(1)) # [B, H, W]
-
- if test_mode or not self.training:
- if task == 'depth':
- # disp_up = disp_up.clamp(min=min_depth, max=max_depth)
-
- if pred_bidir_depth:
- half = disp_up.size(0) // 2
-
- if return_features:
- return disp_up[:half], disp_up[half:], torch.cat((vit_feat, features_left[0]), dim=1), prob
-
- return disp_up[:half], disp_up[half:]
-
- return disp_up
-
- if task == 'depth':
- # upsample to the full resolution to add supervison
- init_disp = F.interpolate(init_disp, scale_factor=4, mode='bilinear', align_corners=True).squeeze(1)
-
- if self.supervise_init_depth:
- depth_preds = [init_disp] + disp_preds
- else:
- depth_preds = disp_preds
-
- if pred_bidir_depth:
- half = depth_preds[0].size(0) // 2
- fwd_depth_preds = [pred[:half] for pred in depth_preds]
- bwd_depth_preds = [pred[half:] for pred in depth_preds]
-
- if return_features:
- return fwd_depth_preds, bwd_depth_preds, torch.cat((vit_feat, features_left[0]), dim=1), prob
-
- return fwd_depth_preds, bwd_depth_preds
-
- return depth_preds
-
- return init_disp, disp_preds
-
-
- def run_hierachical(self, image1, image2, iters=12, test_mode=False, low_memory=False, small_ratio=0.5):
- B,_,H,W = image1.shape
- img1_small = F.interpolate(image1, scale_factor=small_ratio, align_corners=False, mode='bilinear')
- img2_small = F.interpolate(image2, scale_factor=small_ratio, align_corners=False, mode='bilinear')
- padder = InputPadder(img1_small.shape[-2:], divis_by=32, force_square=False)
- img1_small, img2_small = padder.pad(img1_small, img2_small)
- disp_small = self.forward(img1_small, img2_small, test_mode=True, iters=iters, low_memory=low_memory)
- disp_small = padder.unpad(disp_small.float())
- disp_small_up = F.interpolate(disp_small, size=(H,W), mode='bilinear', align_corners=True) * 1/small_ratio
- disp_small_up = disp_small_up.clip(0, None)
-
- padder = InputPadder(image1.shape[-2:], divis_by=32, force_square=False)
- image1, image2, disp_small_up = padder.pad(image1, image2, disp_small_up)
- disp_small_up += padder._pad[0]
- init_disp = F.interpolate(disp_small_up, scale_factor=0.25, mode='bilinear', align_corners=True) * 0.25 # Init disp will be 1/4
- disp = self.forward(image1, image2, iters=iters, test_mode=test_mode, low_memory=low_memory, init_disp=init_disp)
- disp = padder.unpad(disp.float())
- return disp
-
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/core/geometry.py b/optgs/model/encoder/.deprecated/foundationstereo/core/geometry.py
deleted file mode 100644
index eb19acb506021710d7a404a135526489c678d235..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/foundationstereo/core/geometry.py
+++ /dev/null
@@ -1,81 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
-#
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto. Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
-
-
-import torch,pdb,os,sys
-import torch.nn.functional as F
-from core.utils.utils import bilinear_sampler
-code_dir = os.path.dirname(os.path.realpath(__file__))
-sys.path.append(f'{code_dir}/../')
-from Utils import *
-
-class Combined_Geo_Encoding_Volume:
- def __init__(self, init_fmap1, init_fmap2, geo_volume, num_levels=2, dx=None, no_corr=False):
- self.num_levels = num_levels
- self.geo_volume_pyramid = []
- self.init_corr_pyramid = []
- self.dx = dx
- self.no_corr = no_corr
-
- b, c, d, h, w = geo_volume.shape
- geo_volume = geo_volume.permute(0, 3, 4, 1, 2).reshape(b*h*w, c, 1, d).contiguous()
-
- self.geo_volume_pyramid.append(geo_volume)
- for i in range(self.num_levels-1):
- geo_volume = F.avg_pool2d(geo_volume, [1,2], stride=[1,2])
- self.geo_volume_pyramid.append(geo_volume)
-
- if not no_corr:
- # all pairs correlation
- init_corr = Combined_Geo_Encoding_Volume.corr(init_fmap1, init_fmap2)
- b, h, w, _, w2 = init_corr.shape
- init_corr = init_corr.reshape(b*h*w, 1, 1, w2)
- self.init_corr_pyramid.append(init_corr)
-
- for i in range(self.num_levels-1):
- init_corr = F.avg_pool2d(init_corr, [1,2], stride=[1,2])
- self.init_corr_pyramid.append(init_corr)
-
-
- def __call__(self, disp, coords, low_memory=False, no_geo_volume=False):
- b, _, h, w = disp.shape
- self.dx = self.dx.to(disp.device)
- out_pyramid = []
- for i in range(self.num_levels):
- x0 = self.dx + disp.reshape(b*h*w, 1, 1, 1) / 2**i
- y0 = torch.zeros_like(x0)
-
- if not no_geo_volume:
- geo_volume = self.geo_volume_pyramid[i]
- disp_lvl = torch.cat([x0,y0], dim=-1)
- geo_volume = bilinear_sampler(geo_volume, disp_lvl, low_memory=low_memory)
- geo_volume = geo_volume.reshape(b, h, w, -1)
- out_pyramid.append(geo_volume)
-
- if not self.no_corr:
- init_corr = self.init_corr_pyramid[i]
- init_x0 = coords.reshape(b*h*w, 1, 1, 1)/2**i - disp.reshape(b*h*w, 1, 1, 1) / 2**i + self.dx # X on right image
- init_coords_lvl = torch.cat([init_x0,y0], dim=-1)
- init_corr = bilinear_sampler(init_corr, init_coords_lvl, low_memory=low_memory)
- init_corr = init_corr.reshape(b, h, w, -1)
- out_pyramid.append(init_corr)
-
- out_pyramid = torch.cat(out_pyramid, dim=-1)
- return out_pyramid.permute(0, 3, 1, 2).contiguous() #(B,C,H,W)
-
-
- @staticmethod
- def corr(fmap1, fmap2):
- B, D, H, W1 = fmap1.shape
- _, _, _, W2 = fmap2.shape
- fmap1 = fmap1.reshape(B, D, H, W1)
- fmap2 = fmap2.reshape(B, D, H, W2)
- with torch.amp.autocast('cuda', enabled=False):
- corr = torch.einsum('aijk,aijh->ajkh', F.normalize(fmap1.float(), dim=1), F.normalize(fmap2.float(), dim=1))
- corr = corr.reshape(B, H, W1, 1, W2)
- return corr
\ No newline at end of file
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/core/submodule.py b/optgs/model/encoder/.deprecated/foundationstereo/core/submodule.py
deleted file mode 100644
index ec78f411951ca2d21e0e85616fbde64b2819cc3c..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/foundationstereo/core/submodule.py
+++ /dev/null
@@ -1,596 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
-#
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto. Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
-
-
-import torch,pdb,os,sys
-import torch.nn as nn
-import torch.nn.functional as F
-import numpy as np
-from einops import rearrange
-from torch import einsum
-code_dir = os.path.dirname(os.path.realpath(__file__))
-sys.path.append(f'{code_dir}/../')
-from Utils import *
-from flash_attn import flash_attn_qkvpacked_func, flash_attn_func
-
-
-def _is_contiguous(tensor: torch.Tensor) -> bool:
- if torch.jit.is_scripting():
- return tensor.is_contiguous()
- else:
- return tensor.is_contiguous(memory_format=torch.contiguous_format)
-
-
-class LayerNorm2d(nn.LayerNorm):
- r""" https://huggingface.co/spaces/Roll20/pet_score/blob/b258ef28152ab0d5b377d9142a23346f863c1526/lib/timm/models/convnext.py#L85
- LayerNorm for channels_first tensors with 2d spatial dimensions (ie N, C, H, W).
- """
-
- def __init__(self, normalized_shape, eps=1e-6):
- super().__init__(normalized_shape, eps=eps)
-
- def forward(self, x) -> torch.Tensor:
- """
- @x: (B,C,H,W)
- """
- if _is_contiguous(x):
- return F.layer_norm(x.permute(0, 2, 3, 1), self.normalized_shape, self.weight, self.bias, self.eps).permute(0, 3, 1, 2).contiguous()
- else:
- s, u = torch.var_mean(x, dim=1, keepdim=True)
- x = (x - u) * torch.rsqrt(s + self.eps)
- x = x * self.weight[:, None, None] + self.bias[:, None, None]
- return x
-
-
-
-class BasicConv(nn.Module):
-
- def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, bn=True, relu=True, norm='batch', **kwargs):
- super(BasicConv, self).__init__()
-
- self.relu = relu
- self.use_bn = bn
- self.bn = nn.Identity()
- if is_3d:
- if deconv:
- self.conv = nn.ConvTranspose3d(in_channels, out_channels, bias=False, **kwargs)
- else:
- self.conv = nn.Conv3d(in_channels, out_channels, bias=False, **kwargs)
- if self.use_bn:
- if norm=='batch':
- self.bn = nn.BatchNorm3d(out_channels)
- elif norm=='instance':
- self.bn = nn.InstanceNorm3d(out_channels)
- else:
- if deconv:
- self.conv = nn.ConvTranspose2d(in_channels, out_channels, bias=False, **kwargs)
- else:
- self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
- if self.use_bn:
- if norm=='batch':
- self.bn = nn.BatchNorm2d(out_channels)
- elif norm=='instance':
- self.bn = nn.InstanceNorm2d(out_channels)
-
- def forward(self, x):
- x = self.conv(x)
- if self.use_bn:
- x = self.bn(x)
- if self.relu:
- x = nn.LeakyReLU()(x)#, inplace=True)
- return x
-
-
-class Conv3dNormActReduced(nn.Module):
- def __init__(self, C_in, C_out, hidden=None, kernel_size=3, kernel_disp=None, stride=1, norm=nn.BatchNorm3d):
- super().__init__()
- if kernel_disp is None:
- kernel_disp = kernel_size
- if hidden is None:
- hidden = C_out
- self.conv1 = nn.Sequential(
- nn.Conv3d(C_in, hidden, kernel_size=(1,kernel_size,kernel_size), padding=(0, kernel_size//2, kernel_size//2), stride=(1, stride, stride)),
- norm(hidden),
- nn.ReLU(),
- )
- self.conv2 = nn.Sequential(
- nn.Conv3d(hidden, C_out, kernel_size=(kernel_disp, 1, 1), padding=(kernel_disp//2, 0, 0), stride=(stride, 1, 1)),
- norm(C_out),
- nn.ReLU(),
- )
-
-
- def forward(self, x):
- """
- @x: (B,C,D,H,W)
- """
- x = self.conv1(x)
- x = self.conv2(x)
- return x
-
-
-
-
-class ResnetBasicBlock(nn.Module):
- def __init__(self, inplanes, planes, kernel_size=3, stride=1, padding=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=nn.BatchNorm2d, bias=False):
- super().__init__()
- self.norm_layer = norm_layer
- if groups != 1 or base_width != 64:
- raise ValueError('BasicBlock only supports groups=1 and base_width=64')
- if dilation > 1:
- raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
- # Both self.conv1 and self.downsample layers downsample the input when stride != 1
- self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=kernel_size, stride=stride, bias=bias, padding=padding)
- if self.norm_layer is not None:
- self.bn1 = norm_layer(planes)
- self.relu = nn.ReLU(inplace=True)
- self.conv2 = nn.Conv2d(planes, planes, kernel_size=kernel_size, stride=stride, bias=bias, padding=padding)
- if self.norm_layer is not None:
- self.bn2 = norm_layer(planes)
- self.downsample = downsample
- self.stride = stride
-
-
- def forward(self, x):
- identity = x
-
- out = self.conv1(x)
- if self.norm_layer is not None:
- out = self.bn1(out)
- out = self.relu(out)
-
- out = self.conv2(out)
- if self.norm_layer is not None:
- out = self.bn2(out)
-
- if self.downsample is not None:
- identity = self.downsample(x)
- out += identity
- out = self.relu(out)
-
- return out
-
-
-class ResnetBasicBlock3D(nn.Module):
- def __init__(self, inplanes, planes, kernel_size=3, stride=1, padding=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=nn.BatchNorm3d, bias=False):
- super().__init__()
- self.norm_layer = norm_layer
- if groups != 1 or base_width != 64:
- raise ValueError('BasicBlock only supports groups=1 and base_width=64')
- if dilation > 1:
- raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
- self.conv1 = nn.Conv3d(inplanes, planes, kernel_size=kernel_size, stride=stride, bias=bias, padding=padding)
- if self.norm_layer is not None:
- self.bn1 = norm_layer(planes)
- self.relu = nn.ReLU(inplace=True)
- self.conv2 = nn.Conv3d(planes, planes, kernel_size=kernel_size, stride=stride, bias=bias, padding=padding)
- if self.norm_layer is not None:
- self.bn2 = norm_layer(planes)
- self.downsample = downsample
- self.stride = stride
-
-
- def forward(self, x):
- identity = x
-
- out = self.conv1(x)
- if self.norm_layer is not None:
- out = self.bn1(out)
- out = self.relu(out)
-
- out = self.conv2(out)
- if self.norm_layer is not None:
- out = self.bn2(out)
-
- if self.downsample is not None:
- identity = self.downsample(x)
- out += identity
- out = self.relu(out)
-
- return out
-
-
-class FlashMultiheadAttention(nn.Module):
- def __init__(self, embed_dim, num_heads):
- super().__init__()
- self.num_heads = num_heads
- self.embed_dim = embed_dim
- self.head_dim = embed_dim // num_heads
- assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
-
- self.q_proj = nn.Linear(embed_dim, embed_dim)
- self.k_proj = nn.Linear(embed_dim, embed_dim)
- self.v_proj = nn.Linear(embed_dim, embed_dim)
- self.out_proj = nn.Linear(embed_dim, embed_dim)
-
- def forward(self, query, key, value, attn_mask=None, window_size=(-1,-1)):
- """
- @query: (B,L,C)
- """
- B,L,C = query.shape
- Q = self.q_proj(query)
- K = self.k_proj(key)
- V = self.v_proj(value)
-
- Q = Q.view(Q.size(0), Q.size(1), self.num_heads, self.head_dim)
- K = K.view(K.size(0), K.size(1), self.num_heads, self.head_dim)
- V = V.view(V.size(0), V.size(1), self.num_heads, self.head_dim)
-
- attn_output = flash_attn_func(Q, K, V, window_size=window_size) # Replace with actual FlashAttention function
-
- attn_output = attn_output.reshape(B,L,-1)
- output = self.out_proj(attn_output)
-
- return output
-
-
-
-class FlashAttentionTransformerEncoderLayer(nn.Module):
- def __init__(self, embed_dim, num_heads, dim_feedforward, dropout=0.1, act=nn.GELU, norm=nn.LayerNorm):
- super().__init__()
- self.self_attn = FlashMultiheadAttention(embed_dim, num_heads)
- self.act = act()
-
- self.linear1 = nn.Linear(embed_dim, dim_feedforward)
- self.dropout = nn.Dropout(dropout)
- self.linear2 = nn.Linear(dim_feedforward, embed_dim)
-
- self.norm1 = norm(embed_dim)
- self.norm2 = norm(embed_dim)
- self.dropout1 = nn.Dropout(dropout)
- self.dropout2 = nn.Dropout(dropout)
-
- def forward(self, src, src_mask=None, window_size=(-1, -1)):
- src2 = self.self_attn(src, src, src, src_mask, window_size=window_size)
- src = src + self.dropout1(src2)
- src = self.norm1(src)
-
- src2 = self.linear2(self.dropout(self.act(self.linear1(src))))
- src = src + self.dropout2(src2)
- src = self.norm2(src)
-
- return src
-
-
-
-class UpsampleConv(nn.Module):
- def __init__(self, C_in, C_out, is_3d=False, kernel_size=3, bias=True, stride=1, padding=1):
- super().__init__()
- self.is_3d = is_3d
- if is_3d:
- self.conv = nn.Conv3d(C_in, C_out, kernel_size=kernel_size, stride=1, padding=kernel_size//2, bias=bias)
- else:
- self.conv = nn.Conv2d(C_in, C_out, kernel_size=kernel_size, stride=1, padding=kernel_size//2, bias=bias)
-
- def forward(self, x):
- if self.is_3d:
- mode = 'trilinear'
- else:
- mode = 'bilinear'
- x = F.interpolate(x, size=None, scale_factor=2, align_corners=False, mode=mode)
- x = self.conv(x)
- return x
-
-
-
-class Conv2x(nn.Module):
-
- def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, concat=True, keep_concat=True, bn=True, relu=True, keep_dispc=False):
- super(Conv2x, self).__init__()
- self.concat = concat
- self.is_3d = is_3d
- if deconv and is_3d:
- kernel = (4, 4, 4)
- elif deconv:
- kernel = 4
- else:
- kernel = 3
-
- if deconv and is_3d and keep_dispc:
- kernel = (1, 4, 4)
- stride = (1, 2, 2)
- padding = (0, 1, 1)
- self.conv1 = BasicConv(in_channels, out_channels, deconv, is_3d, bn=bn, relu=True, kernel_size=kernel, stride=stride, padding=padding)
- else:
- self.conv1 = BasicConv(in_channels, out_channels, deconv, is_3d, bn=bn, relu=True, kernel_size=kernel, stride=2, padding=1)
-
- if self.concat:
- mul = 2 if keep_concat else 1
- self.conv2 = BasicConv(out_channels*2, out_channels*mul, False, is_3d, bn, relu, kernel_size=3, stride=1, padding=1)
- else:
- self.conv2 = BasicConv(out_channels, out_channels, False, is_3d, bn, relu, kernel_size=3, stride=1, padding=1)
-
- def forward(self, x, rem):
- x = self.conv1(x)
- if x.shape != rem.shape:
- x = F.interpolate(x, size=(rem.shape[-2], rem.shape[-1]), mode='bilinear')
- if self.concat:
- x = torch.cat((x, rem), 1)
- else:
- x = x + rem
- x = self.conv2(x)
- return x
-
-
-class BasicConv_IN(nn.Module):
-
- def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, IN=True, relu=True, **kwargs):
- super(BasicConv_IN, self).__init__()
-
- self.relu = relu
- self.use_in = IN
- if is_3d:
- if deconv:
- self.conv = nn.ConvTranspose3d(in_channels, out_channels, bias=False, **kwargs)
- else:
- self.conv = nn.Conv3d(in_channels, out_channels, bias=False, **kwargs)
- self.IN = nn.InstanceNorm3d(out_channels)
- else:
- if deconv:
- self.conv = nn.ConvTranspose2d(in_channels, out_channels, bias=False, **kwargs)
- else:
- self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
- self.IN = nn.InstanceNorm2d(out_channels)
-
- def forward(self, x):
- x = self.conv(x)
- if self.use_in:
- x = self.IN(x)
- if self.relu:
- x = nn.LeakyReLU()(x)#, inplace=True)
- return x
-
-
-class Conv2x_IN(nn.Module):
-
- def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, concat=True, keep_concat=True, IN=True, relu=True, keep_dispc=False):
- super(Conv2x_IN, self).__init__()
- self.concat = concat
- self.is_3d = is_3d
- if deconv and is_3d:
- kernel = (4, 4, 4)
- elif deconv:
- kernel = 4
- else:
- kernel = 3
-
- if deconv and is_3d and keep_dispc:
- kernel = (1, 4, 4)
- stride = (1, 2, 2)
- padding = (0, 1, 1)
- self.conv1 = BasicConv_IN(in_channels, out_channels, deconv, is_3d, IN=True, relu=True, kernel_size=kernel, stride=stride, padding=padding)
- else:
- self.conv1 = BasicConv_IN(in_channels, out_channels, deconv, is_3d, IN=True, relu=True, kernel_size=kernel, stride=2, padding=1)
-
- if self.concat:
- mul = 2 if keep_concat else 1
- self.conv2 = ResnetBasicBlock(out_channels*2, out_channels*mul, kernel_size=3, stride=1, padding=1, norm_layer=nn.InstanceNorm2d)
- else:
- self.conv2 = BasicConv_IN(out_channels, out_channels, False, is_3d, IN, relu, kernel_size=3, stride=1, padding=1)
-
- def forward(self, x, rem):
- x = self.conv1(x)
- if x.shape != rem.shape:
- x = F.interpolate(x, size=(rem.shape[-2], rem.shape[-1]), mode='bilinear')
- if self.concat:
- x = torch.cat((x, rem), 1)
- else:
- x = x + rem
- x = self.conv2(x)
- return x
-
-
-def groupwise_correlation(fea1, fea2, num_groups):
- B, C, H, W = fea1.shape
- assert C % num_groups == 0, f"C:{C}, num_groups:{num_groups}"
- channels_per_group = C // num_groups
- fea1 = fea1.reshape(B, num_groups, channels_per_group, H, W)
- fea2 = fea2.reshape(B, num_groups, channels_per_group, H, W)
- with torch.cuda.amp.autocast(enabled=False):
- cost = (F.normalize(fea1.float(), dim=2) * F.normalize(fea2.float(), dim=2)).sum(dim=2) #!NOTE Divide first for numerical stability
- assert cost.shape == (B, num_groups, H, W)
- return cost
-
-def build_gwc_volume(refimg_fea, targetimg_fea, maxdisp, num_groups, stride=1):
- """
- @refimg_fea: left image feature
- @targetimg_fea: right image feature
- """
- B, C, H, W = refimg_fea.shape
- volume = refimg_fea.new_zeros([B, num_groups, maxdisp, H, W])
- for i in range(maxdisp):
- if i > 0:
- volume[:, :, i, :, i:] = groupwise_correlation(refimg_fea[:, :, :, i:], targetimg_fea[:, :, :, :-i], num_groups)
- else:
- volume[:, :, i, :, :] = groupwise_correlation(refimg_fea, targetimg_fea, num_groups)
- volume = volume.contiguous()
- return volume
-
-
-
-def build_concat_volume(refimg_fea, targetimg_fea, maxdisp):
- B, C, H, W = refimg_fea.shape
- volume = refimg_fea.new_zeros([B, 2 * C, maxdisp, H, W])
- for i in range(maxdisp):
- if i > 0:
- volume[:, :C, i, :, :] = refimg_fea[:, :, :, :]
- volume[:, C:, i, :, i:] = targetimg_fea[:, :, :, :-i]
- else:
- volume[:, :C, i, :, :] = refimg_fea
- volume[:, C:, i, :, :] = targetimg_fea
- volume = volume.contiguous()
- return volume
-
-
-
-def disparity_regression(x, maxdisp):
- assert len(x.shape) == 4
- disp_values = torch.arange(0, maxdisp, dtype=x.dtype, device=x.device)
- disp_values = disp_values.reshape(1, maxdisp, 1, 1)
- return torch.sum(x * disp_values, 1, keepdim=True)
-
-
-class FeatureAtt(nn.Module):
- def __init__(self, cv_chan, feat_chan):
- super(FeatureAtt, self).__init__()
-
- self.feat_att = nn.Sequential(
- BasicConv(feat_chan, feat_chan//2, kernel_size=1, stride=1, padding=0),
- nn.Conv2d(feat_chan//2, cv_chan, 1)
- )
-
- def forward(self, cv, feat):
- '''
- @cv: cost volume (B,C,D,H,W)
- @feat: (B,C,H,W)
- '''
- feat_att = self.feat_att(feat).unsqueeze(2) #(B,C,1,H,W)
- cv = torch.sigmoid(feat_att)*cv
- return cv
-
-def context_upsample(disp_low, up_weights, task='stereo'):
- """
- @disp_low: (b,1,h,w) 1/4 resolution
- @up_weights: (b,9,4*h,4*w) Image resolution
- """
- b, c, h, w = disp_low.shape
-
- assert task in ['stereo', 'depth']
- if task == 'depth':
- # since we predict inverse depth, we don't want to do zero padding
- disp_unfold = F.unfold(F.pad(disp_low.reshape(b,c,h,w), pad=(1,1,1,1), mode='replicate'),3,1,0).reshape(b,-1,h,w)
- else:
- disp_unfold = F.unfold(disp_low.reshape(b,c,h,w),3,1,1).reshape(b,-1,h,w)
- disp_unfold = F.interpolate(disp_unfold,(h*4,w*4),mode='nearest').reshape(b,9,h*4,w*4)
-
- disp = (disp_unfold*up_weights).sum(1)
-
- return disp
-
-
-
-class PositionalEmbedding(nn.Module):
- def __init__(self, d_model, max_len=512):
- super().__init__()
-
- # Compute the positional encodings once in log space.
- pe = torch.zeros(max_len, d_model).float()
- pe.require_grad = False
-
- position = torch.arange(0, max_len).float().unsqueeze(1) #(N,1)
- div_term = (torch.arange(0, d_model, 2).float() * -(np.log(10000.0) / d_model)).exp()[None]
-
- pe[:, 0::2] = torch.sin(position * div_term) #(N, d_model/2)
- pe[:, 1::2] = torch.cos(position * div_term)
-
- pe = pe.unsqueeze(0)
- self.pe = pe
- # self.register_buffer('pe', pe) #(1, max_len, D)
-
-
- def forward(self, x, resize_embed=False):
- '''
- @x: (B,N,D)
- '''
- self.pe = self.pe.to(x.device).to(x.dtype)
- pe = self.pe
- if pe.shape[1] 0 else None
-
- def forward(self, x):
- input = x
- x = self.dwconv(x)
- x = self.norm(x)
- x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
- x = self.pwconv1(x)
- x = self.act(x)
- x = self.pwconv2(x)
- if self.gamma is not None:
- x = self.gamma * x
- x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
-
- x = input + x
- return x
\ No newline at end of file
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/core/unimatch_geometry.py b/optgs/model/encoder/.deprecated/foundationstereo/core/unimatch_geometry.py
deleted file mode 100644
index edf2c26d7c22ec88e6d3f83505d60b610e494537..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/foundationstereo/core/unimatch_geometry.py
+++ /dev/null
@@ -1,302 +0,0 @@
-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)) # [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].repeat(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 generate_window_grid(h_min, h_max, w_min, w_max, len_h, len_w, device=None):
- assert device is not None
-
- x, y = torch.meshgrid([torch.linspace(w_min, w_max, len_w, device=device),
- torch.linspace(h_min, h_max, len_h, device=device)],
- )
- grid = torch.stack((x, y), -1).transpose(0, 1).float() # [H, W, 2]
-
- return grid
-
-
-def normalize_coords(coords, h, w):
- # coords: [B, H, W, 2]
- c = torch.Tensor([(w - 1) / 2., (h - 1) / 2.]).float().to(coords.device)
- return (coords - c) / c # [-1, 1]
-
-
-def bilinear_sample(img, sample_coords, mode='bilinear', padding_mode='zeros', return_mask=False):
- # img: [B, C, H, W]
- # sample_coords: [B, 2, H, W] in image scale
- if sample_coords.size(1) != 2: # [B, H, W, 2]
- sample_coords = sample_coords.permute(0, 3, 1, 2)
-
- b, _, h, w = sample_coords.shape
-
- # Normalize to [-1, 1]
- x_grid = 2 * sample_coords[:, 0] / (w - 1) - 1
- y_grid = 2 * sample_coords[:, 1] / (h - 1) - 1
-
- grid = torch.stack([x_grid, y_grid], dim=-1) # [B, H, W, 2]
-
- img = F.grid_sample(img, grid, mode=mode, padding_mode=padding_mode, align_corners=True)
-
- if return_mask:
- mask = (x_grid >= -1) & (y_grid >= -1) & (x_grid <= 1) & (y_grid <= 1) # [B, H, W]
-
- return img, mask
-
- return img
-
-
-def flow_warp(feature, flow, mask=False, padding_mode='zeros'):
- b, c, h, w = feature.size()
- assert flow.size(1) == 2
-
- grid = coords_grid(b, h, w).to(flow.device) + flow # [B, 2, H, W]
-
- return bilinear_sample(feature, grid, padding_mode=padding_mode,
- return_mask=mask)
-
-
-def forward_backward_consistency_check(fwd_flow, bwd_flow,
- alpha=0.01,
- beta=0.5
- ):
- # fwd_flow, bwd_flow: [B, 2, H, W]
- # alpha and beta values are following UnFlow (https://arxiv.org/abs/1711.07837)
- assert fwd_flow.dim() == 4 and bwd_flow.dim() == 4
- assert fwd_flow.size(1) == 2 and bwd_flow.size(1) == 2
- flow_mag = torch.norm(fwd_flow, dim=1) + torch.norm(bwd_flow, dim=1) # [B, H, W]
-
- warped_bwd_flow = flow_warp(bwd_flow, fwd_flow) # [B, 2, H, W]
- warped_fwd_flow = flow_warp(fwd_flow, bwd_flow) # [B, 2, H, W]
-
- diff_fwd = torch.norm(fwd_flow + warped_bwd_flow, dim=1) # [B, H, W]
- diff_bwd = torch.norm(bwd_flow + warped_fwd_flow, dim=1)
-
- threshold = alpha * flow_mag + beta
-
- fwd_occ = (diff_fwd > threshold).float() # [B, H, W]
- bwd_occ = (diff_bwd > threshold).float()
-
- return fwd_occ, bwd_occ
-
-
-def back_project(depth, intrinsics):
- # Back project 2D pixel coords to 3D points
- # depth: [B, H, W]
- # intrinsics: [B, 3, 3]
- b, h, w = depth.shape
- grid = coords_grid(b, h, w, homogeneous=True, device=depth.device) # [B, 3, H, W]
-
- intrinsics_inv = torch.inverse(intrinsics) # [B, 3, 3]
-
- points = intrinsics_inv.bmm(grid.view(b, 3, -1)).view(b, 3, h, w) * depth.unsqueeze(1) # [B, 3, H, W]
-
- return points
-
-
-def camera_transform(points_ref, extrinsics_ref=None, extrinsics_tgt=None, extrinsics_rel=None):
- # Transform 3D points from reference camera to target camera
- # points_ref: [B, 3, H, W]
- # extrinsics_ref: [B, 4, 4]
- # extrinsics_tgt: [B, 4, 4]
- # extrinsics_rel: [B, 4, 4], relative pose transform
- b, _, h, w = points_ref.shape
-
- if extrinsics_rel is None:
- extrinsics_rel = torch.bmm(extrinsics_tgt, torch.inverse(extrinsics_ref)) # [B, 4, 4]
-
- points_tgt = torch.bmm(extrinsics_rel[:, :3, :3],
- points_ref.view(b, 3, -1)) + extrinsics_rel[:, :3, -1:] # [B, 3, H*W]
-
- points_tgt = points_tgt.view(b, 3, h, w) # [B, 3, H, W]
-
- return points_tgt
-
-
-def reproject(points_tgt, intrinsics, return_mask=False):
- # reproject to target view
- # points_tgt: [B, 3, H, W]
- # intrinsics: [B, 3, 3]
-
- b, _, h, w = points_tgt.shape
-
- proj_points = torch.bmm(intrinsics, points_tgt.view(b, 3, -1)).view(b, 3, h, w) # [B, 3, H, W]
-
- X = proj_points[:, 0]
- Y = proj_points[:, 1]
- Z = proj_points[:, 2].clamp(min=1e-3)
-
- pixel_coords = torch.stack([X / Z, Y / Z], dim=1).view(b, 2, h, w) # [B, 2, H, W] in image scale
-
- if return_mask:
- # valid mask in pixel space
- mask = (pixel_coords[:, 0] >= 0) & (pixel_coords[:, 0] <= (w - 1)) & (
- pixel_coords[:, 1] >= 0) & (pixel_coords[:, 1] <= (h - 1)) # [B, H, W]
-
- return pixel_coords, mask
-
- return pixel_coords
-
-
-def reproject_coords(depth_ref, intrinsics, extrinsics_ref=None, extrinsics_tgt=None, extrinsics_rel=None,
- return_mask=False):
- # Compute reprojection sample coords
- points_ref = back_project(depth_ref, intrinsics) # [B, 3, H, W]
- points_tgt = camera_transform(points_ref, extrinsics_ref, extrinsics_tgt, extrinsics_rel=extrinsics_rel)
-
- if return_mask:
- reproj_coords, mask = reproject(points_tgt, intrinsics,
- return_mask=return_mask) # [B, 2, H, W] in image scale
-
- return reproj_coords, mask
-
- reproj_coords = reproject(points_tgt, intrinsics,
- return_mask=return_mask) # [B, 2, H, W] in image scale
-
- return reproj_coords
-
-
-def compute_flow_with_depth_pose(depth_ref, intrinsics,
- extrinsics_ref=None, extrinsics_tgt=None, extrinsics_rel=None,
- return_coords=False,
- return_mask=False):
- b, h, w = depth_ref.shape
- coords_init = coords_grid(b, h, w, device=depth_ref.device) # [B, 2, H, W]
-
- if return_mask:
- reproj_coords, mask = reproject_coords(depth_ref, intrinsics, extrinsics_ref, extrinsics_tgt,
- extrinsics_rel=extrinsics_rel,
- return_mask=return_mask) # [B, 2, H, W]
- rigid_flow = reproj_coords - coords_init
-
- return rigid_flow, mask
-
- reproj_coords = reproject_coords(depth_ref, intrinsics, extrinsics_ref, extrinsics_tgt,
- extrinsics_rel=extrinsics_rel,
- return_mask=return_mask) # [B, 2, H, W]
-
- if return_coords:
- return reproj_coords
-
- rigid_flow = reproj_coords - coords_init
-
- return rigid_flow
-
-
-def forward_backward_consistency_check(fwd_flow, bwd_flow,
- alpha=0.01,
- beta=0.5,
- return_flow_diff=False,
- ):
- # fwd_flow, bwd_flow: [B, 2, H, W]
- # alpha and beta values are following UnFlow (https://arxiv.org/abs/1711.07837)
- assert fwd_flow.dim() == 4 and bwd_flow.dim() == 4
- assert fwd_flow.size(1) == 2 and bwd_flow.size(1) == 2
- flow_mag = torch.norm(fwd_flow, dim=1) + torch.norm(bwd_flow, dim=1) # [B, H, W]
-
- warped_bwd_flow = flow_warp(bwd_flow, fwd_flow) # [B, 2, H, W]
- warped_fwd_flow = flow_warp(fwd_flow, bwd_flow) # [B, 2, H, W]
-
- diff_fwd = torch.norm(fwd_flow + warped_bwd_flow, dim=1) # [B, H, W]
- diff_bwd = torch.norm(bwd_flow + warped_fwd_flow, dim=1)
-
- if return_flow_diff:
- return diff_fwd, diff_bwd
-
- threshold = alpha * flow_mag + beta
-
- fwd_occ = (diff_fwd > threshold).float() # [B, H, W]
- bwd_occ = (diff_bwd > threshold).float()
-
- return fwd_occ, bwd_occ
-
-
-def warp_with_depth_pose(feature1, intrinsics, pose, depth,
- padding_mode='zeros',
- return_rigid_flow=False,
- ):
- assert depth.dim() == 3 # [B, H, W]
- sample_coords = reproject_coords(depth,
- intrinsics,
- extrinsics_rel=pose,
- ) # [B, 2, H, W]
-
- sample_coords = sample_coords.permute(0, 2, 3, 1) # [B, H, W, 2]
-
- warped_feature1 = bilinear_sample(feature1, sample_coords,
- padding_mode=padding_mode) # [B, C, H, W]
-
- if return_rigid_flow:
- b, h, w = depth.size()
- coords_init = coords_grid(b, h, w, device=depth.device) # [B, 2, H, W]
- rigid_flow = sample_coords.permute(0, 3, 1, 2) - coords_init
-
- return warped_feature1, rigid_flow
-
- return warped_feature1
-
-
-def warp_with_pose_depth_candidates(feature1, intrinsics, pose, depth,
- padding_mode='zeros',
- rigid_flow_to_subtract=None,
- ):
- # pixel-specific depth candidates, useful for refinement
- # feature1: [B, C, H, W]
- # intrinsics: [B, 3, 3]
- # pose: [B, 4, 4]
- # depth: [B, D, H, W]
- assert intrinsics.size(1) == 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)
-
- # stop gradient
- with torch.no_grad():
- # 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.inverse(intrinsics).bmm(grid.view(b, 3, -1)) # [B, 3, H*W]
- points = torch.bmm(pose[:, :3, :3], points).unsqueeze(2).repeat(
- 1, 1, d, 1) * depth.view(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(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=1e-3) # [B, 2, D, H*W]
-
- if rigid_flow_to_subtract is not None:
- assert rigid_flow_to_subtract.dim() == 4 # [B, 2, H, W]
- assert rigid_flow_to_subtract.size(1) == 2
-
- pixel_coords = pixel_coords - rigid_flow_to_subtract.view(b, 2, h * w).unsqueeze(2)
-
- # 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]
-
- # sample features
- warped_feature = F.grid_sample(feature1, grid.view(b, d * h, w, 2), mode='bilinear',
- padding_mode=padding_mode,
- align_corners=True).view(b, c, d, h, w) # [B, C, D, H, W]
-
- return warped_feature
-
\ No newline at end of file
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/core/unimatch_matching.py b/optgs/model/encoder/.deprecated/foundationstereo/core/unimatch_matching.py
deleted file mode 100644
index 2da593b59ca4e246160e5349b00611f938855bdc..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/foundationstereo/core/unimatch_matching.py
+++ /dev/null
@@ -1,394 +0,0 @@
-import torch
-import torch.nn.functional as F
-
-from .unimatch_geometry import coords_grid, generate_window_grid, normalize_coords
-
-
-def global_correlation_softmax(feature0, feature1,
- pred_bidir_flow=False,
- ):
- # global correlation
- b, c, h, w = feature0.shape
- feature0 = feature0.view(b, c, -1).permute(0, 2, 1) # [B, H*W, C]
- feature1 = feature1.view(b, c, -1) # [B, C, H*W]
-
- correlation = torch.matmul(feature0, feature1).view(b, h, w, h, w) / (c ** 0.5) # [B, H, W, H, W]
-
- # flow from softmax
- init_grid = coords_grid(b, h, w).to(correlation.device) # [B, 2, H, W]
- grid = init_grid.view(b, 2, -1).permute(0, 2, 1) # [B, H*W, 2]
-
- correlation = correlation.view(b, h * w, h * w) # [B, H*W, H*W]
-
- if pred_bidir_flow:
- correlation = torch.cat((correlation, correlation.permute(0, 2, 1)), dim=0) # [2*B, H*W, H*W]
- init_grid = init_grid.repeat(2, 1, 1, 1) # [2*B, 2, H, W]
- grid = grid.repeat(2, 1, 1) # [2*B, H*W, 2]
- b = b * 2
-
- prob = F.softmax(correlation, dim=-1) # [B, H*W, H*W]
-
- correspondence = torch.matmul(prob, grid).view(b, h, w, 2).permute(0, 3, 1, 2) # [B, 2, H, W]
-
- # when predicting bidirectional flow, flow is the concatenation of forward flow and backward flow
- flow = correspondence - init_grid
-
- return flow, prob
-
-
-def local_correlation_softmax(feature0, feature1, local_radius,
- padding_mode='zeros',
- ):
- b, c, h, w = feature0.size()
- coords_init = coords_grid(b, h, w).to(feature0.device) # [B, 2, H, W]
- coords = coords_init.view(b, 2, -1).permute(0, 2, 1) # [B, H*W, 2]
-
- local_h = 2 * local_radius + 1
- local_w = 2 * local_radius + 1
-
- window_grid = generate_window_grid(-local_radius, local_radius,
- -local_radius, local_radius,
- local_h, local_w, device=feature0.device) # [2R+1, 2R+1, 2]
- window_grid = window_grid.reshape(-1, 2).repeat(b, 1, 1, 1) # [B, 1, (2R+1)^2, 2]
- sample_coords = coords.unsqueeze(-2) + window_grid # [B, H*W, (2R+1)^2, 2]
-
- sample_coords_softmax = sample_coords
-
- # exclude coords that are out of image space
- valid_x = (sample_coords[:, :, :, 0] >= 0) & (sample_coords[:, :, :, 0] < w) # [B, H*W, (2R+1)^2]
- valid_y = (sample_coords[:, :, :, 1] >= 0) & (sample_coords[:, :, :, 1] < h) # [B, H*W, (2R+1)^2]
-
- valid = valid_x & valid_y # [B, H*W, (2R+1)^2], used to mask out invalid values when softmax
-
- # normalize coordinates to [-1, 1]
- sample_coords_norm = normalize_coords(sample_coords, h, w) # [-1, 1]
- window_feature = F.grid_sample(feature1, sample_coords_norm,
- padding_mode=padding_mode, align_corners=True
- ).permute(0, 2, 1, 3) # [B, H*W, C, (2R+1)^2]
- feature0_view = feature0.permute(0, 2, 3, 1).view(b, h * w, 1, c) # [B, H*W, 1, C]
-
- corr = torch.matmul(feature0_view, window_feature).view(b, h * w, -1) / (c ** 0.5) # [B, H*W, (2R+1)^2]
-
- # mask invalid locations
- corr[~valid] = -1e9
-
- prob = F.softmax(corr, -1) # [B, H*W, (2R+1)^2]
-
- correspondence = torch.matmul(prob.unsqueeze(-2), sample_coords_softmax).squeeze(-2).view(
- b, h, w, 2).permute(0, 3, 1, 2) # [B, 2, H, W]
-
- flow = correspondence - coords_init
- match_prob = prob
-
- return flow, match_prob
-
-
-def local_correlation_with_flow(feature0, feature1,
- flow,
- local_radius,
- padding_mode='zeros',
- dilation=1,
- ):
- b, c, h, w = feature0.size()
- coords_init = coords_grid(b, h, w).to(feature0.device) # [B, 2, H, W]
- coords = coords_init.view(b, 2, -1).permute(0, 2, 1) # [B, H*W, 2]
-
- local_h = 2 * local_radius + 1
- local_w = 2 * local_radius + 1
-
- window_grid = generate_window_grid(-local_radius, local_radius,
- -local_radius, local_radius,
- local_h, local_w, device=feature0.device) # [2R+1, 2R+1, 2]
- window_grid = window_grid.reshape(-1, 2).repeat(b, 1, 1, 1) # [B, 1, (2R+1)^2, 2]
- sample_coords = coords.unsqueeze(-2) + window_grid * dilation # [B, H*W, (2R+1)^2, 2]
-
- # flow can be zero when using features after transformer
- if not isinstance(flow, float):
- sample_coords = sample_coords + flow.view(
- b, 2, -1).permute(0, 2, 1).unsqueeze(-2) # [B, H*W, (2R+1)^2, 2]
- else:
- assert flow == 0.
-
- # normalize coordinates to [-1, 1]
- sample_coords_norm = normalize_coords(sample_coords, h, w) # [-1, 1]
- window_feature = F.grid_sample(feature1, sample_coords_norm,
- padding_mode=padding_mode, align_corners=True
- ).permute(0, 2, 1, 3) # [B, H*W, C, (2R+1)^2]
- feature0_view = feature0.permute(0, 2, 3, 1).view(b, h * w, 1, c) # [B, H*W, 1, C]
-
- corr = torch.matmul(feature0_view, window_feature).view(b, h * w, -1) / (c ** 0.5) # [B, H*W, (2R+1)^2]
-
- corr = corr.view(b, h, w, -1).permute(0, 3, 1, 2).contiguous() # [B, (2R+1)^2, H, W]
-
- return corr
-
-
-def global_correlation_softmax_stereo(feature0, feature1,
- ):
- # global correlation on horizontal direction
- b, c, h, w = feature0.shape
-
- x_grid = torch.linspace(0, w - 1, w, device=feature0.device) # [W]
-
- feature0 = feature0.permute(0, 2, 3, 1) # [B, H, W, C]
- feature1 = feature1.permute(0, 2, 1, 3) # [B, H, C, W]
-
- correlation = torch.matmul(feature0, feature1) / (c ** 0.5) # [B, H, W, W]
-
- # mask subsequent positions to make disparity positive
- mask = torch.triu(torch.ones((w, w)), diagonal=1).type_as(feature0) # [W, W]
- valid_mask = (mask == 0).unsqueeze(0).unsqueeze(0).repeat(b, h, 1, 1) # [B, H, W, W]
-
- correlation[~valid_mask] = -1e9
-
- prob = F.softmax(correlation, dim=-1) # [B, H, W, W]
-
- correspondence = (x_grid.view(1, 1, 1, w) * prob).sum(-1) # [B, H, W]
-
- # NOTE: unlike flow, disparity is typically positive
- disparity = x_grid.view(1, 1, w).repeat(b, h, 1) - correspondence # [B, H, W]
-
- return disparity.unsqueeze(1), prob # feature resolution
-
-
-def local_correlation_softmax_stereo(feature0, feature1, local_radius,
- ):
- b, c, h, w = feature0.size()
- coords_init = coords_grid(b, h, w).to(feature0.device) # [B, 2, H, W]
- coords = coords_init.view(b, 2, -1).permute(0, 2, 1).contiguous() # [B, H*W, 2]
-
- local_h = 1
- local_w = 2 * local_radius + 1
-
- window_grid = generate_window_grid(0, 0,
- -local_radius, local_radius,
- local_h, local_w, device=feature0.device) # [1, 2R+1, 2]
- window_grid = window_grid.reshape(-1, 2).repeat(b, 1, 1, 1) # [B, 1, (2R+1), 2]
- sample_coords = coords.unsqueeze(-2) + window_grid # [B, H*W, (2R+1), 2]
-
- sample_coords_softmax = sample_coords
-
- # exclude coords that are out of image space
- valid_x = (sample_coords[:, :, :, 0] >= 0) & (sample_coords[:, :, :, 0] < w) # [B, H*W, (2R+1)^2]
- valid_y = (sample_coords[:, :, :, 1] >= 0) & (sample_coords[:, :, :, 1] < h) # [B, H*W, (2R+1)^2]
-
- valid = valid_x & valid_y # [B, H*W, (2R+1)^2], used to mask out invalid values when softmax
-
- # normalize coordinates to [-1, 1]
- sample_coords_norm = normalize_coords(sample_coords, h, w) # [-1, 1]
- window_feature = F.grid_sample(feature1, sample_coords_norm,
- padding_mode='zeros', align_corners=True
- ).permute(0, 2, 1, 3) # [B, H*W, C, (2R+1)]
- feature0_view = feature0.permute(0, 2, 3, 1).contiguous().view(b, h * w, 1, c) # [B, H*W, 1, C]
-
- corr = torch.matmul(feature0_view, window_feature).view(b, h * w, -1) / (c ** 0.5) # [B, H*W, (2R+1)]
-
- # mask invalid locations
- corr[~valid] = -1e9
-
- prob = F.softmax(corr, -1) # [B, H*W, (2R+1)]
-
- correspondence = torch.matmul(prob.unsqueeze(-2),
- sample_coords_softmax).squeeze(-2).view(
- b, h, w, 2).permute(0, 3, 1, 2).contiguous() # [B, 2, H, W]
-
- flow = correspondence - coords_init # flow at feature resolution
- match_prob = prob
-
- flow_x = -flow[:, :1] # [B, 1, H, W]
-
- return flow_x, match_prob
-
-
-def correlation_softmax_depth(feature0, feature1,
- intrinsics,
- pose,
- depth_candidates,
- depth_from_argmax=False,
- pred_bidir_depth=False,
- ):
- b, c, h, w = feature0.size()
- assert depth_candidates.dim() == 4 # [B, D, H, W]
- scale_factor = c ** 0.5
-
- if pred_bidir_depth:
- feature0, feature1 = torch.cat((feature0, feature1), dim=0), torch.cat((feature1, feature0), dim=0)
- intrinsics = intrinsics.repeat(2, 1, 1)
- pose = torch.cat((pose, torch.inverse(pose)), dim=0)
- depth_candidates = depth_candidates.repeat(2, 1, 1, 1)
-
- # depth candidates are actually inverse depth
- warped_feature1 = warp_with_pose_depth_candidates(feature1, intrinsics, pose,
- 1. / depth_candidates,
- ) # [B, C, D, H, W]
-
- correlation = (feature0.unsqueeze(2) * warped_feature1).sum(1) / scale_factor # [B, D, H, W]
-
- match_prob = F.softmax(correlation, dim=1) # [B, D, H, W]
-
- # for cross-task transfer (flow -> depth), extract depth with argmax at test time
- if depth_from_argmax:
- index = torch.argmax(match_prob, dim=1, keepdim=True)
- depth = torch.gather(depth_candidates, dim=1, index=index)
- else:
- depth = (match_prob * depth_candidates).sum(dim=1, keepdim=True) # [B, 1, H, W]
-
- return depth, match_prob
-
-
-def group_correlation_softmax_depth(feature0, feature1,
- intrinsics,
- pose,
- depth_candidates,
- depth_from_argmax=False,
- pred_bidir_depth=False,
- num_groups=8,
- sample_log_depth=False,
- ):
- b, c, h, w = feature0.size()
- d = depth_candidates.size(1)
- assert depth_candidates.dim() == 4 # [B, D, H, W]
- assert c % num_groups == 0, f"c: {c}, num_groups: {num_groups}"
-
- if pred_bidir_depth:
- feature0, feature1 = torch.cat((feature0, feature1), dim=0), torch.cat((feature1, feature0), dim=0)
- intrinsics = intrinsics.repeat(2, 1, 1)
- pose = torch.cat((pose, torch.inverse(pose)), dim=0)
- depth_candidates = depth_candidates.repeat(2, 1, 1, 1)
-
- # depth candidates are actually inverse depth
- with torch.amp.autocast('cuda', enabled=False):
- warped_feature1 = warp_with_pose_depth_candidates(feature1.float(), intrinsics, pose,
- torch.exp(depth_candidates) if sample_log_depth else (1. / depth_candidates),
- grid_sample_disable_cudnn=False,
- ) # [B, C, D, H, W]
-
- channels_per_group = c // num_groups
- warped_feature1_reshape = warped_feature1.reshape(b, num_groups, channels_per_group, d, h, w) # [B, G, C, D, H, W]
- feature0_reshape = feature0.reshape(b, num_groups, channels_per_group, 1, h, w) # [B, G, C, 1, H, W]
-
- with torch.amp.autocast('cuda', enabled=False):
- correlation = (F.normalize(feature0_reshape.float(), dim=2) * F.normalize(warped_feature1_reshape.float(), dim=2)).sum(dim=2) # [B, G, D, H, W]
-
- return correlation, warped_feature1
-
-
-def warp_with_pose_depth_candidates(feature1, intrinsics, pose, depth,
- clamp_min_depth=1e-3,
- grid_sample_disable_cudnn=False,
- ):
- """
- feature1: [B, C, H, W]
- intrinsics: [B, 3, 3]
- pose: [B, 4, 4]
- depth: [B, D, H, W]
- """
-
- assert intrinsics.size(1) == 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)
-
- with torch.no_grad():
- # 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.inverse(intrinsics).bmm(grid.view(b, 3, -1)) # [B, 3, H*W]
- points = torch.bmm(pose[:, :3, :3], points).unsqueeze(2).repeat(
- 1, 1, d, 1) * depth.view(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(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]
-
- # sample features
- # ref: https://github.com/pytorch/pytorch/issues/88380
- with torch.backends.cudnn.flags(enabled=not grid_sample_disable_cudnn):
- warped_feature = F.grid_sample(feature1, 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
-
-
-class CorrBlock:
- def __init__(self, fmap1, fmap2, num_levels=4, radius=4):
- self.num_levels = num_levels
- self.radius = radius
- self.corr_pyramid = []
-
- # all pairs correlation
- corr = CorrBlock.corr(fmap1, fmap2)
-
- batch, h1, w1, dim, h2, w2 = corr.shape
- corr = corr.reshape(batch*h1*w1, dim, h2, w2)
-
- self.corr_pyramid.append(corr)
- for i in range(self.num_levels-1):
- corr = F.avg_pool2d(corr, 2, stride=2)
- self.corr_pyramid.append(corr)
-
- def __call__(self, coords):
- r = self.radius
- coords = coords.permute(0, 2, 3, 1)
- batch, h1, w1, _ = coords.shape
-
- out_pyramid = []
- for i in range(self.num_levels):
- corr = self.corr_pyramid[i]
- dx = torch.linspace(-r, r, 2*r+1, device=coords.device)
- dy = torch.linspace(-r, r, 2*r+1, device=coords.device)
- delta = torch.stack(torch.meshgrid(dy, dx), axis=-1)
-
- centroid_lvl = coords.reshape(batch*h1*w1, 1, 1, 2) / 2**i
- delta_lvl = delta.view(1, 2*r+1, 2*r+1, 2)
- coords_lvl = centroid_lvl + delta_lvl
-
- corr = bilinear_sampler(corr, coords_lvl)
- corr = corr.view(batch, h1, w1, -1)
- out_pyramid.append(corr)
-
- out = torch.cat(out_pyramid, dim=-1)
- return out.permute(0, 3, 1, 2).contiguous().float()
-
- @staticmethod
- def corr(fmap1, fmap2):
- batch, dim, ht, wd = fmap1.shape
- fmap1 = fmap1.view(batch, dim, ht*wd)
- fmap2 = fmap2.view(batch, dim, ht*wd)
-
- with torch.amp.autocast('cuda', enabled=False):
- corr = torch.matmul(F.normalize(fmap1.float(), dim=1).transpose(1, 2), F.normalize(fmap2.float(), dim=1))
-
- corr = corr.view(batch, ht, wd, 1, ht, wd)
-
- # NOTE: normalize first for numerical stability
- # corr = torch.matmul(fmap1.transpose(1,2), fmap2)
- # return corr / torch.sqrt(torch.tensor(dim).float())
- return corr
-
-
-def bilinear_sampler(img, coords, mode='bilinear', mask=False):
- """ Wrapper for grid_sample, uses pixel coordinates """
- H, W = img.shape[-2:]
- xgrid, ygrid = coords.split([1,1], dim=-1)
- xgrid = 2*xgrid/(W-1) - 1
- ygrid = 2*ygrid/(H-1) - 1
-
- grid = torch.cat([xgrid, ygrid], dim=-1)
- img = F.grid_sample(img, grid, align_corners=True)
-
- if mask:
- mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
- return img, mask.float()
-
- return img
-
\ No newline at end of file
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/core/update.py b/optgs/model/encoder/.deprecated/foundationstereo/core/update.py
deleted file mode 100644
index ea4d909927b3963a4d5032e67d18f61f8168d3dd..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/foundationstereo/core/update.py
+++ /dev/null
@@ -1,168 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
-#
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto. Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
-
-
-import torch,pdb,os,sys
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
-from torch import einsum
-code_dir = os.path.dirname(os.path.realpath(__file__))
-sys.path.append(f'{code_dir}/../')
-from core.submodule import *
-from core.extractor import *
-
-class DispHead(nn.Module):
- def __init__(self, input_dim=128, hidden_dim=256, output_dim=1):
- super(DispHead, self).__init__()
- self.conv = nn.Sequential(
- nn.Conv2d(input_dim, input_dim, kernel_size=3, padding=1),
- nn.ReLU(),
- EdgeNextConvEncoder(input_dim, expan_ratio=4, kernel_size=7, norm=None),
- EdgeNextConvEncoder(input_dim, expan_ratio=4, kernel_size=7, norm=None),
- nn.Conv2d(input_dim, output_dim, 3, padding=1),
- )
-
- def forward(self, x):
- return self.conv(x)
-
-class ConvGRU(nn.Module):
- def __init__(self, hidden_dim, input_dim, kernel_size=3):
- super(ConvGRU, self).__init__()
- self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size//2)
- self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size//2)
- self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size//2)
-
- def forward(self, h, cz, cr, cq, *x_list):
- x = torch.cat(x_list, dim=1)
- hx = torch.cat([h, x], dim=1)
- z = torch.sigmoid(self.convz(hx) + cz)
- r = torch.sigmoid(self.convr(hx) + cr)
- q = torch.tanh(self.convq(torch.cat([r*h, x], dim=1)) + cq)
- h = (1-z) * h + z * q
- return h
-
-
-class BasicMotionEncoder(nn.Module):
- def __init__(self, args, ngroup=8):
- super(BasicMotionEncoder, self).__init__()
- self.args = args
- cor_planes = args.corr_levels * (2*args.corr_radius + 1) * (ngroup+1)
- self.convc1 = nn.Conv2d(cor_planes, 256, 1, padding=0)
- self.convc2 = nn.Conv2d(256, 256, 3, padding=1)
- self.convd1 = nn.Conv2d(1, 64, 7, padding=3)
- self.convd2 = nn.Conv2d(64, 64, 3, padding=1)
- self.conv = nn.Conv2d(64+256, 128-1, 3, padding=1)
-
- def forward(self, disp, corr):
- cor = F.relu(self.convc1(corr))
- cor = F.relu(self.convc2(cor))
- disp_ = F.relu(self.convd1(disp))
- disp_ = F.relu(self.convd2(disp_))
-
- cor_disp = torch.cat([cor, disp_], dim=1)
- out = F.relu(self.conv(cor_disp))
- return torch.cat([out, disp], dim=1)
-
-def pool2x(x):
- return F.avg_pool2d(x, 3, stride=2, padding=1)
-
-def pool4x(x):
- return F.avg_pool2d(x, 5, stride=4, padding=1)
-
-def interp(x, dest):
- interp_args = {'mode': 'bilinear', 'align_corners': True}
- return F.interpolate(x, dest.shape[2:], **interp_args)
-
-
-class RaftConvGRU(nn.Module):
- def __init__(self, hidden_dim=128, input_dim=256, kernel_size=3):
- super().__init__()
- self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size // 2)
- self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size // 2)
- self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size // 2)
-
- def forward(self, h, x, hx):
- z = torch.sigmoid(self.convz(hx))
- r = torch.sigmoid(self.convr(hx))
- q = torch.tanh(self.convq(torch.cat([r*h, x], dim=1)))
- h = (1-z) * h + z * q
- return h
-
-
-class SelectiveConvGRU(nn.Module):
- def __init__(self, hidden_dim=128, input_dim=256, small_kernel_size=1, large_kernel_size=3, patch_size=None):
- super(SelectiveConvGRU, self).__init__()
- self.conv0 = nn.Sequential(
- nn.Conv2d(input_dim, input_dim, kernel_size=3, padding=1),
- nn.ReLU(),
- )
- self.conv1 = nn.Sequential(
- nn.Conv2d(input_dim+hidden_dim, input_dim+hidden_dim, kernel_size=3, padding=1),
- nn.ReLU(),
- )
- self.small_gru = RaftConvGRU(hidden_dim, input_dim, small_kernel_size)
- self.large_gru = RaftConvGRU(hidden_dim, input_dim, large_kernel_size)
-
- def forward(self, att, h, *x):
- x = torch.cat(x, dim=1)
- x = self.conv0(x)
- hx = torch.cat([x, h], dim=1)
- hx = self.conv1(hx)
- h = self.small_gru(h, x, hx) * att + self.large_gru(h, x, hx) * (1 - att)
-
- return h
-
-
-class BasicSelectiveMultiUpdateBlock(nn.Module):
- def __init__(self, args, hidden_dim=128, volume_dim=8, depth_head_dim=1):
- super().__init__()
- self.args = args
- self.encoder = BasicMotionEncoder(args, volume_dim)
-
- self.depth_head_dim = depth_head_dim
-
- if args.n_gru_layers == 3:
- self.gru16 = SelectiveConvGRU(hidden_dim, hidden_dim * 2)
- if args.n_gru_layers >= 2:
- self.gru08 = SelectiveConvGRU(hidden_dim, hidden_dim * (args.n_gru_layers == 3) + hidden_dim * 2)
- self.gru04 = SelectiveConvGRU(hidden_dim, hidden_dim * (args.n_gru_layers > 1) + hidden_dim * 2)
- if depth_head_dim > 1:
- # change the name to skip the pretrained weights
- self.depth_head = DispHead(hidden_dim, 256, output_dim=depth_head_dim)
- else:
- self.disp_head = DispHead(hidden_dim, 256)
- self.mask = nn.Sequential(
- nn.Conv2d(128, 64, 3, padding=1),
- nn.ReLU(inplace=True),
- nn.Conv2d(64, 32, 3, padding=1),
- nn.ReLU(inplace=True),
- )
-
- def forward(self, net, inp, corr, disp, att):
- if self.args.n_gru_layers == 3:
- net[2] = self.gru16(att[2], net[2], inp[2], pool2x(net[1]))
- if self.args.n_gru_layers >= 2:
- if self.args.n_gru_layers > 2:
- net[1] = self.gru08(att[1], net[1], inp[1], pool2x(net[0]), interp(net[2], net[1]))
- else:
- net[1] = self.gru08(att[1], net[1], inp[1], pool2x(net[0]))
-
- motion_features = self.encoder(disp, corr)
- motion_features = torch.cat([inp[0], motion_features], dim=1)
- if self.args.n_gru_layers > 1:
- net[0] = self.gru04(att[0], net[0], motion_features, interp(net[1], net[0]))
-
- if self.depth_head_dim > 1:
- delta_disp = self.depth_head(net[0])
- else:
- delta_disp = self.disp_head(net[0])
-
- # scale mask to balence gradients
- mask = .25 * self.mask(net[0])
- return net, mask, delta_disp
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/core/utils/utils.py b/optgs/model/encoder/.deprecated/foundationstereo/core/utils/utils.py
deleted file mode 100644
index 6852481ceeecb126e42999a268e145a189f2faac..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/foundationstereo/core/utils/utils.py
+++ /dev/null
@@ -1,62 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
-#
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto. Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
-
-
-
-import torch,pdb,logging
-import torch.nn.functional as F
-import numpy as np
-from scipy import interpolate
-
-
-class InputPadder:
- """ Pads images such that dimensions are divisible by 8 """
- def __init__(self, dims, mode='sintel', divis_by=8, force_square=False):
- self.ht, self.wd = dims[-2:]
- if force_square:
- max_side = max(self.ht, self.wd)
- pad_ht = ((max_side // divis_by) + 1) * divis_by - self.ht
- pad_wd = ((max_side // divis_by) + 1) * divis_by - self.wd
- else:
- pad_ht = (((self.ht // divis_by) + 1) * divis_by - self.ht) % divis_by
- pad_wd = (((self.wd // divis_by) + 1) * divis_by - self.wd) % divis_by
- if mode == 'sintel':
- self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2]
- else:
- self._pad = [pad_wd//2, pad_wd - pad_wd//2, 0, pad_ht]
-
- def pad(self, *inputs):
- assert all((x.ndim == 4) for x in inputs)
- return [F.pad(x, self._pad, mode='replicate') for x in inputs]
-
- def unpad(self, x):
- assert x.ndim == 4
- ht, wd = x.shape[-2:]
- c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]]
- return x[..., c[0]:c[1], c[2]:c[3]]
-
-
-def bilinear_sampler(img, coords, mode='bilinear', mask=False, low_memory=False):
- """ Wrapper for grid_sample, uses pixel coordinates """
- H, W = img.shape[-2:]
- xgrid, ygrid = coords.split([1,1], dim=-1)
- xgrid = 2*xgrid/(W-1) - 1 # Normalize to [-1,1]
- assert torch.unique(ygrid).numel() == 1 and H == 1 # This is a stereo problem
- grid = torch.cat([xgrid, ygrid], dim=-1).to(img.dtype)
- img = F.grid_sample(img, grid, align_corners=True)
- if mask:
- mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
- return img, mask.float()
- return img
-
-
-def coords_grid(batch, ht, wd):
- coords = torch.meshgrid(torch.arange(ht), torch.arange(wd))
- coords = torch.stack(coords[::-1], dim=0).float()
- return coords[None].repeat(batch, 1, 1, 1)
-
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/depth_anything/LICENSE.txt b/optgs/model/encoder/.deprecated/foundationstereo/depth_anything/LICENSE.txt
deleted file mode 100644
index 261eeb9e9f8b2b4b0d119366dda99c6fd7d35c64..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/foundationstereo/depth_anything/LICENSE.txt
+++ /dev/null
@@ -1,201 +0,0 @@
- Apache License
- Version 2.0, January 2004
- http://www.apache.org/licenses/
-
- TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
-
- 1. Definitions.
-
- "License" shall mean the terms and conditions for use, reproduction,
- and distribution as defined by Sections 1 through 9 of this document.
-
- "Licensor" shall mean the copyright owner or entity authorized by
- the copyright owner that is granting the License.
-
- "Legal Entity" shall mean the union of the acting entity and all
- other entities that control, are controlled by, or are under common
- control with that entity. For the purposes of this definition,
- "control" means (i) the power, direct or indirect, to cause the
- direction or management of such entity, whether by contract or
- otherwise, or (ii) ownership of fifty percent (50%) or more of the
- outstanding shares, or (iii) beneficial ownership of such entity.
-
- "You" (or "Your") shall mean an individual or Legal Entity
- exercising permissions granted by this License.
-
- "Source" form shall mean the preferred form for making modifications,
- including but not limited to software source code, documentation
- source, and configuration files.
-
- "Object" form shall mean any form resulting from mechanical
- transformation or translation of a Source form, including but
- not limited to compiled object code, generated documentation,
- and conversions to other media types.
-
- "Work" shall mean the work of authorship, whether in Source or
- Object form, made available under the License, as indicated by a
- copyright notice that is included in or attached to the work
- (an example is provided in the Appendix below).
-
- "Derivative Works" shall mean any work, whether in Source or Object
- form, that is based on (or derived from) the Work and for which the
- editorial revisions, annotations, elaborations, or other modifications
- represent, as a whole, an original work of authorship. For the purposes
- of this License, Derivative Works shall not include works that remain
- separable from, or merely link (or bind by name) to the interfaces of,
- the Work and Derivative Works thereof.
-
- "Contribution" shall mean any work of authorship, including
- the original version of the Work and any modifications or additions
- to that Work or Derivative Works thereof, that is intentionally
- submitted to Licensor for inclusion in the Work by the copyright owner
- or by an individual or Legal Entity authorized to submit on behalf of
- the copyright owner. For the purposes of this definition, "submitted"
- means any form of electronic, verbal, or written communication sent
- to the Licensor or its representatives, including but not limited to
- communication on electronic mailing lists, source code control systems,
- and issue tracking systems that are managed by, or on behalf of, the
- Licensor for the purpose of discussing and improving the Work, but
- excluding communication that is conspicuously marked or otherwise
- designated in writing by the copyright owner as "Not a Contribution."
-
- "Contributor" shall mean Licensor and any individual or Legal Entity
- on behalf of whom a Contribution has been received by Licensor and
- subsequently incorporated within the Work.
-
- 2. Grant of Copyright License. Subject to the terms and conditions of
- this License, each Contributor hereby grants to You a perpetual,
- worldwide, non-exclusive, no-charge, royalty-free, irrevocable
- copyright license to reproduce, prepare Derivative Works of,
- publicly display, publicly perform, sublicense, and distribute the
- Work and such Derivative Works in Source or Object form.
-
- 3. Grant of Patent License. Subject to the terms and conditions of
- this License, each Contributor hereby grants to You a perpetual,
- worldwide, non-exclusive, no-charge, royalty-free, irrevocable
- (except as stated in this section) patent license to make, have made,
- use, offer to sell, sell, import, and otherwise transfer the Work,
- where such license applies only to those patent claims licensable
- by such Contributor that are necessarily infringed by their
- Contribution(s) alone or by combination of their Contribution(s)
- with the Work to which such Contribution(s) was submitted. If You
- institute patent litigation against any entity (including a
- cross-claim or counterclaim in a lawsuit) alleging that the Work
- or a Contribution incorporated within the Work constitutes direct
- or contributory patent infringement, then any patent licenses
- granted to You under this License for that Work shall terminate
- as of the date such litigation is filed.
-
- 4. Redistribution. You may reproduce and distribute copies of the
- Work or Derivative Works thereof in any medium, with or without
- modifications, and in Source or Object form, provided that You
- meet the following conditions:
-
- (a) You must give any other recipients of the Work or
- Derivative Works a copy of this License; and
-
- (b) You must cause any modified files to carry prominent notices
- stating that You changed the files; and
-
- (c) You must retain, in the Source form of any Derivative Works
- that You distribute, all copyright, patent, trademark, and
- attribution notices from the Source form of the Work,
- excluding those notices that do not pertain to any part of
- the Derivative Works; and
-
- (d) If the Work includes a "NOTICE" text file as part of its
- distribution, then any Derivative Works that You distribute must
- include a readable copy of the attribution notices contained
- within such NOTICE file, excluding those notices that do not
- pertain to any part of the Derivative Works, in at least one
- of the following places: within a NOTICE text file distributed
- as part of the Derivative Works; within the Source form or
- documentation, if provided along with the Derivative Works; or,
- within a display generated by the Derivative Works, if and
- wherever such third-party notices normally appear. The contents
- of the NOTICE file are for informational purposes only and
- do not modify the License. You may add Your own attribution
- notices within Derivative Works that You distribute, alongside
- or as an addendum to the NOTICE text from the Work, provided
- that such additional attribution notices cannot be construed
- as modifying the License.
-
- You may add Your own copyright statement to Your modifications and
- may provide additional or different license terms and conditions
- for use, reproduction, or distribution of Your modifications, or
- for any such Derivative Works as a whole, provided Your use,
- reproduction, and distribution of the Work otherwise complies with
- the conditions stated in this License.
-
- 5. Submission of Contributions. Unless You explicitly state otherwise,
- any Contribution intentionally submitted for inclusion in the Work
- by You to the Licensor shall be under the terms and conditions of
- this License, without any additional terms or conditions.
- Notwithstanding the above, nothing herein shall supersede or modify
- the terms of any separate license agreement you may have executed
- with Licensor regarding such Contributions.
-
- 6. Trademarks. This License does not grant permission to use the trade
- names, trademarks, service marks, or product names of the Licensor,
- except as required for reasonable and customary use in describing the
- origin of the Work and reproducing the content of the NOTICE file.
-
- 7. Disclaimer of Warranty. Unless required by applicable law or
- agreed to in writing, Licensor provides the Work (and each
- Contributor provides its Contributions) on an "AS IS" BASIS,
- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
- implied, including, without limitation, any warranties or conditions
- of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
- PARTICULAR PURPOSE. You are solely responsible for determining the
- appropriateness of using or redistributing the Work and assume any
- risks associated with Your exercise of permissions under this License.
-
- 8. Limitation of Liability. In no event and under no legal theory,
- whether in tort (including negligence), contract, or otherwise,
- unless required by applicable law (such as deliberate and grossly
- negligent acts) or agreed to in writing, shall any Contributor be
- liable to You for damages, including any direct, indirect, special,
- incidental, or consequential damages of any character arising as a
- result of this License or out of the use or inability to use the
- Work (including but not limited to damages for loss of goodwill,
- work stoppage, computer failure or malfunction, or any and all
- other commercial damages or losses), even if such Contributor
- has been advised of the possibility of such damages.
-
- 9. Accepting Warranty or Additional Liability. While redistributing
- the Work or Derivative Works thereof, You may choose to offer,
- and charge a fee for, acceptance of support, warranty, indemnity,
- or other liability obligations and/or rights consistent with this
- License. However, in accepting such obligations, You may act only
- on Your own behalf and on Your sole responsibility, not on behalf
- of any other Contributor, and only if You agree to indemnify,
- defend, and hold each Contributor harmless for any liability
- incurred by, or claims asserted against, such Contributor by reason
- of your accepting any such warranty or additional liability.
-
- END OF TERMS AND CONDITIONS
-
- APPENDIX: How to apply the Apache License to your work.
-
- To apply the Apache License to your work, attach the following
- boilerplate notice, with the fields enclosed by brackets "[]"
- replaced with your own identifying information. (Don't include
- the brackets!) The text should be enclosed in the appropriate
- comment syntax for the file format. We also recommend that a
- file or class name and description of purpose be included on the
- same "printed page" as the copyright notice for easier
- identification within third-party archives.
-
- Copyright [yyyy] [name of copyright owner]
-
- Licensed under the Apache License, Version 2.0 (the "License");
- you may not use this file except in compliance with the License.
- You may obtain a copy of the License at
-
- http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing, software
- distributed under the License is distributed on an "AS IS" BASIS,
- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- See the License for the specific language governing permissions and
- limitations under the License.
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/depth_anything/blocks.py b/optgs/model/encoder/.deprecated/foundationstereo/depth_anything/blocks.py
deleted file mode 100644
index cdf594729fbcc09189ac0ba91b8c95581b80f443..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/foundationstereo/depth_anything/blocks.py
+++ /dev/null
@@ -1,153 +0,0 @@
-import torch.nn as nn
-
-
-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
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/depth_anything/dpt.py b/optgs/model/encoder/.deprecated/foundationstereo/depth_anything/dpt.py
deleted file mode 100644
index 442a04a83f890105f5b149f31b72c4a031319b0c..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/foundationstereo/depth_anything/dpt.py
+++ /dev/null
@@ -1,203 +0,0 @@
-import argparse
-import torch,os,sys,pdb
-import torch.nn as nn
-import torch.nn.functional as F
-code_dir = os.path.dirname(os.path.realpath(__file__))
-sys.path.append(f'{code_dir}/../')
-from depth_anything.blocks import FeatureFusionBlock, _make_scratch
-
-
-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, nclass, in_channels, features=256, use_bn=False, out_channels=[256, 512, 1024, 1024], use_clstoken=False):
- super(DPTHead, self).__init__()
-
- self.nclass = nclass
- self.use_clstoken = use_clstoken
-
- 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)
-
- head_features_1 = features
- head_features_2 = 32
-
- if nclass > 1:
- self.scratch.output_conv = nn.Sequential(
- nn.Conv2d(head_features_1, head_features_1, kernel_size=3, stride=1, padding=1),
- nn.ReLU(True),
- nn.Conv2d(head_features_1, nclass, kernel_size=1, stride=1, padding=0),
- )
- else:
- self.scratch.output_conv1 = nn.Conv2d(head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1)
-
- self.scratch.output_conv2 = nn.Sequential(
- nn.Conv2d(head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1),
- nn.ReLU(True),
- nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0),
- nn.ReLU(True),
- nn.Identity(),
- )
-
- def forward(self, out_features, patch_h, patch_w, return_intermediate=False, patch_size=14):
- out = []
- for i, x in enumerate(out_features):
- if self.use_clstoken:
- x, cls_token = x[0], x[1]
- readout = cls_token.unsqueeze(1).expand_as(x)
- x = self.readout_projects[i](torch.cat((x, readout), -1))
- else:
- x = x[0]
-
- x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
-
- x = self.projects[i](x)
- x = self.resize_layers[i](x)
-
- out.append(x)
-
- layer_1, layer_2, layer_3, layer_4 = out
-
- 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:])
- path_3 = self.scratch.refinenet3(path_4, layer_3_rn, size=layer_2_rn.shape[2:])
- path_2 = self.scratch.refinenet2(path_3, layer_2_rn, size=layer_1_rn.shape[2:])
- path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
-
- out = self.scratch.output_conv1(path_1)
- out = F.interpolate(out, (int(patch_h * patch_size), int(patch_w * patch_size)), mode="bilinear", align_corners=True)
- if return_intermediate:
- depth = self.scratch.output_conv2(out)
- depth = F.relu(depth)
- disp = 1/depth
- disp[depth==0] = 0
- disp = disp/disp.max()
- return out, path_1, path_2, path_3, path_4, disp
-
- else:
- out = self.scratch.output_conv2(out)
- return out
-
-
-class DPT_DINOv2(nn.Module):
- def __init__(self, encoder='vitl', features=256, out_channels=[256, 512, 1024, 1024], use_bn=False, use_clstoken=False, pretrained_dino=False):
- super(DPT_DINOv2, self).__init__()
-
- assert encoder in ['vits', 'vitb', 'vitl']
-
- # in case the Internet connection is not stable, please load the DINOv2 locally
- # if localhub:
- # self.pretrained = torch.hub.load('torchhub/facebookresearch_dinov2_main', 'dinov2_{:}14'.format(encoder), source='local', pretrained=False)
- # else:
- self.pretrained = torch.hub.load('facebookresearch/dinov2', 'dinov2_{:}14'.format(encoder), pretrained=pretrained_dino)
-
-
- dim = self.pretrained.blocks[0].attn.qkv.in_features
-
- self.depth_head = DPTHead(1, dim, features, use_bn, out_channels=out_channels, use_clstoken=use_clstoken)
-
- def forward(self, x):
- h, w = x.shape[-2:]
-
- features = self.pretrained.get_intermediate_layers(x, 4, return_class_token=True)
- patch_size = self.pretrained.patch_size
- patch_h, patch_w = h // patch_size, w // patch_size
- output = self.depth_head(features, patch_h, patch_w, patch_size=patch_size, return_intermediate=True)
- return output
-
-
-class DepthAnything(DPT_DINOv2):
- def __init__(self, config):
- super().__init__(**config)
-
- def forward(self, x):
- h, w = x.shape[-2:]
-
- features = self.pretrained.get_intermediate_layers(x, 4, return_class_token=True)
- patch_size = self.pretrained.patch_size
- patch_h, patch_w = h // patch_size, w // patch_size
- depth = self.depth_head(features, patch_h, patch_w, patch_size=patch_size)
- depth = F.interpolate(depth, size=(h, w), mode="bilinear", align_corners=True)
- depth = F.relu(depth)
-
- return depth.squeeze(1)
-
-
-if __name__ == '__main__':
- parser = argparse.ArgumentParser()
- parser.add_argument(
- "--encoder",
- default="vits",
- type=str,
- choices=["vits", "vitb", "vitl"],
- )
- args = parser.parse_args()
-
- model = DepthAnything.from_pretrained("LiheYoung/depth_anything_{:}14".format(args.encoder))
-
- print(model)
diff --git a/optgs/model/encoder/.deprecated/foundationstereo/depth_anything/util/transform.py b/optgs/model/encoder/.deprecated/foundationstereo/depth_anything/util/transform.py
deleted file mode 100644
index d542fefee7ef4c23d7e4425161902e8a8f1e8b6e..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/foundationstereo/depth_anything/util/transform.py
+++ /dev/null
@@ -1,248 +0,0 @@
-import random
-from PIL import Image, ImageOps, ImageFilter
-import torch
-from torchvision import transforms
-import torch.nn.functional as F
-
-import numpy as np
-import cv2
-import math
-
-
-def apply_min_size(sample, size, image_interpolation_method=cv2.INTER_AREA):
- """Rezise the sample to ensure the given size. Keeps aspect ratio.
-
- Args:
- sample (dict): sample
- size (tuple): image size
-
- Returns:
- tuple: new size
- """
- shape = list(sample["disparity"].shape)
-
- if shape[0] >= size[0] and shape[1] >= size[1]:
- return sample
-
- scale = [0, 0]
- scale[0] = size[0] / shape[0]
- scale[1] = size[1] / shape[1]
-
- scale = max(scale)
-
- shape[0] = math.ceil(scale * shape[0])
- shape[1] = math.ceil(scale * shape[1])
-
- # resize
- sample["image"] = cv2.resize(
- sample["image"], tuple(shape[::-1]), interpolation=image_interpolation_method
- )
-
- sample["disparity"] = cv2.resize(
- sample["disparity"], tuple(shape[::-1]), interpolation=cv2.INTER_NEAREST
- )
- sample["mask"] = cv2.resize(
- sample["mask"].astype(np.float32),
- tuple(shape[::-1]),
- interpolation=cv2.INTER_NEAREST,
- )
- sample["mask"] = sample["mask"].astype(bool)
-
- return tuple(shape)
-
-
-class Resize(object):
- """Resize sample to given size (width, height).
- """
-
- def __init__(
- self,
- width,
- height,
- resize_target=True,
- keep_aspect_ratio=False,
- ensure_multiple_of=1,
- resize_method="lower_bound",
- image_interpolation_method=cv2.INTER_AREA,
- ):
- """Init.
-
- Args:
- width (int): desired output width
- height (int): desired output height
- resize_target (bool, optional):
- True: Resize the full sample (image, mask, target).
- False: Resize image only.
- Defaults to True.
- keep_aspect_ratio (bool, optional):
- True: Keep the aspect ratio of the input sample.
- Output sample might not have the given width and height, and
- resize behaviour depends on the parameter 'resize_method'.
- Defaults to False.
- ensure_multiple_of (int, optional):
- Output width and height is constrained to be multiple of this parameter.
- Defaults to 1.
- resize_method (str, optional):
- "lower_bound": Output will be at least as large as the given size.
- "upper_bound": Output will be at max as large as the given size. (Output size might be smaller than given size.)
- "minimal": Scale as least as possible. (Output size might be smaller than given size.)
- Defaults to "lower_bound".
- """
- self.__width = width
- self.__height = height
-
- self.__resize_target = resize_target
- self.__keep_aspect_ratio = keep_aspect_ratio
- self.__multiple_of = ensure_multiple_of
- self.__resize_method = resize_method
- self.__image_interpolation_method = image_interpolation_method
-
- def constrain_to_multiple_of(self, x, min_val=0, max_val=None):
- y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
-
- if max_val is not None and y > max_val:
- y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
-
- if y < min_val:
- y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
-
- return y
-
- def get_size(self, width, height):
- # determine new height and width
- scale_height = self.__height / height
- scale_width = self.__width / width
-
- if self.__keep_aspect_ratio:
- if self.__resize_method == "lower_bound":
- # scale such that output size is lower bound
- if scale_width > scale_height:
- # fit width
- scale_height = scale_width
- else:
- # fit height
- scale_width = scale_height
- elif self.__resize_method == "upper_bound":
- # scale such that output size is upper bound
- if scale_width < scale_height:
- # fit width
- scale_height = scale_width
- else:
- # fit height
- scale_width = scale_height
- elif self.__resize_method == "minimal":
- # scale as least as possbile
- if abs(1 - scale_width) < abs(1 - scale_height):
- # fit width
- scale_height = scale_width
- else:
- # fit height
- scale_width = scale_height
- else:
- raise ValueError(
- f"resize_method {self.__resize_method} not implemented"
- )
-
- if self.__resize_method == "lower_bound":
- new_height = self.constrain_to_multiple_of(
- scale_height * height, min_val=self.__height
- )
- new_width = self.constrain_to_multiple_of(
- scale_width * width, min_val=self.__width
- )
- elif self.__resize_method == "upper_bound":
- new_height = self.constrain_to_multiple_of(
- scale_height * height, max_val=self.__height
- )
- new_width = self.constrain_to_multiple_of(
- scale_width * width, max_val=self.__width
- )
- elif self.__resize_method == "minimal":
- new_height = self.constrain_to_multiple_of(scale_height * height)
- new_width = self.constrain_to_multiple_of(scale_width * width)
- else:
- raise ValueError(f"resize_method {self.__resize_method} not implemented")
-
- return (new_width, new_height)
-
- def __call__(self, sample):
- width, height = self.get_size(
- sample["image"].shape[1], sample["image"].shape[0]
- )
-
- # resize sample
- sample["image"] = cv2.resize(
- sample["image"],
- (width, height),
- interpolation=self.__image_interpolation_method,
- )
-
- if self.__resize_target:
- if "disparity" in sample:
- sample["disparity"] = cv2.resize(
- sample["disparity"],
- (width, height),
- interpolation=cv2.INTER_NEAREST,
- )
-
- if "depth" in sample:
- sample["depth"] = cv2.resize(
- sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST
- )
-
- if "semseg_mask" in sample:
- # sample["semseg_mask"] = cv2.resize(
- # sample["semseg_mask"], (width, height), interpolation=cv2.INTER_NEAREST
- # )
- sample["semseg_mask"] = F.interpolate(torch.from_numpy(sample["semseg_mask"]).float()[None, None, ...], (height, width), mode='nearest').numpy()[0, 0]
-
- if "mask" in sample:
- sample["mask"] = cv2.resize(
- sample["mask"].astype(np.float32),
- (width, height),
- interpolation=cv2.INTER_NEAREST,
- )
- # sample["mask"] = sample["mask"].astype(bool)
-
- # print(sample['image'].shape, sample['depth'].shape)
- return sample
-
-
-class NormalizeImage(object):
- """Normlize image by given mean and std.
- """
-
- def __init__(self, mean, std):
- self.__mean = mean
- self.__std = std
-
- def __call__(self, sample):
- sample["image"] = (sample["image"] - self.__mean) / self.__std
-
- return sample
-
-
-class PrepareForNet(object):
- """Prepare sample for usage as network input.
- """
-
- def __init__(self):
- pass
-
- def __call__(self, sample):
- image = np.transpose(sample["image"], (2, 0, 1))
- sample["image"] = np.ascontiguousarray(image).astype(np.float32)
-
- if "mask" in sample:
- sample["mask"] = sample["mask"].astype(np.float32)
- sample["mask"] = np.ascontiguousarray(sample["mask"])
-
- if "depth" in sample:
- depth = sample["depth"].astype(np.float32)
- sample["depth"] = np.ascontiguousarray(depth)
-
- if "semseg_mask" in sample:
- sample["semseg_mask"] = sample["semseg_mask"].astype(np.float32)
- sample["semseg_mask"] = np.ascontiguousarray(sample["semseg_mask"])
-
- return sample
diff --git a/optgs/model/encoder/.deprecated/raft_backbone.py b/optgs/model/encoder/.deprecated/raft_backbone.py
deleted file mode 100644
index e4b3cb73c5ff6041f1bb0211cb3591d35579c7f1..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/.deprecated/raft_backbone.py
+++ /dev/null
@@ -1,211 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class ResidualBlock(nn.Module):
- def __init__(self, in_planes, planes, norm_fn='group', stride=1):
- super(ResidualBlock, self).__init__()
-
- self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride)
- self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)
- self.relu = nn.ReLU(inplace=True)
-
- num_groups = planes // 8
-
- if norm_fn == 'group':
- self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
- self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
- if not stride == 1:
- self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
-
- elif norm_fn == 'batch':
- self.norm1 = nn.BatchNorm2d(planes)
- self.norm2 = nn.BatchNorm2d(planes)
- if not stride == 1:
- self.norm3 = nn.BatchNorm2d(planes)
-
- elif norm_fn == 'instance':
- self.norm1 = nn.InstanceNorm2d(planes)
- self.norm2 = nn.InstanceNorm2d(planes)
- if not stride == 1:
- self.norm3 = nn.InstanceNorm2d(planes)
-
- elif norm_fn == 'none':
- self.norm1 = nn.Sequential()
- self.norm2 = nn.Sequential()
- if not stride == 1:
- self.norm3 = nn.Sequential()
-
- if stride == 1:
- 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 BottleneckBlock(nn.Module):
- def __init__(self, in_planes, planes, norm_fn='group', stride=1):
- super(BottleneckBlock, self).__init__()
-
- self.conv1 = nn.Conv2d(in_planes, planes//4, kernel_size=1, padding=0)
- self.conv2 = nn.Conv2d(planes//4, planes//4, kernel_size=3, padding=1, stride=stride)
- self.conv3 = nn.Conv2d(planes//4, planes, kernel_size=1, padding=0)
- self.relu = nn.ReLU(inplace=True)
-
- num_groups = planes // 8
-
- if norm_fn == 'group':
- self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes//4)
- self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes//4)
- self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
- if not stride == 1:
- self.norm4 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
-
- elif norm_fn == 'batch':
- self.norm1 = nn.BatchNorm2d(planes//4)
- self.norm2 = nn.BatchNorm2d(planes//4)
- self.norm3 = nn.BatchNorm2d(planes)
- if not stride == 1:
- self.norm4 = nn.BatchNorm2d(planes)
-
- elif norm_fn == 'instance':
- self.norm1 = nn.InstanceNorm2d(planes//4)
- self.norm2 = nn.InstanceNorm2d(planes//4)
- self.norm3 = nn.InstanceNorm2d(planes)
- if not stride == 1:
- self.norm4 = nn.InstanceNorm2d(planes)
-
- elif norm_fn == 'none':
- self.norm1 = nn.Sequential()
- self.norm2 = nn.Sequential()
- self.norm3 = nn.Sequential()
- if not stride == 1:
- self.norm4 = nn.Sequential()
-
- if stride == 1:
- self.downsample = None
-
- else:
- self.downsample = nn.Sequential(
- nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm4)
-
-
- def forward(self, x):
- y = x
- y = self.relu(self.norm1(self.conv1(y)))
- y = self.relu(self.norm2(self.conv2(y)))
- y = self.relu(self.norm3(self.conv3(y)))
-
- if self.downsample is not None:
- x = self.downsample(x)
-
- return self.relu(x+y)
-
-class RAFTBasicEncoder(nn.Module):
- def __init__(self, output_dim=256, norm_fn='instance', dropout=0.0):
- super(RAFTBasicEncoder, self).__init__()
- self.norm_fn = norm_fn
-
- if self.norm_fn == 'group':
- self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64)
-
- elif self.norm_fn == 'batch':
- self.norm1 = nn.BatchNorm2d(64)
-
- elif self.norm_fn == 'instance':
- self.norm1 = nn.InstanceNorm2d(64)
-
- elif self.norm_fn == 'none':
- self.norm1 = nn.Sequential()
-
- self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
- self.relu1 = nn.ReLU(inplace=True)
-
- self.in_planes = 64
- self.layer1 = self._make_layer(64, stride=1)
- self.layer2 = self._make_layer(96, stride=2)
- self.layer3 = self._make_layer(128, stride=2)
-
- # output convolution
- self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1)
-
- self.dropout = None
- if dropout > 0:
- self.dropout = nn.Dropout2d(p=dropout)
-
- 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):
- layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
- layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
- layers = (layer1, layer2)
-
- self.in_planes = dim
- return nn.Sequential(*layers)
-
-
- def forward(self, x):
-
- # if input is list, combine batch dimension
- is_list = isinstance(x, tuple) or isinstance(x, list)
- if is_list:
- batch_dim = x[0].shape[0]
- x = torch.cat(x, dim=0)
-
- x = self.conv1(x)
- x = self.norm1(x)
- x = self.relu1(x)
-
- x = self.layer1(x)
- x = self.layer2(x)
- x = self.layer3(x)
-
- x = self.conv2(x)
-
- if self.training and self.dropout is not None:
- x = self.dropout(x)
-
- if is_list:
- x = torch.split(x, [batch_dim, batch_dim], dim=0)
-
- return x
-
-
-
-if __name__ == '__main__':
- model = RAFTBasicEncoder().cuda()
- x = torch.randn(2, 3, 64, 128).cuda()
- y = model(x)
- print(y.shape)
-
- weights = '/users/hxu/projects/optgs/pretrained/raft_models/raft-sintel.pth'
- weights = torch.load(weights)
- # remove 'module.fnet.' in the keys
- new_state_dict = {}
- for k, v in weights.items():
- if 'module.fnet.' in k:
- new_key = k.replace('module.fnet.', '') if k.startswith('module.fnet.') else k
- new_state_dict[new_key] = v
- model.load_state_dict(new_state_dict)
-
diff --git a/optgs/model/encoder/__init__.py b/optgs/model/encoder/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/optgs/model/encoder/depth_anything_v2/__init__.py b/optgs/model/encoder/depth_anything_v2/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/optgs/model/encoder/depth_anything_v2/util/__init__.py b/optgs/model/encoder/depth_anything_v2/util/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/optgs/model/encoder/lvsm/__init__.py b/optgs/model/encoder/lvsm/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/optgs/model/encoder/lvsm/layer.py b/optgs/model/encoder/lvsm/layer.py
deleted file mode 100644
index 5b8170a2dd2f52ffc51a9f81c5dce7fbc3ca50d7..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/lvsm/layer.py
+++ /dev/null
@@ -1,340 +0,0 @@
-import torch
-import torch.nn as nn
-from einops import rearrange
-
-import warnings
-import torch.nn.functional as F
-
-
-USE_FLASH_ATTENTION3 = True
-try:
- from flash_attn_interface import flash_attn_func
- FA3_AVAILABLE = True
- warnings.warn('flash attention 3 is available (LVSM)')
-except ImportError:
- FA3_AVAILABLE = False
- warnings.warn('flash attention 3 is not available (LVSM)')
-
-
-try:
- import xformers.ops as xops
- XFORMERS_AVAILABLE = True
-except ImportError:
- XFORMERS_AVAILABLE = False
- warnings.warn('xformers is not available (LVSM)')
- # raise ImportError("Please install xformers to use flashatt v2")
-
-
-def init_weights(module, std=0.02):
- """Initialize weights for linear and embedding layers.
-
- Args:
- module: Module to initialize
- std: Standard deviation for normal initialization
- """
- if isinstance(module, (nn.Linear, nn.Embedding)):
- torch.nn.init.normal_(module.weight, mean=0.0, std=std)
- if isinstance(module, nn.Linear) and module.bias is not None:
- torch.nn.init.zeros_(module.bias)
-
-
-
-# src: https://github.com/pytorch/benchmark/blob/main/torchbenchmark/models/llama/model.py#L28
-class RMSNorm(nn.Module):
- def __init__(self, dim: int, eps: float = 1e-5):
- super().__init__()
- self.eps = eps
- self.weight = nn.Parameter(torch.ones(dim))
-
- def _norm(self, x):
- return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
-
- def forward(self, x):
- output = self._norm(x.float()).type_as(x)
-
- return output * self.weight.type_as(x)
-
-
-
-class MLP(nn.Module):
- """
- Multi-Layer Perceptron block.
- Reference: https://github.com/facebookresearch/dino/blob/7c446df5b9f45747937fb0d72314eb9f7b66930a/vision_transformer.py#L49-L65
- """
-
- def __init__(
- self,
- dim,
- mlp_ratio=4,
- bias=False,
- dropout=0.0,
- activation=nn.GELU,
- mlp_dim=None,
- ):
- """
- Args:
- dim: Input dimension
- mlp_ratio: Multiplier for hidden dimension
- bias: Whether to use bias in linear layers
- dropout: Dropout probability
- activation: Activation function
- mlp_dim: Optional explicit hidden dimension (overrides mlp_ratio)
- """
- super().__init__()
- hidden_dim = mlp_dim if mlp_dim is not None else int(dim * mlp_ratio)
-
- self.mlp = nn.Sequential(
- nn.Linear(dim, hidden_dim, bias=bias),
- activation(),
- nn.Linear(hidden_dim, dim, bias=bias),
- nn.Dropout(dropout),
- )
-
- def forward(self, x):
- return self.mlp(x)
-
-
-
-class QK_Norm_SelfAttention(nn.Module):
- """
- Self-attention with optional Q-K normalization.
- Reference: https://github.com/facebookresearch/dino/blob/7c446df5b9f45747937fb0d72314eb9f7b66930a/vision_transformer.py#L68-L92
- """
-
- def __init__(
- self,
- dim,
- head_dim,
- qkv_bias=False,
- fc_bias=True,
- attn_dropout=0.0,
- fc_dropout=0.0,
- use_qk_norm=True,
- ):
- """
- Args:
- dim: Input dimension
- head_dim: Dimension of each attention head
- qkv_bias: Whether to use bias in QKV projection
- fc_bias: Whether to use bias in output projection
- attn_dropout: Dropout probability for attention weights
- fc_dropout: Dropout probability for output projection
- use_qk_norm: Whether to use Q-K normalization
- We use flash attention V2 for efficiency.
- """
- super().__init__()
- assert dim % head_dim == 0, f"Token dimension {dim} should be divisible by head dimension {head_dim}"
-
- self.dim = dim
- self.head_dim = head_dim
- self.num_heads = dim // head_dim
- self.attn_dropout = attn_dropout
- self.use_qk_norm = use_qk_norm
-
- self.to_qkv = nn.Linear(dim, 3 * dim, bias=qkv_bias)
- self.fc = nn.Linear(dim, dim, bias=fc_bias)
- self.attn_fc_dropout = nn.Dropout(fc_dropout)
-
- # Optional Q-K normalization
- if self.use_qk_norm:
- self.q_norm = RMSNorm(head_dim)
- self.k_norm = RMSNorm(head_dim)
-
- def forward(self, x, attn_bias=None):
- """
- Args:
- x: Input tensor of shape (batch, seq_len, dim)
- attn_bias: Optional attention bias mask
-
- Returns:
- Output tensor of shape (batch, seq_len, dim)
- """
- q, k, v = self.to_qkv(x).chunk(3, dim=-1)
-
- q, k, v = (rearrange(t, "b l (nh dh) -> b l nh dh", dh=self.head_dim) for t in (q, k, v))
-
- # Apply qk normalization if enabled
- if self.use_qk_norm:
- q = self.q_norm(q)
- k = self.k_norm(k)
-
- if USE_FLASH_ATTENTION3 and FA3_AVAILABLE:
- x = flash_attn_func(q, k, v)[0]
- elif XFORMERS_AVAILABLE:
- x = xops.memory_efficient_attention(
- q, k, v,
- attn_bias=attn_bias,
- p=self.attn_dropout if self.training else 0.0,
- op=(xops.fmha.flash.FwOp, xops.fmha.flash.BwOp),
- )
- else:
- # use pytorch's built-in attention
- q = q.permute(0, 2, 1, 3).contiguous() # [B, H, L, C]
- k = k.permute(0, 2, 1, 3).contiguous()
- v = v.permute(0, 2, 1, 3).contiguous()
- x = F.scaled_dot_product_attention(q, k, v)
- x = x.permute(0, 2, 1, 3).contiguous() # [B, L, H, C]
-
- x = rearrange(x, "b l nh dh -> b l (nh dh)")
- x = self.attn_fc_dropout(self.fc(x))
-
- return x
-
-
-
-
-class SubsetAttention(nn.Module):
- """Attention that can attend to subsets of queries or keys/values."""
-
- def __init__(
- self,
- dim,
- head_dim,
- qkv_bias=False,
- attn_dropout=0.0,
- fc_bias=False,
- fc_dropout=0.0,
- use_qk_norm=False
- ):
- """
- Args:
- dim: Input dimension
- head_dim: Dimension of each attention head
- qkv_bias: Whether to use bias in QKV projection
- attn_dropout: Dropout probability for attention weights
- fc_bias: Whether to use bias in output projection
- fc_dropout: Dropout probability for output projection
- use_qk_norm: Whether to use Q-K normalization
- We use flash attention V2 for efficiency.
- """
- super().__init__()
- assert dim % head_dim == 0, f"Token dimension {dim} should be divisible by head dimension {head_dim}"
-
- self.dim = dim
- self.head_dim = head_dim
- self.num_heads = dim // head_dim
- self.attn_dropout = attn_dropout
- self.use_qk_norm = use_qk_norm
-
- # Projections
- self.to_qkv = nn.Linear(dim, 3 * dim, bias=qkv_bias)
- self.fc = nn.Linear(dim, dim, bias=fc_bias)
- self.attn_fc_dropout = nn.Dropout(fc_dropout)
-
- # Optional Q-K normalization
- if self.use_qk_norm:
- self.q_norm = RMSNorm(head_dim)
- self.k_norm = RMSNorm(head_dim)
-
- def forward(self, x, subset_kv_size=None, subset_q_size=None):
- """
- Args:
- x: Input tensor of shape (batch, seq_len, dim)
- subset_kv_size: If provided, only attend to tokens after this index in KV
- subset_q_size: If provided, only compute attention for queries up to this index
-
- Returns:
- Output tensor of shape (batch, seq_len, dim)
- """
- # Only one subset parameter can be provided
- assert not (subset_kv_size is not None and subset_q_size is not None), \
- "Only one of subset_kv_size or subset_q_size can be provided"
-
- q, k, v = self.to_qkv(x).chunk(3, dim=-1)
-
- q, k, v = (rearrange(t, "b l (nh dh) -> b l nh dh", dh=self.head_dim) for t in (q, k, v))
-
- if self.use_qk_norm:
- q = self.q_norm(q)
- k = self.k_norm(k)
-
- # Handle subset attention cases
- if subset_kv_size is not None and subset_kv_size < k.shape[1]:
- # Attend to subset of key/value tokens
- k_subset = k[:, subset_kv_size:, :, :].contiguous()
- v_subset = v[:, subset_kv_size:, :, :].contiguous()
-
- x = xops.memory_efficient_attention(
- q, k_subset, v_subset,
- attn_bias=None,
- p=self.attn_dropout if self.training else 0.0,
- op=(xops.fmha.flash.FwOp, xops.fmha.flash.BwOp),
- )
- elif subset_q_size is not None and subset_q_size < q.shape[1]:
- # Only compute attention for subset of query tokens
- q_subset = q[:, :subset_q_size, :, :].contiguous()
-
- x = xops.memory_efficient_attention(
- q_subset, k, v,
- attn_bias=None,
- p=self.attn_dropout if self.training else 0.0,
- op=(xops.fmha.flash.FwOp, xops.fmha.flash.BwOp),
- )
- else:
- # Regular attention for all tokens
- x = xops.memory_efficient_attention(
- q, k, v,
- attn_bias=None,
- p=self.attn_dropout if self.training else 0.0,
- op=(xops.fmha.flash.FwOp, xops.fmha.flash.BwOp),
- )
-
- x = rearrange(x, "b l nh dh -> b l (nh dh)")
-
- # Final projection
- x = self.attn_fc_dropout(self.fc(x))
-
- return x
-
-
-
-
-class QK_Norm_TransformerBlock(nn.Module):
- """
- Standard transformer block with pre-normalization architecture.
- Reference: https://github.com/facebookresearch/dino/blob/7c446df5b9f45747937fb0d72314eb9f7b66930a/vision_transformer.py#L95-L113
- """
-
- def __init__(
- self,
- dim,
- head_dim,
- ln_bias=False,
- attn_qkv_bias=False,
- attn_dropout=0.0,
- attn_fc_bias=False,
- attn_fc_dropout=0.0,
- mlp_ratio=4,
- mlp_bias=False,
- mlp_dropout=0.0,
- use_qk_norm=True,
- ):
- super().__init__()
- self.norm1 = nn.LayerNorm(dim, bias=ln_bias)
- self.attn = QK_Norm_SelfAttention(
- dim=dim,
- head_dim=head_dim,
- qkv_bias=attn_qkv_bias,
- fc_bias=attn_fc_bias,
- attn_dropout=attn_dropout,
- fc_dropout=attn_fc_dropout,
- use_qk_norm=use_qk_norm,
- )
-
- self.norm2 = nn.LayerNorm(dim, bias=ln_bias)
- self.mlp = MLP(
- dim=dim,
- mlp_ratio=mlp_ratio,
- bias=mlp_bias,
- dropout=mlp_dropout,
- )
-
-
- def forward(self, x):
- x = x + self.attn(self.norm1(x))
- x = x + self.mlp(self.norm2(x))
- return x
-
-
-
-
\ No newline at end of file
diff --git a/optgs/model/encoder/lvsm/transformer.py b/optgs/model/encoder/lvsm/transformer.py
deleted file mode 100644
index 2ce07e1e46be9b7e068dc561336098e93c68fd2b..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/lvsm/transformer.py
+++ /dev/null
@@ -1,62 +0,0 @@
-import torch
-import torch.nn as nn
-
-from .layer import QK_Norm_TransformerBlock, init_weights
-
-
-class LVSMTransformer(nn.Module):
- def __init__(self,
- dim=768,
- d_head=64,
- n_layer=24,
- special_init=True,
- depth_init=True,
- use_qk_norm=True,
- ):
- super().__init__()
-
- # Create transformer blocks
- self.transformer_blocks = [
- QK_Norm_TransformerBlock(
- dim, d_head, use_qk_norm=use_qk_norm
- ) for _ in range(n_layer)
- ]
-
- # Apply special initialization if configured
- if special_init:
- for idx, block in enumerate(self.transformer_blocks):
- if depth_init:
- weight_init_std = 0.02 / (2 * (idx + 1)) ** 0.5
- else:
- weight_init_std = 0.02 / (2 * n_layer) ** 0.5
- block.apply(lambda module: init_weights(module, weight_init_std))
- else:
- for block in self.transformer_blocks:
- block.apply(init_weights)
-
- self.transformer_blocks = nn.ModuleList(self.transformer_blocks)
-
-
- def forward(self, x):
-
- for blk in self.transformer_blocks:
- x = blk(x)
-
- return x
-
-
-
-if __name__ == '__main__':
- device = torch.device('cuda')
- model = LVSMTransformer().to(device)
-
- x = torch.randn(2, 64, 768).to(device)
-
- with torch.autocast('cuda', dtype=torch.bfloat16):
- y = model(x)
-
- print(y.shape)
-
-
-
-
diff --git a/optgs/model/encoder/point_transformer/__init__.py b/optgs/model/encoder/point_transformer/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/optgs/model/encoder/point_transformer/debug.py b/optgs/model/encoder/point_transformer/debug.py
deleted file mode 100644
index 3f2314d4dd172c5b2a67c5e567fce28088702e40..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/point_transformer/debug.py
+++ /dev/null
@@ -1,59 +0,0 @@
-import torch
-import time
-
-
-from layer import KNNAttention, TransformerBlock, PlainPointTransformer, SubsampleBlock
-
-
-device = torch.device('cuda')
-
-def test():
- bs = 4
- npts = 1024
- len_xyz = 3
- feat_dims = 64
- num_classes = 23
- coord = torch.rand(bs * npts, len_xyz).cuda()
- feat = torch.rand(bs * npts, feat_dims).cuda()
- offset = [npts * i for i in range(1, bs + 1)]
- offset = torch.tensor(offset).cuda()
-
- # data_dict = dict(
- # coord = coord,
- # feat = feat,
- # offset = offset
- # )
-
- # model = PointTransformerSeg26().cuda()
-
- # model = KNNAttention(feat_dims, num_samples=16).cuda()
-
- # model = TransformerBlock(feat_dims).cuda()
-
- # model = PlainPointTransformer(feat_dims, num_blocks=2).cuda()
- model = SubsampleBlock(feat_dims, feat_dims).cuda()
-
- print(model)
-
-
- # count time
- # count = 100
-
- # torch.cuda.synchronize()
- # start = time.time()
-
- # for _ in range(count):
- # out = model((coord, feat, offset))
-
- # torch.cuda.synchronize()
- # print(time.time() - start)
-
-
- out = model((coord, feat, offset))
- print(out[0].shape)
-
- # print(out.shape)
-
-
-
-test()
diff --git a/optgs/model/encoder/point_transformer/layer.py b/optgs/model/encoder/point_transformer/layer.py
deleted file mode 100644
index 148a57edbbbcf725037d3b57ee423dc138c25ca9..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/point_transformer/layer.py
+++ /dev/null
@@ -1,1491 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils.checkpoint
-
-import pointops
-from pointops import grouping, grouping2
-from einops import rearrange
-import time
-
-from ..unimatch.dinov2.layers.block import Block as MultiViewBlock
-from ..unimatch.utils import mv_feature_add_position
-from ..unimatch.mv_transformer import MultiViewFeatureTransformer
-
-
-USE_PYTORCH_ATTN = False
-USE_FLASH_ATTN3 = False
-
-# try:
-# from flash_attn_interface import flash_attn_func
-# FA3_AVAILABLE = True
-# warnings.warn('flash attention 3 is available (point attn)')
-
-# except ImportError:
-# FA3_AVAILABLE = False
-# warnings.warn('flash attention 3 is not available (point attn)')
-
-
-class KNNAttention(nn.Module):
- # TODO: multi-head
- def __init__(self, channels, knn_samples=16, no_rpe=True,
- qk_norm=False,
- num_heads=1,
- proj_channels=None,
- use_fused=False,
- ):
- super().__init__()
- self.proj_channels = proj_channels
-
- self.knn_samples = knn_samples
- self.no_rpe = no_rpe
- self.num_heads = num_heads
- assert self.num_heads == 1
-
- self.use_fused = use_fused
- if use_fused:
- try:
- import sys
- from optgs.paths import PROJECT_DIR
- sys.path.append(str(PROJECT_DIR / "submodules"))
- from fused_knn_attn import fused_knn_attention, FUSED_KNN_ATTN_CUDA_AVAILABLE
- self._fused_knn_attention = fused_knn_attention
- if not FUSED_KNN_ATTN_CUDA_AVAILABLE:
- import warnings
- warnings.warn(
- "Fused KNN attention CUDA extension not available, "
- "using PyTorch fallback (still avoids [N,K,C] intermediates)"
- )
- except ImportError:
- import warnings
- warnings.warn(
- "fused_knn_attn package not found, falling back to unfused attention"
- )
- self.use_fused = False
-
- self.qk_norm = qk_norm
- if qk_norm:
- self.q_norm = nn.RMSNorm(channels)
- self.k_norm = nn.RMSNorm(channels)
-
- if self.proj_channels is not None:
- self.qkv = nn.Linear(channels, self.proj_channels * 3, bias=False)
- self.proj = nn.Linear(self.proj_channels, channels)
- else:
- self.qkv = nn.Linear(channels, channels * 3, bias=False)
- self.proj = nn.Linear(channels, channels)
-
- if not self.no_rpe:
- self.rpe = nn.Sequential(
- nn.Linear(3, 32),
- nn.GELU(),
- nn.Linear(32, 1)
- )
-
-
- def forward(self, pxo, knn_idx=None):
- # [N, 3], [N, C], [B]
- p, x, o = pxo
- c = x.size(1)
-
- if self.proj_channels is not None:
- c = self.proj_channels
-
- assert c % self.num_heads == 0
- head_dim = c // self.num_heads
- scale_factor = head_dim ** -0.5
-
- qkv = self.qkv(x) # [N, 3*C]
- x_q, x_k, x_v = torch.chunk(qkv, chunks=3, dim=-1) # each [N, C]
-
- # ---- Fused path: gather + attention in one kernel ----
- if self.use_fused and self.no_rpe:
- # Ensure we have KNN indices
- if knn_idx is None:
- knn_idx, _ = pointops.knn_query(
- self.knn_samples, p, o, p, o
- )
-
- # qk_norm: RMSNorm normalizes each C-dim vector independently,
- # so applying before gather is equivalent to applying after gather.
- if self.qk_norm:
- x_q = self.q_norm(x_q)
- x_k = self.k_norm(x_k)
-
- out = self._fused_knn_attention(
- x_q.contiguous(), x_k.contiguous(), x_v.contiguous(),
- knn_idx.contiguous(), scale_factor
- )
- out = self.proj(out)
- return out
-
- # ---- Original unfused path ----
- # # [N, K, C], [N, K]
- # x_k, idx = pointops.knn_query_and_group(
- # x_k.contiguous(), p, o, new_xyz=p, new_offset=o,
- # idx=knn_idx,
- # nsample=self.knn_samples, with_xyz=False
- # ) # [N, K, C]
- #
- # # [N, K, C]
- # x_v, _ = pointops.knn_query_and_group(
- # x_v.contiguous(),
- # p,
- # o,
- # new_xyz=p,
- # new_offset=o,
- # idx=idx,
- # nsample=self.knn_samples,
- # with_xyz=False,
- # )
-
- # ---- Initial improved version ----
- x_kv = torch.cat([x_k, x_v], dim=-1) # [N, 2C/3]
- x_kv_query, _ = pointops.knn_query_and_group(
- x_kv.contiguous(), p, o, new_xyz=p, new_offset=o,
- idx=knn_idx, nsample=self.knn_samples, with_xyz=False
- ) # [N, K, 2C/3]
- x_k, x_v = torch.chunk(x_kv_query, chunks=2, dim=-1)
-
- # [N, K, 3], [N, K, C]
- # NOTE: without xyz in knn
- # p_r, x_k = x_k[:, :, :3], x_k[:, :, 3:]
-
- # [N, 1, K]
- assert self.no_rpe
- if not self.no_rpe:
- rpe = self.rpe(p_r).permute(0, 2, 1)
- else:
- rpe = 0
-
- if self.qk_norm:
- x_q = self.q_norm(x_q)
- x_k = self.k_norm(x_k)
-
- n, k, c = x_k.shape
-
- # attention
- if USE_PYTORCH_ATTN:
- out = F.scaled_dot_product_attention(
- x_q.view(n, 1, c),
- x_k.view(n, k, c),
- x_v.view(n, k, c),
- ).reshape(n, c) # [N, C]
-
- elif (USE_FLASH_ATTN3 and FA3_AVAILABLE and self.no_rpe):
- # no relative pos enc
- out = flash_attn_func(
- x_q.view(n, 1, self.num_heads, head_dim).to(torch.bfloat16),
- x_k.view(n, k, self.num_heads, head_dim).to(torch.bfloat16),
- x_v.view(n, k, self.num_heads, head_dim).to(torch.bfloat16),
- )[0].reshape(n, c).float() # [N, C]
- else:
- # [N, 1, K]
- scores = torch.matmul(x_q.unsqueeze(1), x_k.permute(0, 2, 1)) * scale_factor + rpe
- # [N, C]
- out = torch.matmul(torch.softmax(scores, dim=2), x_v).squeeze(1)
-
- out = self.proj(out)
-
- return out
-
-
-class MLP(nn.Module):
- def __init__(
- self,
- channels,
- act="gelu",
- ):
- super().__init__()
-
- expansion = 4
-
- self.fc1 = nn.Linear(channels, channels * expansion)
- if act is None or act in ['none', 'identity']:
- self.act = nn.Identity()
- elif act == 'gelu':
- self.act = nn.GELU()
- elif act == 'tanh':
- self.act = nn.Tanh()
- else:
- raise ValueError(f"unsupported activation {act}")
- self.fc2 = nn.Linear(channels * expansion, channels)
-
- def forward(self, x):
- x = self.fc1(x)
- x = self.act(x)
- x = self.fc2(x)
- return x
-
-
-class TransformerBlock(nn.Module):
- def __init__(self, channels, knn_samples=16, post_norm=False,
- no_rpe=False,
- no_attn=False,
- no_norm=False,
- act="gelu",
- qk_norm=False,
- norm_pt_block=False,
- num_heads=1,
- attn_proj_channels=None,
- use_fused_attn=False,
- ):
- super().__init__()
- self.post_norm = post_norm
- self.no_attn = no_attn
- self.norm_pt_block = norm_pt_block
-
- if no_norm:
- self.norm1 = nn.Identity()
- self.norm2 = nn.Identity()
- else:
- self.norm1 = nn.LayerNorm(channels)
- self.norm2 = nn.LayerNorm(channels)
-
- if self.no_attn:
- self.linear = nn.Linear(channels, channels)
- else:
- self.attn = KNNAttention(channels, knn_samples=knn_samples, no_rpe=no_rpe,
- qk_norm=qk_norm,
- num_heads=num_heads,
- proj_channels=attn_proj_channels,
- use_fused=use_fused_attn,
- )
- self.mlp = MLP(channels, act=act)
-
- if self.norm_pt_block:
- self.norm3 = nn.LayerNorm(channels)
-
- def forward(self, pxo, knn_idx=None):
- p, x, o = pxo
-
- if self.post_norm:
- if self.no_attn:
- x = x + self.norm1(self.linear(x))
- else:
- x = x + self.norm1(self.attn((p, x, o), knn_idx=knn_idx))
- x = x + self.norm2(self.mlp(x))
- else:
- if self.no_attn:
- x = x + self.linear(self.norm1(x))
- else:
- x = x + self.attn((p, self.norm1(x), o), knn_idx=knn_idx)
- x = x + self.mlp(self.norm2(x))
-
- if self.norm_pt_block:
- x = self.norm3(x)
-
- return x
-
-
-class FPSSubsample(nn.Module):
- def __init__(self, in_planes, out_planes, stride=4, nsample=16,
- agg_func='attn',
- subsample_method='fps',
- return_idx=False,
- fps_num_samples=None,
- attn_channels=64,
- ):
- super().__init__()
-
- assert stride > 0
-
- self.agg_func = agg_func
- self.subsample_method = subsample_method
- self.knn_samples = nsample
- self.return_idx = return_idx
-
- self.stride, self.nsample = stride, nsample
-
- if fps_num_samples is not None:
- self.nsample = fps_num_samples
-
- # if stride != 1:
- # # xyz + feature
- # # self.linear = nn.Linear(3 + in_planes, out_planes, bias=not post_norm)
- # # only feature
- # # TODO: attention aggregation
- # if agg_func == 'maxpool':
- # self.agg = nn.MaxPool1d(nsample)
- # elif agg_func == 'avgpool':
- # self.agg = nn.AvgPool1d(nsample)
- # else:
- # raise ValueError(f"unsupported agg_func {agg_func}")
-
- # fewer channels to save memory
- assert agg_func in ['attn', 'avgpool']
- if self.agg_func == 'attn':
- self.q = nn.Linear(in_planes, attn_channels, bias=False)
- self.k = nn.Linear(in_planes, attn_channels, bias=False)
- self.v = nn.Linear(in_planes, attn_channels, bias=False)
-
- self.proj = nn.Linear(attn_channels, out_planes, bias=True)
- self.residual = nn.Linear(in_planes, out_planes, bias=True)
- else:
- self.proj = nn.Linear(in_planes, out_planes, bias=True)
-
- def forward(self, pxo):
- p, x, o = pxo # (n, 3), (n, c), (b)
- if self.stride != 1:
- if self.subsample_method == 'density':
- assert False # not well tested
- n_o, count = [o[0].item() // self.stride], o[0].item() // self.stride
- for i in range(1, o.shape[0]):
- count += (o[i].item() - o[i - 1].item()) // self.stride
- n_o.append(count)
- n_o = torch.tensor(n_o, dtype=torch.int32, device=x.device)
-
- # [N, K, C+3]
- x_k, _ = pointops.knn_query_and_group(
- x.contiguous(), p, o, new_xyz=p, new_offset=o, nsample=self.knn_samples, with_xyz=True
- )
-
- p_r = x_k[:, :, 0:3]
- density = torch.mean(torch.norm(p_r, dim=-1), dim=-1) # [N]
-
- # TODO: normalize the distance
- weights = (density - density.min()) / (density.max() - density.min() + 1e-6)
- # weights = density
-
- # weights = 1.0 / (density + 1e-6) # Inverse density weighting
-
- # to batch
- lists = [weights[:o[0]]]
- for i in range(o.shape[0] - 1):
- lists.append(weights[o[i]:o[i+1]])
-
- weights = torch.stack(lists, dim=0) # [B, N]
-
- weights = weights / weights.sum(dim=1, keepdim=True) # Normalize weights
-
- # Sample points based on weights
- batch = n_o.shape[0]
- num_samples = o[0].item() // self.stride
- sampled_indices = torch.stack([
- torch.multinomial(weights[b], num_samples, replacement=False)
- for b in range(batch)
- ], dim=0) # (B, num_samples)
-
- idx = rearrange(sampled_indices, "b n -> (b n)")
-
- point_list = [p[:o[0]]]
- for i in range(o.shape[0] - 1):
- point_list.append(p[o[i]:o[i+1], :])
-
- points = torch.stack(point_list, dim=0) # [B, N, 3]
-
- # Gather sampled points
- sampled_points = torch.gather(points, 1, sampled_indices.unsqueeze(-1).expand(-1, -1, 3))
-
- # print(sampled_points.shape) # [B, M, 3]
-
- sampled_points = rearrange(sampled_points, "b m c -> (b m) c")
-
- # average pooling
- # TODO: try others
- x = x_k.mean(dim=1) # [N, C]
- x_list = [x[:o[0]]]
- for i in range(o.shape[0] - 1):
- x_list.append(x[o[i]:o[i+1], :])
- x = torch.stack(x_list, dim=0) # [B, N, C]
-
- # Gather sampled points
- x = torch.gather(x, 1, sampled_indices.unsqueeze(-1).expand(-1, -1, x.size(-1)))
- x = rearrange(x, "b n c -> (b n) c")
-
- # TODO: do we need to add residual to x here?
- # use the index to subsample the initial features
- x = self.proj(x)
-
- p, o = sampled_points, n_o
- elif self.subsample_method in ['fps', 'grid']:
- n_o, count = [o[0].item() // self.stride], o[0].item() // self.stride
- for i in range(1, o.shape[0]):
- count += (o[i].item() - o[i - 1].item()) // self.stride
- n_o.append(count)
- n_o = torch.tensor(n_o, dtype=torch.int32, device=x.device)
-
- if self.subsample_method == 'fps':
- idx = pointops.farthest_point_sampling(p, o, n_o) # (m)
- else:
- # uniform sampling: sanity check
- # first reshape to V, H, W, then do grid sampling
- # Generate grid indices
- # TODO: grid sample in the image space
- idx = torch.arange(0, p.size(0), self.stride).to(x.device)
-
- n_p = p[idx.long(), :] # (m, 3)
- x_subsample = x[idx.long(), :] # [M, C]
- if self.agg_func == 'attn':
- x_q = self.q(x_subsample) # [M, C]
- # [M, K, C]
- x_k = self.k(x) # [N, C]
- else:
- x_k = x
-
- x_k, knn_idx = pointops.knn_query_and_group(
- x_k,
- p,
- offset=o,
- new_xyz=n_p,
- new_offset=n_o,
- nsample=self.nsample,
- with_xyz=False, # remove xyz
- )
-
- if self.agg_func == 'attn':
- x_v = self.v(x)
- x_v, _ = pointops.knn_query_and_group(
- x_v,
- p,
- offset=o,
- new_xyz=n_p,
- new_offset=n_o,
- idx=knn_idx,
- nsample=self.nsample,
- with_xyz=False, # remove xyz
- )
-
- # attention
- # x_q: [M, C], x_k: [M, K, C], x_v: [M, K, C]
- scale_factor = x_q.shape[-1] ** -0.5
-
- # [M, 1, K]
- # no relative posenc
- scores = torch.matmul(x_q.unsqueeze(1), x_k.permute(0, 2, 1)) * scale_factor
- # [M, C]
- x = torch.matmul(torch.softmax(scores, dim=2), x_v).squeeze(1)
-
- # if self.agg_func in ['avgpool', 'maxpool']:
- # x = self.agg(x.transpose(1, 2).contiguous()).squeeze(-1) # (m, c)
- # else:
- # raise NotImplementedError
-
- # add residual to x here?
- # use the index to subsample the initial features
- x = self.residual(x_subsample) + self.proj(x)
- else:
- x = x_k.mean(dim=1)
- x = self.proj(x)
-
- p, o = n_p, n_o
-
- else:
- raise ValueError(f"unsupported subsampling method {self.subsample_method}")
- else:
- # add residual to x here?
- x = x + self.proj(x)
-
- idx = torch.arange(0, p.size(0)).to(x.device)
-
- if self.return_idx:
- return [p, x, o], idx
- return [p, x, o]
-
-
-class SubsampleBlock(nn.Module):
- def __init__(self, in_channels, out_channels, stride=4, knn_samples=16, post_norm=False,
- agg_func='attn',
- subsample_method='fps',
- return_idx=False,
- fps_num_samples=None,
- attn_proj_channels=None,
- ):
- super().__init__()
-
- assert not post_norm
-
- self.return_idx = return_idx
-
- self.post_norm = post_norm
- self.norm1 = nn.LayerNorm(in_channels)
- self.fps = FPSSubsample(in_channels, out_channels, stride=stride, nsample=knn_samples,
- agg_func=agg_func,
- subsample_method=subsample_method,
- return_idx=return_idx,
- fps_num_samples=fps_num_samples,
- attn_channels=attn_proj_channels,
- )
-
- self.norm2 = nn.LayerNorm(out_channels)
- self.mlp = MLP(out_channels)
-
- def forward(self, pxo):
-
- # pre norm
- p, x, o = pxo
- x = self.norm1(x)
-
- if self.return_idx:
- pxo, idx = self.fps([p, x, o])
- else:
- pxo = self.fps([p, x, o])
-
- p, x, o = pxo
-
- x = x + self.mlp(self.norm2(x))
-
- if self.return_idx:
- return [p, x, o], idx
-
- return [p, x, o]
-
-
-class SkipConnect(nn.Module):
- def __init__(self, in_channels, out_channels):
- super().__init__()
- self.proj1 = nn.Linear(out_channels, out_channels)
- self.proj2 = nn.Linear(in_channels, out_channels)
- self.proj3 = nn.Linear(out_channels, out_channels)
-
- def forward(self, pxo1, pxo2):
- p1, x1, o1 = pxo1
- p2, x2, o2 = pxo2
-
- # TODO: support half precision
- with torch.amp.autocast(device_type='cuda', enabled=True, dtype=torch.float32):
- x = self.proj1(x1) + pointops.interpolation2(
- p2, p1, self.proj2(x2), o2, o1
- )
-
- x = self.proj3(x)
-
- return x
-
-
-
-class PlainPointTransformer(nn.Module):
- def __init__(self, channels, knn_samples=16, num_blocks=4, post_norm=False,
- no_rpe=False,
- no_attn=False,
- no_norm=False,
- act="gelu",
- qk_norm=False,
- norm_pt_block=False,
- num_heads=1,
- attn_proj_channels=None,
- cache_knn_idx=None,
- knn_idx_update_every=1,
- with_mv_attn=False,
- with_mv_attn_lowres=False,
- mv_attn_first=False,
- no_mv_attn=False,
- conv_with_norm=False,
- mv_shuffle_attn=False,
- with_pos_enc=False,
- shuffle_attn_no_norm=False,
- mv_unimatch_attn=False,
- use_checkpointing=False,
- init_use_checkpointing=False,
- use_fused_attn=False,
- ):
- super().__init__()
-
- self.cache_knn_idx = cache_knn_idx
- self.knn_idx_update_every = knn_idx_update_every
- self.knn_samples = knn_samples
- self.use_checkpointing = use_checkpointing
- self.init_use_checkpointing = init_use_checkpointing
-
- self.with_mv_attn = with_mv_attn
- self.with_mv_attn_lowres = with_mv_attn_lowres
- if with_pos_enc:
- assert mv_shuffle_attn
-
- self.blocks = nn.ModuleList()
- for _ in range(num_blocks):
- self.blocks.append(TransformerBlock(channels, knn_samples=knn_samples,
- post_norm=post_norm,
- no_rpe=no_rpe,
- no_attn=no_attn,
- no_norm=no_norm,
- act=act,
- qk_norm=qk_norm,
- norm_pt_block=norm_pt_block,
- num_heads=num_heads,
- attn_proj_channels=attn_proj_channels,
- use_fused_attn=use_fused_attn,
- ))
-
- # multi-view attention
- if self.with_mv_attn:
- self.mv_blocks = nn.ModuleList()
- for _ in range(num_blocks):
- # if mv_shuffle_attn:
- if self.with_mv_attn_lowres:
- self.mv_blocks.append(
- MultViewLowresAttn(
- channels,
- )
- )
- else:
- self.mv_blocks.append(
- MultiViewBlock(
- channels,
- num_heads=4,
- )
- )
- # elif mv_unimatch_attn:
- # self.mv_blocks.append(
- # MultViewUniMatchAttn(
- # channels,
- # )
- # )
- # else:
- # self.mv_blocks.append(
- # MultViewUnetAttn(channels,
- # no_mv_attn=no_mv_attn,
- # conv_with_norm=conv_with_norm,
- # )
- # )
-
- def forward(self, pxo, iter=0, b=None, v=None, h=None, w=None):
- p, x, o = pxo
- # compute knn idx here only once and pass it to the model
- # the positions are not changed inside the blocks
- if self.cache_knn_idx is None or (iter % self.knn_idx_update_every) == 0:
- knn_idx, _ = pointops.knn_query(self.knn_samples, p, o, p, o)
- self.cache_knn_idx = knn_idx
- # print(knn_idx.float().mean().item())
- else:
- knn_idx = self.cache_knn_idx
-
- if self.with_mv_attn:
- assert b is not None and v is not None and h is not None and w is not None
- if self.use_checkpointing:
- raise NotImplementedError
-
- for i in range(len(self.blocks)):
- # knn attention
- x = self.blocks[i]([p, x, o], knn_idx=knn_idx)
- # global multi-view attention
- x = rearrange(x, "(b v h w) c -> b (v h w) c", b=b, v=v, h=h, w=w)
- if self.with_mv_attn_lowres:
- x = self.mv_blocks[i](x, v=v, h=h, w=w)
- # # TODO: hard-coded for now
- # if x.size(1) == 8 * 512 // 4 * 960 // 4:
- # x = self.mv_blocks[i](x, v=8, h=512 // 4, w=960 // 4)
- # elif x.size(1) == 8 * 256 // 4 * 448 // 4:
- # x = self.mv_blocks[i](x, v=8, h=256 // 4, w=448 // 4)
- # else:
- # raise ValueError(f"unsupported input size {x.size(1)} for multi-view attention")
- # # print(x.shape)
- else:
- x = self.mv_blocks[i](x)
- # x = x.squeeze(0)
- x = rearrange(x, "b (v h w) c -> (b v h w) c",
- b=b, v=v, h=h, w=w)
- else:
- for blk in self.blocks:
- if self.init_use_checkpointing:
- # checkpointing the inital reconstruction model
- # NOTE: cannot cache knn_idx here, otherwise index out error
- def custom_forward(p, x, o):
- return blk((p, x, o), knn_idx=None) # knn_idx is closed over
- x = torch.utils.checkpoint.checkpoint(custom_forward, p, x, o)
- else:
- x = blk((p, x, o), knn_idx=knn_idx)
-
- return x
-
-
-class MultViewUnetAttn(nn.Module):
- def __init__(self, channels, no_mv_attn=False, conv_with_norm=False):
- super().__init__()
-
- self.conv_with_norm = conv_with_norm
-
- self.down1 = nn.Conv2d(channels, channels, 3, 2, 1)
- self.down2 = nn.Conv2d(channels, channels, 3, 2, 1)
-
- self.up2 = nn.Conv2d(channels, channels, 3, 1, 1)
- self.up1 = nn.Conv2d(channels, channels, 3, 1, 1)
-
- self.attn = MultiViewBlock(channels, 4, no_attn=no_mv_attn)
-
- if self.conv_with_norm:
- self.norm1 = nn.LayerNorm(channels)
- self.norm2 = nn.LayerNorm(channels)
- self.norm3 = nn.LayerNorm(channels)
- self.norm4 = nn.LayerNorm(channels)
-
- def forward(self, x):
- v = 8
- h = 256 // 4
- w = 448 // 4
- b = 1
- assert x.size(0) == b * v * h * w
- residual = x
- x = rearrange(x, "(b v h w) c -> (b v) c h w", b=b, v=v, h=h, w=w)
- x1 = self.down1(x) # 1/2
- if self.conv_with_norm:
- x1 = self.norm1(x1.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
- x2 = self.down2(x1) # 1/4
- if self.conv_with_norm:
- x2 = self.norm2(x2.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
- x2 = rearrange(x2, "(b v) c h w -> b (v h w) c", b=b, v=v)
- x2 = self.attn(x2) # 1/4
- x2 = rearrange(x2, "b (v h w) c -> (b v) c h w", b=b, v=v, h=h//4, w=w//4)
- x2 = self.up2(x1 + F.interpolate(x2, scale_factor=2, mode='bilinear', align_corners=True)) # 1/2
- if self.conv_with_norm:
- x2 = self.norm3(x2.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
- x = self.up1(x + F.interpolate(x2, scale_factor=2, mode='bilinear', align_corners=True)) # 1
- if self.conv_with_norm:
- x = self.norm4(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
-
- x = rearrange(x, "(b v) c h w -> (b v h w) c", b=b, v=v)
-
- x = residual + x
-
- return x
-
-
-class MultViewShuffleAttn(nn.Module):
- def __init__(self, channels, no_mv_attn=False, with_pos_enc=False, shuffle_attn_no_norm=False):
- super().__init__()
-
- self.down_factor = 4
- self.with_pos_enc = with_pos_enc
-
- self.proj1 = nn.Linear(channels * self.down_factor ** 2, channels)
- if shuffle_attn_no_norm:
- self.norm1 = nn.Identity()
- else:
- self.norm1 = nn.LayerNorm(channels)
-
- self.proj2 = nn.Linear(channels, channels * self.down_factor ** 2)
-
- if shuffle_attn_no_norm:
- self.norm2 = nn.Identity()
- else:
- self.norm2 = nn.LayerNorm(channels * self.down_factor ** 2)
-
- self.conv = nn.Conv2d(channels, channels, 3, 1, 1)
-
- if no_mv_attn:
- self.attn = nn.Identity()
- else:
- self.attn = MultiViewBlock(channels, 4, no_attn=no_mv_attn)
-
- def forward(self, x):
- v = 8
- h = 256 // 4
- w = 448 // 4
- b = 1
- assert x.size(0) == b * v * h * w
- residual = x
- x = rearrange(x, "(b v h w) c -> (b v) c h w", b=b, v=v, h=h, w=w)
-
- # TODO: add positional encoding to x
- if self.with_pos_enc:
- x = mv_feature_add_position(x, attn_splits=1, feature_channels=x.size(1))
- # print(x.shape)
-
- x = F.pixel_unshuffle(x, self.down_factor)
-
- x = rearrange(x, "(b v) c h w -> b (v h w) c", b=b)
- x = self.proj1(x)
- x = self.norm1(x)
-
- x = self.attn(x)
-
- x = self.proj2(x)
- x = self.norm2(x)
-
- x = rearrange(x, "b (v h w) c -> (b v) c h w", b=b, v=v, h=h // self.down_factor, w=w // self.down_factor)
- x = F.pixel_shuffle(x, self.down_factor)
- x = self.conv(x)
- x = rearrange(x, "(b v) c h w -> (b v h w) c", b=b, v=v)
- x = x + residual
-
- return x
-
-
-class MultViewLowresAttn(nn.Module):
- def __init__(self, channels, no_mv_attn=False, with_pos_enc=False, shuffle_attn_no_norm=False,
- down_factor=4,
- attn_proj_channels=None,
- ):
- super().__init__()
-
- self.down_factor = down_factor
- self.with_pos_enc = with_pos_enc
-
- self.attn_proj_channels = attn_proj_channels
-
- if attn_proj_channels:
- ori_channels = channels
- self.proj0 = nn.Linear(channels, attn_proj_channels)
- channels = attn_proj_channels
-
- if self.down_factor == 8:
- down_factor = 4
- else:
- down_factor = self.down_factor
-
- self.proj1 = nn.Linear(channels * down_factor ** 2, channels)
- if shuffle_attn_no_norm:
- self.norm1 = nn.Identity()
- else:
- self.norm1 = nn.LayerNorm(channels)
-
- self.proj2 = nn.Linear(channels, channels * down_factor ** 2)
-
- if shuffle_attn_no_norm:
- self.norm2 = nn.Identity()
- else:
- self.norm2 = nn.LayerNorm(channels * down_factor ** 2)
-
- self.conv = nn.Conv2d(channels, channels, 3, 1, 1)
-
- if attn_proj_channels:
- self.proj3 = nn.Linear(channels, ori_channels)
-
- if no_mv_attn:
- self.attn = nn.Identity()
- else:
- num_heads = 1 if self.attn_proj_channels else 4
- self.attn = MultiViewBlock(channels, num_heads, no_attn=no_mv_attn)
-
- def forward(self, x, v=None, h=None, w=None, y=None):
- if y is not None:
- return self.forward_cross_attn(x, y, v, h, w)
- residual = x
- if self.attn_proj_channels:
- x = self.proj0(x)
-
- x = rearrange(x, "b (v h w) c -> (b v) c h w", v=v, h=h, w=w)
-
- # TODO: add positional encoding to x
- if self.with_pos_enc:
- x = mv_feature_add_position(x, attn_splits=1, feature_channels=x.size(1))
- # print(x.shape)
-
- if self.down_factor == 8:
- # bilinear to half first to save channels
- x = F.interpolate(x, scale_factor=0.5, mode='bilinear', align_corners=True)
- down_factor = 4
- else:
- down_factor = self.down_factor
-
- x = F.pixel_unshuffle(x, down_factor)
-
- x = rearrange(x, "(b v) c h w -> b (v h w) c", v=v)
- x = self.proj1(x)
- x = self.norm1(x)
-
- x = self.attn(x)
-
- x = self.proj2(x)
- x = self.norm2(x)
-
- x = rearrange(x, "b (v h w) c -> (b v) c h w", v=v, h=h // self.down_factor, w=w // self.down_factor)
- x = F.pixel_shuffle(x, down_factor)
- x = self.conv(x)
- if self.down_factor == 8:
- # bilinear to full
- x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True)
- x = rearrange(x, "(b v) c h w -> b (v h w) c", v=v)
- if self.attn_proj_channels:
- x = self.proj3(x)
- x = x + residual
-
- return x
-
- def forward_cross_attn(self, x, y, v=None, h=None, w=None):
- residual = x
- if self.attn_proj_channels:
- x = self.proj0(x)
-
- assert y is not None
- y = rearrange(y, "b (v h w) c -> (b v) c h w", h=h, w=w) # different v with x
- num_cross_view = y.shape[0] // x.shape[0]
-
- x = rearrange(x, "b (v h w) c -> (b v) c h w", v=v, h=h, w=w)
-
- # TODO: add positional encoding to x
- if self.with_pos_enc:
- x = mv_feature_add_position(x, attn_splits=1, feature_channels=x.size(1))
- # print(x.shape)
-
- if self.down_factor == 8:
- # bilinear to half first to save channels
- x = F.interpolate(x, scale_factor=0.5, mode='bilinear', align_corners=True)
- y = F.interpolate(y, scale_factor=0.5, mode='bilinear', align_corners=True)
- down_factor = 4
- else:
- down_factor = self.down_factor
-
- x = F.pixel_unshuffle(x, down_factor)
- y = F.pixel_unshuffle(y, down_factor)
-
- x = rearrange(x, "(b v) c h w -> b (v h w) c", v=v)
- y = rearrange(y, "(b v) c h w -> b (v h w) c", v=num_cross_view)
- x = self.proj1(x)
- x = self.norm1(x)
-
- y = self.proj1(y)
- y = self.norm1(y)
-
- # x_tmp = self.attn(x)
-
- # print((x - y).abs().max().item())
-
- x = self.attn(x, y)
-
- # there will be slight diff for self and cross attn caused by flash3
- # print((x_tmp - x).abs().max().item())
-
- x = self.proj2(x)
- x = self.norm2(x)
-
- x = rearrange(x, "b (v h w) c -> (b v) c h w", v=v, h=h // self.down_factor, w=w // self.down_factor)
- x = F.pixel_shuffle(x, down_factor)
- x = self.conv(x)
- if self.down_factor == 8:
- # bilinear to full
- x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True)
- x = rearrange(x, "(b v) c h w -> b (v h w) c", v=v)
- if self.attn_proj_channels:
- x = self.proj3(x)
- x = x + residual
-
- return x
-
-
-
-class GaussianErrorCrossAttn(nn.Module):
- def __init__(self, gaussian_channels,
- error_channels,
- model_channels=256,
- no_mv_attn=False, with_pos_enc=False, shuffle_attn_no_norm=False,
- down_factor=4,
- attn_proj_channels=None,
- num_heads=4,
- with_mlp=False,
- ):
- super().__init__()
-
- self.num_heads = num_heads
- self.model_channels = model_channels
- self.down_factor = down_factor
- self.with_mlp = with_mlp
-
- # self.q_norm = nn.LayerNorm(gaussian_channels)
- self.q_proj = nn.Linear(gaussian_channels, model_channels)
-
- kv_channels = error_channels * (down_factor ** 2)
- # self.kv_norm = nn.LayerNorm(kv_channels)
- self.kv_proj = nn.Linear(kv_channels, 2 * model_channels)
-
- # self.out_proj = nn.Linear(model_channels, gaussian_channels)
- # concat
- self.out_proj = nn.Linear(model_channels + gaussian_channels, gaussian_channels)
-
- if with_mlp:
- self.mlp_norm = nn.LayerNorm(gaussian_channels)
- self.mlp = MLP(gaussian_channels)
-
-
- def forward(self, gaussian, error, v=None, h=None, w=None, mask=None):
- # [B, VHW, C]
- residual = gaussian
- b = gaussian.size(0)
-
- # x = self.q_norm(gaussian)
- x = gaussian
- q = self.q_proj(x) # [B, VHW, C]
-
- # spatial reshape to save computation
- error = rearrange(error, "b (v h w) c -> (b v) c h w", v=v, h=h, w=w)
- error = F.pixel_unshuffle(error, self.down_factor)
- error = rearrange(error, "(b v) c h w -> b (v h w) c", v=v)
- # error = self.kv_norm(error)
-
- kv = self.kv_proj(error)
- k, v = kv.chunk(2, dim=-1) # [B, VHW, C]
-
- # attention
- c = self.model_channels
- head_dim = c // self.num_heads
-
- # [B, N, C] → [B, num_heads, N, head_dim]
- def reshape(x):
- return x.view(b, -1, self.num_heads, head_dim).transpose(1, 2) # [B, H, N, D]
-
- q = reshape(q)
- k = reshape(k)
- v = reshape(v)
-
- # Fast fused attention
- out = F.scaled_dot_product_attention(q, k, v)
-
- # [B, H, N, D] → [B, N, C]
- out = out.transpose(1, 2).contiguous().view(b, -1, c)
-
- # return self.out_proj(out)
-
- # out = residual + self.out_proj(out)
- # concat
- out = self.out_proj(torch.cat([out, gaussian], dim=-1))
-
- # if self.with_mlp:
- # out = out + self.mlp(self.mlp_norm(out))
-
- return out
-
-
-
-
-class MultViewUniMatchAttn(nn.Module):
- def __init__(self, channels, no_mv_attn=False, with_pos_enc=False, shuffle_attn_no_norm=False):
- super().__init__()
-
- self.attn = MultiViewFeatureTransformer(num_layers=1,
- d_model=channels,
- )
-
- def forward(self, x, v=None, h=None, w=None):
- residual = x
- x = rearrange(x, "b (v h w) c -> (b v) c h w", v=v, h=h, w=w)
-
- attn_splits = 4
-
- # add pos enc
- x = mv_feature_add_position(x, attn_splits, feature_channels=x.size(1))
- x = rearrange(x, "(b v) c h w -> b v c h w", v=v)
-
- x_list = list(torch.unbind(x, dim=1))
-
- x_list = self.attn(x_list, attn_splits)
-
- x = torch.stack(x_list, dim=1)
-
- x = rearrange(x, "b v c h w -> b (v h w) c")
-
- return x
-
-
-
-class MultiScalePointTransformer(nn.Module):
- def __init__(self, channels, knn_samples=16, post_norm=False,
- no_rpe=True,
- no_attn=False,
- qk_norm=False,
- norm_pt_block=False,
- num_heads=1,
- num_scales=3,
- stride=4,
- downsample_agg_func='attn',
- subsample_method='fps',
- fps_num_samples=None,
- attn_proj_channels=None,
- ):
- super().__init__()
-
- self.blocks = nn.ModuleList()
- # knn 4 at 1
- self.blocks.append(TransformerBlock(channels, knn_samples=4,
- post_norm=post_norm,
- no_rpe=no_rpe,
- no_attn=no_attn,
- qk_norm=qk_norm,
- norm_pt_block=norm_pt_block,
- num_heads=num_heads,
- attn_proj_channels=attn_proj_channels,
- ))
-
- for i in range(num_scales - 2, -1, -1):
- # knn 8 at 1/4
- # knn 16 at 1/16
- self.blocks.append(TransformerBlock(channels * (2 ** i), knn_samples= knn_samples // (2 ** i),
- post_norm=post_norm,
- no_rpe=no_rpe,
- no_attn=no_attn,
- qk_norm=qk_norm,
- norm_pt_block=norm_pt_block,
- num_heads=num_heads,
- attn_proj_channels=attn_proj_channels,
- ))
-
- self.down_blocks = nn.ModuleList()
- for i in range(num_scales - 1):
- self.down_blocks.append(
- SubsampleBlock(
- channels * (2 ** i), channels * (2 ** (i + 1)),
- stride=stride,
- knn_samples=knn_samples // (2 ** (num_scales - 1 - i)),
- subsample_method=subsample_method,
- agg_func=downsample_agg_func,
- fps_num_samples=fps_num_samples,
- attn_proj_channels=attn_proj_channels,
- )
- )
-
- self.down_agg = nn.ModuleList()
- for i in range(num_scales - 1):
- self.down_agg.append(
- TransformerBlock(channels * (2 ** (i + 1)), knn_samples=knn_samples // (2 ** (num_scales - 1 - i)),
- post_norm=post_norm,
- no_rpe=no_rpe,
- no_attn=no_attn,
- qk_norm=qk_norm,
- norm_pt_block=norm_pt_block,
- num_heads=num_heads,
- attn_proj_channels=attn_proj_channels,
- )
- )
-
- self.skip_blocks = nn.ModuleList()
- for i in range(num_scales - 1, 0, -1):
- self.skip_blocks.append(
- SkipConnect(
- channels * (2 ** i),
- channels * (2 ** (i - 1))
- )
- )
-
- def forward(self, pxo):
- x1 = self.blocks[0](pxo) # 1
- p1, o1 = pxo[0], pxo[2]
- p2, x2, o2 = self.down_blocks[0]([p1, x1, o1]) # 1/4
- x2 = self.down_agg[0]([p2, x2, o2]) # 1/4
- p3, x3, o3 = self.down_blocks[1]([p2, x2, o2]) # 1/16
- x3 = self.down_agg[1]([p3, x3, o3]) # 1/16
-
- x4 = self.skip_blocks[0]([p2, x2, o2], [p3, x3, o3]) # 1/4
- p4, o4 = p2, o2
- x4 = self.blocks[1]([p4, x4, o4])
- x5 = self.skip_blocks[1]([p1, x1, o1], [p4, x4, o4]) # 1
- p5, o5 = p1, o1
- x5 = self.blocks[2]([p5, x5, o5])
-
- return x5
-
-
-class PointLinearWrapper(nn.Module):
- def __init__(self, in_channels, out_channels):
- super().__init__()
-
- self.linear = nn.Linear(in_channels, out_channels)
-
- def forward(self, pxo, b=None, v=None, h=None, w=None):
- p, x, o = pxo
- x = self.linear(x)
-
- return [p, x, o]
-
-
-class SwiGLUFFN(nn.Module):
- def __init__(
- self,
- in_features: int,
- hidden_features: int | None = None,
- out_features: int | None = None,
- 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):
- x12 = self.w12(x)
- x1, x2 = x12.chunk(2, dim=-1)
- hidden = F.silu(x1) * x2
- return self.w3(hidden)
-
-
-def test_fps():
- model = FPSSubsample(256, 256,
- fps_num_samples=16,
- subsample_method='fps',
- ).cuda()
- print(model)
-
- # FPS is significantly slower than grid with many points
-
- c = 256
- b, n = 2, 40480
-
- x = torch.randn(b, n, c).cuda()
- offset = torch.tensor([n * (i + 1) for i in range(b)]).to(x.device)
- p = torch.randn(b, n, 3).cuda()
- pxo = [p.view(-1, 3), x.view(-1, c), offset]
- y = model(pxo)
- print(y[1].shape)
-
- count = 100
-
- for _ in range(5):
- model(pxo)
-
- torch.cuda.synchronize()
- start = time.time()
-
- for i in range(count):
- model(pxo)
-
- torch.cuda.synchronize()
- print(time.time() - start)
-
-def test_knn_query_and_group():
- c = 256
- # b, n = 2, 80480
- b, n = 8, 57344
- knn_samples = 16
-
- x = torch.randn(b, n, c).cuda()
- offset = torch.tensor([n * (i + 1) for i in range(b)]).to(x.device)
- o = offset
- p = torch.randn(b, n, 3).cuda()
- p = p.view(-1, 3)
-
- knn_idx, _ = pointops.knn_query(knn_samples, p, o, p, o)
-
- print(knn_idx.shape)
-
- c_qkv = 192
- qkv = torch.randn(b*n, c_qkv).cuda()
- T = 1000
-
- # chunk first, then query twice
- torch.cuda.synchronize()
- start_time = time.time()
- for _ in range(T):
- x_q, x_k, x_v = torch.chunk(qkv, chunks=3, dim=-1)
- x_k_query, idx = pointops.knn_query_and_group(
- x_k.contiguous(), p, o, new_xyz=p, new_offset=o,
- idx=knn_idx,
- nsample=knn_samples, with_xyz=False
- ) # [N, K, C/3]
- x_v_query, _ = pointops.knn_query_and_group(
- x_v.contiguous(),
- p,
- o,
- new_xyz=p,
- new_offset=o,
- idx=idx,
- nsample=knn_samples,
- with_xyz=False,
- )
- torch.cuda.synchronize()
- end_time = time.time()
- print(f"KNN query and group time: {(end_time - start_time) / T * 1000:.2f} ms")
-
- # query first, then chunk
- torch.cuda.synchronize()
- start_time = time.time()
- for _ in range(T):
- x_qkv_query = pointops.knn_query_and_group(
- qkv.contiguous(), p, o, new_xyz=p, new_offset=o,
- idx=knn_idx,
- nsample=knn_samples, with_xyz=False
- )[0] # [N, K, C*3]
- x_q, x_k, x_v = torch.chunk(x_qkv_query, chunks=3, dim=-1)
- torch.cuda.synchronize()
- end_time = time.time()
- print(f"KNN query and group time: {(end_time - start_time) / T * 1000:.2f} ms")
-
- # chunk first, then query once
- torch.cuda.synchronize()
- start_time = time.time()
- for _ in range(T):
- x_q, x_k, x_v = torch.chunk(qkv, chunks=3, dim=-1)
- x_kv = torch.cat([x_k, x_v], dim=-1) # [N, 2C/3]
- x_kv_query = pointops.knn_query_and_group(
- x_kv.contiguous(), p, o, new_xyz=p, new_offset=o,
- idx=knn_idx, nsample=knn_samples, with_xyz=False
- )[0] # [N, K, 2C/3]
- x_k_query, x_v_query = torch.chunk(x_kv_query, 2, dim=-1)
- torch.cuda.synchronize()
- end_time = time.time()
- print(f"KNN query and group time: {(end_time - start_time) / T * 1000:.2f} ms")
-
-def test_knn():
- c = 256
- b, n = 2, 80480
- model = KNNAttention(channels=c,
- # proj_feature=64,
- ).cuda()
- print(model)
-
- x = torch.randn(b, n, c).cuda()
- offset = torch.tensor([n * (i + 1) for i in range(b)]).to(x.device)
- p = torch.randn(b, n, 3).cuda()
- pxo = [p.view(-1, 3), x.view(-1, c), offset]
- y = model(pxo)
- print(y.shape)
-
- count = 100
-
- for _ in range(5):
- model(pxo)
-
- torch.cuda.synchronize()
- start = time.time()
-
- for i in range(count):
- model(pxo)
-
- torch.cuda.synchronize()
- print(time.time() - start)
-
-
-def test_faiss_knn():
- # cannot install faiss unfortunately
- # TODO: maybe implement a sliding window knn search later
- c = 256
- b, n = 2, 80480
- knn_samples = 16
-
- x = torch.randn(b, n, c).cuda()
- offset = torch.tensor([n * (i + 1) for i in range(b)]).to(x.device)
- o = offset
- p = torch.randn(b, n, 3).cuda()
- p = p.view(-1, 3)
- # pxo = [p.view(-1, 3), x.view(-1, c), offset]
-
- # print(p.shape, o.shape)
- # print(o)
-
- knn_idx, _ = pointops.knn_query(knn_samples, p, o, p, o)
-
- print(knn_idx.shape)
-
- count = 100
-
- for _ in range(5):
- pointops.knn_query(knn_samples, p, o, p, o)
-
- torch.cuda.synchronize()
- start = time.time()
-
- for i in range(count):
- pointops.knn_query(knn_samples, p, o, p, o)
-
- torch.cuda.synchronize()
- print(time.time() - start)
-
-
-def count_parameters(model):
- return sum(p.numel() for p in model.parameters() if p.requires_grad)
-
-
-def test_mlp():
- b, n, c = 2, 40240, 256
- model = MLP(c).cuda()
- x = torch.randn(b, n, c).cuda()
-
- # model = SwiGLUFFN(c, c * 3).cuda()
-
- print('parameters:', count_parameters(model))
-
- x = x.to(torch.bfloat16)
- model.to(dtype=torch.bfloat16)
-
- with torch.autocast('cuda', enabled=True, dtype=torch.bfloat16):
- y = model(x)
- print(y.shape)
-
- count = 100
-
- for _ in range(5):
- model(x)
-
- torch.cuda.synchronize()
- start = time.time()
-
- for i in range(count):
- model(x)
-
- torch.cuda.synchronize()
- print(time.time() - start)
-
-
-def test_mv_block():
- c = 256
- num_heads = 4
- model = MultiViewBlock(c, num_heads).cuda()
- x = torch.rand(2, 256, c).cuda()
-
- print(model)
-
- y = model(x)
-
- print(y.shape)
-
-
-def test_cross_attn():
- c = 256
- v, h, w = 8, 64, 128
- num_heads = 4
- model = GaussianErrorCrossAttn(512, c, c).cuda()
- x = torch.rand(2, v * h * w, 512).cuda()
- y = torch.rand(2, v * h * w, c).cuda()
-
- print(model)
-
- y = model(x, y, v=v, h=h, w=w)
-
- print(x.shape, y.shape)
-
-
-def test_grouping():
- c = 256
- # b, n = 2, 80480
- b, n = 1, 57344
- knn_samples = 16
-
- x = torch.randn(b, n, c).cuda()
- offset = torch.tensor([n * (i + 1) for i in range(b)]).to(x.device)
- o = offset
- p = torch.randn(b, n, 3).cuda()
- p = p.view(-1, 3)
-
- knn_idx, _ = pointops.knn_query(knn_samples, p, o, p, o)
-
- print(knn_idx.shape)
-
- c_qkv = 192
- qkv = torch.randn(b*n, c_qkv).cuda()
- x_q, x_k, x_v = torch.chunk(qkv, chunks=3, dim=-1)
- x_kv = torch.cat([x_k, x_v], dim=-1) # [N, 2C/3]
-
- m, nsample, c = knn_idx.shape[0], knn_idx.shape[1], x_kv.shape[1]
- feat = torch.cat([x_kv, torch.zeros([1, c]).to(x_kv.device)], dim=0)
- T = 1000
-
- torch.cuda.synchronize()
- start_time = time.time()
- for _ in range(T):
- grouping(idx=knn_idx, feat=x_kv, xyz=p, new_xyz=p, with_xyz=False)
- # grouping_idx = feat[knn_idx.view(-1).long(), :].view(
- # m, nsample, c
- # ) # (m, num_sample, c)
- torch.cuda.synchronize()
- end_time = time.time()
- # print(f"Grouping via indexing: {(end_time - start_time) / T * 1000:.2f} ms")
- print(f"grouping pytorch: {(end_time - start_time) / T * 1000:.2f} ms")
-
- torch.cuda.synchronize()
- start_time = time.time()
- for _ in range(T):
- grouping2(x_kv, knn_idx)
- # grouping_embed = torch.nn.functional.embedding(knn_idx, feat) # [m,num_sample,c]
- torch.cuda.synchronize()
- end_time = time.time()
- # print(f"Grouping via embedding: {(end_time - start_time) / T * 1000:.2f} ms")
- print(f"grouping cuda: {(end_time - start_time) / T * 1000:.2f} ms")
-
-if __name__ == '__main__':
- # test_fps()
- # test_knn()
- # test_mlp()
- # test_mv_block()
- # test_cross_attn()
- # test_faiss_knn()
- # test_knn_query_and_group()
- test_grouping()
-
diff --git a/optgs/model/encoder/unimatch/__init__.py b/optgs/model/encoder/unimatch/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/optgs/model/encoder/unimatch/dinov2/__init__.py b/optgs/model/encoder/unimatch/dinov2/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/optgs/model/encoder/unimatch/geometry.py b/optgs/model/encoder/unimatch/geometry.py
deleted file mode 100644
index d28984173fd998f444a5f9e832674a5b1c505eef..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/unimatch/geometry.py
+++ /dev/null
@@ -1,306 +0,0 @@
-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)) # [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].repeat(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 generate_window_grid(h_min, h_max, w_min, w_max, len_h, len_w, device=None):
- assert device is not None
-
- x, y = torch.meshgrid([torch.linspace(w_min, w_max, len_w, device=device),
- torch.linspace(h_min, h_max, len_h, device=device)],
- )
- grid = torch.stack((x, y), -1).transpose(0, 1).float() # [H, W, 2]
-
- return grid
-
-
-def normalize_coords(coords, h, w):
- # coords: [B, H, W, 2]
- c = torch.Tensor([(w - 1) / 2., (h - 1) / 2.]).float().to(coords.device)
- return (coords - c) / c # [-1, 1]
-
-
-def bilinear_sample(img, sample_coords, mode='bilinear', padding_mode='zeros', return_mask=False):
- # img: [B, C, H, W]
- # sample_coords: [B, 2, H, W] in image scale
- if sample_coords.size(1) != 2: # [B, H, W, 2]
- sample_coords = sample_coords.permute(0, 3, 1, 2)
-
- b, _, h, w = sample_coords.shape
-
- # Normalize to [-1, 1]
- x_grid = 2 * sample_coords[:, 0] / (w - 1) - 1
- y_grid = 2 * sample_coords[:, 1] / (h - 1) - 1
-
- grid = torch.stack([x_grid, y_grid], dim=-1) # [B, H, W, 2]
-
- img = F.grid_sample(img, grid, mode=mode, padding_mode=padding_mode, align_corners=True)
-
- if return_mask:
- mask = (x_grid >= -1) & (y_grid >= -1) & (x_grid <= 1) & (y_grid <= 1) # [B, H, W]
-
- return img, mask
-
- return img
-
-
-def flow_warp(feature, flow, mask=False, padding_mode='zeros'):
- b, c, h, w = feature.size()
- assert flow.size(1) == 2
-
- grid = coords_grid(b, h, w).to(flow.device) + flow # [B, 2, H, W]
-
- return bilinear_sample(feature, grid, padding_mode=padding_mode,
- return_mask=mask)
-
-
-def forward_backward_consistency_check(fwd_flow, bwd_flow,
- alpha=0.01,
- beta=0.5
- ):
- # fwd_flow, bwd_flow: [B, 2, H, W]
- # alpha and beta values are following UnFlow (https://arxiv.org/abs/1711.07837)
- assert fwd_flow.dim() == 4 and bwd_flow.dim() == 4
- assert fwd_flow.size(1) == 2 and bwd_flow.size(1) == 2
- flow_mag = torch.norm(fwd_flow, dim=1) + torch.norm(bwd_flow, dim=1) # [B, H, W]
-
- warped_bwd_flow = flow_warp(bwd_flow, fwd_flow) # [B, 2, H, W]
- warped_fwd_flow = flow_warp(fwd_flow, bwd_flow) # [B, 2, H, W]
-
- diff_fwd = torch.norm(fwd_flow + warped_bwd_flow, dim=1) # [B, H, W]
- diff_bwd = torch.norm(bwd_flow + warped_fwd_flow, dim=1)
-
- threshold = alpha * flow_mag + beta
-
- fwd_occ = (diff_fwd > threshold).float() # [B, H, W]
- bwd_occ = (diff_bwd > threshold).float()
-
- return fwd_occ, bwd_occ
-
-
-def back_project(depth, intrinsics):
- # Back project 2D pixel coords to 3D points
- # depth: [B, H, W]
- # intrinsics: [B, 3, 3]
- b, h, w = depth.shape
- grid = coords_grid(b, h, w, homogeneous=True, device=depth.device) # [B, 3, H, W]
-
- intrinsics_inv = torch.inverse(intrinsics) # [B, 3, 3]
-
- points = intrinsics_inv.bmm(grid.view(b, 3, -1)).view(b, 3, h, w) * depth.unsqueeze(1) # [B, 3, H, W]
-
- return points
-
-
-def camera_transform(points_ref, extrinsics_ref=None, extrinsics_tgt=None, extrinsics_rel=None):
- # Transform 3D points from reference camera to target camera
- # points_ref: [B, 3, H, W]
- # extrinsics_ref: [B, 4, 4]
- # extrinsics_tgt: [B, 4, 4]
- # extrinsics_rel: [B, 4, 4], relative pose transform
- b, _, h, w = points_ref.shape
-
- if extrinsics_rel is None:
- extrinsics_rel = torch.bmm(extrinsics_tgt, torch.inverse(extrinsics_ref)) # [B, 4, 4]
-
- points_tgt = torch.bmm(extrinsics_rel[:, :3, :3],
- points_ref.view(b, 3, -1)) + extrinsics_rel[:, :3, -1:] # [B, 3, H*W]
-
- points_tgt = points_tgt.view(b, 3, h, w) # [B, 3, H, W]
-
- return points_tgt
-
-
-def reproject(points_tgt, intrinsics, return_mask=False):
- # reproject to target view
- # points_tgt: [B, 3, H, W]
- # intrinsics: [B, 3, 3]
-
- b, _, h, w = points_tgt.shape
-
- proj_points = torch.bmm(intrinsics, points_tgt.view(b, 3, -1)).view(b, 3, h, w) # [B, 3, H, W]
-
- X = proj_points[:, 0]
- Y = proj_points[:, 1]
- Z = proj_points[:, 2].clamp(min=1e-3)
-
- pixel_coords = torch.stack([X / Z, Y / Z], dim=1).view(b, 2, h, w) # [B, 2, H, W] in image scale
-
- if return_mask:
- # valid mask in pixel space
- mask = (pixel_coords[:, 0] >= 0) & (pixel_coords[:, 0] <= (w - 1)) & (
- pixel_coords[:, 1] >= 0) & (pixel_coords[:, 1] <= (h - 1)) # [B, H, W]
-
- return pixel_coords, mask
-
- return pixel_coords
-
-
-def reproject_coords(depth_ref, intrinsics, extrinsics_ref=None, extrinsics_tgt=None, extrinsics_rel=None,
- return_mask=False):
- # Compute reprojection sample coords
- points_ref = back_project(depth_ref, intrinsics) # [B, 3, H, W]
- points_tgt = camera_transform(points_ref, extrinsics_ref, extrinsics_tgt, extrinsics_rel=extrinsics_rel)
-
- if return_mask:
- reproj_coords, mask = reproject(points_tgt, intrinsics,
- return_mask=return_mask) # [B, 2, H, W] in image scale
-
- return reproj_coords, mask
-
- reproj_coords = reproject(points_tgt, intrinsics,
- return_mask=return_mask) # [B, 2, H, W] in image scale
-
- return reproj_coords
-
-
-def compute_flow_with_depth_pose(depth_ref, intrinsics,
- extrinsics_ref=None, extrinsics_tgt=None, extrinsics_rel=None,
- return_mask=False):
- b, h, w = depth_ref.shape
- coords_init = coords_grid(b, h, w, device=depth_ref.device) # [B, 2, H, W]
-
- if return_mask:
- reproj_coords, mask = reproject_coords(depth_ref, intrinsics, extrinsics_ref, extrinsics_tgt,
- extrinsics_rel=extrinsics_rel,
- return_mask=return_mask) # [B, 2, H, W]
- rigid_flow = reproj_coords - coords_init
-
- return rigid_flow, mask
-
- reproj_coords = reproject_coords(depth_ref, intrinsics, extrinsics_ref, extrinsics_tgt,
- extrinsics_rel=extrinsics_rel,
- return_mask=return_mask) # [B, 2, H, W]
-
- rigid_flow = reproj_coords - coords_init
-
- return rigid_flow
-
-
-def forward_backward_consistency_check(fwd_flow, bwd_flow,
- alpha=0.01,
- beta=0.5,
- return_flow_diff=False,
- ):
- # fwd_flow, bwd_flow: [B, 2, H, W]
- # alpha and beta values are following UnFlow (https://arxiv.org/abs/1711.07837)
- assert fwd_flow.dim() == 4 and bwd_flow.dim() == 4
- assert fwd_flow.size(1) == 2 and bwd_flow.size(1) == 2
- flow_mag = torch.norm(fwd_flow, dim=1) + torch.norm(bwd_flow, dim=1) # [B, H, W]
-
- warped_bwd_flow = flow_warp(bwd_flow, fwd_flow) # [B, 2, H, W]
- warped_fwd_flow = flow_warp(fwd_flow, bwd_flow) # [B, 2, H, W]
-
- diff_fwd = torch.norm(fwd_flow + warped_bwd_flow, dim=1) # [B, H, W]
- diff_bwd = torch.norm(bwd_flow + warped_fwd_flow, dim=1)
-
- if return_flow_diff:
- return diff_fwd, diff_bwd
-
- threshold = alpha * flow_mag + beta
-
- fwd_occ = (diff_fwd > threshold).float() # [B, H, W]
- bwd_occ = (diff_bwd > threshold).float()
-
- return fwd_occ, bwd_occ
-
-
-def warp_with_depth_pose(feature1, intrinsics, pose, depth,
- padding_mode='zeros',
- return_rigid_flow=False,
- return_mask=False,
- ):
- assert depth.dim() == 3 # [B, H, W]
- sample_coords = reproject_coords(depth,
- intrinsics,
- extrinsics_rel=pose,
- return_mask=return_mask,
- ) # [B, 2, H, W]
-
- if return_mask:
- sample_coords, mask = sample_coords
-
- sample_coords = sample_coords.permute(0, 2, 3, 1) # [B, H, W, 2]
-
- warped_feature1 = bilinear_sample(feature1, sample_coords,
- padding_mode=padding_mode) # [B, C, H, W]
-
- if return_mask:
- return warped_feature1, mask
-
- if return_rigid_flow:
- b, h, w = depth.size()
- coords_init = coords_grid(b, h, w, device=depth.device) # [B, 2, H, W]
- rigid_flow = sample_coords.permute(0, 3, 1, 2) - coords_init
-
- return warped_feature1, rigid_flow
-
- return warped_feature1
-
-
-def warp_with_pose_depth_candidates(feature1, intrinsics, pose, depth,
- padding_mode='zeros',
- rigid_flow_to_subtract=None,
- ):
- # pixel-specific depth candidates, useful for refinement
- # feature1: [B, C, H, W]
- # intrinsics: [B, 3, 3]
- # pose: [B, 4, 4]
- # depth: [B, D, H, W]
- assert intrinsics.size(1) == 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)
-
- # stop gradient
- with torch.no_grad():
- # 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.inverse(intrinsics).bmm(grid.view(b, 3, -1)) # [B, 3, H*W]
- points = torch.bmm(pose[:, :3, :3], points).unsqueeze(2).repeat(
- 1, 1, d, 1) * depth.view(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(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=MIN_DEPTH) # [B, 2, D, H*W]
-
- if rigid_flow_to_subtract is not None:
- assert rigid_flow_to_subtract.dim() == 4 # [B, 2, H, W]
- assert rigid_flow_to_subtract.size(1) == 2
-
- pixel_coords = pixel_coords - rigid_flow_to_subtract.view(b, 2, h * w).unsqueeze(2)
-
- # 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]
-
- # sample features
- warped_feature = F.grid_sample(feature1, grid.view(b, d * h, w, 2), mode='bilinear',
- padding_mode=padding_mode,
- align_corners=True).view(b, c, d, h, w) # [B, C, D, H, W]
-
- return warped_feature
-
\ No newline at end of file
diff --git a/optgs/model/encoder/unimatch/ldm_unet/__init__.py b/optgs/model/encoder/unimatch/ldm_unet/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/optgs/model/encoder/visualization/__init__.py b/optgs/model/encoder/visualization/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/optgs/model/encoder/visualization/encoder_visualizer.py b/optgs/model/encoder/visualization/encoder_visualizer.py
deleted file mode 100644
index 5af485b01faed25fbe3557f650d0c3809b216b6d..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/visualization/encoder_visualizer.py
+++ /dev/null
@@ -1,25 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import Generic, TypeVar
-
-from jaxtyping import Float
-from torch import Tensor
-
-T_cfg = TypeVar("T_cfg")
-T_encoder = TypeVar("T_encoder")
-
-
-class EncoderVisualizer(ABC, Generic[T_cfg, T_encoder]):
- cfg: T_cfg
- encoder: T_encoder
-
- def __init__(self, cfg: T_cfg, encoder: T_encoder) -> None:
- self.cfg = cfg
- self.encoder = encoder
-
- @abstractmethod
- def visualize(
- self,
- context: dict,
- global_step: int,
- ) -> dict[str, Float[Tensor, "3 _ _"]]:
- pass
diff --git a/optgs/model/encoder/visualization/encoder_visualizer_depthsplat.py b/optgs/model/encoder/visualization/encoder_visualizer_depthsplat.py
deleted file mode 100644
index 96b043e34b75337fdc2c06bda8d848fd3d74db91..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/visualization/encoder_visualizer_depthsplat.py
+++ /dev/null
@@ -1,527 +0,0 @@
-from pathlib import Path
-from random import randrange
-from typing import Optional
-
-import numpy as np
-import torch
-import wandb
-from einops import rearrange, reduce, repeat
-from jaxtyping import Bool, Float
-from torch import Tensor
-
-from ....dataset.data_types import BatchedViews
-from ....misc.heterogeneous_pairings import generate_heterogeneous_index
-from ....visualization.annotation import add_label
-from ....visualization.color_map import apply_color_map, apply_color_map_to_image
-from ....visualization.colors import get_distinct_color
-from ....visualization.drawing.lines import draw_lines
-from ....visualization.drawing.points import draw_points
-from ....visualization.layout import add_border, hcat, vcat
-# from ...ply_export import export_ply
-from ..encoder_depthsplat import EncoderDepthSplat
-# from ..epipolar.epipolar_sampler import EpipolarSampling
-from .encoder_visualizer import EncoderVisualizer
-from .encoder_visualizer_depthsplat_cfg import EncoderVisualizerDepthSplatCfg
-
-
-def box(
- image: Float[Tensor, "3 height width"],
-) -> Float[Tensor, "3 new_height new_width"]:
- return add_border(add_border(image), 1, 0)
-
-
-class EncoderVisualizerDepthSplat(
- EncoderVisualizer[EncoderVisualizerDepthSplatCfg, EncoderDepthSplat]
-):
- def visualize(
- self,
- context: BatchedViews,
- global_step: int,
- ) -> dict[str, Float[Tensor, "3 _ _"]]:
- # Short-circuit execution when using mvsplat.
- return {}
-
- visualization_dump = {}
-
- softmax_weights = []
-
- def hook(module, input, output):
- softmax_weights.append(output)
-
- # Register hooks to grab attention.
- handles = [
- layer[0].fn.attend.register_forward_hook(hook)
- for layer in self.encoder.epipolar_transformer.transformer.layers
- ]
-
- result = self.encoder.forward(
- context,
- global_step,
- visualization_dump=visualization_dump,
- deterministic=True,
- )
-
- # De-register hooks.
- for handle in handles:
- handle.remove()
-
- softmax_weights = torch.stack(softmax_weights)
-
- # Generate high-resolution context images that can be drawn on.
- context_images = context["image"]
- _, _, _, h, w = context_images.shape
- length = min(h, w)
- min_resolution = self.cfg.min_resolution
- scale_multiplier = (min_resolution + length - 1) // length
- if scale_multiplier > 1:
- context_images = repeat(
- context_images,
- "b v c h w -> b v c (h rh) (w rw)",
- rh=scale_multiplier,
- rw=scale_multiplier,
- )
-
- # This is kind of hacky for now, since we're using it for short experiments.
- if self.cfg.export_ply and wandb.run is not None:
- name = wandb.run._name.split(" ")[0]
- ply_path = Path(f"outputs/gaussians/{name}/{global_step:0>6}.ply")
- export_ply(
- context["extrinsics"][0, 0],
- result.means[0],
- visualization_dump["scales"][0],
- visualization_dump["rotations"][0],
- result.harmonics[0],
- result.opacities[0],
- ply_path,
- )
-
- return {
- "attention": self.visualize_attention(
- context_images,
- visualization_dump["sampling"],
- softmax_weights,
- ),
- "epipolar_samples": self.visualize_epipolar_samples(
- context_images,
- visualization_dump["sampling"],
- ),
- "epipolar_color_samples": self.visualize_epipolar_color_samples(
- context_images,
- context,
- ),
- "gaussians": self.visualize_gaussians(
- context["image"],
- result.opacities,
- result.covariances,
- result.harmonics[..., 0], # Just visualize DC component.
- ),
- "overlaps": self.visualize_overlaps(
- context["image"],
- visualization_dump["sampling"],
- visualization_dump.get("is_monocular", None),
- ),
- "depth": self.visualize_depth(
- context,
- visualization_dump["depth"],
- ),
- }
-
- def visualize_attention(
- self,
- context_images: Float[Tensor, "batch view 3 height width"],
- sampling: None,
- attention: Float[Tensor, "layer bvr head 1 sample"],
- ) -> Float[Tensor, "3 vis_height vis_width"]:
- device = context_images.device
-
- # Pick a random batch element, view, and other view.
- b, v, ov, r, s, _ = sampling.xy_sample.shape
- rb = randrange(b)
- rv = randrange(v)
- rov = randrange(ov)
- num_samples = self.cfg.num_samples
- rr = np.random.choice(r, num_samples, replace=False)
- rr = torch.tensor(rr, dtype=torch.int64, device=device)
-
- # Visualize the rays in the ray view.
- ray_view = draw_points(
- context_images[rb, rv],
- sampling.xy_ray[rb, rv, rr],
- 0,
- radius=4,
- x_range=(0, 1),
- y_range=(0, 1),
- )
- ray_view = draw_points(
- ray_view,
- sampling.xy_ray[rb, rv, rr],
- [get_distinct_color(i) for i, _ in enumerate(rr)],
- radius=3,
- x_range=(0, 1),
- y_range=(0, 1),
- )
-
- # Visualize attention in the sample view.
- attention = rearrange(
- attention, "l (b v r) hd () s -> l b v r hd s", b=b, v=v, r=r
- )
- attention = attention[:, rb, rv, rr, :, :]
- num_layers, _, hd, _ = attention.shape
-
- vis = []
- for il in range(num_layers):
- vis_layer = []
- for ihd in range(hd):
- # Create colors according to attention.
- color = [get_distinct_color(i) for i, _ in enumerate(rr)]
- color = torch.tensor(color, device=attention.device)
- color = rearrange(color, "r c -> r () c")
- attn = rearrange(attention[il, :, ihd], "r s -> r s ()")
- color = rearrange(attn * color, "r s c -> (r s ) c")
-
- # Draw the alternating bucket lines.
- vis_layer_head = draw_lines(
- context_images[rb, self.encoder.sampler.index_v[rv, rov]],
- rearrange(
- sampling.xy_sample_near[rb, rv, rov, rr], "r s xy -> (r s) xy"
- ),
- rearrange(
- sampling.xy_sample_far[rb, rv, rov, rr], "r s xy -> (r s) xy"
- ),
- color,
- 3,
- cap="butt",
- x_range=(0, 1),
- y_range=(0, 1),
- )
- vis_layer.append(vis_layer_head)
- vis.append(add_label(vcat(*vis_layer), f"Layer {il}"))
- vis = add_label(add_border(add_border(hcat(*vis)), 1, 0), "Keys & Values")
- vis = add_border(hcat(add_label(ray_view, "ray_view"), vis, align="top"))
- return vis
-
- def visualize_depth(
- self,
- context: BatchedViews,
- multi_depth: Float[Tensor, "batch view height width surface spp"],
- ) -> Float[Tensor, "3 vis_width vis_height"]:
- multi_vis = []
- *_, srf, _ = multi_depth.shape
- for i in range(srf):
- depth = multi_depth[..., i, :]
- depth = depth.mean(dim=-1)
-
- # Compute relative depth and disparity.
- near = rearrange(context["near"], "b v -> b v () ()")
- far = rearrange(context["far"], "b v -> b v () ()")
- relative_depth = (depth - near) / (far - near)
- relative_disparity = 1 - (1 / depth - 1 / far) / (1 / near - 1 / far)
-
- relative_depth = apply_color_map_to_image(relative_depth, "turbo")
- relative_depth = vcat(*[hcat(*x) for x in relative_depth])
- relative_depth = add_label(relative_depth, "Depth")
- relative_disparity = apply_color_map_to_image(relative_disparity, "turbo")
- relative_disparity = vcat(*[hcat(*x) for x in relative_disparity])
- relative_disparity = add_label(relative_disparity, "Disparity")
- multi_vis.append(add_border(hcat(relative_depth, relative_disparity)))
-
- return add_border(vcat(*multi_vis))
-
- def visualize_overlaps(
- self,
- context_images: Float[Tensor, "batch view 3 height width"],
- sampling: None,
- is_monocular: Optional[Bool[Tensor, "batch view height width"]] = None,
- ) -> Float[Tensor, "3 vis_width vis_height"]:
- device = context_images.device
- b, v, _, h, w = context_images.shape
- green = torch.tensor([0.235, 0.706, 0.294], device=device)[..., None, None]
- rb = randrange(b)
- valid = sampling.valid[rb].float()
- ds = self.encoder.cfg.epipolar_transformer.downscale
- valid = repeat(
- valid,
- "v ov (h w) -> v ov c (h rh) (w rw)",
- c=3,
- h=h // ds,
- w=w // ds,
- rh=ds,
- rw=ds,
- )
-
- if is_monocular is not None:
- is_monocular = is_monocular[rb].float()
- is_monocular = repeat(is_monocular, "v h w -> v c h w", c=3, h=h, w=w)
-
- # Select context images in grid.
- context_images = context_images[rb]
- index, _ = generate_heterogeneous_index(v)
- valid = valid * (green + context_images[index]) / 2
-
- vis = vcat(*(hcat(im, hcat(*v)) for im, v in zip(context_images, valid)))
- vis = add_label(vis, "Context Overlaps")
-
- if is_monocular is not None:
- vis = hcat(vis, add_label(vcat(*is_monocular), "Monocular?"))
-
- return add_border(vis)
-
- def visualize_gaussians(
- self,
- context_images: Float[Tensor, "batch view 3 height width"],
- opacities: Float[Tensor, "batch vrspp"],
- covariances: Float[Tensor, "batch vrspp 3 3"],
- colors: Float[Tensor, "batch vrspp 3"],
- ) -> Float[Tensor, "3 vis_height vis_width"]:
- b, v, _, h, w = context_images.shape
- rb = randrange(b)
- context_images = context_images[rb]
- opacities = repeat(
- opacities[rb], "(v h w spp) -> spp v c h w", v=v, c=3, h=h, w=w
- )
- colors = rearrange(colors[rb], "(v h w spp) c -> spp v c h w", v=v, h=h, w=w)
-
- # Color-map Gaussian covariawnces.
- det = covariances[rb].det()
- det = apply_color_map(det / det.max(), "inferno")
- det = rearrange(det, "(v h w spp) c -> spp v c h w", v=v, h=h, w=w)
-
- return add_border(
- hcat(
- add_label(box(hcat(*context_images)), "Context"),
- add_label(box(vcat(*[hcat(*x) for x in opacities])), "Opacities"),
- add_label(
- box(vcat(*[hcat(*x) for x in (colors * opacities)])), "Colors"
- ),
- add_label(box(vcat(*[hcat(*x) for x in colors])), "Colors (Raw)"),
- add_label(box(vcat(*[hcat(*x) for x in det])), "Determinant"),
- )
- )
-
- def visualize_probabilities(
- self,
- context_images: Float[Tensor, "batch view 3 height width"],
- sampling: None,
- pdf: Float[Tensor, "batch view ray sample"],
- ) -> Float[Tensor, "3 vis_height vis_width"]:
- device = context_images.device
-
- # Pick a random batch element, view, and other view.
- b, v, ov, r, _, _ = sampling.xy_sample.shape
- rb = randrange(b)
- rv = randrange(v)
- rov = randrange(ov)
- num_samples = self.cfg.num_samples
- rr = np.random.choice(r, num_samples, replace=False)
- rr = torch.tensor(rr, dtype=torch.int64, device=device)
- colors = [get_distinct_color(i) for i, _ in enumerate(rr)]
- colors = torch.tensor(colors, dtype=torch.float32, device=device)
-
- # Visualize the rays in the ray view.
- ray_view = draw_points(
- context_images[rb, rv],
- sampling.xy_ray[rb, rv, rr],
- 0,
- radius=4,
- x_range=(0, 1),
- y_range=(0, 1),
- )
- ray_view = draw_points(
- ray_view,
- sampling.xy_ray[rb, rv, rr],
- colors,
- radius=3,
- x_range=(0, 1),
- y_range=(0, 1),
- )
-
- # Visualize probabilities in the sample view.
- pdf = pdf[rb, rv, rr]
- pdf = rearrange(pdf, "r s -> r s ()")
- colors = rearrange(colors, "r c -> r () c")
- sample_view = draw_lines(
- context_images[rb, self.encoder.sampler.index_v[rv, rov]],
- rearrange(sampling.xy_sample_near[rb, rv, rov, rr], "r s xy -> (r s) xy"),
- rearrange(sampling.xy_sample_far[rb, rv, rov, rr], "r s xy -> (r s) xy"),
- rearrange(pdf * colors, "r s c -> (r s) c"),
- 6,
- cap="butt",
- x_range=(0, 1),
- y_range=(0, 1),
- )
-
- # Visualize rescaled probabilities in the sample view.
- pdf_magnified = pdf / reduce(pdf, "r s () -> r () ()", "max")
- sample_view_magnified = draw_lines(
- context_images[rb, self.encoder.sampler.index_v[rv, rov]],
- rearrange(sampling.xy_sample_near[rb, rv, rov, rr], "r s xy -> (r s) xy"),
- rearrange(sampling.xy_sample_far[rb, rv, rov, rr], "r s xy -> (r s) xy"),
- rearrange(pdf_magnified * colors, "r s c -> (r s) c"),
- 6,
- cap="butt",
- x_range=(0, 1),
- y_range=(0, 1),
- )
-
- return add_border(
- hcat(
- add_label(ray_view, "Rays"),
- add_label(sample_view, "Samples"),
- add_label(sample_view_magnified, "Samples (Magnified PDF)"),
- )
- )
-
- def visualize_epipolar_samples(
- self,
- context_images: Float[Tensor, "batch view 3 height width"],
- sampling: None,
- ) -> Float[Tensor, "3 vis_height vis_width"]:
- device = context_images.device
-
- # Pick a random batch element, view, and other view.
- b, v, ov, r, s, _ = sampling.xy_sample.shape
- rb = randrange(b)
- rv = randrange(v)
- rov = randrange(ov)
- num_samples = self.cfg.num_samples
- rr = np.random.choice(r, num_samples, replace=False)
- rr = torch.tensor(rr, dtype=torch.int64, device=device)
-
- # Visualize the rays in the ray view.
- ray_view = draw_points(
- context_images[rb, rv],
- sampling.xy_ray[rb, rv, rr],
- 0,
- radius=4,
- x_range=(0, 1),
- y_range=(0, 1),
- )
- ray_view = draw_points(
- ray_view,
- sampling.xy_ray[rb, rv, rr],
- [get_distinct_color(i) for i, _ in enumerate(rr)],
- radius=3,
- x_range=(0, 1),
- y_range=(0, 1),
- )
-
- # Visualize the samples and epipolar lines in the sample view.
- # First, draw the epipolar line in black.
- sample_view = draw_lines(
- context_images[rb, self.encoder.sampler.index_v[rv, rov]],
- sampling.xy_sample_near[rb, rv, rov, rr, 0],
- sampling.xy_sample_far[rb, rv, rov, rr, -1],
- 0,
- 5,
- cap="butt",
- x_range=(0, 1),
- y_range=(0, 1),
- )
-
- # Create an alternating line color for the buckets.
- color = repeat(
- torch.tensor([0, 1], device=device),
- "ab -> r (s ab) c",
- r=len(rr),
- s=(s + 1) // 2,
- c=3,
- )
- color = rearrange(color[:, :s], "r s c -> (r s) c")
-
- # Draw the alternating bucket lines.
- sample_view = draw_lines(
- sample_view,
- rearrange(sampling.xy_sample_near[rb, rv, rov, rr], "r s xy -> (r s) xy"),
- rearrange(sampling.xy_sample_far[rb, rv, rov, rr], "r s xy -> (r s) xy"),
- color,
- 3,
- cap="butt",
- x_range=(0, 1),
- y_range=(0, 1),
- )
-
- # Draw the sample points.
- sample_view = draw_points(
- sample_view,
- rearrange(sampling.xy_sample[rb, rv, rov, rr], "r s xy -> (r s) xy"),
- 0,
- radius=4,
- x_range=(0, 1),
- y_range=(0, 1),
- )
- sample_view = draw_points(
- sample_view,
- rearrange(sampling.xy_sample[rb, rv, rov, rr], "r s xy -> (r s) xy"),
- [get_distinct_color(i // s) for i in range(s * len(rr))],
- radius=3,
- x_range=(0, 1),
- y_range=(0, 1),
- )
-
- return add_border(
- hcat(add_label(ray_view, "Ray View"), add_label(sample_view, "Sample View"))
- )
-
- def visualize_epipolar_color_samples(
- self,
- context_images: Float[Tensor, "batch view 3 height width"],
- context: BatchedViews,
- ) -> Float[Tensor, "3 vis_height vis_width"]:
- device = context_images.device
-
- sampling = self.encoder.sampler(
- context["image"],
- context["extrinsics"],
- context["intrinsics"],
- context["near"],
- context["far"],
- )
-
- # Pick a random batch element, view, and other view.
- b, v, ov, r, s, _ = sampling.xy_sample.shape
- rb = randrange(b)
- rv = randrange(v)
- rov = randrange(ov)
- num_samples = self.cfg.num_samples
- rr = np.random.choice(r, num_samples, replace=False)
- rr = torch.tensor(rr, dtype=torch.int64, device=device)
-
- # Visualize the rays in the ray view.
- ray_view = draw_points(
- context_images[rb, rv],
- sampling.xy_ray[rb, rv, rr],
- 0,
- radius=4,
- x_range=(0, 1),
- y_range=(0, 1),
- )
- ray_view = draw_points(
- ray_view,
- sampling.xy_ray[rb, rv, rr],
- [get_distinct_color(i) for i, _ in enumerate(rr)],
- radius=3,
- x_range=(0, 1),
- y_range=(0, 1),
- )
-
- # Visualize the samples and in the sample view.
- sample_view = draw_points(
- context_images[rb, self.encoder.sampler.index_v[rv, rov]],
- rearrange(sampling.xy_sample[rb, rv, rov, rr], "r s xy -> (r s) xy"),
- [get_distinct_color(i // s) for i in range(s * len(rr))],
- radius=4,
- x_range=(0, 1),
- y_range=(0, 1),
- )
- sample_view = draw_points(
- sample_view,
- rearrange(sampling.xy_sample[rb, rv, rov, rr], "r s xy -> (r s) xy"),
- rearrange(sampling.features[rb, rv, rov, rr], "r s c -> (r s) c"),
- radius=3,
- x_range=(0, 1),
- y_range=(0, 1),
- )
-
- return add_border(
- hcat(add_label(ray_view, "Ray View"), add_label(sample_view, "Sample View"))
- )
diff --git a/optgs/model/encoder/visualization/encoder_visualizer_depthsplat_cfg.py b/optgs/model/encoder/visualization/encoder_visualizer_depthsplat_cfg.py
deleted file mode 100644
index 2f0831572770802d97e23f7cc1fa3daad0cfef41..0000000000000000000000000000000000000000
--- a/optgs/model/encoder/visualization/encoder_visualizer_depthsplat_cfg.py
+++ /dev/null
@@ -1,10 +0,0 @@
-from dataclasses import dataclass
-
-# This is in a separate file to avoid circular imports.
-
-
-@dataclass
-class EncoderVisualizerDepthSplatCfg:
- num_samples: int
- min_resolution: int
- export_ply: bool
diff --git a/optgs/model/ply_export.py b/optgs/model/ply_export.py
index 36ea93af81d80fc39fabdbfb1f6bd2a73a6a9940..5c9aa24365115c0487d2979c63977344f6a0abc4 100644
--- a/optgs/model/ply_export.py
+++ b/optgs/model/ply_export.py
@@ -62,7 +62,6 @@ def export_ply(
def save_gaussian_ply(
gaussians: Gaussians | GaussiansModule,
save_path,
- save_all_gaussians=True, # no trim
):
"""
Save Gaussians to a .ply file for visualization.
@@ -72,9 +71,6 @@ def save_gaussian_ply(
"""
- if not save_all_gaussians:
- raise NotImplementedError("Not implemented yet.")
-
if isinstance(gaussians, GaussiansModule):
# no batch dimension
@@ -138,9 +134,10 @@ def load_gaussians_ply(path, max_sh_degree=3) -> Gaussians:
#
if len(extra_f_names) == 0:
- # loaded ply has no SH coefficients
- # TODO: does this mean that features_dc probably encodes RGB which needs to be converted to SH0?
- # all other features are zero
+ # SH degree 0: only the DC band is present, no higher-order coefficients.
+ # f_dc still holds the SH DC term (SH0), not raw RGB -- the 3DGS .ply
+ # convention stores SH0, and the colour is recovered as rgb = 0.5 + C0 * SH0.
+ # So f_dc is read directly into harmonics[..., 0]; the remaining bands are zero.
print("Loaded PLY has no SH coefficients, only DC features.")
features_extra = np.zeros((xyz.shape[0], 3, (max_sh_degree + 1) ** 2 - 1))
diff --git a/optgs/model/types.py b/optgs/model/types.py
index f26de4c957cad3e5262a6c68c72ab0a44e7467f4..79a6b560024520c5839210485564cca4464ad165 100644
--- a/optgs/model/types.py
+++ b/optgs/model/types.py
@@ -1,6 +1,14 @@
+"""Core `Gaussians` dataclass shared across the whole pipeline.
+
+`Gaussians` holds one scene's primitives — means, SH harmonics, opacities, scales, rotations — plus the
+optional per-Gaussian fields the optimizer attaches during refinement (gradients, deltas, visibility,
+optimization state) and a `nr_valid` count for padded scenes. Helpers (`to`, `clone`, `__getitem__`)
+move it across devices and batch dimensions. Initializers produce it, the optimizer refines it, and the
+decoder renders it.
+"""
+
from dataclasses import dataclass, fields
-import torch
from jaxtyping import Float, Bool, Int64, BFloat16
from torch import Tensor
@@ -73,6 +81,15 @@ class Gaussians:
new_tensors[field.name] = None
return Gaussians(**new_tensors)
+ def select_valid(self) -> "Gaussians":
+ """Return the subset of valid (non-padding) gaussians selected by `sel`.
+
+ Returns self unchanged when `sel` is None.
+ """
+ if self.sel is None:
+ return self
+ return self[:, self.sel]
+
def sample_subset(self, sampled_indices) -> "Gaussians":
""" Randomly sample a subset of gaussians. """
total_gaussians = self.means.shape[1]
diff --git a/optgs/paths.py b/optgs/paths.py
index 89f1d77b6d5807d056a6be4182f7cc1f22331aea..a1997a90f23fdba0b0cd20dff8f02443f2625950 100644
--- a/optgs/paths.py
+++ b/optgs/paths.py
@@ -63,8 +63,15 @@ if _IS_SOURCE_CHECKOUT:
except OSError:
pass
-if sys.gettrace() is not None:
- print("Running in debug mode")
- DEBUG = True
+# DEBUG enables debug-only code paths (shorter runs, extra logging, "DEBUG"
+# run names). Force it on or off with OPTGS_DEBUG (1/0, true/false); otherwise
+# it auto-enables under any debugger that installs a trace function (PyCharm,
+# pdb, debugpy).
+_debug_env = os.environ.get("OPTGS_DEBUG")
+if _debug_env is not None:
+ DEBUG = _debug_env.strip().lower() in ("1", "true", "yes", "on")
else:
- DEBUG = False
+ DEBUG = sys.gettrace() is not None
+
+if DEBUG:
+ print("Running in debug mode")
diff --git a/optgs/scene_trainer/adc/base.py b/optgs/scene_trainer/adc/base.py
index 69f075a4c3252b23ea1cf0ad3de2e36c3ae83ac1..0d4cda5119c603b6f771fca995b971f850408587 100644
--- a/optgs/scene_trainer/adc/base.py
+++ b/optgs/scene_trainer/adc/base.py
@@ -1,6 +1,5 @@
from dataclasses import dataclass
-from typing import Any, Literal, Callable
-import torch
+from typing import Literal, Callable
from torch import Tensor
from optgs.scene_trainer.gaussian_module import GaussiansModule
from optgs.model.types import Gaussians
@@ -48,7 +47,7 @@ class BaseStrategyCfg:
relocate_copy_state: bool = False
# MCMC noise: cap on the scales used for the noise covariance (does NOT affect rendered scales).
- # Needed because knn_based saturates clamp_refine_max_scale, producing covariances orders of
+ # Needed because knn_based saturates clamp_max_scale, producing covariances orders of
# magnitude larger than vanilla's Adam-evolved scales. The resulting MCMC noise overflows the
# renderer's tile-binning math and causes silent CUDA OOB. Rule of thumb: ~scene_scale / 5.
noise_scale_cap: float = 1.0
diff --git a/optgs/scene_trainer/adc/mcmc.py b/optgs/scene_trainer/adc/mcmc.py
index 03c9fc12b4a2b59d5a26f2b34d0627e3270781f6..a0019c6a86e2cf7ba9df58a47d40729c001a8935 100644
--- a/optgs/scene_trainer/adc/mcmc.py
+++ b/optgs/scene_trainer/adc/mcmc.py
@@ -1,19 +1,16 @@
-from dataclasses import dataclass
-from typing import Any, Callable, Literal, Optional, Tuple
import math
-import torch
+from dataclasses import dataclass
+from typing import Any, Callable, Optional, Tuple
+
import numpy as np
-from jaxtyping import Float, Bool
-from torch import Tensor
+import torch
import torch.nn.functional as F
+from torch import Tensor
+
from optgs.model.types import Gaussians
-import torch
-from typeguard import value
-from optgs.scene_trainer.gaussian_module import GaussiansModule
-from optgs.scene_trainer.optimizer.layer import AdamInputSmoothing
-import matplotlib.pyplot as plt
from optgs.scene_trainer.adc.base import BaseStrategyCfg, GenericStrategyState
from optgs.scene_trainer.adc.base import _replace_objects, _add_to_objects
+from optgs.scene_trainer.gaussian_module import GaussiansModule
def _make_lazy_cuda_func(name: str) -> Callable:
@@ -25,17 +22,18 @@ def _make_lazy_cuda_func(name: str) -> Callable:
return call_cuda
+
class _QuatScaleToCovarPreci(torch.autograd.Function):
"""Converts quaternions and scales to covariance and precision matrices."""
@staticmethod
def forward(
- ctx,
- quats: Tensor, # [..., 4],
- scales: Tensor, # [..., 3],
- compute_covar: bool = True,
- compute_preci: bool = True,
- triu: bool = False,
+ ctx,
+ quats: Tensor, # [..., 4],
+ scales: Tensor, # [..., 3],
+ compute_covar: bool = True,
+ compute_preci: bool = True,
+ triu: bool = False,
) -> Tuple[Optional[Tensor], Optional[Tensor]]:
covars, precis = _make_lazy_cuda_func("quat_scale_to_covar_preci_fwd")(
quats, scales, compute_covar, compute_preci, triu
@@ -65,6 +63,7 @@ class _QuatScaleToCovarPreci(torch.autograd.Function):
)
return v_quats, v_scales, None, None, None
+
@dataclass
class McmcStrategyState(GenericStrategyState):
# Add MCMC specific state variables here
@@ -76,29 +75,31 @@ class McmcStrategyState(GenericStrategyState):
@classmethod
def initialize(cls, device: torch.device) -> "McmcStrategyState":
-
+
# from gsplat
n_max = 51
binoms = torch.zeros((n_max, n_max))
for n in range(n_max):
for k in range(n + 1):
binoms[n, k] = math.comb(n, k)
-
+
return cls(
binoms=binoms.to(device)
)
+
def update_mcmc_strategy_state(
- adc_state: McmcStrategyState
+ adc_state: McmcStrategyState
) -> None:
"""Updates adc_state in place."""
pass
+
@torch.no_grad()
def inject_noise_to_position(
- gaussians: Gaussians | GaussiansModule,
- scaler: float,
- scale_cap: float = 1.0,
+ gaussians: Gaussians | GaussiansModule,
+ scaler: float,
+ scale_cap: float = 1.0,
):
if isinstance(gaussians, GaussiansModule):
raise NotImplementedError("noise injection not implemented for GaussiansModule")
@@ -118,13 +119,13 @@ def inject_noise_to_position(
# activate
opacities = torch.sigmoid(opacities) # [G]
scales = torch.exp(scales) # [G, 3]
-
+
def _quat_scale_to_covar_preci(
- quats: Tensor, # [..., 4],
- scales: Tensor, # [..., 3],
- compute_covar: bool = True,
- compute_preci: bool = True,
- triu: bool = False,
+ quats: Tensor, # [..., 4],
+ scales: Tensor, # [..., 3],
+ compute_covar: bool = True,
+ compute_preci: bool = True,
+ triu: bool = False,
) -> Tuple[Optional[Tensor], Optional[Tensor]]:
"""Converts quaternions and scales to covariance and precision matrices.
@@ -154,7 +155,7 @@ def inject_noise_to_position(
return covars if compute_covar else None, precis if compute_preci else None
# Cap the scales used for the noise covariance only — does NOT change the rendered Gaussian
- # scales. knn_based's network saturates clamp_refine_max_scale, producing covariances orders
+ # scales. knn_based's network saturates clamp_max_scale, producing covariances orders
# of magnitude larger than vanilla's Adam-evolved scales; the resulting noise overflows the
# renderer's tile-binning math and causes a silent CUDA OOB downstream. See BaseStrategyCfg.
scales_for_noise = scales.clamp(max=scale_cap)
@@ -170,9 +171,9 @@ def inject_noise_to_position(
return 1 / (1 + torch.exp(-k * (x - x0)))
noise = (
- torch.randn_like(means)
- * (op_sigmoid(1 - opacities)).unsqueeze(-1)
- * scaler
+ torch.randn_like(means)
+ * (op_sigmoid(1 - opacities)).unsqueeze(-1)
+ * scaler
)
noise = torch.einsum("bij,bj->bi", covars, noise)
@@ -182,6 +183,7 @@ def inject_noise_to_position(
# replace the original leaf tensor with a requires_grad=False view (created
# inside @torch.no_grad), breaking gradient flow on the next iteration.
+
@torch.no_grad()
def _multinomial_sample(weights: Tensor, n: int, replacement: bool = True) -> Tensor:
"""Sample from a distribution using torch.multinomial or numpy.random.choice.
@@ -200,7 +202,7 @@ def _multinomial_sample(weights: Tensor, n: int, replacement: bool = True) -> Te
"""
num_elements = weights.size(0)
- if num_elements <= 2**24:
+ if num_elements <= 2 ** 24:
# Use torch.multinomial for elements within the limit
return torch.multinomial(weights, n, replacement=replacement)
else:
@@ -215,12 +217,13 @@ def _multinomial_sample(weights: Tensor, n: int, replacement: bool = True) -> Te
# Return the sampled indices on the original device
return sampled_idxs.to(weights.device)
+
@torch.no_grad()
def _compute_relocation(
- opacities: Tensor, # [N]
- scales: Tensor, # [N, 3]
- ratios: Tensor, # [N]
- binoms: Tensor, # [n_max, n_max]
+ opacities: Tensor, # [N]
+ scales: Tensor, # [N, 3]
+ ratios: Tensor, # [N]
+ binoms: Tensor, # [n_max, n_max]
) -> Tuple[Tensor, Tensor]:
"""Compute new Gaussians from a set of old Gaussians.
@@ -259,11 +262,11 @@ def _compute_relocation(
def relocate(
- gaussians: Gaussians | GaussiansModule,
- smoothers: dict[str, Any],
- adc_state: McmcStrategyState,
- min_opacity: float,
- copy_state: bool = False,
+ gaussians: Gaussians | GaussiansModule,
+ smoothers: dict[str, Any],
+ adc_state: McmcStrategyState,
+ min_opacity: float,
+ copy_state: bool = False,
) -> int:
"""Relocates Gaussians based on MCMC strategy.
@@ -275,7 +278,7 @@ def relocate(
Returns:
int: Number of relocated Gaussians.
"""
-
+
if isinstance(gaussians, GaussiansModule):
raise NotImplementedError("noise injection not implemented for GaussiansModule")
@@ -298,13 +301,13 @@ def relocate(
# activate
opacities = torch.sigmoid(opacities) # [G]
scales = torch.exp(scales) # [G, 3]
-
+
dead_mask = opacities <= min_opacity
n_gs = dead_mask.sum().item()
if n_gs > 0:
# Inplace relocate some dead Gaussians to the lives ones.
n_relocated = int(n_gs)
-
+
dead_indices = dead_mask.nonzero(as_tuple=True)[0]
alive_indices = (~dead_mask).nonzero(as_tuple=True)[0]
n = len(dead_indices)
@@ -332,7 +335,7 @@ def relocate(
# replace values (batch dim = 0, Gaussian dim = 1)
gaussians.means[0, dead_indices] = means[sampled_idxs]
- gaussians.scales[0, dead_indices] = new_scales # relocated scale (deactivated if needed)
+ gaussians.scales[0, dead_indices] = new_scales # relocated scale (deactivated if needed)
gaussians.opacities[0, dead_indices] = new_opacities # relocated opacity (deactivated if needed)
gaussians.rotations_unnorm[0, dead_indices] = rotations_unnorm[sampled_idxs]
gaussians.harmonics[0, dead_indices] = harmonics[sampled_idxs]
@@ -344,12 +347,13 @@ def relocate(
return n_relocated
+
def add_new(
- gaussians: Gaussians | GaussiansModule,
- smoothers: dict[str, Any],
- adc_state: McmcStrategyState,
- cap_max: int,
- min_opacity: float
+ gaussians: Gaussians | GaussiansModule,
+ smoothers: dict[str, Any],
+ adc_state: McmcStrategyState,
+ cap_max: int,
+ min_opacity: float
) -> int:
"""Adds new Gaussians based on MCMC strategy.
@@ -361,7 +365,7 @@ def add_new(
Returns:
int: Number of new Gaussians added.
"""
-
+
if isinstance(gaussians, GaussiansModule):
raise NotImplementedError("noise injection not implemented for GaussiansModule")
@@ -384,14 +388,14 @@ def add_new(
# activate
opacities = torch.sigmoid(opacities) # [G]
scales = torch.exp(scales) # [G, 3]
-
+
current_n_points = means.shape[0]
n_target = min(cap_max, int(1.05 * current_n_points))
n_gs = max(0, n_target - current_n_points)
if n_gs > 0:
# add new
n_new = int(n_gs)
-
+
eps = torch.finfo(torch.float32).eps
probs = opacities.flatten()
sampled_idxs = _multinomial_sample(probs, n_gs, replacement=True)
@@ -402,7 +406,7 @@ def add_new(
binoms=adc_state.binoms,
)
new_opacities = torch.clamp(new_opacities, max=1.0 - eps, min=min_opacity)
-
+
# deactivate new opacities/scales for storage if needed
if not gaussians.stores_activated:
new_opacities = torch.logit(new_opacities.clamp(min=eps, max=1.0 - eps))
@@ -413,11 +417,11 @@ def add_new(
new_harmonics = harmonics[sampled_idxs]
# append to existing Gaussians (batch dim = 0, Gaussian dim = 1)
- gaussians.means = torch.cat([gaussians.means, new_means.unsqueeze(0)], dim=1)
- gaussians.scales = torch.cat([gaussians.scales, new_scales.unsqueeze(0)], dim=1)
- gaussians.opacities = torch.cat([gaussians.opacities, new_opacities.unsqueeze(0)], dim=1)
+ gaussians.means = torch.cat([gaussians.means, new_means.unsqueeze(0)], dim=1)
+ gaussians.scales = torch.cat([gaussians.scales, new_scales.unsqueeze(0)], dim=1)
+ gaussians.opacities = torch.cat([gaussians.opacities, new_opacities.unsqueeze(0)], dim=1)
gaussians.rotations_unnorm = torch.cat([gaussians.rotations_unnorm, new_rotations_unnorm.unsqueeze(0)], dim=1)
- gaussians.harmonics = torch.cat([gaussians.harmonics, new_harmonics.unsqueeze(0)], dim=1)
+ gaussians.harmonics = torch.cat([gaussians.harmonics, new_harmonics.unsqueeze(0)], dim=1)
if gaussians.rotations is not None:
gaussians.rotations = torch.cat(
[gaussians.rotations, F.normalize(new_rotations_unnorm, dim=-1).unsqueeze(0)], dim=1
@@ -429,18 +433,19 @@ def add_new(
# add new entries to smoothers state
_add_to_objects(n_new, smoothers)
-
+
return n_new
+
@torch.no_grad()
def apply_mcmc_strategy(
- cfg: BaseStrategyCfg,
- step: int,
- gaussians: Gaussians | GaussiansModule,
- adc_state: McmcStrategyState,
- smoothers: dict[str, Any],
- lr: float,
- # zero_t: bool = False
+ cfg: BaseStrategyCfg,
+ step: int,
+ gaussians: Gaussians | GaussiansModule,
+ adc_state: McmcStrategyState,
+ smoothers: dict[str, Any],
+ lr: float,
+ # zero_t: bool = False
) -> tuple[int, int, int, float | None, float | None]:
"""Applies MCMC strategy to the given Gaussian distributions.
@@ -452,18 +457,18 @@ def apply_mcmc_strategy(
smoothers (dict[str, Any]): Optimizer smoothers.
lr (float): Learning rate for "means" attribute of the GS.
"""
-
+
if isinstance(gaussians, GaussiansModule):
raise NotImplementedError("cloning not implemented for GaussiansModule")
-
+
# Densification and Pruning
nr_cloned, nr_splitted, nr_pruned = 0, 0, 0
-
+
# check if should densify/prune
if (
- step > cfg.refine_start_iter
- and step % cfg.refine_every == 0
- and step % cfg.reset_every >= cfg.pause_refine_after_reset
+ step > cfg.refine_start_iter
+ and step % cfg.refine_every == 0
+ and step % cfg.reset_every >= cfg.pause_refine_after_reset
):
# teleport dead GSs to positions of alive ones
n_relocated_gs = relocate(gaussians, smoothers, adc_state, cfg.min_opacity, copy_state=cfg.relocate_copy_state)
@@ -472,13 +477,13 @@ def apply_mcmc_strategy(
n_new_gs = 0
if step < cfg.refine_stop_iter:
n_new_gs = add_new(gaussians, smoothers, adc_state, cfg.cap_max, cfg.min_opacity)
-
+
torch.cuda.empty_cache()
-
+
print(
f"MCMC @ iter {step}: n_relocated {n_relocated_gs}, n_new_gs {n_new_gs}, total now {gaussians.means.shape[1]}"
)
-
+
# add noise to GSs
inject_noise_to_position(
gaussians, scaler=lr * cfg.noise_lr, scale_cap=cfg.noise_scale_cap
diff --git a/optgs/scene_trainer/adc/vanilla.py b/optgs/scene_trainer/adc/vanilla.py
index 366a877bd3c75ee595f42478406842626cc2c118..8b4fa0104da742270ce1b8fadd6a51f3d4276622 100644
--- a/optgs/scene_trainer/adc/vanilla.py
+++ b/optgs/scene_trainer/adc/vanilla.py
@@ -108,7 +108,7 @@ def prune(
if isinstance(gaussians, GaussiansModule):
raise NotImplementedError("pruning not implemented for GaussiansModule")
- # TODO Naama: check if we can avoid squeezing and unsqueezing
+ # NOTE: ADC runs per-Gaussian at batch size 1, so the batch dim is squeezed off here and added back after.
# get all params and remove batch dim
means = gaussians.means.squeeze(0) # [G, 3]
scales = gaussians.scales.squeeze(0) # [G, 3]
diff --git a/optgs/scene_trainer/common/gaussian_adapter.py b/optgs/scene_trainer/common/gaussian_adapter.py
index 27757f76e015e31ef756c607176a8fae369b21a8..99c8bc1ba065e100462d57af8873cd3a7fc29a87 100644
--- a/optgs/scene_trainer/common/gaussian_adapter.py
+++ b/optgs/scene_trainer/common/gaussian_adapter.py
@@ -89,7 +89,6 @@ class GaussianAdapter(nn.Module):
scales = scales * gaussian_scale_depth[..., None] * multiplier[..., None]
elif point_cloud is not None:
- # TODO: try other activations
if self.cfg.softplus_scale:
scales = torch.clamp(F.softplus(scales - self.cfg.exp_scale_bias), max=self.cfg.gaussian_scale_max)
else:
@@ -107,7 +106,6 @@ class GaussianAdapter(nn.Module):
multiplier = self.get_scale_multiplier(intrinsics, pixel_size)
scales = scales * depths[..., None] * multiplier[..., None]
- # TODO: avoid nan when using exp scale?
scales = torch.clamp(scales, min=self.cfg.clamp_min_scale)
assert input_images is not None
diff --git a/optgs/scene_trainer/common/gaussians.py b/optgs/scene_trainer/common/gaussians.py
index a55bc7fbb8e6c481706219a1cb8e1d52b8b63a82..937c69264ceed2f20e3d54a589e328aad744c636 100644
--- a/optgs/scene_trainer/common/gaussians.py
+++ b/optgs/scene_trainer/common/gaussians.py
@@ -4,38 +4,6 @@ from jaxtyping import Float
from torch import Tensor
-# def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
-# """
-# Convert rotations given as quaternions to rotation matrices.
-
-# Args:
-# quaternions: quaternions with real part first,
-# as tensor of shape (..., 4).
-
-# Returns:
-# Rotation matrices as tensor of shape (..., 3, 3).
-# """
-# r, i, j, k = torch.unbind(quaternions, -1) # wxyz
-# # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
-# two_s = 2.0 / (quaternions * quaternions).sum(-1)
-
-# 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 o.reshape(quaternions.shape[:-1] + (3, 3))
-
-
def rotation_matrix_to_quaternion_xyzw(R: torch.Tensor, eps: float = 1e-12) -> torch.Tensor:
"""
Convert rotation matrices to quaternions in xyzw format.
diff --git a/optgs/scene_trainer/gaussian_module.py b/optgs/scene_trainer/gaussian_module.py
index 6465a1fb5b2757a49d8e2b52248a672c5488d8bf..ece8836df47c9d5b9d34bdbb5c0127f69183f0c0 100644
--- a/optgs/scene_trainer/gaussian_module.py
+++ b/optgs/scene_trainer/gaussian_module.py
@@ -5,7 +5,13 @@ from torch import Tensor, nn
from optgs.model.types import Gaussians
from optgs.scene_trainer.common.gaussians import build_covariance
-
+# Two Gaussian representations coexist:
+# Gaussians (model/types.py) — a plain dataclass whose fields are tensors; the type that flows
+# through the whole init → optimize → render pipeline.
+# GaussiansModule (here) — an nn.Module that wraps the same parameters as nn.Parameter, so a
+# standard torch optimizer (Adam/SGD) can optimize them directly. Built only for the test-time
+# post-processing path (postprocessing.py); converted to/from Gaussians via the helpers below.
+# The ADC strategies do not support it (they raise NotImplementedError on GaussiansModule).
class GaussiansModule(nn.Module):
def __init__(
self,
@@ -71,9 +77,6 @@ class GaussiansModule(nn.Module):
_register_param("sh0", sh0)
if shN is not None:
_register_param("shN", shN)
-
- for name, param in self.named_parameters():
- print(f"Registered parameter: {name}, shape: {param.shape}, dtype: {param.dtype}, min: {param.min()}, max: {param.max()}, requires_grad: {param.requires_grad}")
@property
def scales(self):
@@ -105,13 +108,6 @@ class GaussiansModule(nn.Module):
rotation_xyzw = self.rotations
covariances = self.covariance_activation(self.scales, rotation_xyzw) # [G, 3, 3]
return covariances
-
- def reset_opacity(self, optimizer):
- opacities_old = self.opacity_activation(self.opacities_raw)
- opacities_raw_new = self.inverse_opacity_activation(torch.min(opacities_old, torch.ones_like(opacities_old)*0.01), eps=1e-6)
- # optimizable_tensors = self.replace_tensor_to_optimizer(optimizer, opacities_raw_new, "opacity")
- # self.opacities_raw = optimizable_tensors["opacity"]
-
def gaussians2module(gaussians: Gaussians, device: torch.device) -> GaussiansModule:
diff --git a/optgs/scene_trainer/initializer/initializer.py b/optgs/scene_trainer/initializer/initializer.py
index e867432c43fb7fa5fd99d842cd0b16a44b56e4ac..a3fa65ba62b4e2c1c85e953aa8e5ba84bff93bbe 100644
--- a/optgs/scene_trainer/initializer/initializer.py
+++ b/optgs/scene_trainer/initializer/initializer.py
@@ -18,6 +18,41 @@ class InitializerOutput:
depths: list[torch.Tensor] | torch.Tensor | None = None
target_render: DecoderOutput | None = None
context_render: DecoderOutput | None = None
+ # View indices used when rendering a subset (training); None means all views were rendered.
+ target_render_index: torch.Tensor | None = None
+ context_render_index: torch.Tensor | None = None
+
+ def get_render(self, which: str) -> DecoderOutput | None:
+ if which == "target":
+ return self.target_render
+ elif which == "context":
+ return self.context_render
+ else:
+ raise ValueError(f"Unknown which: {which}, should be 'target' or 'context'")
+
+ def set_render(self, which: str, value: DecoderOutput) -> None:
+ if which == "target":
+ self.target_render = value
+ elif which == "context":
+ self.context_render = value
+ else:
+ raise ValueError(f"Unknown which: {which}, should be 'target' or 'context'")
+
+ def get_render_index(self, which: str) -> torch.Tensor | None:
+ if which == "target":
+ return self.target_render_index
+ elif which == "context":
+ return self.context_render_index
+ else:
+ raise ValueError(f"Unknown which: {which}, should be 'target' or 'context'")
+
+ def set_render_index(self, which: str, value: torch.Tensor | None) -> None:
+ if which == "target":
+ self.target_render_index = value
+ elif which == "context":
+ self.context_render_index = value
+ else:
+ raise ValueError(f"Unknown which: {which}, should be 'target' or 'context'")
@dataclass
@@ -37,6 +72,17 @@ class InitializerCfg:
train_fixed_gaussians_num: int | None
eval_fixed_gaussians_num: int | None
+ def get_gaussian_param_num(self) -> int:
+ """Per-Gaussian parameter count of this initializer's own representation:
+ scale(3) + rotation(4) + SH(3*sh_d) + opacity(1) + position encoding (get_position_param_num)."""
+ return 3 + 4 + 3 * self.get_sh_d() + 1 + self.get_position_param_num()
+
+ def get_position_param_num(self) -> int:
+ """Size of this initializer's position encoding. Most initializers place Gaussians at explicit
+ 3D positions (3). Initializers that encode position differently (resplat: a 2D pixel offset
+ added to a per-pixel depth) override this."""
+ return 3
+
@dataclass
class NonlearnedInitializerCfg(InitializerCfg):
pass
@@ -47,7 +93,7 @@ class LearnedInitializerCfg(InitializerCfg):
@dataclass
-class PerPixelInitializerCfg(InitializerCfg):
+class PerPixelInitializerCfg(LearnedInitializerCfg):
latent_gs: bool
latent_downsample: int
@@ -59,7 +105,10 @@ class Initializer(nn.Module, ABC, Generic[T]):
super().__init__()
self.cfg = cfg
- def preprocessing(self, batch, train_cfg) -> None:
+ def eval_preprocessing(self, batch, train_cfg) -> None:
+ """Eval/validation-only batch prep (depth-range override + optional scale prediction),
+ applied in-place before the initializer runs. Training does not call this. The
+ universal patch-crop data shim is a separate, always-applied step (see MetaTrainer)."""
pass
@property
diff --git a/optgs/scene_trainer/initializer/initializer_colmap.py b/optgs/scene_trainer/initializer/initializer_colmap.py
index 30e6ed9745b05ee1c108de1e6e89b811df515a6d..40e587121ecadc6d37b43a00dad120b7458daf07 100644
--- a/optgs/scene_trainer/initializer/initializer_colmap.py
+++ b/optgs/scene_trainer/initializer/initializer_colmap.py
@@ -36,12 +36,6 @@ class InitializerColmapCfg(NonlearnedInitializerCfg):
points3d_ply_filename: Optional[str] # if set, loads points from this PLY file (relative to scene dir) instead of COLMAP binary
override_dataset_poses: bool # if true, overrides the dataset poses with the COLMAP poses (after applying T_world transform)
- def get_gaussian_param_num(self):
- # calculate the number of parameters per Gaussian
- sh_d = self.get_sh_d()
- init_gaussian_param_num = 3 + 4 + 3 * sh_d + 2 + 1
- return init_gaussian_param_num
-
def get_sh_d(self):
sh_d = (self.sh_degree + 1) ** 2
return sh_d
@@ -112,7 +106,6 @@ class InitializerColmap(NonlearnedInitializer[InitializerColmapCfg]):
points = parser.points # (N, 3) float32
points_rgb = parser.points_rgb # (N, 3) uint8
- # TODO Patricia: Fix permission denied
# Write atomically with a temp file that already ends in .npz.
try:
tmp_path = ''
@@ -275,7 +268,11 @@ class InitializerColmap(NonlearnedInitializer[InitializerColmapCfg]):
rgbs = F.pad(rgbs, (0, 0, 0, pad), value=0.0)
scales = F.pad(scales, (0, 0, 0, pad), value=1e-10)
opacities = F.pad(opacities, (0, pad), value=1e-10)
- # TODO Naama: might be a problem if we don't freeze zero-grad gaussians
+ # NOTE: These are invisible (zero-opacity) dummy Gaussians, added only so that every GPU optimizes
+ # the same number of Gaussians; nr_valid (set below) is the true count. They must stay
+ # invisible: if the optimizer updates them they gain opacity and become real Gaussians. The
+ # knn optimizer keeps them frozen by updating only Gaussians with nonzero gradient (the dummies
+ # are invisible, so their gradient is zero) — see KnnBasedOptimizer._forward_impl.
points_dict = {
"xyz": xyz,
diff --git a/optgs/scene_trainer/initializer/initializer_edgs.py b/optgs/scene_trainer/initializer/initializer_edgs.py
index 057f37305a14e21422f78e704c67607b79732bcb..fea9f56a1985c813c427986b85c5a0ce44dc26d7 100644
--- a/optgs/scene_trainer/initializer/initializer_edgs.py
+++ b/optgs/scene_trainer/initializer/initializer_edgs.py
@@ -2,13 +2,12 @@ from dataclasses import dataclass
from typing import Literal, Optional
from optgs.dataset.data_types import BatchedViews
-import numpy as np
import torch
import math
import torch.nn.functional as F
from pathlib import Path
from optgs.experimental.edgs.init import init_gaussians_with_corr
-from optgs.experimental.initializers_utils import knn, points_to_gaussians
+from optgs.experimental.initializers_utils import points_to_gaussians
from optgs.model.types import Gaussians
from optgs.scene_trainer.common.gaussian_adapter import build_covariance
from optgs.scene_trainer.initializer.initializer import InitializerOutput, NonlearnedInitializer, NonlearnedInitializerCfg
@@ -23,13 +22,7 @@ class InitializerEdgsCfg(NonlearnedInitializerCfg):
roma_model_type: str
sample_init_gaussians: int # if >0, randomly sample this many gaussians from the initialized set
-
- def get_gaussian_param_num(self):
- # calculate the number of parameters per Gaussian
- sh_d = self.get_sh_d()
- # TODO Naama: check where this is used, and if it is needed
- init_gaussian_param_num = 3 + 4 + 3 * sh_d + 2 + 1
- return init_gaussian_param_num
+ cache_dir: str # root dir for the per-scene points_dict cache (built once per scene by RoMa)
def get_sh_d(self):
sh_d = (self.sh_degree + 1) ** 2
@@ -44,10 +37,15 @@ class InitializerEdgs(NonlearnedInitializer[InitializerEdgsCfg]):
self,
context: BatchedViews,
visualization_dump: Optional[dict] = None,
- cached_data_path: Optional[Path] = None,
+ scene: Optional[list[str]] = None,
**kwargs
) -> InitializerOutput:
+ # Per-scene cache for the (expensive RoMa-correspondence) points_dict.
+ # NOTE: keyed only by scene name, not by the selected context views — re-evaluating
+ # the same scene with a different context set will hit a stale cache.
+ cached_data_path = Path(self.cfg.cache_dir) / scene[0] if scene is not None else None
+
device = context["extrinsics"].device
# unpack context (batch_dim = 1)
diff --git a/optgs/scene_trainer/initializer/initializer_ply.py b/optgs/scene_trainer/initializer/initializer_ply.py
index 7827531515839e7cd05de5c8d5da2350c9fbcf1f..e215543e8d0e8da1ba97e8a050204de0a6a8afca 100644
--- a/optgs/scene_trainer/initializer/initializer_ply.py
+++ b/optgs/scene_trainer/initializer/initializer_ply.py
@@ -4,13 +4,10 @@ from typing import Literal, Optional
from optgs.dataset.data_types import BatchedViews
-import torch
-import torch.nn.functional as F
# from optgs.dataset.colmap.utils import Parser
# from optgs.experimental.initializers_utils import knn, points_to_gaussians
from optgs.scene_trainer.common.gaussian_adapter import build_covariance
-from optgs.model.types import Gaussians
from optgs.model.ply_export import load_gaussians_ply
from optgs.scene_trainer.initializer.initializer import NonlearnedInitializer, InitializerOutput, NonlearnedInitializerCfg
@@ -26,12 +23,6 @@ class InitializerPlyCfg(NonlearnedInitializerCfg):
# dl3dv_settings: bool
ply_filename: str = "gaussians.ply" # relative path under the scene dir; can include subdirs, e.g. "iteration_20000/point_cloud.ply"
- def get_gaussian_param_num(self):
- # calculate the number of parameters per Gaussian
- sh_d = self.get_sh_d()
- init_gaussian_param_num = 3 + 4 + 3 * sh_d + 2 + 1
- return init_gaussian_param_num
-
def get_sh_d(self):
sh_d = (self.sh_degree + 1) ** 2
return sh_d
diff --git a/optgs/scene_trainer/initializer/initializer_random.py b/optgs/scene_trainer/initializer/initializer_random.py
index 281c85e6ef6a00acf824a113965473286d7075bd..31c39f7d732b9aa5153bc75784651906f8b4ffcf 100644
--- a/optgs/scene_trainer/initializer/initializer_random.py
+++ b/optgs/scene_trainer/initializer/initializer_random.py
@@ -1,6 +1,5 @@
from dataclasses import dataclass
-from typing import Literal, Optional
-from pathlib import Path
+from typing import Literal
import torch
import torch.nn.functional as F
@@ -8,7 +7,7 @@ from optgs.dataset.data_types import BatchedViews
from optgs.scene_trainer.common.gaussian_adapter import build_covariance
from optgs.model.types import Gaussians
from optgs.experimental.initializers_utils import knn, points_to_gaussians
-from optgs.scene_trainer.initializer.initializer import NonlearnedInitializer, InitializerOutput, InitializerCfg, NonlearnedInitializerCfg
+from optgs.scene_trainer.initializer.initializer import NonlearnedInitializer, InitializerOutput, NonlearnedInitializerCfg
from optgs.dataset.camera_datasets.camera import get_scene_scale
@@ -21,12 +20,6 @@ class InitializerRandomCfg(NonlearnedInitializerCfg):
init_opacity: float
sh_degree: int
- def get_gaussian_param_num(self):
- # calculate the number of parameters per Gaussian
- sh_d = self.get_sh_d()
- init_gaussian_param_num = 3 + 4 + 3 * sh_d + 2 + 1
- return init_gaussian_param_num
-
def get_sh_d(self):
sh_d = (self.sh_degree + 1) ** 2
return sh_d
diff --git a/optgs/scene_trainer/initializer/initializer_resplat.py b/optgs/scene_trainer/initializer/initializer_resplat.py
index 5527a572956aeea558733ff4d01d3d74e1f63359..830a37e890056306ef6e416193e0c6259c1817d2 100644
--- a/optgs/scene_trainer/initializer/initializer_resplat.py
+++ b/optgs/scene_trainer/initializer/initializer_resplat.py
@@ -13,28 +13,17 @@ from optgs.dataset.shims.patch_shim import apply_patch_shim
from optgs.geometry.projection import sample_image_grid, get_world_rays
from optgs.misc.general_utils import rotate_quats
from optgs.misc.io import FrequencyScheduler
-from optgs.model.encoder.layer import BasicBlock
-from optgs.model.encoder.unimatch.dpt_head import DPTHead
-from optgs.model.encoder.unimatch.feature_upsampler import ResizeConvFeatureUpsampler
-from optgs.model.encoder.unimatch.ldm_unet.unet import UNetModel
-from optgs.model.encoder.unimatch.mv_unimatch import MultiViewUniMatch
-from optgs.model.encoder.visualization.encoder_visualizer_depthsplat_cfg import EncoderVisualizerDepthSplatCfg
+from optgs.model.backbones.layer import BasicBlock
+from optgs.model.backbones.unimatch.dpt_head import DPTHead
+from optgs.model.backbones.unimatch.feature_upsampler import ResizeConvFeatureUpsampler
+from optgs.model.backbones.unimatch.ldm_unet.unet import UNetModel
+from optgs.model.backbones.unimatch.mv_unimatch import MultiViewUniMatch
from optgs.model.types import Gaussians
from optgs.scene_trainer.common.gaussian_adapter import GaussianAdapter, GaussianAdapterCfg, build_covariance, RGB2SH
from optgs.scene_trainer.initializer.initializer import InitializerOutput, LearnedInitializer, PerPixelInitializerCfg
try:
- from optgs.model.encoder.point_transformer.layer import (PlainPointTransformer, SubsampleBlock, PointLinearWrapper,
- MultiScalePointTransformer,
- MultViewLowresAttn, MultViewUniMatchAttn,
- GaussianErrorCrossAttn)
-except:
- pass
-
-from optgs.model.encoder.lvsm.transformer import LVSMTransformer
-
-try:
- from simple_knn._C import distCUDA2
+ from optgs.model.backbones.point_transformer.layer import PlainPointTransformer
except:
pass
@@ -45,7 +34,6 @@ class ResplatInitializerCfg(PerPixelInitializerCfg):
d_feature: int
num_depth_candidates: int
num_surfaces: int
- visualizer: EncoderVisualizerDepthSplatCfg
gaussian_adapter: GaussianAdapterCfg
gaussians_per_pixel: int
unimatch_weights_path: str | None
@@ -89,19 +77,13 @@ class ResplatInitializerCfg(PerPixelInitializerCfg):
pt_head_conv: bool
pt_head_concat_img: bool
pt_head_channels: int | None
- multi_scale_pt: bool
attn_proj_channels: int | None
- fps_num_samples: int | None
knn_samples: int
post_norm: bool
no_rpe: bool
no_knn_attn: bool
num_blocks: int
- pt_downsample: int
- fps_agg_func: str
- subsample_method: str
add_pt_residual: bool
- pt_pred_residual_position: bool # based on the inital point cloud from depth, predict additional residual
latent_dpt_upsampler: bool
latent_dpt_upsampler_no_concat: bool
light_dpt_feature: bool
@@ -119,10 +101,6 @@ class ResplatInitializerCfg(PerPixelInitializerCfg):
unet_gaussian_regressor: bool
resnet_gaussian_regressor: bool
- # lvsm gaussian regressor
- lvsm_gaussian_regressor: bool
- lvsm_layers: int
-
sample_log_depth: bool
bilinear_upsample_depth: bool
no_upsample_depth: bool
@@ -211,17 +189,10 @@ class ResplatInitializerCfg(PerPixelInitializerCfg):
def get_gaussian_param_num(self):
# predict gaussian parameters: scale, q, sh, offset, opacity
# d_in: (scale, q, sh)
- sh_d = self.get_sh_d()
- init_gaussian_param_num = 3 + 4 + 3 * sh_d + 2 + 1
+ # base params (scale, q, sh, opacity) plus the 2D pixel offset (via get_position_param_num)
+ init_gaussian_param_num = super().get_gaussian_param_num()
if self.no_pixel_offset:
init_gaussian_param_num -= 2
- if self.pt_downsample > 0:
- # no pixel offset
- init_gaussian_param_num -= 2
- if self.pt_pred_residual_position:
- # based on the inital point cloud from depth, predict additional residual
- # without pixel offset on 2d
- init_gaussian_param_num = init_gaussian_param_num + 3 - 2
# multiple gaussians per latent
if self.init_gaussian_multiple > 1:
# we use the point cloud unprojected from higher resolution depth map as center
@@ -230,6 +201,10 @@ class ResplatInitializerCfg(PerPixelInitializerCfg):
init_gaussian_param_num *= self.init_gaussian_multiple
return init_gaussian_param_num
+ def get_position_param_num(self) -> int:
+ # resplat encodes position as a 2D pixel offset added to a per-pixel depth, not an explicit 3D point.
+ return 2
+
def get_sh_d(self):
sh_d = (self.gaussian_adapter.sh_degree + 1) ** 2
return sh_d
@@ -326,15 +301,6 @@ class ResplatInitializer(LearnedInitializer[ResplatInitializerCfg]):
BasicBlock(channels, channels),
]
- elif self.cfg.lvsm_gaussian_regressor:
- modules = [
- nn.Linear(in_channels, channels),
- nn.LayerNorm(channels),
- nn.GELU(),
- LVSMTransformer(channels,
- n_layer=self.cfg.lvsm_layers)
- ]
-
else:
# conv regressor
modules = [
@@ -357,52 +323,21 @@ class ResplatInitializer(LearnedInitializer[ResplatInitializerCfg]):
channels = self.cfg.pt_head_channels
self.proj = nn.Linear(in_channels, channels)
- if self.cfg.multi_scale_pt:
- self.pt = MultiScalePointTransformer(channels,
- self.cfg.knn_samples,
- downsample_agg_func=self.cfg.fps_agg_func,
- subsample_method=self.cfg.subsample_method,
- fps_num_samples=self.cfg.fps_num_samples,
- attn_proj_channels=self.cfg.attn_proj_channels,
- )
- else:
- self.pt = PlainPointTransformer(channels, self.cfg.knn_samples,
- post_norm=self.cfg.post_norm,
- no_rpe=self.cfg.no_rpe,
- no_attn=self.cfg.no_knn_attn,
- num_blocks=self.cfg.num_blocks,
- num_heads=self.cfg.pt_heads,
- attn_proj_channels=self.cfg.attn_proj_channels,
- use_checkpointing=self.cfg.use_checkpointing,
- init_use_checkpointing=self.cfg.init_use_checkpointing,
- with_mv_attn=self.cfg.init_pt_with_mv_attn,
- with_mv_attn_lowres=self.cfg.init_pt_with_mv_attn_lowres,
- )
+ self.pt = PlainPointTransformer(channels, self.cfg.knn_samples,
+ post_norm=self.cfg.post_norm,
+ no_rpe=self.cfg.no_rpe,
+ no_attn=self.cfg.no_knn_attn,
+ num_blocks=self.cfg.num_blocks,
+ num_heads=self.cfg.pt_heads,
+ attn_proj_channels=self.cfg.attn_proj_channels,
+ use_checkpointing=self.cfg.use_checkpointing,
+ init_use_checkpointing=self.cfg.init_use_checkpointing,
+ with_mv_attn=self.cfg.init_pt_with_mv_attn,
+ with_mv_attn_lowres=self.cfg.init_pt_with_mv_attn_lowres,
+ )
out_channels = channels
- # point downsample
- if self.cfg.pt_downsample > 0:
- num_downsample = int(np.log2(self.cfg.pt_downsample))
-
- if num_downsample == 0:
- stride = 1
- else:
- stride = 2
-
- assert num_downsample == 1, f"unsupported num_downsample: {num_downsample}"
-
- self.pt_down = SubsampleBlock(channels, out_channels=channels * 2,
- stride=stride,
- knn_samples=self.cfg.knn_samples,
- post_norm=self.cfg.post_norm,
- agg_func=self.cfg.fps_agg_func,
- subsample_method=self.cfg.subsample_method,
- )
-
- out_channels = channels * 2
-
- # TODO: add more pt blocks after downsampling
if self.cfg.pt_head_concat_img:
# concat to the initial image and features
@@ -675,16 +610,8 @@ class ResplatInitializer(LearnedInitializer[ResplatInitializerCfg]):
features,
), dim=1)
- if self.cfg.lvsm_gaussian_regressor:
- h, w = concat.shape[-2:]
- tmp = rearrange(concat, "(b v) c h w -> b (v h w) c", b=b, v=v)
- with torch.autocast('cuda', dtype=torch.bfloat16):
- out = self.gaussian_regressor(tmp)
-
- out = rearrange(out, "b (v h w) c -> (b v) c h w", b=b, v=v, h=h, w=w)
- else:
- with torch.amp.autocast(device_type='cuda', enabled=self.cfg.use_amp, dtype=torch.bfloat16):
- out = self.gaussian_regressor(concat)
+ with torch.amp.autocast(device_type='cuda', enabled=self.cfg.use_amp, dtype=torch.bfloat16):
+ out = self.gaussian_regressor(concat)
if self.cfg.latent_gs:
concat = [out, img_unshuffle, features, match_prob]
@@ -744,10 +671,6 @@ class ResplatInitializer(LearnedInitializer[ResplatInitializerCfg]):
condition_features = rearrange(out, "(bv h w) c -> bv c h w", h=h, w=w)
- if self.cfg.pt_downsample > 0:
- out, fps_idx = self.pt_down((point_cloud, out, offset))
- # [N, 3]
- point_cloud, out, offset = out
with torch.amp.autocast(device_type='cuda', enabled=self.cfg.pt_head_amp, dtype=torch.bfloat16):
if self.cfg.pt_head_concat_img:
@@ -767,22 +690,9 @@ class ResplatInitializer(LearnedInitializer[ResplatInitializerCfg]):
if self.cfg.pt_head_conv:
out = rearrange(out, "(b v) c h w -> (b v h w) c", b=b, v=v)
- if self.cfg.pt_downsample > 0:
- # [N, C]
- gaussians = out
- else:
- if self.cfg.pt_pred_residual_position:
- # TODO: add intermediate supervision to the initial point cloud
- # TODO: multiple scale factor to the delta position to make it more stable
- # residual position
- point_cloud = point_cloud + out[..., -3:] # [BVHW, 3]
-
- # remaining gaussians
- out = out[..., :-3]
-
- point_cloud = rearrange(point_cloud, "(b v h w) c -> b v (h w) () () c", b=b, v=v, h=h, w=w)
+ point_cloud = rearrange(point_cloud, "(b v h w) c -> b v (h w) () () c", b=b, v=v, h=h, w=w)
- gaussians = rearrange(out, "(b v h w) c -> (b v) c h w", b=b, h=h, w=w)
+ gaussians = rearrange(out, "(b v h w) c -> (b v) c h w", b=b, h=h, w=w)
else:
with torch.amp.autocast(device_type='cuda', enabled=self.cfg.use_amp, dtype=torch.bfloat16):
@@ -795,7 +705,6 @@ class ResplatInitializer(LearnedInitializer[ResplatInitializerCfg]):
if self.cfg.init_gaussian_multiple > 1:
# hard coded for now
if self.cfg.init_gaussian_multiple == 4:
- # TODO: try avgpooling downsampling depth
if self.cfg.latent_downsample == 4:
# resize full resolution depth
depths = F.interpolate(depth, scale_factor=0.5, mode='bilinear', align_corners=True)
@@ -806,7 +715,6 @@ class ResplatInitializer(LearnedInitializer[ResplatInitializerCfg]):
else:
raise NotImplementedError
elif self.cfg.init_gaussian_multiple == 16:
- # TODO: try avgpooling downsampling depth
if self.cfg.latent_downsample == 4:
depths = depth
elif self.cfg.latent_downsample == 8:
@@ -822,414 +730,170 @@ class ResplatInitializer(LearnedInitializer[ResplatInitializerCfg]):
else:
depths = rearrange(depth, "b v h w -> b v (h w) () ()")
- if self.cfg.pt_downsample > 0:
+ gaussians = rearrange(gaussians, "(b v) c h w -> b v c h w", b=b, v=v)
- # split batch
- assert offset.shape[0] == b
+ # [B, V, H*W, 84]
+ raw_gaussians = rearrange(
+ gaussians, "b v c h w -> b v (h w) c")
- if self.cfg.latent_gs:
- sh_input_images = rearrange(context["image"], "b v c h w -> (b v) c h w")
- if self.cfg.latent_gs_img_interp == 'bicubic':
- sh_input_images = F.interpolate(sh_input_images, scale_factor=1. / self.cfg.latent_downsample,
- mode='bicubic', align_corners=True)
- elif self.cfg.latent_gs_img_interp == 'area':
- sh_input_images = F.interpolate(sh_input_images, scale_factor=1. / self.cfg.latent_downsample,
- mode='area')
- elif self.cfg.latent_gs_img_interp == 'softmax':
- sh_input_images = self.softmax_downsample(sh_input_images)
- else:
- raise NotImplementedError
+ assert len(depth_preds) == 1, "num_scales must be 1; multi-scale depth supervision is not supported"
- h, w = sh_input_images.shape[-2:]
+ # [B, V, H*W, C]
+ repeat = self.cfg.init_gaussian_multiple
+ num_sh = self.gaussian_adapter.d_sh
- sh_input_images = rearrange(sh_input_images, "(b v) c h w -> b v c h w", b=b, v=v)
- else:
- sh_input_images = context["image"]
-
- sh_input_images = rearrange(sh_input_images, "b v c h w -> (b v h w) c")
-
- # subsample with fps index
- sh_input_images = sh_input_images[fps_idx.long(), :] # [N, 3]
-
- # extrinsics
- extrinsics_all = rearrange(repeat(context["extrinsics"], "b v i j -> b v h w i j", h=h, w=w),
- "b v h w i j -> (b v h w) i j"
- )
- extrinsics_all = extrinsics_all[fps_idx.long(), :, :] # [N, 4, 4]
-
- point_list = [point_cloud[:offset[0]]]
- gaussian_list = [gaussians[:offset[0]]]
- sh_img_list = [sh_input_images[:offset[0]]]
- extrinsics_list = [extrinsics_all[:offset[0]]]
-
- for i in range(b - 1):
- point_list.append(point_cloud[offset[i]:offset[i + 1]])
- gaussian_list.append(gaussians[offset[i]:offset[i + 1]])
- sh_img_list.append(sh_input_images[offset[i]:offset[i + 1]])
- extrinsics_list.append(extrinsics_all[offset[i]:offset[i + 1]])
-
- point_cloud = torch.stack(point_list, dim=0) # [B, N, 3]
- gaussians = torch.stack(gaussian_list, dim=0) # [B, N, C]
- sh_imgs = torch.stack(sh_img_list, dim=0) # [B, N, 3]
- extrinsics_all = torch.stack(extrinsics_list, dim=0) # [B, N, 4, 4]
-
- # point_cloud = [point_cloud[offset[i]:offset[i+1]] for i in range(b)]
- # point_cloud = torch.stack(point_cloud, dim=0) # [B, N, 3]
- # gaussians = [gaussians[offset[i]:offset[i+1]] for i in range(b)]
- # gaussians = torch.stack(gaussians, dim=0) # [B, N, 3]
-
- opacities = gaussians[..., 0].sigmoid() # [B, N]
-
- gaussians = self.gaussian_adapter.forward(
- extrinsics=extrinsics_all,
- intrinsics=None,
- coordinates=None,
- depths=None,
- opacities=opacities,
- raw_gaussians=gaussians[..., 1:],
- image_shape=None,
- point_cloud=point_cloud,
- input_images=sh_imgs,
+ if self.cfg.no_pixel_offset:
+ rotations_unnorm, scales, opacities_raw, sh = raw_gaussians.split(
+ [4 * repeat, 3 * repeat, 1 * repeat, 3 * num_sh * repeat],
+ dim=-1,
)
-
- gaussians = rearrange(gaussians, "(b v) c h w -> b v c h w", b=b, v=v)
-
- # [B, V, H*W, 84]
- raw_gaussians = rearrange(
- gaussians, "b v c h w -> b v (h w) c")
-
- assert len(depth_preds) == 1, "num_scales must be 1; multi-scale depth supervision is not supported"
-
- # [B, V, H*W, C]
- repeat = self.cfg.init_gaussian_multiple
- num_sh = self.gaussian_adapter.d_sh
-
- if self.cfg.no_pixel_offset:
- rotations_unnorm, scales, opacities_raw, sh = raw_gaussians.split(
- [4 * repeat, 3 * repeat, 1 * repeat, 3 * num_sh * repeat],
- dim=-1,
- )
- else:
- rotations_unnorm, scales, opacities_raw, offset, sh = raw_gaussians.split(
- [4 * repeat, 3 * repeat, 1 * repeat, 2 * repeat, 3 * num_sh * repeat],
- dim=-1,
- )
-
- latent_h, latent_w = gaussians.shape[-2:]
-
- if repeat > 1:
- # reshape all the gaussian parameters
- if True or self.cfg.latent_new_reshape:
- # this works
- r = int(np.sqrt(repeat))
- rotations_unnorm = rearrange(rotations_unnorm, "b v (h w) (c x y) -> b v (h x w y) c",
- h=latent_h, w=latent_w, x=r, y=r)
- scales = rearrange(scales, "b v (h w) (c x y) -> b v (h x w y) c", h=latent_h, w=latent_w, x=r,
- y=r)
- opacities_raw = rearrange(opacities_raw, "b v (h w) (c x y) -> b v (h x w y) c", h=latent_h,
- w=latent_w, x=r, y=r)
- offset = rearrange(offset, "b v (h w) (c x y) -> b v (h x w y) c", h=latent_h, w=latent_w, x=r,
- y=r)
- sh = rearrange(sh, "b v (h w) (c x y) -> b v (h x w y) c", h=latent_h, w=latent_w, x=r, y=r)
- else:
- # doesn't work
- rotations_unnorm = rearrange(rotations_unnorm, "b v hw (k c) -> b v (hw k) c", k=repeat)
- scales = rearrange(scales, "b v hw (k c) -> b v (hw k) c", k=repeat)
- opacities_raw = rearrange(opacities_raw, "b v hw (k c) -> b v (hw k) c", k=repeat)
- offset = rearrange(offset, "b v hw (k c) -> b v (hw k) c", k=repeat)
- sh = rearrange(sh, "b v hw (k c) -> b v (hw k) c", k=repeat)
-
- opacities = opacities_raw.sigmoid() # [B, V, H*W*K, 1]
-
- if self.cfg.latent_downsample == 4 and self.cfg.init_gaussian_multiple == 4:
- scale_factor = 2
- elif self.cfg.latent_downsample == 2 and self.cfg.init_gaussian_multiple == 4:
- scale_factor = 2
- elif self.cfg.latent_downsample == 4 and self.cfg.init_gaussian_multiple == 16:
- scale_factor = 4
- elif self.cfg.latent_downsample == 8 and self.cfg.init_gaussian_multiple == 4:
- scale_factor = 2
- elif self.cfg.latent_downsample == 8 and self.cfg.init_gaussian_multiple == 16:
- scale_factor = 4
- else:
- scale_factor = 1
-
- h, w = latent_h * scale_factor, latent_w * scale_factor
-
- # unproject depth
- xy_ray, _ = sample_image_grid((h, w), device) # [H, W, 2] in [0, 1]
- xy_ray = rearrange(xy_ray, "h w xy -> (h w) () xy") # [H*W, 1, 2]
-
- if self.cfg.no_pixel_offset:
- offset_xy = torch.ones_like(raw_gaussians[..., :2]).unsqueeze(-2).to(
- raw_gaussians.device) * 0.5 # [B, V, H*W, 1, 2]
- else:
- offset_xy = offset.sigmoid().unsqueeze(-2) # [B, V, H*W, 1, 2]
-
- pixel_size = 1 / \
- torch.tensor((w, h), dtype=torch.float32, device=device)
- # [H*W, 1, 2]
- if self.cfg.deform_sample_depth and not self.cfg.deform_sample_depth_debug:
- # (offset_xy - 0.5) in -0.5 to 0.5, without multiplying by pixel size such that the points can move in the image space
- xy_ray = (xy_ray + (offset_xy - 0.5)).clamp(min=0., max=1.)
- else:
- xy_ray = xy_ray + (offset_xy - 0.5) * pixel_size
-
- if self.cfg.deform_sample_depth:
- # use low-res xy_ray to sample full-res depth
-
- sample_grid = rearrange(xy_ray, "b v (h w) c xy -> (b v) h w (c xy)", h=h, w=w) # in [0, 1]
- # to [-1, 1]
- sample_grid = 2 * (sample_grid - 0.5) # [BV, h, w, 2]
-
- fullres_depth = rearrange(depth, "b v h w -> (b v) () h w") # [BV, 1, H, W]
- sampled_depth = F.grid_sample(fullres_depth, sample_grid, mode='bilinear', align_corners=True,
- padding_mode="border") # [BV, 1, h, w]
- # reshape
- depths = rearrange(sampled_depth, "(b v) () h w -> b v (h w) () ()", b=b, v=v, h=h, w=w)
-
- if self.cfg.latent_gs:
- sh_input_images = rearrange(context["image"], "b v c h w -> (b v) c h w")
- if self.cfg.latent_downsample == 4 and self.cfg.init_gaussian_multiple == 4:
- sh_input_images = F.interpolate(sh_input_images, scale_factor=0.5, mode='area')
- elif self.cfg.latent_downsample == 4 and self.cfg.init_gaussian_multiple == 16:
- pass
- elif self.cfg.latent_downsample == 2 and self.cfg.init_gaussian_multiple == 4:
- pass
- elif self.cfg.latent_downsample == 8 and self.cfg.init_gaussian_multiple == 4:
- sh_input_images = F.interpolate(sh_input_images, scale_factor=0.25, mode='area')
- elif self.cfg.latent_downsample == 8 and self.cfg.init_gaussian_multiple == 16:
- sh_input_images = F.interpolate(sh_input_images, scale_factor=0.5, mode='area')
- else:
- sh_input_images = F.interpolate(sh_input_images, scale_factor=1. / self.cfg.latent_downsample,
- mode='area')
-
- sh_input_images = rearrange(sh_input_images, "(b v) c h w -> b v c h w", b=b, v=v)
-
- else:
- sh_input_images = context["image"]
-
- assert len(depth_preds) == 1, "num_scales must be 1; multi-scale depth supervision is not supported"
-
- # build gaussians
- # scale
- scales = torch.clamp(F.softplus(scales - self.cfg.gaussian_adapter.exp_scale_bias),
- min=self.cfg.gaussian_adapter.clamp_min_scale,
- max=self.cfg.gaussian_adapter.gaussian_scale_max
- )
-
- # Normalize the quaternion features to yield a valid quaternion.
- # rotations = rotations_unnorm / (rotations_unnorm.norm(dim=-1, keepdim=True) + 1e-8)
-
- # Convert rotations to world-space
- c2w_rotations = context["extrinsics"][..., :3, :3].unsqueeze(2) # [B, V, 1, 3, 3]
- rotations = rotate_quats(c2w_rotations, rotations_unnorm)
- rotations_unnorm = rotations.clone()
- # Create world-space covariance matrices.
- covariances = build_covariance(scale=scales, rotation_xyzw=rotations) # [B, V, H*W, 3, 3]
-
- # means
- # [B, V, H*W, 1, 2]
- # xy_ray = xy_ray.unsqueeze(0).unsqueeze(0).repeat(b, v, 1, 1, 1)
- origins, directions = get_world_rays(xy_ray,
- context["extrinsics"].unsqueeze(2).unsqueeze(2),
- context["intrinsics"].unsqueeze(2).unsqueeze(2))
- means = origins + directions * depths
-
- # sh: [B, V, HW, 3, SH]
- sh = rearrange(sh, "... (xyz d_sh) -> ... xyz d_sh", xyz=3).clone()
- # sh = sh.broadcast_to((*opacities.shape, 3, self.gaussian_adapter.d_sh)).clone()
-
- # [B, V, H*W, 3]
- sh_input_images = rearrange(sh_input_images, "b v c h w -> b v (h w) c")
- # init sh with input images
- sh[..., 0] = sh[..., 0] + RGB2SH(sh_input_images)
-
- gaussians = Gaussians(
- means=rearrange(means, "b v r spp xyz -> b (v r spp) xyz"),
- covariances=rearrange(covariances, "b v r i j -> b (v r) i j"),
- harmonics=rearrange(sh, "b v r c d_sh -> b (v r) c d_sh"),
- opacities=rearrange(opacities, "b v r spp -> b (v r spp)"),
- scales=rearrange(scales, "b v r xyz -> b (v r) xyz"),
- rotations=rearrange(rotations, "b v r wxyz -> b (v r) wxyz"), # in wxyz format
- rotations_unnorm=rearrange(rotations_unnorm, "b v r wxyz -> b (v r) wxyz") # in wxyz format
+ else:
+ rotations_unnorm, scales, opacities_raw, offset, sh = raw_gaussians.split(
+ [4 * repeat, 3 * repeat, 1 * repeat, 2 * repeat, 3 * num_sh * repeat],
+ dim=-1,
)
+ latent_h, latent_w = gaussians.shape[-2:]
+
+ if repeat > 1:
+ # reshape all the gaussian parameters
+ if True or self.cfg.latent_new_reshape:
+ # this works
+ r = int(np.sqrt(repeat))
+ rotations_unnorm = rearrange(rotations_unnorm, "b v (h w) (c x y) -> b v (h x w y) c",
+ h=latent_h, w=latent_w, x=r, y=r)
+ scales = rearrange(scales, "b v (h w) (c x y) -> b v (h x w y) c", h=latent_h, w=latent_w, x=r,
+ y=r)
+ opacities_raw = rearrange(opacities_raw, "b v (h w) (c x y) -> b v (h x w y) c", h=latent_h,
+ w=latent_w, x=r, y=r)
+ offset = rearrange(offset, "b v (h w) (c x y) -> b v (h x w y) c", h=latent_h, w=latent_w, x=r,
+ y=r)
+ sh = rearrange(sh, "b v (h w) (c x y) -> b v (h x w y) c", h=latent_h, w=latent_w, x=r, y=r)
+ else:
+ # doesn't work
+ rotations_unnorm = rearrange(rotations_unnorm, "b v hw (k c) -> b v (hw k) c", k=repeat)
+ scales = rearrange(scales, "b v hw (k c) -> b v (hw k) c", k=repeat)
+ opacities_raw = rearrange(opacities_raw, "b v hw (k c) -> b v (hw k) c", k=repeat)
+ offset = rearrange(offset, "b v hw (k c) -> b v (hw k) c", k=repeat)
+ sh = rearrange(sh, "b v hw (k c) -> b v (hw k) c", k=repeat)
+
+ opacities = opacities_raw.sigmoid() # [B, V, H*W*K, 1]
+
+ if self.cfg.latent_downsample == 4 and self.cfg.init_gaussian_multiple == 4:
+ scale_factor = 2
+ elif self.cfg.latent_downsample == 2 and self.cfg.init_gaussian_multiple == 4:
+ scale_factor = 2
+ elif self.cfg.latent_downsample == 4 and self.cfg.init_gaussian_multiple == 16:
+ scale_factor = 4
+ elif self.cfg.latent_downsample == 8 and self.cfg.init_gaussian_multiple == 4:
+ scale_factor = 2
+ elif self.cfg.latent_downsample == 8 and self.cfg.init_gaussian_multiple == 16:
+ scale_factor = 4
else:
- gaussians = rearrange(gaussians, "(b v) c h w -> b v c h w", b=b, v=v)
+ scale_factor = 1
- # [B, V, H*W, 84]
- raw_gaussians = rearrange(
- gaussians, "b v c h w -> b v (h w) c")
+ h, w = latent_h * scale_factor, latent_w * scale_factor
- assert len(depth_preds) == 1, "num_scales must be 1; multi-scale depth supervision is not supported"
+ # unproject depth
+ xy_ray, _ = sample_image_grid((h, w), device)
+ xy_ray = rearrange(xy_ray, "h w xy -> (h w) () xy")
- # [B, V, H*W, C]
- repeat = self.cfg.init_gaussian_multiple
- num_sh = self.gaussian_adapter.d_sh
+ if self.cfg.no_pixel_offset:
+ offset_xy = torch.ones_like(raw_gaussians[..., :2]).unsqueeze(-2).to(
+ raw_gaussians.device) * 0.5 # [B, V, H*W, 1, 2]
+ else:
+ offset_xy = offset.sigmoid().unsqueeze(-2) # [B, V, H*W, 1, 2]
+
+ pixel_size = 1 / \
+ torch.tensor((w, h), dtype=torch.float32, device=device)
+ # [H*W, 1, 2]
+ if self.cfg.deform_sample_depth and not self.cfg.deform_sample_depth_debug:
+ # (offset_xy - 0.5) in -0.5 to 0.5, without multiplying by pixel size such that the points can move in the image space
+ xy_ray = (xy_ray + (offset_xy - 0.5)).clamp(min=0., max=1.)
+ else:
+ xy_ray = xy_ray + (offset_xy - 0.5) * pixel_size
- if self.cfg.no_pixel_offset:
- rotations_unnorm, scales, opacities_raw, sh = raw_gaussians.split(
- [4 * repeat, 3 * repeat, 1 * repeat, 3 * num_sh * repeat],
- dim=-1,
- )
- else:
- rotations_unnorm, scales, opacities_raw, offset, sh = raw_gaussians.split(
- [4 * repeat, 3 * repeat, 1 * repeat, 2 * repeat, 3 * num_sh * repeat],
- dim=-1,
- )
+ if self.cfg.deform_sample_depth:
+ # use low-res xy_ray to sample full-res depth
- latent_h, latent_w = gaussians.shape[-2:]
-
- if repeat > 1:
- # reshape all the gaussian parameters
- if True or self.cfg.latent_new_reshape:
- # this works
- r = int(np.sqrt(repeat))
- rotations_unnorm = rearrange(rotations_unnorm, "b v (h w) (c x y) -> b v (h x w y) c",
- h=latent_h, w=latent_w, x=r, y=r)
- scales = rearrange(scales, "b v (h w) (c x y) -> b v (h x w y) c", h=latent_h, w=latent_w, x=r,
- y=r)
- opacities_raw = rearrange(opacities_raw, "b v (h w) (c x y) -> b v (h x w y) c", h=latent_h,
- w=latent_w, x=r, y=r)
- offset = rearrange(offset, "b v (h w) (c x y) -> b v (h x w y) c", h=latent_h, w=latent_w, x=r,
- y=r)
- sh = rearrange(sh, "b v (h w) (c x y) -> b v (h x w y) c", h=latent_h, w=latent_w, x=r, y=r)
- else:
- # doesn't work
- rotations_unnorm = rearrange(rotations_unnorm, "b v hw (k c) -> b v (hw k) c", k=repeat)
- scales = rearrange(scales, "b v hw (k c) -> b v (hw k) c", k=repeat)
- opacities_raw = rearrange(opacities_raw, "b v hw (k c) -> b v (hw k) c", k=repeat)
- offset = rearrange(offset, "b v hw (k c) -> b v (hw k) c", k=repeat)
- sh = rearrange(sh, "b v hw (k c) -> b v (hw k) c", k=repeat)
+ sample_grid = rearrange(xy_ray, "b v (h w) c xy -> (b v) h w (c xy)", h=h, w=w) # in [0, 1]
+ # to [-1, 1]
+ sample_grid = 2 * (sample_grid - 0.5) # [BV, h, w, 2]
- opacities = opacities_raw.sigmoid() # [B, V, H*W*K, 1]
+ fullres_depth = rearrange(depth, "b v h w -> (b v) () h w") # [BV, 1, H, W]
+ sampled_depth = F.grid_sample(fullres_depth, sample_grid, mode='bilinear', align_corners=True,
+ padding_mode="border") # [BV, 1, h, w]
+ # reshape
+ depths = rearrange(sampled_depth, "(b v) () h w -> b v (h w) () ()", b=b, v=v, h=h, w=w)
+ if self.cfg.latent_gs:
+ sh_input_images = rearrange(context["image"], "b v c h w -> (b v) c h w")
if self.cfg.latent_downsample == 4 and self.cfg.init_gaussian_multiple == 4:
- scale_factor = 2
- elif self.cfg.latent_downsample == 2 and self.cfg.init_gaussian_multiple == 4:
- scale_factor = 2
+ sh_input_images = F.interpolate(sh_input_images, scale_factor=0.5, mode='area')
elif self.cfg.latent_downsample == 4 and self.cfg.init_gaussian_multiple == 16:
- scale_factor = 4
+ pass
+ elif self.cfg.latent_downsample == 2 and self.cfg.init_gaussian_multiple == 4:
+ pass
elif self.cfg.latent_downsample == 8 and self.cfg.init_gaussian_multiple == 4:
- scale_factor = 2
+ sh_input_images = F.interpolate(sh_input_images, scale_factor=0.25, mode='area')
elif self.cfg.latent_downsample == 8 and self.cfg.init_gaussian_multiple == 16:
- scale_factor = 4
- else:
- scale_factor = 1
-
- h, w = latent_h * scale_factor, latent_w * scale_factor
-
- # unproject depth
- xy_ray, _ = sample_image_grid((h, w), device)
- xy_ray = rearrange(xy_ray, "h w xy -> (h w) () xy")
-
- if self.cfg.no_pixel_offset:
- offset_xy = torch.ones_like(raw_gaussians[..., :2]).unsqueeze(-2).to(
- raw_gaussians.device) * 0.5 # [B, V, H*W, 1, 2]
- else:
- offset_xy = offset.sigmoid().unsqueeze(-2) # [B, V, H*W, 1, 2]
-
- pixel_size = 1 / \
- torch.tensor((w, h), dtype=torch.float32, device=device)
- # [H*W, 1, 2]
- if self.cfg.deform_sample_depth and not self.cfg.deform_sample_depth_debug:
- # (offset_xy - 0.5) in -0.5 to 0.5, without multiplying by pixel size such that the points can move in the image space
- xy_ray = (xy_ray + (offset_xy - 0.5)).clamp(min=0., max=1.)
+ sh_input_images = F.interpolate(sh_input_images, scale_factor=0.5, mode='area')
else:
- xy_ray = xy_ray + (offset_xy - 0.5) * pixel_size
+ sh_input_images = F.interpolate(sh_input_images, scale_factor=1. / self.cfg.latent_downsample,
+ mode='area')
- if self.cfg.deform_sample_depth:
- # use low-res xy_ray to sample full-res depth
+ sh_input_images = rearrange(sh_input_images, "(b v) c h w -> b v c h w", b=b, v=v)
- sample_grid = rearrange(xy_ray, "b v (h w) c xy -> (b v) h w (c xy)", h=h, w=w) # in [0, 1]
- # to [-1, 1]
- sample_grid = 2 * (sample_grid - 0.5) # [BV, h, w, 2]
-
- fullres_depth = rearrange(depth, "b v h w -> (b v) () h w") # [BV, 1, H, W]
- sampled_depth = F.grid_sample(fullres_depth, sample_grid, mode='bilinear', align_corners=True,
- padding_mode="border") # [BV, 1, h, w]
- # reshape
- depths = rearrange(sampled_depth, "(b v) () h w -> b v (h w) () ()", b=b, v=v, h=h, w=w)
-
- if self.cfg.latent_gs:
- sh_input_images = rearrange(context["image"], "b v c h w -> (b v) c h w")
- if self.cfg.latent_downsample == 4 and self.cfg.init_gaussian_multiple == 4:
- sh_input_images = F.interpolate(sh_input_images, scale_factor=0.5, mode='area')
- elif self.cfg.latent_downsample == 4 and self.cfg.init_gaussian_multiple == 16:
- pass
- elif self.cfg.latent_downsample == 2 and self.cfg.init_gaussian_multiple == 4:
- pass
- elif self.cfg.latent_downsample == 8 and self.cfg.init_gaussian_multiple == 4:
- sh_input_images = F.interpolate(sh_input_images, scale_factor=0.25, mode='area')
- elif self.cfg.latent_downsample == 8 and self.cfg.init_gaussian_multiple == 16:
- sh_input_images = F.interpolate(sh_input_images, scale_factor=0.5, mode='area')
- else:
- sh_input_images = F.interpolate(sh_input_images, scale_factor=1. / self.cfg.latent_downsample,
- mode='area')
-
- sh_input_images = rearrange(sh_input_images, "(b v) c h w -> b v c h w", b=b, v=v)
-
- else:
- sh_input_images = context["image"]
-
- assert len(depth_preds) == 1, "num_scales must be 1; multi-scale depth supervision is not supported"
-
- # build gaussians
- # scale
- scales = torch.clamp(F.softplus(scales - self.cfg.gaussian_adapter.exp_scale_bias),
- min=self.cfg.gaussian_adapter.clamp_min_scale,
- max=self.cfg.gaussian_adapter.gaussian_scale_max
- )
-
- # Convert rotations to world-space
- c2w_rotations = context["extrinsics"][..., :3, :3].unsqueeze(2) # [B, V, 1, 3, 3]
- # Here quaternions follow the xyzw format (scalar last)
- rotations = rotate_quats(c2w_rotations, rotations_unnorm)
- rotations_unnorm = rotations.clone()
- # Create world-space covariance matrices.
- covariances = build_covariance(scale=scales, rotation_xyzw=rotations) # [B, V, H*W, 3, 3]
-
- # means
- # [B, V, H*W, 1, 2]
- origins, directions = get_world_rays(xy_ray,
- context["extrinsics"].unsqueeze(2).unsqueeze(2),
- context["intrinsics"].unsqueeze(2).unsqueeze(2))
- means = origins + directions * depths
-
- # sh: [B, V, HW, 3, SH]
- sh = rearrange(sh, "... (xyz d_sh) -> ... xyz d_sh", xyz=3).clone()
-
- # [B, V, H*W, 3]
- sh_input_images = rearrange(sh_input_images, "b v c h w -> b v (h w) c")
- # init sh with input images
- sh[..., 0] = sh[..., 0] + RGB2SH(sh_input_images)
-
- gaussians = Gaussians(
- means=rearrange(means, "b v r spp xyz -> b (v r spp) xyz"),
- covariances=rearrange(covariances, "b v r i j -> b (v r) i j"),
- harmonics=rearrange(sh, "b v r c d_sh -> b (v r) c d_sh"),
- opacities=rearrange(opacities, "b v r spp -> b (v r spp)"),
- scales=rearrange(scales, "b v r xyz -> b (v r) xyz"),
- rotations=rearrange(rotations, "b v r wxyz -> b (v r) wxyz"),
- rotations_unnorm=rearrange(rotations_unnorm, "b v r wxyz -> b (v r) wxyz")
- )
+ else:
+ sh_input_images = context["image"]
+
+ assert len(depth_preds) == 1, "num_scales must be 1; multi-scale depth supervision is not supported"
+
+ # build gaussians
+ # scale
+ scales = torch.clamp(F.softplus(scales - self.cfg.gaussian_adapter.exp_scale_bias),
+ min=self.cfg.gaussian_adapter.clamp_min_scale,
+ max=self.cfg.gaussian_adapter.gaussian_scale_max
+ )
+
+ # Convert rotations to world-space
+ c2w_rotations = context["extrinsics"][..., :3, :3].unsqueeze(2) # [B, V, 1, 3, 3]
+ # Here quaternions follow the xyzw format (scalar last)
+ rotations = rotate_quats(c2w_rotations, rotations_unnorm)
+ rotations_unnorm = rotations.clone()
+ # Create world-space covariance matrices.
+ covariances = build_covariance(scale=scales, rotation_xyzw=rotations) # [B, V, H*W, 3, 3]
+
+ # means
+ # [B, V, H*W, 1, 2]
+ origins, directions = get_world_rays(xy_ray,
+ context["extrinsics"].unsqueeze(2).unsqueeze(2),
+ context["intrinsics"].unsqueeze(2).unsqueeze(2))
+ means = origins + directions * depths
+
+ # sh: [B, V, HW, 3, SH]
+ sh = rearrange(sh, "... (xyz d_sh) -> ... xyz d_sh", xyz=3).clone()
+
+ # [B, V, H*W, 3]
+ sh_input_images = rearrange(sh_input_images, "b v c h w -> b v (h w) c")
+ # init sh with input images
+ sh[..., 0] = sh[..., 0] + RGB2SH(sh_input_images)
+
+ gaussians = Gaussians(
+ means=rearrange(means, "b v r spp xyz -> b (v r spp) xyz"),
+ covariances=rearrange(covariances, "b v r i j -> b (v r) i j"),
+ harmonics=rearrange(sh, "b v r c d_sh -> b (v r) c d_sh"),
+ opacities=rearrange(opacities, "b v r spp -> b (v r spp)"),
+ scales=rearrange(scales, "b v r xyz -> b (v r) xyz"),
+ rotations=rearrange(rotations, "b v r wxyz -> b (v r) wxyz"),
+ rotations_unnorm=rearrange(rotations_unnorm, "b v r wxyz -> b (v r) wxyz")
+ )
# Dump visualizations if needed.
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
)
- # if self.cfg.pt_downsample > 0:
- # visualization_dump["scales"] = gaussians.scales
- # visualization_dump["rotations"] = gaussians.rotations
- # else:
- # 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"
- # )
if self.cfg.return_depth:
# return depth prediction for supervision
@@ -1280,22 +944,6 @@ class ResplatInitializer(LearnedInitializer[ResplatInitializerCfg]):
batch["target"]["near"] = batch["target"]["depth"].min(dim=3)[0].min(dim=2)[0].clamp(min=0.01)
batch["target"]["far"] = batch["target"]["depth"].max(dim=3)[0].max(dim=2)[0].clamp(max=1000.)
- def update_depth_range_from_disparity(self, batch):
- b, v, _, h, w = batch["context"]["image"].shape
- # TODO: support multi-view later
- assert v == 2
- assert self.decoder.cfg.scale_invariant is False
- w = batch["context"]["image"].shape[-1]
- # compute the depth range based on disparity range
- dist = (batch["context"]["extrinsics"][:, 0, :3, 3] - batch["context"]["extrinsics"][:, 1, :3, 3]).norm(
- dim=1, keepdim=True)
- focal = batch["context"]["intrinsics"][:, :, 0, 0] * w
- min_depth = dist * focal / self.train_cfg.max_disparity
- max_depth = dist * focal / self.train_cfg.min_disparity
- batch["context"]["near"] = min_depth
- batch["context"]["far"] = max_depth
- # TODO: also update target near and far
-
def predict_scale(self, batch):
context = batch["context"]
# [B, V, H, W]
@@ -1354,13 +1002,10 @@ class ResplatInitializer(LearnedInitializer[ResplatInitializerCfg]):
batch["context"]["extrinsics"][:, :, :3, -1] /= norm_factor.view(b, 1, 1)
batch["target"]["extrinsics"][:, :, :3, -1] /= norm_factor.view(b, 1, 1)
- def preprocessing(self, batch, train_cfg):
+ def eval_preprocessing(self, batch, train_cfg):
# use gt depth range instead of a fixed one
if train_cfg.use_gt_depth_range:
self.update_gt_depth_range(batch)
- # compute depth range from camera distance and disparity range
- if train_cfg.depth_range_from_disparity:
- self.update_depth_range_from_disparity(batch)
# use a pretrained depth model to predict scale
if self.cfg.predict_scale:
diff --git a/optgs/scene_trainer/optimizer/__init__.py b/optgs/scene_trainer/optimizer/__init__.py
index 05769793054a3f8e6fc2482279e1ca9ceee0ef24..dc32d41ea6f9947f2dd604ee3cae65ebba455132 100644
--- a/optgs/scene_trainer/optimizer/__init__.py
+++ b/optgs/scene_trainer/optimizer/__init__.py
@@ -1,26 +1,15 @@
from .optimizer import Optimizer, OptimizerCfg
-from .optimizer_learn2splat import Learn2SplatOptimizer
+from .optimizer_adam import AdamOptimizerCfg, AdamOptimizer
from .optimizer_knn_based import KnnBasedOptimizer, KnnBasedOptimizerCfg
+from .optimizer_learn2splat import Learn2SplatOptimizer
from .optimizer_resplat import ResplatOptimizerV1, ResplatOptimizerV2
-from .optimizer_adam import AdamOptimizerCfg, AdamOptimizer
SceneOptimizerCfg = KnnBasedOptimizerCfg | AdamOptimizerCfg
-def extract_opt_params(cfg):
- opt_params = cfg.__dict__.copy()
- opt_params.pop("enabled", None)
- opt_params.pop("name", None)
- opt_params.pop("steps", None)
- opt_params.pop("scheduler", None)
- opt_params.pop("scheduler_warm_up_ratio", None)
- opt_params.pop("compute_metrics_every", None)
- return opt_params
-
-
SCENE_OPTIMIZERS = {
"none": None,
- "depthsplat": KnnBasedOptimizer,
+ "knn_based": KnnBasedOptimizer,
"resplat_v1": ResplatOptimizerV1,
"resplat_v2": ResplatOptimizerV2,
"clogs": Learn2SplatOptimizer, # TODO (release): remove
diff --git a/optgs/scene_trainer/optimizer/layer.py b/optgs/scene_trainer/optimizer/layer.py
index 641a68f3c6f64075121f344e41e13e8656e74a29..b387828442cc731338f46e1cd46ee70ef5649528 100644
--- a/optgs/scene_trainer/optimizer/layer.py
+++ b/optgs/scene_trainer/optimizer/layer.py
@@ -1,18 +1,40 @@
+import functools
+
import torch
from jaxtyping import Float
from torch import nn as nn, Tensor
+def _lazy_compile(**compile_kwargs):
+ """Like ``@torch.compile`` but defers the (heavy) dynamo/inductor import
+ until the function is first called. Applying ``torch.compile`` at import
+ eagerly loads the compile stack (~3.5s); this keeps it off the import path
+ and skips it entirely for runs that never call the wrapped function."""
+ def decorator(fn):
+ compiled = None
+
+ @functools.wraps(fn)
+ def wrapper(*args, **kwargs):
+ nonlocal compiled
+ if compiled is None:
+ compiled = torch.compile(fn, **compile_kwargs)
+ return compiled(*args, **kwargs)
+
+ return wrapper
+
+ return decorator
+
+
class SlicedG3RNorm(nn.Module):
def __init__(self, num_features, input_slice, eps=1e-8):
"""
Apply G3R normalization to a slice of the input features.
- Devide the input by the maximum absolute value in each channel.
+ Divide the input by the maximum absolute value in each channel.
Args:
num_features (int): Total number of features (channels).
- input_slice (slice): Size of each slice to normalize independently.
+ input_slice (slice): Channel slice to normalize independently.
eps (float): Small constant to prevent division by zero.
"""
super().__init__()
@@ -44,50 +66,6 @@ class SlicedG3RNorm(nn.Module):
return x
-class SlicedBatchNorm1d(nn.Module):
- def __init__(self, num_features, input_slice, eps=1e-8, affine=False, track_running_stats=True):
- """
- Apply normalization independently to a slice of the input features.
-
- Args:
- num_features (int): Total number of features (channels).
- input_slice (slice): Size of each slice to normalize independently.
- eps (float): Small constant to prevent division by zero.
- affine (bool): Whether to include learnable scale and bias per slice.
- """
- super().__init__()
- self.input_slice = input_slice
- dummy = torch.zeros(1, num_features)
- chunk = dummy[:, input_slice]
-
- self.slice_size = chunk.shape[-1]
- self.eps = eps
-
- # Create a BatchNorm1d module for each slice
- self.slice_norm = nn.BatchNorm1d(self.slice_size, eps=eps, affine=affine,
- track_running_stats=track_running_stats)
-
- def forward(self, x):
- """
- Args:
- x (Tensor): Shape (B, C) where C = num_features
- Returns:
- Tensor: Same shape, only subset of channels normalized
- """
- B, C = x.shape
-
- # Split input into the slice to normalize and the rest
- chunk = x[:, self.input_slice]
-
- # Apply normalization to the selected slice
- chunk = self.slice_norm(chunk)
-
- # Replace the normalized slice back into the original input
- x = x.clone()
- x[:, self.input_slice] = chunk
- return x
-
-
class CustomGroupNorm(nn.Module):
def __init__(self, group_sizes, eps=1e-8, affine=True):
"""
@@ -110,7 +88,7 @@ class CustomGroupNorm(nn.Module):
def forward(self, x):
"""
Args:
- x (Tensor): Shape (B, C, H, W)
+ x (Tensor): Shape (B, C)
Returns:
Tensor: Same shape, group-wise normalized
"""
@@ -148,7 +126,7 @@ def slice_length(s, dim):
max(0, (start - stop + (-step - 1)) // -step)
-@torch.compile(dynamic=True)
+@_lazy_compile(dynamic=True)
def _adam_smooth_unmasked(m, v, t, chunk, beta1, beta2, eps) -> Tensor:
"""Fused moment update + bias-corrected output for the unmasked path."""
m.lerp_(chunk, 1 - beta1)
@@ -160,7 +138,7 @@ def _adam_smooth_unmasked(m, v, t, chunk, beta1, beta2, eps) -> Tensor:
return m.div(bias1).div_(denom)
-@torch.compile(dynamic=True)
+@_lazy_compile(dynamic=True)
def _adam_smooth_masked(m, v, t, sel, chunk, beta1, beta2, eps) -> Tensor:
"""Fused moment update + bias-corrected output for the masked path."""
m_sel = m[sel].lerp_(chunk, 1 - beta1)
@@ -291,7 +269,7 @@ class AdamInputSmoothing(nn.Module):
def zero_out(self, zero_t=False) -> None:
"""Zero out the moments. Called when resetting gaussians opacities."""
assert not self.is_reset(), (
- "Cannot extend state that has not been initialized. Call forward() at least once first."
+ "Cannot zero out state that has not been initialized. Call forward() at least once first."
)
self.m = torch.zeros_like(self.m)
self.v = torch.zeros_like(self.v)
@@ -301,7 +279,7 @@ class AdamInputSmoothing(nn.Module):
def replace(self, from_indices: Tensor, dest_indices: Tensor, zero_t=False) -> None:
"""Replace the internal state to duplicate entries at the specified indices."""
assert not self.is_reset(), (
- "Cannot extend state that has not been initialized. Call forward() at least once first."
+ "Cannot replace state that has not been initialized. Call forward() at least once first."
)
self.m[dest_indices] = self.m[from_indices]
diff --git a/optgs/scene_trainer/optimizer/lr_scheduler.py b/optgs/scene_trainer/optimizer/lr_scheduler.py
index a6064dbf7f2d63a00cef3e75d5e595db85e0409e..5437ba0416ac9687341341577e83f4381f5f1e64 100644
--- a/optgs/scene_trainer/optimizer/lr_scheduler.py
+++ b/optgs/scene_trainer/optimizer/lr_scheduler.py
@@ -1,5 +1,6 @@
from dataclasses import dataclass
+from optgs.misc.general_utils import get_expon_lr_func
from optgs.scene_trainer.optimizer.optimizer_utils import Number3DGSCfg, Bool3DGSCfg
@@ -17,7 +18,15 @@ class DDIMSchedulerCfg(SchedulerCfg):
s: float | int
-LrSchedulerCfgType = DDIMSchedulerCfg | SchedulerCfg
+@dataclass
+class ExponSchedulerCfg(SchedulerCfg):
+ lr_final: float | int # final LR the schedule decays to (in _base=1 units)
+ lr_delay_steps: int # 0 disables the warm-up ramp
+ lr_delay_mult: float | int
+ max_steps: int # step at which lr_final is reached; = total optimization steps
+
+
+LrSchedulerCfgType = DDIMSchedulerCfg | ExponSchedulerCfg | SchedulerCfg
class Scheduler[SchedulerCfg]:
@@ -36,11 +45,6 @@ class Scheduler[SchedulerCfg]:
raise NotImplementedError
-class DummyScheduler(Scheduler[SchedulerCfg]):
- def get_lr(self, t: int, param: str) -> float | int:
- return 1
-
-
class DDIMCosineScheduler(Scheduler[DDIMSchedulerCfg]):
def scheduler_fn(self, t: int, base_lr: float | int) -> float | int:
# Implement DDIM Cosine scheduling logic here
@@ -52,9 +56,27 @@ class DDIMCosineScheduler(Scheduler[DDIMSchedulerCfg]):
return lr
+class ExponScheduler(Scheduler[ExponSchedulerCfg]):
+ """Log-linear (exponential) LR decay, as used by the 3DGS means schedule."""
+
+ def scheduler_fn(self, t: int, base_lr: float | int) -> float | int:
+ # base_lr is the per-param LR (lr_data.param = _base * _param); lr_final carries the same
+ # _base multiplier so the whole schedule scales uniformly with _base. The schedule is
+ # linear in its endpoints, so a caller can pull out scene-extent scaling by multiplying
+ # the result (see AdamOptimizer means LR).
+ return get_expon_lr_func(
+ lr_init=base_lr,
+ lr_final=self.cfg.lr_final * self.cfg.lr_data.base,
+ lr_delay_steps=self.cfg.lr_delay_steps,
+ lr_delay_mult=self.cfg.lr_delay_mult,
+ max_steps=self.cfg.max_steps,
+ )(t)
+
+
SCHEDULERS = {
"none": Scheduler,
"ddim": DDIMCosineScheduler,
+ "expon": ExponScheduler,
}
def get_scheduler(cfg: SchedulerCfg) -> Scheduler:
diff --git a/optgs/scene_trainer/optimizer/optimizer.py b/optgs/scene_trainer/optimizer/optimizer.py
index bb1f2030a85c1799732ab343ccba19e550e41642..f9ed156f7d48eb9ce299aa5096750b88274ab1a3 100644
--- a/optgs/scene_trainer/optimizer/optimizer.py
+++ b/optgs/scene_trainer/optimizer/optimizer.py
@@ -1,3 +1,13 @@
+"""Base classes and shared types for the scene optimizers.
+
+Defines the optimizer interface that `SceneTrainer` calls: `Optimizer` (the abstract base, with the
+per-iteration timing/benchmarking and ADC plumbing) and `LearnedOptimizer` (adds the network-based
+machinery). Also defines the pipeline data types passed in and out each iteration: `OptimizerInput`,
+`OptimizerOutput`, `OptimizerState`, `OptimizerPreviousOutput`, and the `OptimizerCfg` base config. The
+learned optimizer used in the paper is `KnnBasedOptimizer` (optimizer_knn_based.py); the 3DGS Adam
+baseline is `AdamOptimizer` (optimizer_adam.py).
+"""
+
from abc import ABC
from dataclasses import dataclass, field
from pathlib import Path
@@ -6,9 +16,8 @@ import torch
from matplotlib import pyplot as plt
from torch import nn
from torch import Tensor
-import numpy as np
import os
-from optgs.dataset.camera_datasets.camera import get_scene_scale
+from optgs.misc.benchmarker import Benchmarker
from optgs.misc.io import FrequencyScheduler
from optgs.dataset.data_types import BatchedViews
from optgs.model.decoder import Decoder
@@ -16,7 +25,6 @@ from optgs.model.decoder.decoder import DecoderOutput
from optgs.model.types import Gaussians
from optgs.scene_trainer.adc.base import BaseStrategyCfg
from optgs.scene_trainer.initializer.initializer import InitializerOutput
-from optgs.scene_trainer.optimizer.layer import AdamState
from optgs.scene_trainer.initializer import InitializerCfg
from optgs.misc.detaching_cpu_list import DetachingCPUList
from optgs.scene_trainer.optimizer.lr_scheduler import LrSchedulerCfgType, get_scheduler
@@ -28,9 +36,6 @@ if TYPE_CHECKING:
@dataclass
class OptimizerState:
- state: torch.Tensor | None = None
- init_state: torch.Tensor | None = None # state at the beginning of the optimization
- adam_state: AdamState | None = None
adc_state: Any = None # VanillaStrategyState | McmcStrategyState | None
@@ -38,6 +43,48 @@ class OptimizerState:
class OptimizerPreviousOutput:
gaussians: Gaussians
state: OptimizerState | None = None
+ # Optional pre-computed renders of `gaussians`, used by the splice in
+ # SceneTrainer.get_optimized_gaussians to put init/resumed renders at position 0
+ # of optimizer_output's render lists. Mirrors the same fields on InitializerOutput
+ # so the splice can treat both prev_output types uniformly.
+ target_render: DecoderOutput | None = None
+ context_render: DecoderOutput | None = None
+
+ # View indices used when rendering a subset (training); None means all views were rendered.
+ target_render_index: "torch.Tensor | None" = None
+ context_render_index: "torch.Tensor | None" = None
+
+ def get_render(self, which: str) -> "DecoderOutput | None":
+ if which == "target":
+ return self.target_render
+ elif which == "context":
+ return self.context_render
+ else:
+ raise ValueError(f"Unknown which: {which}, should be 'target' or 'context'")
+
+ def set_render(self, which: str, value: "DecoderOutput") -> None:
+ if which == "target":
+ self.target_render = value
+ elif which == "context":
+ self.context_render = value
+ else:
+ raise ValueError(f"Unknown which: {which}, should be 'target' or 'context'")
+
+ def get_render_index(self, which: str) -> "torch.Tensor | None":
+ if which == "target":
+ return self.target_render_index
+ elif which == "context":
+ return self.context_render_index
+ else:
+ raise ValueError(f"Unknown which: {which}, should be 'target' or 'context'")
+
+ def set_render_index(self, which: str, value: "torch.Tensor | None") -> None:
+ if which == "target":
+ self.target_render_index = value
+ elif which == "context":
+ self.context_render_index = value
+ else:
+ raise ValueError(f"Unknown which: {which}, should be 'target' or 'context'")
@dataclass
@@ -45,12 +92,10 @@ class OptimizerInput:
context: BatchedViews
renderer: Decoder
prev_output: InitializerOutput | OptimizerPreviousOutput
- num_refine: int
iter_batch_size: int | None
target: BatchedViews | None = None
context_remain: dict | None = None
debug_dict: dict | None = None
- additional_info: tuple | None = None
@property
def device(self) -> torch.device:
@@ -59,7 +104,6 @@ class OptimizerInput:
@dataclass
class OptimizerOutput:
- # TODO Naama: should we add here iterations?
gaussian_list: DetachingCPUList[Gaussians]
t: int | None = None
T: int | None = None
@@ -100,12 +144,12 @@ class OptimizerOutput:
class OptimizerCfg:
# subset optimization flags
- no_refine_mean: bool
- no_refine_scale: bool
- no_refine_rotation: bool
- no_refine_opacity: bool
- no_refine_sh0: bool
- no_refine_shN: bool
+ freeze_mean: bool
+ freeze_scale: bool
+ freeze_rotation: bool
+ freeze_opacity: bool
+ freeze_sh0: bool
+ freeze_shN: bool
# lr scheduler
lr_scheduler: LrSchedulerCfgType
@@ -131,14 +175,14 @@ class OptimizerCfg:
@property
def optimize_all(self):
- # All the no_refine_* are False
+ # All the freeze_* are False
return not any([
- self.no_refine_mean,
- self.no_refine_scale,
- self.no_refine_rotation,
- self.no_refine_opacity,
- self.no_refine_sh0,
- self.no_refine_shN,
+ self.freeze_mean,
+ self.freeze_scale,
+ self.freeze_rotation,
+ self.freeze_opacity,
+ self.freeze_sh0,
+ self.freeze_shN,
])
@@ -153,49 +197,87 @@ class Optimizer(nn.Module, ABC, Generic[T]):
self.cfg = cfg
self.save_every = save_every
- # for timing
- self.iter_start = torch.cuda.Event(enable_timing=True)
- self.iter_end = torch.cuda.Event(enable_timing=True)
- # decoder_event_start/end bracket only the rendering-for-gradients call inside
- # apply_one_update_step, letting us split iter_time into decoder vs optimizer.
- self.decoder_event_start = torch.cuda.Event(enable_timing=True)
- self.decoder_event_end = torch.cuda.Event(enable_timing=True)
+ # Per-iteration timing. The optimizer brackets each iteration with benchmarker.time("iter") and
+ # the render inside it with benchmarker.time("decoder"); these use CUDA events with a deferred
+ # sync, so the loop is not stalled every step. iter/decoder/optimizer_time_log read back from it.
+ self.benchmarker = Benchmarker()
# scene_start_event_start/end bracket optimizer.on_scene_start() (KNN, Adam init).
# Read after the post-loop cuda.synchronize() in scene_trainer.get_optimized_gaussians.
self.scene_start_event_start = torch.cuda.Event(enable_timing=True)
self.scene_start_event_end = torch.cuda.Event(enable_timing=True)
- # Init logs for densification/pruning
- self.radii_max_log = []
- self.grads_max_log = []
- self.nr_cloned_log = []
- self.nr_splitted_log = []
- self.nr_pruned_log = []
- self.nr_gaussians_log = []
- self.iter_time_log = [] # total ms per iteration
- self.decoder_time_log = [] # ms spent in rendering-for-gradients per iteration
- self.optimizer_time_log = [] # ms spent in update step (iter_time - decoder_time)
- self.scene_start_ms = 0.0 # ms for on_scene_start (KNN lookup, Adam state init)
- self.nr_nonzero_grad_log = []
-
# LR scheduler
self.scheduler = get_scheduler(self.cfg.lr_scheduler)
def forward(self, i, optimizer_input: OptimizerInput, optimizer_output: OptimizerOutput, **kwargs) -> OptimizerOutput:
return self._forward_impl(i, optimizer_input, optimizer_output, **kwargs)
- def _record_iter_timing(self) -> None:
- """Record per-iteration timing into iter/decoder/optimizer_time_log.
- Call right after the timed region; iter_start must already be recorded."""
- self.iter_end.record()
- torch.cuda.synchronize()
- elapsed_time = self.iter_start.elapsed_time(self.iter_end)
- self.iter_time_log.append(elapsed_time)
- decoder_ms = self.decoder_event_start.elapsed_time(self.decoder_event_end)
- self.decoder_time_log.append(decoder_ms)
- self.optimizer_time_log.append(elapsed_time - decoder_ms)
+ @property
+ def iter_time_log(self) -> list[float]:
+ """Total ms per optimization iteration (from benchmarker.time("iter"))."""
+ return self.benchmarker.execution_times["iter"]
+
+ @property
+ def decoder_time_log(self) -> list[float]:
+ """Ms spent rendering-for-gradients per iteration (from benchmarker.time("decoder"))."""
+ return self.benchmarker.execution_times["decoder"]
+
+ @property
+ def optimizer_time_log(self) -> list[float]:
+ """Ms spent in the update step per iteration (iter_time - decoder_time)."""
+ return [it - dec for it, dec in zip(self.iter_time_log, self.decoder_time_log)]
+
+ # Per-iteration stats, stored in `self.benchmarker` next to the timings (one entry per step).
+ # Counts are ints, max-radii/grads are floats, so the element type is float | int.
+ @property
+ def nr_gaussians_log(self) -> list[float | int]:
+ return self.benchmarker.execution_times["gaussians"]
+
+ @property
+ def nr_nonzero_grad_log(self) -> list[float | int]:
+ return self.benchmarker.execution_times["nonzero_grads"]
+
+ @property
+ def nr_cloned_log(self) -> list[float | int]:
+ return self.benchmarker.execution_times["cloned"]
+
+ @property
+ def nr_splitted_log(self) -> list[float | int]:
+ return self.benchmarker.execution_times["splitted"]
+
+ @property
+ def nr_pruned_log(self) -> list[float | int]:
+ return self.benchmarker.execution_times["pruned"]
+
+ @property
+ def radii_max_log(self) -> list[float | int]:
+ return self.benchmarker.execution_times["radii_max"]
+
+ @property
+ def grads_max_log(self) -> list[float | int]:
+ return self.benchmarker.execution_times["grads_max"]
+
+ @property
+ def scene_start_ms(self) -> float:
+ """Ms for on_scene_start (Adam/state init, preprocessing); one value per scene."""
+ times = self.benchmarker.execution_times["scene_start"]
+ return times[-1] if times else 0.0
+
+ @staticmethod
+ def _append_info(optimizer_output: "OptimizerOutput", gaussians, **info_lists) -> None:
+ """Append per-iteration debug info at test time (shared by all optimizers).
+
+ Saves `gaussians`, then appends each keyword value (e.g. deltas / grads /
+ normalized_grads / states_norms / learning_rates) to optimizer_output.info[key].
+ Each optimizer prepares its own already-CPU values before calling.
+ """
+ optimizer_output.gaussian_list.append(gaussians, detach_and_cpu=True)
+ assert optimizer_output.info is not None
+ for key, value in info_lists.items():
+ optimizer_output.info.setdefault(key, []).append(value)
def on_scene_start(self, optimizer_input: OptimizerInput) -> None:
+ self.benchmarker.clear_history() # isolate this scene's per-iteration timing
self._on_scene_start_impl(optimizer_input)
def _on_scene_start_impl(self, optimizer_input: OptimizerInput) -> None:
@@ -222,17 +304,8 @@ class Optimizer(nn.Module, ABC, Generic[T]):
pass
def reset_logs(self):
- self.radii_max_log = []
- self.grads_max_log = []
- self.nr_cloned_log = []
- self.nr_splitted_log = []
- self.nr_pruned_log = []
- self.nr_gaussians_log = []
- self.iter_time_log = []
- self.decoder_time_log = []
- self.optimizer_time_log = []
- self.scene_start_ms = 0.0
- self.nr_nonzero_grad_log = []
+ # All per-scene logs (timings, stats, scene_start) live in self.benchmarker; one clear resets them.
+ self.benchmarker.clear_history()
@staticmethod
def initialize_adc_state(cfg: OptimizerCfg, optimizer_input: OptimizerInput) -> None:
@@ -274,12 +347,18 @@ class Optimizer(nn.Module, ABC, Generic[T]):
full_context: BatchedViews,
full_target: BatchedViews,
) -> None:
- """Render and append post-update context+target views.
-
- Renders every iteration during training (so per-step renders can feed the meta-loss);
- otherwise renders only when save_every fires for the given tag. The per-iter subset
- (optimizer_input.context/target) is used in training when sampling indices exist,
- otherwise the full views.
+ """Render and append post-update context+target views into optimizer_output's render lists.
+
+ When and what is rendered:
+ - Training: render every iteration so per-step renders can feed the meta-loss.
+ If `opt_batch_size < V` (subset sampling), render only that per-iter subset
+ to match the views the optimization step saw.
+ - Test/eval (self.training=False): render only when save_every fires for this tag,
+ and ALWAYS use the full V views — even if opt_batch_size < V drove the
+ optimization step on a subset. This is the invariant the test-time
+ save/score paths in MetaTrainer._eval_and_save rely on: render_list[k]
+ is always [1, V, ...], so it lines up 1:1 with batch[tag]["index"][0]
+ (the full-V frame indices) used for filename labelling and metric GT.
"""
for tag, full, iter_views in (
("context", full_context, optimizer_input.context),
@@ -288,6 +367,8 @@ class Optimizer(nn.Module, ABC, Generic[T]):
if not (self.training or self.save_every(i + 1, tag=tag)):
continue
index_list = optimizer_output.get_index_list(tag)
+ # Subset rendering is training-only — at test time we always re-render the full V
+ # views so downstream save/score paths see uniform [1, V, ...] tensors.
subset = iter_views if (index_list and self.training) else full
render_output = optimizer_input.renderer.forward_batch_subset(
updated_gaussians,
@@ -347,17 +428,11 @@ class Optimizer(nn.Module, ABC, Generic[T]):
lr=lr
)
- self.nr_cloned_log.append(nr_cloned)
- self.nr_splitted_log.append(nr_splitted)
- self.nr_pruned_log.append(nr_pruned)
- if max_radii is not None:
- self.radii_max_log.append(max_radii)
- else:
- self.radii_max_log.append(0.0)
- if max_grad2d is not None:
- self.grads_max_log.append(max_grad2d)
- else:
- self.grads_max_log.append(0.0)
+ self.benchmarker.record("cloned", nr_cloned)
+ self.benchmarker.record("splitted", nr_splitted)
+ self.benchmarker.record("pruned", nr_pruned)
+ self.benchmarker.record("radii_max", max_radii if max_radii is not None else 0.0)
+ self.benchmarker.record("grads_max", max_grad2d if max_grad2d is not None else 0.0)
def plot_info(self, step, output_path: Path | None = None, scene_name: str | None = None) -> None:
@@ -380,7 +455,7 @@ class Optimizer(nn.Module, ABC, Generic[T]):
if len(self.nr_cloned_log) == len(self.iter_time_log):
data.append((range(len(self.iter_time_log)), self.nr_cloned_log, "Cloned"))
if len(self.nr_splitted_log) == len(self.iter_time_log):
- data.append((range(len(self.iter_time_log)), self.nr_splitted_log, "Splitted"))
+ data.append((range(len(self.iter_time_log)), self.nr_splitted_log, "Split"))
if len(self.nr_pruned_log) == len(self.iter_time_log):
data.append((range(len(self.iter_time_log)), self.nr_pruned_log, "Pruned"))
@@ -404,17 +479,13 @@ class Optimizer(nn.Module, ABC, Generic[T]):
ax.set_title(f"{label} Gaussians", fontsize=13, pad=5)
# show x-axis ticks on all plots
ax.tick_params(axis='x', which='both', bottom=True, top=False, labelbottom=True)
- # set y-axis vmin to 0
- # ax.set_ylim(bottom=0)
-
+
# Shared x-axis label
axes[-1].set_xlabel("Iteration", fontsize=12)
# Improve layout
plt.tight_layout()
plt.subplots_adjust(hspace=0.3)
- #
- # module_name = self.__class__.__name__.lower()
-
+
# Save and close
save_path = save_path / f"stats_{step}.png"
plt.savefig(save_path, dpi=300, bbox_inches='tight')
diff --git a/optgs/scene_trainer/optimizer/optimizer_adam.py b/optgs/scene_trainer/optimizer/optimizer_adam.py
index 2632db990e489de54d4f7f521949a2ff8f7f5b69..35445ee2d314cca8da3898c4750d913166dc5658 100644
--- a/optgs/scene_trainer/optimizer/optimizer_adam.py
+++ b/optgs/scene_trainer/optimizer/optimizer_adam.py
@@ -6,7 +6,6 @@ import torch
from torch import Tensor
from optgs.dataset.data_types import BatchedViews
-from optgs.misc.general_utils import get_expon_lr_func
from optgs.misc.io import FrequencyScheduler
from optgs.model.decoder.decoder import Decoder
from optgs.model.types import Gaussians
@@ -33,17 +32,8 @@ class AdamOptimizerCfg(OptimizerCfg):
eps: float
weight_decay: float
- # learning rates
- base_lr: int | float
- means_lr_init: float
- means_lr_final: float
- means_lr_delay_mult: float
- means_lr_max_steps: int # should be equal to total optimization steps
- scales_lr: float
- rotations_lr: float
- opacities_lr: float
- sh0s_lr: float
- shNs_lr: float # 20 times less as sh0s_lr in original paper
+ # Per-param learning rates and the means decay schedule come from the inherited
+ # lr_scheduler cfg (lr_data / apply_scheduler + the expon scheduler params).
def update(self, initializer_cfg: InitializerCfg):
pass
@@ -102,13 +92,11 @@ class AdamOptimizer(NonlearnedOptimizer[AdamOptimizerCfg]):
scene_scale = torch.ones(1, 1, device=device)
scene_scale = scene_scale[0].item()
- # initialize learning rate scheduler for means
- self.means_lr_scheduler = get_expon_lr_func(
- lr_init=self.cfg.means_lr_init * scene_scale,
- lr_final=self.cfg.means_lr_final * scene_scale,
- lr_delay_mult=self.cfg.means_lr_delay_mult,
- max_steps=self.cfg.means_lr_max_steps
- )
+ # Means LR follows the lr_scheduler's expon decay, scaled by scene extent. The schedule is
+ # linear in its endpoints, so multiplying its output by scene_scale matches scaling both
+ # endpoints (the original 3DGS recipe). Exposed as means_lr_scheduler so the shared ADC
+ # path (MCMC noise injection) can read the same per-step means LR.
+ self.means_lr_scheduler = lambda step: self.scheduler.get_lr(step, "means") * scene_scale
def on_scene_end(self) -> None:
super().on_scene_end()
@@ -126,106 +114,70 @@ class AdamOptimizer(NonlearnedOptimizer[AdamOptimizerCfg]):
**kwargs
) -> OptimizerOutput:
- # Timing
- self.iter_start.record()
-
- # Unpack
- iter_context: BatchedViews = optimizer_input.context
- target: BatchedViews = optimizer_input.target
- renderer: Decoder = optimizer_input.renderer
- b, v, _, h, w = iter_context["image"].shape
- assert b == 1, "Batch size > 1 not supported for post-processing"
-
- # Log number of gaussians
- self.nr_gaussians_log.append(
- optimizer_input.prev_output.gaussians.means.shape[1]
- )
-
- # One optimization step
- res = self.apply_one_update_step(i, optimizer_input, optimizer_output, sh_degree=kwargs.get("sh_degree", None))
- gaussians: Gaussians = res[0]
- meta_for_adc: dict = res[1]
- updates: dict[str, Tensor] = res[2]
- grads_raw: dict[str, Tensor] = res[3]
- normalized_grads: dict[str, Tensor] = res[4]
- learning_rates: dict[str, float] = res[5]
-
- # Densification and Pruning
- if self.cfg.any_adc:
- # Apply ADC
- self.apply_adc(
- i=i, v=v, h=h, w=w,
- adc_state=optimizer_input.prev_output.state.adc_state,
- gaussians=gaussians,
- meta=meta_for_adc,
- object_dict_to_adjust=self.smoothers
- )
- # ADC changes N → cached buffers are invalid; re-make tensors as fresh leaves.
- # torch.cat (used by add_new/relocate) produces a non-leaf even with requires_grad=True,
- # so .grad is never populated by backward(). detach() cuts the grad_fn first.
- buf_nr_gaussians = self._meta_bufs['N']
- actual_nr_gaussians = gaussians.means.shape[1]
- if buf_nr_gaussians != actual_nr_gaussians:
- self._meta_bufs.clear()
- # TODO Naama: need to think if the detach is necessary (was added during mcmc implementation)
- gaussians.means = gaussians.means.detach().requires_grad_(True)
- gaussians.scales = gaussians.scales.detach().requires_grad_(True)
- gaussians.rotations_unnorm = gaussians.rotations_unnorm.detach().requires_grad_(True)
- gaussians.opacities = gaussians.opacities.detach().requires_grad_(True)
- gaussians.harmonics = gaussians.harmonics.detach().requires_grad_(True)
-
- # Timing
- self._record_iter_timing()
-
- # TODO Naama: we can log stats with save_every, but need to change stuff later.
- # Log stats — guard with save_every
- if grads_raw is not None: # and self.save_every(i + 1, tag="info"):
+ with self.benchmarker.time("iter"):
+ # Unpack
+ iter_context: BatchedViews = optimizer_input.context
+ target: BatchedViews = optimizer_input.target
+ renderer: Decoder = optimizer_input.renderer
+ b, v, _, h, w = iter_context["image"].shape
+ assert b == 1, "Batch size > 1 not supported for post-processing"
+
+ # Log number of gaussians
+ self.benchmarker.record("gaussians", optimizer_input.prev_output.gaussians.means.shape[1])
+
+ # One optimization step
+ res = self._apply_step(i, optimizer_input, optimizer_output, sh_degree=kwargs.get("sh_degree", None))
+ gaussians: Gaussians = res[0]
+ meta_for_adc: dict = res[1]
+ updates: dict[str, Tensor] = res[2]
+ grads_raw: dict[str, Tensor] = res[3]
+ normalized_grads: dict[str, Tensor] = res[4]
+ learning_rates: dict[str, float] = res[5]
+
+ # Densification and Pruning
+ if self.cfg.any_adc:
+ # Apply ADC
+ self.apply_adc(
+ i=i, v=v, h=h, w=w,
+ adc_state=optimizer_input.prev_output.state.adc_state,
+ gaussians=gaussians,
+ meta=meta_for_adc,
+ object_dict_to_adjust=self.smoothers
+ )
+ # ADC changes N → cached buffers are invalid; re-make tensors as fresh leaves.
+ buf_nr_gaussians = self._meta_bufs['N']
+ actual_nr_gaussians = gaussians.means.shape[1]
+ if buf_nr_gaussians != actual_nr_gaussians:
+ self._meta_bufs.clear()
+ # ADC rebuilt these tensors. Re-make them as clean leaves that require grad —
+ # the invariant the i==0 setup establishes and that the next step's in-place
+ # updates and autograd.grad both assume.
+ gaussians.means = gaussians.means.detach().requires_grad_(True)
+ gaussians.scales = gaussians.scales.detach().requires_grad_(True)
+ gaussians.rotations_unnorm = gaussians.rotations_unnorm.detach().requires_grad_(True)
+ gaussians.opacities = gaussians.opacities.detach().requires_grad_(True)
+ gaussians.harmonics = gaussians.harmonics.detach().requires_grad_(True)
+
+ # Log stats — one entry per step (consumer indexes nr_nonzero_grad_log by step index)
+ if grads_raw is not None:
G = grads_raw["means"].shape[0]
nonzero_grads = [(g.reshape(G, -1) != 0).any(dim=-1) for g in grads_raw.values() if g is not None]
nonzero_grads = torch.stack(nonzero_grads) # [num_params, G]
nonzero_grads = nonzero_grads.any(dim=0) # [G]
- self.nr_nonzero_grad_log.append(nonzero_grads.sum().item())
+ self.benchmarker.record("nonzero_grads", nonzero_grads.sum().item())
# Save updated gaussians (for next iteration)
optimizer_input.prev_output.gaussians = gaussians
# Info
if self.save_every(i + 1, tag="info"):
-
- # save gaussians
- optimizer_output.gaussian_list.append(gaussians, detach_and_cpu=True, save_to_disk=False, no_cache=False)
-
- # Save delta stats
- assert optimizer_output.info is not None
-
- # log deltas
- if "deltas" not in optimizer_output.info:
- optimizer_output.info["deltas"] = []
- optimizer_output.info["deltas"].append({k: v.cpu() for k, v in updates.items() if v is not None})
-
- # log gradients
- if "grads" not in optimizer_output.info:
- optimizer_output.info["grads"] = []
- optimizer_output.info["grads"].append({k: v.cpu() for k, v in grads_raw.items() if v is not None})
-
- # log normalized gradients
- if "normalized_grads" not in optimizer_output.info:
- optimizer_output.info["normalized_grads"] = []
- optimizer_output.info["normalized_grads"].append(
- {k: v.cpu() for k, v in normalized_grads.items() if v is not None})
-
- # log learning rates
- if "learning_rates" not in optimizer_output.info:
- optimizer_output.info["learning_rates"] = []
- optimizer_output.info["learning_rates"].append(learning_rates)
-
- # Check if output_path in kwargs
- output_path = kwargs.get("output_path", None)
- scene_name = kwargs.get("scene_name", None)
-
- # Plot stats
- # if self.cfg.any_adc:
- # self.plot_info(i, output_path=output_path, scene_name=scene_name)
+ self._append_info(
+ optimizer_output, gaussians,
+ deltas={k: v.cpu() for k, v in updates.items() if v is not None},
+ grads={k: v.cpu() for k, v in grads_raw.items() if v is not None},
+ normalized_grads={k: v.cpu() for k, v in normalized_grads.items() if v is not None},
+ learning_rates=learning_rates,
+ )
# Post-update context + target renders
self._save_post_update_renders(
@@ -236,7 +188,7 @@ class AdamOptimizer(NonlearnedOptimizer[AdamOptimizerCfg]):
# Optimizer output is being changed in place, but for clarity we return it
return optimizer_output
- def apply_one_update_step(
+ def _apply_step(
self, i, optimizer_input: OptimizerInput, optimizer_output: OptimizerOutput, sh_degree: int | None = None
) -> tuple[Gaussians, dict | None, dict, dict[str, Tensor], dict[str, Tensor], dict[str, float]]:
@@ -264,23 +216,16 @@ class AdamOptimizer(NonlearnedOptimizer[AdamOptimizerCfg]):
# assert gaussians does not store activated values
assert not gaussians.stores_activated, "Gaussians must not store activated values."
- # learning rates
- # TODO Naama: use current cfg field lr_scheduler, which also defines the lr per param
+ # learning rates — all from the lr_scheduler cfg (lr_data per-param values, with the
+ # expon decay applied only where apply_scheduler is set, i.e. means). Means is additionally
+ # scaled by scene extent via means_lr_scheduler.
assert self.means_lr_scheduler is not None, "means_lr_scheduler is not initialized"
- means_lr = self.means_lr_scheduler(i) * self.cfg.base_lr
- scales_lr = self.cfg.scales_lr * self.cfg.base_lr
- rotations_lr = self.cfg.rotations_lr * self.cfg.base_lr
- opacities_lr = self.cfg.opacities_lr * self.cfg.base_lr
- sh0s_lr = self.cfg.sh0s_lr * self.cfg.base_lr
- shNs_lr = self.cfg.shNs_lr * self.cfg.base_lr
-
- # scale learning rates by number of views in the batch
- # means_lr *= v
- # scales_lr *= v
- # rotations_lr *= v
- # opacities_lr *= v
- # sh0s_lr *= v
- # shNs_lr *= v
+ means_lr = self.means_lr_scheduler(i)
+ scales_lr = self.scheduler.get_lr(i, "scales")
+ rotations_lr = self.scheduler.get_lr(i, "rotations")
+ opacities_lr = self.scheduler.get_lr(i, "opacities")
+ sh0s_lr = self.scheduler.get_lr(i, "sh0")
+ shNs_lr = self.scheduler.get_lr(i, "shN")
assert (
iter_context["extrinsics"].shape[0] == iter_context["extrinsics"].shape[0] == 1
@@ -293,23 +238,22 @@ class AdamOptimizer(NonlearnedOptimizer[AdamOptimizerCfg]):
opacities_raw = gaussians.opacities # [B, N]
shs = gaussians.harmonics # [B, N, 3, sh_d]
- self.decoder_event_start.record()
- loss, grads_raw, meta_for_adc = calc_input_gradients(
- iter_context,
- means,
- scales_raw,
- rotations_unnorm,
- opacities_raw,
- shs,
- renderer,
- need_2d_grads=self.cfg.need_2d_grads,
- chunk_size=self.cfg.input_gradients_chunk_size,
- any_adc=self.cfg.any_adc,
- sh_degree=sh_degree,
- meta_bufs=self._meta_bufs,
- opacity_reg_lambda=self.cfg.opacity_reg_lambda,
- )
- self.decoder_event_end.record()
+ with self.benchmarker.time("decoder"):
+ loss, grads_raw, meta_for_adc = calc_input_gradients(
+ iter_context,
+ means,
+ scales_raw,
+ rotations_unnorm,
+ opacities_raw,
+ shs,
+ renderer,
+ need_2d_grads=self.cfg.need_2d_grads,
+ chunk_size=self.cfg.input_gradients_chunk_size,
+ any_adc=self.cfg.any_adc,
+ sh_degree=sh_degree,
+ meta_bufs=self._meta_bufs,
+ opacity_reg_lambda=self.cfg.opacity_reg_lambda,
+ )
# get updates from adam optimizer
grads_raw = squeeze_grad_dict(grads_raw)
@@ -320,25 +264,25 @@ class AdamOptimizer(NonlearnedOptimizer[AdamOptimizerCfg]):
# Batch delta computation for contiguous params with _foreach_mul to reduce kernel launches.
# no_refine flags are handled by excluding the param from the batch (delta stays None).
_grad_lr_pairs = [
- (grads_adam["means"], -means_lr, self.cfg.no_refine_mean),
- (grads_adam["scales"], -scales_lr, self.cfg.no_refine_scale),
- (grads_adam["rotations"], -rotations_lr, self.cfg.no_refine_rotation),
- (grads_adam["opacities"], -opacities_lr, self.cfg.no_refine_opacity),
+ (grads_adam["means"], -means_lr, self.cfg.freeze_mean),
+ (grads_adam["scales"], -scales_lr, self.cfg.freeze_scale),
+ (grads_adam["rotations"], -rotations_lr, self.cfg.freeze_rotation),
+ (grads_adam["opacities"], -opacities_lr, self.cfg.freeze_opacity),
]
_active_grads = [g for g, lr, skip in _grad_lr_pairs if not skip]
_active_lrs = [lr for g, lr, skip in _grad_lr_pairs if not skip]
_active_deltas = torch._foreach_mul(_active_grads, _active_lrs) if _active_grads else []
_delta_iter = iter(_active_deltas)
- delta_means = next(_delta_iter) if not self.cfg.no_refine_mean else None
- delta_scales_raw = next(_delta_iter) if not self.cfg.no_refine_scale else None
- delta_rotations_unnorm = next(_delta_iter) if not self.cfg.no_refine_rotation else None
- delta_opacities_raw = next(_delta_iter) if not self.cfg.no_refine_opacity else None
+ delta_means = next(_delta_iter) if not self.cfg.freeze_mean else None
+ delta_scales_raw = next(_delta_iter) if not self.cfg.freeze_scale else None
+ delta_rotations_unnorm = next(_delta_iter) if not self.cfg.freeze_rotation else None
+ delta_opacities_raw = next(_delta_iter) if not self.cfg.freeze_opacity else None
# SH deltas stay separate (non-contiguous slice views)
- delta_sh0s = None if self.cfg.no_refine_sh0 else -sh0s_lr * grads_adam["sh0s"]
+ delta_sh0s = None if self.cfg.freeze_sh0 else -sh0s_lr * grads_adam["sh0s"]
delta_shNs = None
- if grads_adam["shNs"] is not None and not self.cfg.no_refine_shN:
+ if grads_adam["shNs"] is not None and not self.cfg.freeze_shN:
delta_shNs = -shNs_lr * grads_adam["shNs"]
# step — batch contiguous params with _foreach_add_ to reduce kernel launches;
diff --git a/optgs/scene_trainer/optimizer/optimizer_knn_based.py b/optgs/scene_trainer/optimizer/optimizer_knn_based.py
index 974235fb4cd4bcf830c2559a13a7401bd455d065..2679e01ed4f3e8317fa53db42fbfbe56ed1ebe66 100644
--- a/optgs/scene_trainer/optimizer/optimizer_knn_based.py
+++ b/optgs/scene_trainer/optimizer/optimizer_knn_based.py
@@ -1,5 +1,4 @@
import math
-import random
from dataclasses import dataclass
from typing import Literal, Optional, Any
@@ -10,14 +9,12 @@ import torchvision.transforms as T
from einops import rearrange
from torch import nn, Tensor
-from optgs.dataset.data_types import BatchedExample, DataShim
from optgs.dataset.data_types import BatchedViews
-from optgs.dataset.shims.patch_shim import apply_patch_shim
from optgs.geometry.projection import project, sample_image_grid
from optgs.misc.general_utils import SkipBatchException
from optgs.misc.io import FrequencyScheduler
from optgs.model.decoder.decoder import Decoder
-from optgs.model.encoder.layer import ResNetFeatureWarpper
+from optgs.model.backbones.layer import ResNetFeatureWarpper
from optgs.model.types import Gaussians
from optgs.scene_trainer.common.gaussian_adapter import build_covariance
from optgs.scene_trainer.initializer import InitializerCfg, InitializerColmapCfg, InitializerEdgsCfg, \
@@ -31,18 +28,13 @@ from optgs.scene_trainer.optimizer.optimizer_utils import unpack_gaussians, \
get_visibility_contribution_from_gaussian_obj
try:
- from optgs.model.encoder.point_transformer.layer import (PlainPointTransformer, SubsampleBlock, PointLinearWrapper,
- MultiScalePointTransformer,
- MultViewLowresAttn)
+ from optgs.model.backbones.point_transformer.layer import (PlainPointTransformer,
+ PointLinearWrapper,
+ MultViewLowresAttn)
except:
pass
-try:
- from simple_knn._C import distCUDA2
-except:
- pass
-
-from optgs.scene_trainer.optimizer.layer import CustomGroupNorm, AdamInputSmoothing, SlicedG3RNorm
+from optgs.scene_trainer.optimizer.layer import CustomGroupNorm, AdamInputSmoothing, SlicedG3RNorm, AdamState
from optgs.scene_trainer.initializer.initializer import InitializerOutput
from optgs.scene_trainer.optimizer.time_embed import get_embedder, TimeEncodingWrapper
@@ -58,6 +50,16 @@ from optgs.scene_trainer.optimizer.optimizer_utils import (
_IMAGENET_NORM = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+# Gaussian parameter groups the delta head predicts, paired with their key in
+# get_gaussian_param_sizes. Order matches pack_gaussians / split_delta_gaussians.
+_GAUSSIAN_GROUPS = [
+ ("means", "means"),
+ ("scales", "scales"),
+ ("rotations", "quats"),
+ ("opacities", "opacities"),
+ ("shs", "shs"),
+]
+
@dataclass
class KnnBasedOptimizerCfg(OptimizerCfg):
@@ -66,19 +68,18 @@ class KnnBasedOptimizerCfg(OptimizerCfg):
no_render_error: bool
input_error_shallow_resnet_feature: bool
input_error_resnet_feature_layers: int
- refine_sh_only: bool
- num_basic_refine_blocks: int
- num_refine_blocks: int
+ num_basic_blocks: int
+ num_blocks: int
concat_init_state: bool # always concat init state during updates
replace_init_state: bool # always use the init state during updates
state_channels: int
- refine_block_rmsnorm: bool
- refine_block_layernorm: bool
+ block_rmsnorm: bool
+ block_layernorm: bool
pt_qk_norm: bool
norm_pt_block: bool
- refine_gaussian_multiple: int # predict more gaussian residuals based on the previous gaussian center
- refine_residual_init_state: bool # add residual connection in the prediction head to the inital state
- clamp_refine_max_scale: float
+ delta_gaussian_multiple: int # predict more gaussian residuals based on the previous gaussian center
+ residual_init_state: bool # add residual connection in the prediction head to the inital state
+ clamp_max_scale: float
clamp_min_scale: float | int
clamp_min_raw_scales: float | int
clamp_max_raw_scales: float | int
@@ -90,24 +91,23 @@ class KnnBasedOptimizerCfg(OptimizerCfg):
clamp_max_shs: float | int
clamp_shs_soft: bool
+ # TODO (release): remove gaussian_head_multiple (kept for next-project research; unused here, always 1)
gaussian_head_multiple: int # use multiple non-weight sharing heads to predict multiple gaussians
- gradient_update_scale: float | int
+ residual_grad_scale: float | int
input_gradient_with_ssim_loss: bool
- update_attn_proj_channels: int | None
- update_no_knn_attn: bool
- update_no_tran_block_norm: bool
- update_tran_block_act: str | None
+ attn_proj_channels: int | None
+ no_knn_attn: bool
+ no_tran_block_norm: bool
+ tran_block_act: str | None
multi_gaussian_scale_smaller: bool
- refine_condition_pt_feature: bool
- reinit_gaussian_when_refine_multiple: bool
- refine_same_num_points: bool # when init_gaussian_multiple > 1, refine directly works on it instead of subsampling points
+ condition_pt_feature: bool
+ same_num_points: bool # when init_gaussian_multiple > 1, refine directly works on it instead of subsampling points
- refine_knn_samples: int
- refine_multi_scale_pt: bool
+ knn_samples: int
+ max_active_gaussians: int
# KNN
use_fused_attn: bool
- prune_invisible_gaussians: bool
knn_idx_update_every: int
# point transformer
@@ -145,14 +145,14 @@ class KnnBasedOptimizerCfg(OptimizerCfg):
input_error_mv_attn_blocks: int
# refine global attention
- refine_with_mv_attn: bool
- refine_with_mv_attn_lowres: bool
- refine_no_mv_attn: bool # remove only the attn
+ with_mv_attn: bool
+ with_mv_attn_lowres: bool
+ no_mv_attn: bool # remove only the attn
mv_attn_conv_with_norm: bool # unet-attn conv with norm
- refine_mv_shuffle_attn: bool # use pixel shuffle to save computation instead of unet
- refine_mv_attn_with_pos_enc: bool
- refine_shuffle_attn_no_norm: bool
- refine_mv_unimatch_attn: bool
+ mv_shuffle_attn: bool # use pixel shuffle to save computation instead of unet
+ mv_attn_with_pos_enc: bool
+ shuffle_attn_no_norm: bool
+ mv_unimatch_attn: bool
# input gradients
input_gradient: bool
@@ -163,24 +163,15 @@ class KnnBasedOptimizerCfg(OptimizerCfg):
input_gradient_loss_reduction: str
scale_residual_grads: bool
- # sliding window
- window_local_refine: bool # refine each local window separately and then combine all windows
- window_global_refine: bool # refine all windows together
- window_local_global_refine: bool # first refine each window seprately, and then refine all windows together
-
- # sliding window update instead of update all gaussians together
- update_window_size: int
- local_gaussian_render: bool
-
# time encoding
use_time_encoding: bool
time_encoding_max_steps: int
- train_global_update_only: bool
+ train_global_only: bool
# random size refine
# update more for low resolution, less for high
- random_update_with_size: bool
+ random_step_with_size: bool
# amp
use_amp: bool
@@ -190,9 +181,6 @@ class KnnBasedOptimizerCfg(OptimizerCfg):
use_checkpointing: bool
recurrent_use_checkpointing: bool
- # Debugging
- debug_refine_update_module: bool
-
# Normalizing input
input_gradient_normalize: bool
input_gradient_normalize_type: str
@@ -213,53 +201,61 @@ class KnnBasedOptimizerCfg(OptimizerCfg):
# Preprocessing the init gaussians
scale_initial_opacities: float | int
- # Experimental
+ # TODO (release): remove experimental delta override cluster
experimental_run: bool
experimental_update: Bool3DGSCfg
experimental_use_grads: bool
experimental_use_norm_grads: Bool3DGSCfg
experimental_lr: Number3DGSCfg
# Deactivate gaussians
+ # TODO (release): wire in or remove (only used in experimental_get_visible_gaussian_mask)
local_prune_zero_radii: bool
local_prune_low_weights: bool
local_prune_low_weights_thresh: float | int
+
update_only_nonzero_grad: bool
# update learn residual state
residual_state: bool
# Update head
- update_head_layer_num: int
- update_head_concat_img: bool
- update_head_act: str | None # update_head activation to predict the deltas
- update_head_final_act: str | None # final activation in the update_head
- update_head_hidden_dim_matches: str # rebuttal or submission version
+ delta_head_layer_num: int
+ delta_head_concat_img: bool
+ delta_head_act: str | None # delta_head activation to predict the deltas
+ delta_head_final_act: str | None # final activation in the delta_head
+ delta_head_hidden_dim_matches: str # rebuttal or submission version
- update_head_scale_mag: bool # predict deltas as scale * 0.01 * jnp.exp(mag * 0.01)
- update_head_scalar_scale: bool # predict deltas as scalar * delta / norm(delta)
- update_head_scalar_scale_act: str # activation for the scalar scale output
+ delta_head_scalar_scale: bool # predict deltas as scalar * delta / norm(delta)
+ delta_head_scalar_scale_act: str # activation for the scalar scale output
# Per-parameter-group update head (Feature A)
- update_head_per_param_heads: bool # separate heads per param group, each with own normalize+scale
- update_head_per_param_hidden_dim: int # hidden dim for per-param heads (SH head gets 2x)
- # Per-parameter scalar scales (Feature B) — requires update_head_scalar_scale=true
- update_head_per_param_scales: bool # per-group scalar scales instead of one global scalar
+ delta_head_per_param_heads: bool # separate heads per param group, each with own normalize+scale
+ delta_head_per_param_hidden_dim: int # hidden dim for per-param heads (SH head gets 2x)
+ # Per-parameter scalar scales (Feature B) — requires delta_head_scalar_scale=true
+ delta_head_per_param_scales: bool # per-group scalar scales instead of one global scalar
# Config from initializer
sh_d: int | None
- init_gaussian_param_num: int | None = None
init_sh_d: int | None = None
# Fow initialization from feed forward, gaussians are aligned with pixels.
init_gaussian_multiple: int | None = None
latent_downsample: int | None = None
- delta_adam_combine_step: int = 0 # combine deltas and adam updates
+ @property
+ def scales_clamp_lims(self) -> tuple:
+ """Scale clamp limits in the space scales are refined in.
+
+ With opt_scales_before_act scales are refined in log space, so the raw (log-space)
+ limits apply; otherwise the activation-space limits apply.
+ """
+ if self.opt_scales_before_act:
+ return self.clamp_min_raw_scales, self.clamp_max_raw_scales
+ return self.clamp_min_scale, self.clamp_max_scale
def update(self, initializer_cfg: InitializerCfg):
""" Update the optimizer config based on the initializer config"""
# General settings
- self.init_gaussian_param_num = initializer_cfg.get_gaussian_param_num()
self.init_sh_d = initializer_cfg.get_sh_d()
if self.sh_d is None:
# get sh_d from initializer if not set
@@ -271,7 +267,7 @@ class KnnBasedOptimizerCfg(OptimizerCfg):
self.init_gaussian_multiple = initializer_cfg.init_gaussian_multiple
# update proj channels
- if self.refine_condition_pt_feature:
+ if self.condition_pt_feature:
self.condition_channels = initializer_cfg.gaussian_regressor_channels
else:
self.condition_channels = initializer_cfg.get_pt_in_channels()
@@ -281,9 +277,8 @@ class KnnBasedOptimizerCfg(OptimizerCfg):
InitializerPointcloudCfg)):
# Since pixels and gaussians are not alligned, we can not use pixel attributes
assert not self.input_error, "The error calculation assumes per pixel gaussians"
- assert not self.update_head_concat_img
+ assert not self.delta_head_concat_img
assert not self.input_alpha
- assert not self.local_gaussian_render, "The local rendering assumes per view gaussians"
assert self.init_state_wo_features, "Colmap initializer does not have point features, init_state_wo_features must be set to True"
@@ -293,47 +288,69 @@ class KnnBasedOptimizerCfg(OptimizerCfg):
raise ValueError(f"Unsupported initializer config type: {type(initializer_cfg)}")
-class KnnBasedOptimizerState:
- # TODO Naama: OptimizerState class already exists
- def __init__(self, state: torch.Tensor):
- self.state = state
+@dataclass
+class KnnBasedOptimizerState(OptimizerState):
+ """OptimizerState subclass for KNN-based optimizers.
+
+ Holds the per-Gaussian learned state vectors and adds ADC mutation methods
+ (called by adc/base.py helpers) that keep them in sync with clone/split/prune.
+ """
+ state: torch.Tensor | None = None
+ init_state: torch.Tensor | None = None # scene-start state; None when concat/replace/residual flags are all off
+ adam_state: AdamState | None = None
def clone(self, clone_mask: torch.Tensor, zero_t: bool) -> None:
- cloned_state = self.state[clone_mask]
- if zero_t:
- cloned_state = torch.zeros_like(cloned_state)
- self.state = torch.cat([self.state, cloned_state], dim=0)
-
- def split(self, split_mask, num_splits: int, zero_t: bool) -> None:
- states_to_split = self.state[split_mask]
- split_states = states_to_split.chunk(num_splits, dim=0)
- new_states = []
- for i in range(num_splits):
- if zero_t:
- new_states.append(torch.zeros_like(split_states[i]))
- else:
- new_states.append(split_states[i])
- self.state = torch.cat([self.state, *new_states], dim=0)
+ if self.state is not None:
+ cloned = torch.zeros_like(self.state[clone_mask]) if zero_t else self.state[clone_mask]
+ self.state = torch.cat([self.state, cloned], dim=0)
+ if self.init_state is not None:
+ # NOTE: cloned Gaussians inherit the parent's init_state (or zeros when zero_t).
+ # Whether a clone should keep, reset, or re-derive init_state is an open design question.
+ cloned_init = torch.zeros_like(self.init_state[clone_mask]) if zero_t else self.init_state[clone_mask]
+ self.init_state = torch.cat([self.init_state, cloned_init], dim=0)
+
+ def split(self, split_mask: torch.Tensor, num_splits: int, zero_t: bool) -> None:
+ if self.state is not None:
+ chunks = self.state[split_mask].chunk(num_splits, dim=0)
+ new_states = [torch.zeros_like(c) if zero_t else c for c in chunks]
+ self.state = torch.cat([self.state, *new_states], dim=0)
+ if self.init_state is not None:
+ # NOTE: split Gaussians inherit the parent's init_state (or zeros when zero_t).
+ # Whether a split should keep, reset, or re-derive init_state is an open design question.
+ init_chunks = self.init_state[split_mask].chunk(num_splits, dim=0)
+ new_init = [torch.zeros_like(c) if zero_t else c for c in init_chunks]
+ self.init_state = torch.cat([self.init_state, *new_init], dim=0)
def replace(self, from_indices: torch.Tensor, dest_indices: torch.Tensor, zero_t: bool) -> None:
- if zero_t:
- self.state[dest_indices] = 0.0
- else:
- self.state[dest_indices] = self.state[from_indices]
+ if self.state is not None:
+ if zero_t:
+ self.state[dest_indices] = 0.0
+ else:
+ self.state[dest_indices] = self.state[from_indices]
+ if self.init_state is not None:
+ if zero_t:
+ self.init_state[dest_indices] = 0.0
+ else:
+ self.init_state[dest_indices] = self.init_state[from_indices]
def prune(self, prune_mask: torch.Tensor) -> None:
- self.state = self.state[~prune_mask]
+ if self.state is not None:
+ self.state = self.state[~prune_mask]
+ if self.init_state is not None:
+ self.init_state = self.init_state[~prune_mask]
def add(self, num_new: int) -> None:
if num_new <= 0:
return
- device = self.state.device
- dtype = self.state.dtype
- input_dim = self.state.shape[1:]
- self.state = torch.cat([self.state, torch.zeros((num_new, *input_dim), device=device, dtype=dtype)], dim=0)
-
- def extend(self, num_new):
- self.add(num_new)
+ if self.state is not None:
+ zeros = torch.zeros((num_new, *self.state.shape[1:]), device=self.state.device, dtype=self.state.dtype)
+ self.state = torch.cat([self.state, zeros], dim=0)
+ # NOTE: newly added Gaussians get a zero init_state (they have no parent to inherit from).
+ # Whether they should instead re-derive init_state is an open design question.
+ if self.init_state is not None:
+ zeros_init = torch.zeros((num_new, *self.init_state.shape[1:]),
+ device=self.init_state.device, dtype=self.init_state.dtype)
+ self.init_state = torch.cat([self.init_state, zeros_init], dim=0)
class Abs(nn.Module):
@@ -361,6 +378,37 @@ def get_activation_cls(activation: Optional[str] = None):
class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
+ """Learned optimizer that refines a scene's Gaussians with a point-transformer (the core method).
+
+ Replaces hand-tuned gradient descent (Adam + heuristics) with a network that predicts, at each
+ iteration, how every Gaussian should change. The scene is treated as a point cloud (one point per
+ Gaussian) and the network reasons over local neighborhoods via KNN attention, so an update to one
+ Gaussian is informed by its spatial neighbors.
+
+ Each Gaussian carries a learned latent state vector alongside its parameters (means, scales,
+ rotations, opacity, SH). The state is what the network reads and writes across iterations; the
+ Gaussian parameters are updated by adding predicted deltas.
+
+ One optimization step (`_apply_step`, run per iteration by `_forward_impl`):
+ 1. Input signal (`prepare_input_signal`): render the context views and measure how well they match
+ the ground-truth inputs. The signal is the per-Gaussian rendering gradients
+ (`_calc_input_gradients` — a backprop of the reconstruction loss, used only as a feature, never
+ to update weights) and/or the render error in feature space. This tells the network "what is
+ wrong" with the current Gaussians.
+ 2. Update (point transformer): KNN attention over the point cloud mixes each Gaussian's state with
+ its neighbors', conditioned on the input signal, producing an updated per-Gaussian state.
+ 3. Delta head: maps the updated state to per-parameter deltas, which are added to the Gaussians
+ (with clamping/activation handling per parameter group).
+ 4. Optional adaptive density control (clone/split/prune), keeping the state in sync with the
+ changing Gaussian count.
+
+ `optimizer_preprocessing` runs once at scene start to seed the latent state from the initializer's
+ Gaussians. This class implements both ReSplat and Learn2Splat through different settings in
+ `KnnBasedOptimizerCfg`: ReSplat (`resplat_v1.yaml`) drives the update with feature-space render
+ error and multi-view attention, while Learn2Splat (`learn2splat.yaml`) drives it with normalized
+ per-Gaussian gradients. The many `_build_*` / `_calc_*` helpers wire up these settings.
+ """
+
OPTIMIZER_NAME = "knn_based"
OPTIMIZER_NAME_ALIASES: tuple[str, ...] = ()
@@ -371,7 +419,7 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
super().__init__(cfg, save_every)
if self.cfg.residual_state:
- assert not self.cfg.refine_residual_init_state
+ assert not self.cfg.residual_init_state
# State channel
self.state_channels = self.cfg.state_channels
@@ -383,36 +431,36 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
self.time_encoder_fn = None
self.time_embedding_dim = 0
- # update_proj
+ # state_proj
if not self.cfg.init_state_wo_features:
- self.update_proj = nn.Conv2d(self.cfg.condition_channels, self.state_channels, 1)
+ self.state_proj = nn.Conv2d(self.cfg.condition_channels, self.state_channels, 1)
channels, in_channels, update_gaussian_param_num, out_channels, error_features_channels = (
- self.define_update_channels(self.cfg.init_gaussian_param_num))
+ self._define_channels())
self.error_features_channels = error_features_channels
self.gaussian_param_num = out_channels
if self.cfg.input_error:
- self.update_feature = self.get_input_error_feature_extractor()
+ self.error_feature_extractor = self.get_input_error_feature_extractor()
if self.cfg.input_error_add_rgb_feature:
if self.cfg.init_gaussian_multiple == 4: # re10k
- self.update_rgb_error_proj = nn.Sequential(
+ self.rgb_error_proj = nn.Sequential(
nn.Linear(3, error_features_channels),
nn.LayerNorm(error_features_channels)
)
else:
- self.update_rgb_error_proj = nn.Sequential(
+ self.rgb_error_proj = nn.Sequential(
nn.Linear(3 * self.cfg.latent_downsample ** 2, error_features_channels),
nn.LayerNorm(error_features_channels)
)
- self.update_input_norm = self.get_update_input_norm(in_channels)
- self.update_module = self.get_update_module(channels, in_channels)
+ self.input_norm = self._build_input_norm(in_channels)
+ self.point_transformer = self._build_point_transformer(channels, in_channels)
# predict multiple gaussians
- out_channels = out_channels * self.cfg.refine_gaussian_multiple
+ out_channels = out_channels * self.cfg.delta_gaussian_multiple
- if not self.cfg.refine_same_num_points:
+ if not self.cfg.same_num_points:
out_channels = out_channels * self.cfg.init_gaussian_multiple
# make sure the input size of the gaussian head is updated accordingly
@@ -420,24 +468,24 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
channels += self.time_embedding_dim
# Compute per-param group dims (needed by per_param_heads and per_param_scales)
- if self.cfg.update_head_per_param_heads or self.cfg.update_head_per_param_scales:
+ if self.cfg.delta_head_per_param_heads or self.cfg.delta_head_per_param_scales:
self._per_param_group_dims = self._compute_per_param_group_dims(out_channels)
# Scaling state for update head
if self.cfg.predict_state_scale:
- self.state_scale_head = self.get_state_scale_head(in_channels)
+ self.state_scale_head = self._build_state_scale_head(in_channels)
- self.update_head = self.get_update_head(in_channels, channels, out_channels)
+ self.delta_head = self._build_delta_head(in_channels, channels, out_channels)
# multiple gaussian heads to predict multiple gaussians
if self.cfg.gaussian_head_multiple > 1:
- self.update_head_list = self.get_update_head_list(channels, out_channels)
+ self.extra_gaussian_heads = self._build_extra_gaussian_heads(channels, out_channels)
# Define error calculation
# add global attention to the render error
if self.cfg.input_error and self.cfg.input_error_mv_attn:
assert self.cfg.input_error_resnet_feature
- self.update_error_attn = nn.ModuleList([
+ self.error_mv_attn = nn.ModuleList([
MultViewLowresAttn(error_features_channels)
for _ in range(self.cfg.input_error_mv_attn_blocks)
])
@@ -458,47 +506,51 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
@property
def adc_object_dict_to_adjust(self):
if self.cfg.any_adc:
- object_dict: dict[str, Any] = {"depthsplat_state": None}
+ object_dict: dict[str, Any] = {"optimizer_state": None}
# For ADC
if self.cfg.input_gradient_normalize and self.cfg.input_gradient_normalize_type == "adam":
- object_dict.update(self.update_input_norm.subgroups_view(self.param_slices))
+ object_dict.update(self.input_norm.subgroups_view(self.param_slices))
else:
return None
return object_dict
+ def _frozen_param_groups(self) -> set:
+ """Gaussian parameter groups the delta head does not predict (frozen).
+
+ Any combination of freeze_mean/scale/rotation/opacity is supported. The knn
+ optimizer treats SH as a single group, so SH is frozen only when both freeze_sh0
+ and freeze_shN are set — a partial SH freeze is not expressible here.
+ """
+ frozen = set()
+ if self.cfg.freeze_mean:
+ frozen.add("means")
+ if self.cfg.freeze_scale:
+ frozen.add("scales")
+ if self.cfg.freeze_rotation:
+ frozen.add("rotations")
+ if self.cfg.freeze_opacity:
+ frozen.add("opacities")
+ assert self.cfg.freeze_sh0 == self.cfg.freeze_shN, \
+ "knn optimizer treats SH as one group; freeze_sh0 and freeze_shN must match"
+ if self.cfg.freeze_sh0 and self.cfg.freeze_shN:
+ frozen.add("shs")
+ return frozen
+
def _compute_per_param_group_dims(self, out_channels):
"""Compute per-parameter-group output dimensions from total out_channels.
- Returns a dict {group_name: dim} in the same order as split_delta_gaussians.
- Accounts for no_refine_rotation, no_refine_mean, refine_sh_only, and multipliers.
+ Returns a dict {group_name: dim} in the same order as split_delta_gaussians,
+ omitting frozen groups, with the delta/init-gaussian multipliers folded in.
"""
-
- # TODO Naama: allow combination of no_refine_*
p = get_gaussian_param_sizes(self.cfg.sh_d)
+ excluded = self._frozen_param_groups()
- all_params = [
- ("means", "means"),
- ("scales", "scales"),
- ("rotations", "quats"),
- ("opacities", "opacities"),
- ("shs", "shs"),
- ]
-
- if self.cfg.refine_sh_only:
- excluded = {"means", "scales", "rotations", "opacities"}
- elif self.cfg.no_refine_rotation:
- excluded = {"rotations"}
- elif self.cfg.no_refine_mean:
- excluded = {"means"}
- else:
- excluded = set()
-
- multiplier = self.cfg.refine_gaussian_multiple
- if not self.cfg.refine_same_num_points:
+ multiplier = self.cfg.delta_gaussian_multiple
+ if not self.cfg.same_num_points:
multiplier *= self.cfg.init_gaussian_multiple
- group_dims = {name: p[key] * multiplier for name, key in all_params if name not in excluded}
+ group_dims = {name: p[key] * multiplier for name, key in _GAUSSIAN_GROUPS if name not in excluded}
assert sum(group_dims.values()) == out_channels, (
f"Per-param group dims {dict(group_dims)} sum={sum(group_dims.values())} != out_channels={out_channels}"
@@ -511,11 +563,11 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
Each head: Linear(channels, hidden) -> act -> Linear(hidden, dim+1)
The +1 is a per-group scalar scale. Each head independently normalizes + scales.
"""
- act_cls = get_activation_cls(self.cfg.update_head_act)
- hidden_dim = self.cfg.update_head_per_param_hidden_dim
+ act_cls = get_activation_cls(self.cfg.delta_head_act)
+ hidden_dim = self.cfg.delta_head_per_param_hidden_dim
# Set up scale activation (shared across all per-param heads)
- scale_act_name = self.cfg.update_head_scalar_scale_act
+ scale_act_name = self.cfg.delta_head_scalar_scale_act
init_bias_map = {'softplus': -1, 'relu': 1e-8, 'abs': 1e-8}
if scale_act_name not in init_bias_map:
raise ValueError(f"Unsupported scalar_scale_act: {scale_act_name}")
@@ -528,7 +580,7 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
h = hidden_dim * 2 if name == "shs" else hidden_dim
layers = [nn.Linear(channels, h), act_cls()]
- for _ in range(self.cfg.update_head_layer_num - 2):
+ for _ in range(self.cfg.delta_head_layer_num - 2):
layers += [nn.Linear(h, h), act_cls()]
layers.append(nn.Linear(h, dim + 1)) # +1 for scalar scale
@@ -544,45 +596,40 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
return heads
- def get_update_head(self, in_channels, channels, out_channels):
- update_head_activation_cls = get_activation_cls(self.cfg.update_head_act)
- final_head_activation_cls = get_activation_cls(self.cfg.update_head_final_act)
+ def _build_delta_head(self, in_channels, channels, out_channels):
+ update_head_activation_cls = get_activation_cls(self.cfg.delta_head_act)
+ final_head_activation_cls = get_activation_cls(self.cfg.delta_head_final_act)
# skip connection to the image color
- if self.cfg.update_head_concat_img:
+ if self.cfg.delta_head_concat_img:
channels += 3 * (self.cfg.latent_downsample ** 2)
# Feature A: per-parameter-group heads (early return — builds ModuleDict instead of Sequential)
- if self.cfg.update_head_per_param_heads:
- assert not self.cfg.update_head_scale_mag, "update_head_scale_mag not supported with per_param_heads"
- assert not self.cfg.update_head_per_param_scales, "per_param_heads already includes per-group scales"
+ if self.cfg.delta_head_per_param_heads:
+ assert not self.cfg.delta_head_per_param_scales, "per_param_heads already includes per-group scales"
return self._build_per_param_heads(channels, out_channels)
- # predict delta = scale * 0.01 * jnp.exp(mag * 0.01)
- if self.cfg.update_head_scale_mag:
- out_channels = out_channels * 2
-
- if self.cfg.update_head_scalar_scale:
- if self.cfg.update_head_per_param_scales:
+ if self.cfg.delta_head_scalar_scale:
+ if self.cfg.delta_head_per_param_scales:
# Feature B: one scalar scale per parameter group
out_channels = out_channels + len(self._per_param_group_dims)
else:
out_channels = out_channels + 1
# Determine hidden layer size
- # TODO: update_head_hidden_dim_source should be "output" (out_channels).
+ # TODO (release): update_head_hidden_dim_source should be "output" (out_channels).
# Using "input" currently as default to reproduce rebuttal results.
- if self.cfg.update_head_hidden_dim_matches == "input":
- hidden_dim = channels # rebuttal version
+ if self.cfg.delta_head_hidden_dim_matches == "input":
+ hidden_dim = channels # cvpr rebuttal version (current default)
else:
- hidden_dim = out_channels # submitted version
+ hidden_dim = out_channels # cvpr submitted version
# Build update head
layers_list = [
nn.Linear(channels, hidden_dim),
update_head_activation_cls()
]
- for i in range(self.cfg.update_head_layer_num - 2):
+ for i in range(self.cfg.delta_head_layer_num - 2):
layers_list += [
nn.Linear(hidden_dim, hidden_dim),
update_head_activation_cls(),
@@ -592,19 +639,19 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
nn.Linear(hidden_dim, out_channels),
final_head_activation_cls()
]
- update_head = nn.Sequential(*layers_list)
+ delta_head = nn.Sequential(*layers_list)
# init the delta as 0
- nn.init.zeros_(update_head[-2].weight)
+ nn.init.zeros_(delta_head[-2].weight)
if final_head_activation_cls == torch.nn.Sigmoid:
desired_init_delta = 0.005
bias = math.log(desired_init_delta / (1 - desired_init_delta)) # ~= -4.6
- nn.init.constant_(update_head[-2].bias, bias)
+ nn.init.constant_(delta_head[-2].bias, bias)
else:
- nn.init.zeros_(update_head[-2].bias)
+ nn.init.zeros_(delta_head[-2].bias)
# Scalar scale output
- if self.cfg.update_head_scalar_scale:
+ if self.cfg.delta_head_scalar_scale:
# Set the initial scale to very low number, to get the gradients flow
init_bias_map = {
'softplus': -1,
@@ -612,30 +659,30 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
'abs': 1e-8,
}
- act_name = self.cfg.update_head_scalar_scale_act
+ act_name = self.cfg.delta_head_scalar_scale_act
if act_name not in init_bias_map:
raise ValueError(f"Unsupported scalar_scale_out_act: {act_name}")
# Initialize bias for scale output(s)
- if self.cfg.update_head_per_param_scales:
+ if self.cfg.delta_head_per_param_scales:
num_groups = len(self._per_param_group_dims)
for i in range(num_groups):
- nn.init.constant_(update_head[-2].bias[-(num_groups - i)], init_bias_map[act_name])
+ nn.init.constant_(delta_head[-2].bias[-(num_groups - i)], init_bias_map[act_name])
else:
- nn.init.constant_(update_head[-2].bias[-1], init_bias_map[act_name])
+ nn.init.constant_(delta_head[-2].bias[-1], init_bias_map[act_name])
# Create activation
act_class = get_activation_cls(act_name)
self.scale_act = act_class(beta=1) if act_name == 'softplus' else act_class()
- return update_head
+ return delta_head
- def get_update_head_list(self, channels, out_channels):
- update_head_activation = get_activation_cls(self.cfg.update_head_act)
- final_head_activation = get_activation_cls(self.cfg.final_head_act)
- update_head_list = nn.ModuleList()
+ def _build_extra_gaussian_heads(self, channels, out_channels):
+ update_head_activation = get_activation_cls(self.cfg.delta_head_act)
+ final_head_activation = get_activation_cls(self.cfg.delta_head_final_act)
+ extra_gaussian_heads = nn.ModuleList()
for i in range(self.cfg.gaussian_head_multiple - 1):
- update_head_list.append(
+ extra_gaussian_heads.append(
nn.Sequential(
nn.Linear(channels, channels),
update_head_activation(),
@@ -645,12 +692,12 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
)
# init the delta as 0
- nn.init.zeros_(update_head_list[i][-2].weight)
- nn.init.zeros_(update_head_list[i][-2].bias)
+ nn.init.zeros_(extra_gaussian_heads[i][-2].weight)
+ nn.init.zeros_(extra_gaussian_heads[i][-2].bias)
- return update_head_list
+ return extra_gaussian_heads
- def get_update_input_norm(self, in_channels):
+ def _build_input_norm(self, in_channels):
if self.cfg.input_gradient_normalize:
assert self.cfg.input_gradient, "for now we only normalize when using gradient as input"
if self.cfg.input_gradient_normalize_type == 'layer':
@@ -669,57 +716,45 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
else:
return nn.Identity()
- def get_update_module(self, channels, in_channels):
- if not self.cfg.debug_refine_update_module:
- return None
-
- if self.cfg.refine_multi_scale_pt:
- update_module = nn.Sequential(
- PointLinearWrapper(in_channels, channels),
- MultiScalePointTransformer(channels,
- self.cfg.refine_knn_samples,
- subsample_method=self.cfg.subsample_method,
- attn_proj_channels=self.cfg.update_attn_proj_channels,
- )
- )
- else:
- update_module = nn.Sequential(
- PointLinearWrapper(in_channels, channels),
- PlainPointTransformer(channels, self.cfg.refine_knn_samples,
- num_blocks=self.cfg.num_basic_refine_blocks,
- qk_norm=self.cfg.pt_qk_norm,
- norm_pt_block=self.cfg.norm_pt_block,
- num_heads=self.cfg.pt_heads,
- no_rpe=True,
- no_attn=self.cfg.update_no_knn_attn,
- no_norm=self.cfg.update_no_tran_block_norm,
- act=self.cfg.update_tran_block_act,
- attn_proj_channels=self.cfg.update_attn_proj_channels,
- with_mv_attn=self.cfg.refine_with_mv_attn,
- with_mv_attn_lowres=self.cfg.refine_with_mv_attn_lowres,
- no_mv_attn=self.cfg.refine_no_mv_attn,
- conv_with_norm=self.cfg.mv_attn_conv_with_norm,
- mv_shuffle_attn=self.cfg.refine_mv_shuffle_attn,
- with_pos_enc=self.cfg.refine_mv_attn_with_pos_enc,
- shuffle_attn_no_norm=self.cfg.refine_shuffle_attn_no_norm,
- mv_unimatch_attn=self.cfg.refine_mv_unimatch_attn,
- use_checkpointing=self.cfg.use_checkpointing,
- use_fused_attn=self.cfg.use_fused_attn,
- knn_idx_update_every=self.cfg.knn_idx_update_every
- )
- )
+ def _build_point_transformer(self, channels, in_channels):
+ point_transformer = nn.Sequential(
+ PointLinearWrapper(in_channels, channels),
+ PlainPointTransformer(channels, self.cfg.knn_samples,
+ num_blocks=self.cfg.num_basic_blocks,
+ qk_norm=self.cfg.pt_qk_norm,
+ norm_pt_block=self.cfg.norm_pt_block,
+ num_heads=self.cfg.pt_heads,
+ no_rpe=True,
+ no_attn=self.cfg.no_knn_attn,
+ no_norm=self.cfg.no_tran_block_norm,
+ act=self.cfg.tran_block_act,
+ attn_proj_channels=self.cfg.attn_proj_channels,
+ with_mv_attn=self.cfg.with_mv_attn,
+ with_mv_attn_lowres=self.cfg.with_mv_attn_lowres,
+ no_mv_attn=self.cfg.no_mv_attn,
+ conv_with_norm=self.cfg.mv_attn_conv_with_norm,
+ mv_shuffle_attn=self.cfg.mv_shuffle_attn,
+ with_pos_enc=self.cfg.mv_attn_with_pos_enc,
+ shuffle_attn_no_norm=self.cfg.shuffle_attn_no_norm,
+ mv_unimatch_attn=self.cfg.mv_unimatch_attn,
+ use_checkpointing=self.cfg.use_checkpointing,
+ use_fused_attn=self.cfg.use_fused_attn,
+ knn_idx_update_every=self.cfg.knn_idx_update_every
+ )
+ )
# Init normalization layers
if self.cfg.input_normalize_state:
- for block in update_module[1].blocks:
+ for block in point_transformer[1].blocks:
nn.init.zeros_(block.norm1.bias)
nn.init.zeros_(block.norm2.bias)
nn.init.ones_(block.norm1.weight)
nn.init.ones_(block.norm2.weight)
- return update_module
+ return point_transformer
- def get_state_scale_head(self, in_channels):
+ @staticmethod
+ def _build_state_scale_head(in_channels):
state_scale_head = nn.Sequential(
nn.Linear(in_channels, in_channels // 2),
nn.ReLU(),
@@ -733,21 +768,14 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
return state_scale_head
- def define_update_channels(self, init_gaussian_param_num):
- if self.cfg.init_gaussian_multiple > 1:
- gaussian_param_num = init_gaussian_param_num // self.cfg.init_gaussian_multiple
- else:
- gaussian_param_num = init_gaussian_param_num
-
- # no pixel offset
- gaussian_param_num -= 2
-
- # update position
- gaussian_param_num += 3
-
- # SHs
- if self.cfg.sh_d != self.cfg.init_sh_d:
- gaussian_param_num += 3 * (self.cfg.sh_d - self.cfg.init_sh_d)
+ def _define_channels(self):
+ # The optimizer refines a fixed per-Gaussian representation, independent of the initializer:
+ # means/3D position(3), scales(3), rotation(4), opacity(1), SH(3*sh_d). Whatever position
+ # encoding the initializer used (e.g. resplat's 2D pixel offset) has already been unprojected
+ # to a 3D point by the time gaussians reach the optimizer. Sized from the same per-group sizes
+ # the freeze logic uses below, so the total and the per-group counts cannot drift.
+ param_sizes = get_gaussian_param_sizes(self.cfg.sh_d)
+ gaussian_param_num = sum(param_sizes.values())
# Get error channels
if self.cfg.input_error:
@@ -764,7 +792,7 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
# final input channels
input_signal_channels = gradient_channels + error_channels
- if self.cfg.refine_same_num_points:
+ if self.cfg.same_num_points:
in_channels = (gaussian_param_num
+ self.state_channels
+ input_signal_channels)
@@ -777,8 +805,9 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
in_channels += self.state_channels
out_channels = gaussian_param_num
- if self.cfg.no_refine_mean:
- out_channels -= 3
+ for name, key in _GAUSSIAN_GROUPS:
+ if name in self._frozen_param_groups():
+ out_channels -= param_sizes[key]
channels = self.state_channels
if self.cfg.input_alpha:
# pixel shuffle the alpha channel to the latent resolution
@@ -812,6 +841,11 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
return error_channels, error_feature_channels
def optimizer_preprocessing(self, optimizer_input: OptimizerInput, from_init: bool) -> None:
+ """Prepare the initializer's Gaussians and seed the per-Gaussian latent state at scene start.
+
+ On a fresh scene (from_init) scales the initial opacities, pads/truncates SH to sh_d, and builds
+ the latent state vector; on a buffer resume it only refreshes the state, leaving init_state intact.
+ """
if self.cfg.input_error_remain_context or self.cfg.input_error_merge_remain_context:
assert self.cfg.input_error_cache_resnet_feature
@@ -822,7 +856,6 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
# Prepare Gaussians
if from_init:
# Scale initial opacities (in normal scale)
- # TODO Naama: add option to reset opacities and randomly reset/scale opacities of intermidiate updates
opacities = optimizer_input.prev_output.gaussians.opacities # post activation, in [0, 1]
scaled_opacities = opacities * self.cfg.scale_initial_opacities # default to 1.0
optimizer_input.prev_output.gaussians.opacities = scaled_opacities
@@ -838,11 +871,6 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
shs = F.pad(shs, (0, pad), "constant", 0)
optimizer_input.prev_output.gaussians.harmonics = shs
- # Right now, this does not do anything, since we do not use windows
- local_window_update, test_window_size, window_end, window_start = self.get_window_size(v)
- optimizer_input.additional_info = local_window_update, test_window_size, window_end, window_start
- self.update_gaussians_for_window(v, h, w, optimizer_input)
-
# Prepare state
# Gaussians dimensions
n = optimizer_input.prev_output.gaussians.means.shape[1]
@@ -852,27 +880,14 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
# Set everything so that the optimizer isn't aware whether it's a new scene
# Convert InitializerOutput to OptimizerPreviousOutput
optimizer_input.prev_output = OptimizerPreviousOutput(gaussians=optimizer_input.prev_output.gaussians,
- state=OptimizerState())
+ state=KnnBasedOptimizerState())
optimizer_input.prev_output.state.state = vector_state
# init_state captures the scene-start state used by some experiments;
# only set it on a fresh scene so replay-buffer resumes preserve the original value.
- if from_init:
+ # Left None when none of the three init_state flags are active.
+ if from_init and (self.cfg.concat_init_state or self.cfg.replace_init_state or self.cfg.residual_init_state):
optimizer_input.prev_output.state.init_state = vector_state
- def update_gaussians_for_window(self, v, h, w, optimizer_input):
- # Get window parameters and set gaussians accordingly
- local_window_update, test_window_size, window_end, window_start = optimizer_input.additional_info
-
- if local_window_update and self.cfg.local_gaussian_render:
- init_gaussians = optimizer_input.prev_output.gaussians
- # select subset of gaussians
- init_gaussians_subset = select_gaussian_subset(init_gaussians, window_start, window_end,
- v=v,
- h=h // self.cfg.latent_downsample,
- w=w // self.cfg.latent_downsample,
- )
- optimizer_input.prev_output.gaussians = init_gaussians_subset
-
def _forward_impl(
self,
i: int,
@@ -882,81 +897,61 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
full_target: BatchedViews,
**kwargs
) -> OptimizerOutput:
-
- # Timing
- self.iter_start.record()
-
- # Unpack
- iter_context: BatchedViews = optimizer_input.context
- target: BatchedViews = optimizer_input.target
- renderer: Decoder = optimizer_input.renderer
- b, v, _, h, w = iter_context["image"].shape
- assert b == 1, "Batch size > 1 not supported for post-processing"
-
- # Log number of gaussians
- self.nr_gaussians_log.append(
- optimizer_input.prev_output.gaussians.means.shape[1]
- )
-
- # One optimization step
- res = self.apply_one_update_step(
- i, optimizer_input, optimizer_output
- )
- updated_gaussians: Gaussians = res[0]
- state: Tensor = res[1]
- meta_for_adc: dict = res[2]
- updates: dict[str, Tensor] = res[3]
- grads_raw: Tensor | None = res[4]
- normalized_grads: Tensor | None = res[5]
- scaled_state: Tensor | None = res[6]
- gaussians_sel: Tensor | None = res[7]
-
- # Timing
- self._record_iter_timing()
+ """Run one optimization iteration: apply the update step, then ADC, then save renders/stats."""
+
+ with self.benchmarker.time("iter"):
+ # Unpack
+ iter_context: BatchedViews = optimizer_input.context
+ target: BatchedViews = optimizer_input.target
+ renderer: Decoder = optimizer_input.renderer
+ b, v, _, h, w = iter_context["image"].shape
+ assert b == 1, "Batch size > 1 not supported for post-processing"
+
+ # Log number of gaussians
+ self.benchmarker.record("gaussians", optimizer_input.prev_output.gaussians.means.shape[1])
+
+ # One optimization step
+ res = self._apply_step(
+ i, optimizer_input, optimizer_output
+ )
+ updated_gaussians: Gaussians = res[0]
+ state: Tensor = res[1]
+ meta_for_adc: dict = res[2]
+ updates: dict[str, Tensor] = res[3]
+ grads_raw: Tensor | None = res[4]
+ normalized_grads: Tensor | None = res[5]
+ updated_scaled_state: Tensor | None = res[6]
+ gaussians_sel: Tensor | None = res[7]
# Log stats
if grads_raw is not None:
- grads = grads_raw # [B, G, D]
- nonzero_grads = (grads != 0).any(-1) # [B, G]
- # Filter out strictly zero gradients for logging
- grads = grads[nonzero_grads].unsqueeze(0) # [1, N_nonzero, D]
+ nonzero_grads = (grads_raw != 0).any(-1) # [B, G]
assert nonzero_grads.shape[0] == 1
- self.nr_nonzero_grad_log.append(nonzero_grads[0].sum().item())
-
- # Local ADC
- # if optimizer_output.t == 500:
- # weight_vis_contribution, _ = get_visibility_contribution_from_gaussian_obj(iter_context, updated_gaussians) # [N]
- # prune_mask = weight_vis_contribution < 5
- # print(f"Pruning {torch.sum(prune_mask)} gaussians out of {prune_mask.shape[0]} at iteration {i}")
- # updated_gaussians = updated_gaussians[:, ~prune_mask]
- # state = state[~prune_mask]
- # if self.cfg.normalize_update_input and self.cfg.normalize_update_input_type == "adam":
- # if not self.update_input_norm.is_reset():
- # self.update_input_norm.prune(prune_mask)
+ self.benchmarker.record("nonzero_grads", nonzero_grads[0].sum().item())
# Densification and Pruning
if self.cfg.any_adc:
n_before_adc = updated_gaussians.means.shape[1]
- # Prepare objects to adjust during ADC
+ # Prepare objects to adjust during ADC.
+ # Write the current state back so ADC mutation methods see it.
+ optimizer_input.prev_output.state.state = state
object_dict = self.adc_object_dict_to_adjust
- object_dict["depthsplat_state"] = KnnBasedOptimizerState(state)
- object_dict["depthsplat_init_state"] = KnnBasedOptimizerState(optimizer_input.prev_output.state.init_state)
+ object_dict["optimizer_state"] = optimizer_input.prev_output.state
- # Apply ADC
+ # Apply ADC — mutates state (and init_state if set) in place
self.apply_adc(
i=i, v=v, h=h, w=w, adc_state=optimizer_input.prev_output.state.adc_state,
gaussians=updated_gaussians, meta=meta_for_adc, object_dict_to_adjust=object_dict
)
- # Update state after ADC
- state = object_dict["depthsplat_state"].state
- optimizer_input.prev_output.state.init_state = object_dict["depthsplat_init_state"].state
+ # Read back the (possibly grown/pruned) state tensor
+ state = optimizer_input.prev_output.state.state
- del object_dict["depthsplat_state"]
+ del object_dict["optimizer_state"]
if self.cfg.input_gradient_normalize and self.cfg.input_gradient_normalize_type == "adam":
- self.update_input_norm.aggregate_from_subgroups(object_dict, self.param_slices)
+ self.input_norm.aggregate_from_subgroups(object_dict, self.param_slices)
# If N changed (add_new grew the population), stale KNN caches in the
# point transformer modules would index out-of-bounds on the next forward
@@ -969,87 +964,15 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
optimizer_input.prev_output.state.state = state
if self.cfg.input_gradient_normalize_type == "adam":
- optimizer_input.prev_output.state.adam_state = self.update_input_norm.get_state()
+ optimizer_input.prev_output.state.adam_state = self.input_norm.get_state()
if self.training:
optimizer_output.gaussian_list.append(updated_gaussians)
- # Info
+ # Save per-iteration debug info (test time only)
if not self.training and self.save_every(i + 1, tag="info"):
- # TODO Naama: review and refactor
-
- # save guassians
- optimizer_output.gaussian_list.append(updated_gaussians, detach_and_cpu=True, save_to_disk=False)
-
- # Save delta stats
- assert optimizer_output.info is not None
-
- # log updates
-
- # unpack shs
- shs = updates.pop("shs") # [1, N, 3*sh_d]
- assert shs.shape[0] == 1, "Batch size > 1 not supported"
- shs = shs.squeeze(0) # [N, 3*sh_d]
- shs = rearrange(shs, "n (c x) -> n c x", c=3, x=self.cfg.sh_d) # [N, 3, sh_d]
- updates["sh0s"] = shs[..., 0:1]
- if self.cfg.sh_d > 1:
- updates["shNs"] = shs[..., 1:]
- else:
- updates["shNs"] = None
-
- # log deltas
- if "deltas" not in optimizer_output.info:
- optimizer_output.info["deltas"] = []
- optimizer_output.info["deltas"].append(
- {k: v.squeeze(0).cpu() if v is not None else None for k, v in updates.items()})
-
- # Split each vector grad into gaussians components
- if grads_raw is not None:
- if gaussians_sel is not None:
- # Restore the zero gradients for tracking
- b, g_valid, d = grads_raw.shape
- g = state.shape[0]
- grads_raw_full = torch.zeros((b, g, d))
- normalized_grads_full = torch.zeros((b, g, d))
- grads_raw_full[:, gaussians_sel, :] = grads_raw.cpu()
- normalized_grads_full[:, gaussians_sel, :] = normalized_grads.cpu()
- grads_raw = grads_raw_full
- normalized_grads = normalized_grads_full
-
- grads_raw: dict[str, Tensor] = split_grads(grads_raw.cpu(), self.cfg)
-
- # Split each vector normalized_grads into gaussians components
- if normalized_grads is not None:
- normalized_grads: dict[str, Tensor] = split_grads(normalized_grads.cpu(), self.cfg)
-
- assert grads_raw["means"].shape == normalized_grads["means"].shape, \
- f"Shape mismatch between grads and normalized_grads: {grads_raw['means'].shape} vs {normalized_grads['means'].shape}"
-
- # log states
- if scaled_state is not None:
- if "states_norms" not in optimizer_output.info:
- optimizer_output.info["states_norms"] = []
- state_norm = torch.norm(scaled_state, dim=-1) # [B, N]
- optimizer_output.info["states_norms"].append(state_norm.cpu())
-
- # log gradients
- if "grads" not in optimizer_output.info:
- optimizer_output.info["grads"] = []
- optimizer_output.info["grads"].append(grads_raw)
-
- # log normalized gradients
- if "normalized_grads" not in optimizer_output.info:
- optimizer_output.info["normalized_grads"] = []
- optimizer_output.info["normalized_grads"].append(normalized_grads)
-
- # Check if output_path in kwargs
- output_path = kwargs.get("output_path", None)
- scene_name = kwargs.get("scene_name", None)
-
- if self.cfg.any_adc:
- pass
- # Plot stats
- # self.plot_info(i, output_path=output_path, scene_name=scene_name)
+ self._save_info_stats(optimizer_output, updated_gaussians, updates, grads_raw,
+ normalized_grads, updated_scaled_state, gaussians_sel, state)
# Post-update context + target renders
self._save_post_update_renders(
@@ -1060,29 +983,69 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
# Optimizer output is being changed in place, but for clarity we return it
return optimizer_output
- def apply_one_update_step(
+ def _save_info_stats(self, optimizer_output, updated_gaussians, updates, grads_raw,
+ normalized_grads, updated_scaled_state, gaussians_sel, state):
+ """Prepare per-iteration debug stats (deltas, gradients, state norms) and save them.
+
+ knn-specific prep: unpack shs into sh0s/shNs, split the flat gradient vectors into
+ per-param dicts, and restore zeroed-out (excluded) Gaussians; then hand the
+ already-CPU values to the shared Optimizer._append_info.
+ """
+ # unpack shs into sh0s / shNs
+ shs = updates.pop("shs") # [1, N, 3*sh_d]
+ assert shs.shape[0] == 1, "Batch size > 1 not supported"
+ shs = rearrange(shs.squeeze(0), "n (c x) -> n c x", c=3, x=self.cfg.sh_d) # [N, 3, sh_d]
+ updates["sh0s"] = shs[..., 0:1]
+ updates["shNs"] = shs[..., 1:] if self.cfg.sh_d > 1 else None
+
+ info = {"deltas": {k: v.squeeze(0).cpu() if v is not None else None for k, v in updates.items()}}
+
+ if grads_raw is not None:
+ if gaussians_sel is not None:
+ # Restore the zeroed-out (excluded) Gaussians so the logged tensors are full-size
+ b, g_valid, d = grads_raw.shape
+ g = state.shape[0]
+ grads_raw_full = torch.zeros((b, g, d))
+ normalized_grads_full = torch.zeros((b, g, d))
+ grads_raw_full[:, gaussians_sel, :] = grads_raw.cpu()
+ normalized_grads_full[:, gaussians_sel, :] = normalized_grads.cpu()
+ grads_raw = grads_raw_full
+ normalized_grads = normalized_grads_full
+
+ info["grads"] = split_grads(grads_raw.cpu(), self.cfg)
+ if normalized_grads is not None:
+ info["normalized_grads"] = split_grads(normalized_grads.cpu(), self.cfg)
+ assert info["grads"]["means"].shape == info["normalized_grads"]["means"].shape, \
+ f"Shape mismatch between grads and normalized_grads: " \
+ f"{info['grads']['means'].shape} vs {info['normalized_grads']['means'].shape}"
+ else:
+ info["normalized_grads"] = None
+ if updated_scaled_state is not None:
+ info["states_norms"] = torch.norm(updated_scaled_state, dim=-1).cpu() # [B, N]
+
+ self._append_info(optimizer_output, updated_gaussians, **info)
+
+ def _apply_step(
self,
i,
optimizer_input: OptimizerInput,
optimizer_output: OptimizerOutput
) -> tuple[Gaussians, Tensor, dict, dict[str, Tensor], Tensor | None, Tensor | None, Tensor | None, Tensor | None]:
+ """Compute one update: input signal -> point-transformer state update -> delta head -> new Gaussians.
+
+ Returns the updated Gaussians and latent state plus the intermediates ADC and the debug logging need.
+ """
# Unpacking
context = optimizer_input.context
- target = optimizer_input.target
- renderer = optimizer_input.renderer
debug_dict = optimizer_input.debug_dict
- num_refine = optimizer_input.num_refine
gaussians = optimizer_input.prev_output.gaussians # Gaussian object of [B, N, C]
state = optimizer_input.prev_output.state.state # [N, C]
init_state = optimizer_input.prev_output.state.init_state # [N, C]
- local_window_update, test_window_size, window_end, window_start = optimizer_input.additional_info
- # Get input signal for the optimizer model (erros/gradients)
- self.decoder_event_start.record()
- input_signal, gaussian_grads_raw, gaussian_grads, grad_sign, context_render_output, means2d_grads = (
- self.prepare_input_signal(context, i, gaussians, local_window_update, renderer, window_end,
- window_start, num_refine)
- )
- self.decoder_event_end.record()
+ # Get input signal for the optimizer model (errors/gradients)
+ with self.benchmarker.time("decoder"):
+ input_signal, gaussian_grads_raw, gaussian_grads, grad_sign, context_render_output, means2d_grads = (
+ self.prepare_input_signal(context, gaussians, optimizer_input.renderer)
+ )
# Preparing meta for ADC
if means2d_grads is not None:
@@ -1093,49 +1056,45 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
"means_2d_grads": means2d_grads, # [B, V, N, 2]
}
+ # Some initializers pad the scene with extra, invisible (zero-opacity) Gaussians so that every
+ # GPU optimizes the same number of Gaussians (ColmapInitializer, train_fixed_gaussians_num; the
+ # true count is gaussians.nr_valid). The optimizer predicts an update for every Gaussian, so
+ # updating one of these dummies would raise its opacity and turn it into a real Gaussian. Because
+ # the dummies are invisible their gradient is zero, so updating only nonzero-gradient Gaussians
+ # (update_only_nonzero_grad=True, done by filter_gaussians_for_optimizer_step) leaves them as they
+ # are. Every training config does this; the check below rejects the combination we don't support.
+ has_padding = 0 <= gaussians.nr_valid < gaussians.means.shape[1]
+ if has_padding and not self.cfg.update_only_nonzero_grad:
+ raise ValueError(
+ f"update_only_nonzero_grad=False with padded init Gaussians (nr_valid="
+ f"{gaussians.nr_valid} < N={gaussians.means.shape[1]}): the invisible padding Gaussians "
+ f"would be optimized and turn into real ones. Set update_only_nonzero_grad=True (it skips "
+ f"them, since their gradient is zero), or use an initializer without padding."
+ )
+
# Handle zero gradient gaussians
# We either prune them, or exclude them from the input/output update
if self.cfg.update_only_nonzero_grad and gaussian_grads is not None:
gaussian_grads, gaussian_grads_raw, gaussians, grad_sign, init_state, input_signal, state = (
- self.handle_zero_grad_gaussians(
- context,
- context_render_output,
- gaussian_grads,
- gaussian_grads_raw,
- gaussians,
- grad_sign,
- init_state,
- input_signal,
- means2d_grads,
- meta_for_adc,
- optimizer_input,
- state)
+ self.filter_gaussians_for_optimizer_step(gaussian_grads, gaussian_grads_raw, gaussians, grad_sign,
+ init_state, input_signal, state)
)
- # For training, if the number of active gaussians is too high, skip this batch
- # TODO Naama: maybe sampling?
+ # During training, skip a scene whose active-Gaussian count is outside the workable
+ # range: too many to fit in memory, or too few for the KNN neighborhood.
+ # NOTE: over-cap scenes are discarded wholesale; subsampling down to
+ # max_active_gaussians instead is a possible future alternative (would change
+ # training dynamics, so it needs validation).
active_gaussians_num = state.shape[0]
if self.training:
- if active_gaussians_num > 100_000:
+ if active_gaussians_num > self.cfg.max_active_gaussians:
print(f"Skipping batch at iteration {i} with {active_gaussians_num} active gaussians.")
raise SkipBatchException()
- if active_gaussians_num < self.cfg.refine_knn_samples:
+ if active_gaussians_num < self.cfg.knn_samples:
print(
- f"Skipping batch at iteration {i} with only {active_gaussians_num} active gaussians (need >= {self.cfg.refine_knn_samples}).")
+ f"Skipping batch at iteration {i} with only {active_gaussians_num} active gaussians (need >= {self.cfg.knn_samples}).")
raise SkipBatchException()
- # Training only: save the rendering of initialization for logging
- # Will not be used for loss calculation
- # TODO Naama: this cause to many confusion. Pull it out of this function
- if self.training and i == 0:
- # Append context images initialization
- assert context_render_output is not None
- optimizer_output.context_render_list.append(context_render_output, detach_and_cpu=False)
-
- # render target images initialization
- target_render_output = renderer.forward_batch_subset(gaussians, target)
- optimizer_output.target_render_list.append(target_render_output, detach_and_cpu=False)
-
# Unpack Gaussians
means, scales, rotations_unnorm, opacities_raw, shs = unpack_gaussians(
gaussians,
@@ -1143,7 +1102,7 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
opacities_logit=True,
opacities_unsqueeze=True,
detach=True, # stop gradient of last predictions
- scales_lims=(self.cfg.clamp_min_scale, self.cfg.clamp_refine_max_scale),
+ scales_lims=self.cfg.scales_clamp_lims,
raw_opacities_lims=(self.cfg.clamp_min_raw_opacities, self.cfg.clamp_max_raw_opacities)
)
@@ -1158,58 +1117,44 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
self.debug_reprojection_error(means, debug_dict, context, i, latent_h, latent_w)
# prepare pt input
- point_cloud, tmp_batch_size = self.get_point_cloud(latent_h, latent_w, local_window_update, means,
- test_window_size, v)
+ point_cloud, tmp_batch_size = self.get_point_cloud(latent_h, latent_w, means, v)
# Create offset directly on device to avoid CPU-GPU transfer
offset = torch.arange(1, b + 1, device=state.device, dtype=torch.long) * tmp_batch_size
# reshape
- tmp_gaussian = self.reshape_gaussians_to_nc(latent_h, latent_w, gaussians_concat, v) # [B, N, C] --> [BN, C]
+ gaussians_flat = self.reshape_gaussians_to_nc(latent_h, latent_w, gaussians_concat, v) # [B, N, C] --> [BN, C]
# add global attention to exchange info across views
if self.cfg.input_error_mv_attn:
- input_signal = self.apply_global_attn(b, h, input_signal, latent_h,
- latent_w, local_window_update, test_window_size, v, w)
+ input_signal = self.apply_global_attn(h, input_signal, latent_h, latent_w, v, w)
- tmp_input_signal = input_signal.reshape(-1,
- input_signal.shape[-1]) # [B, N, C] --> [BN, C] - faster than rearrange
- tmp_input_signal = self.append_to_input_signal(b, context, context_render_output, tmp_input_signal, v)
+ input_signal_flat = input_signal.reshape(-1,
+ input_signal.shape[
+ -1]) # [B, N, C] --> [BN, C] - faster than rearrange
+ input_signal_flat = self.append_to_input_signal(b, context, context_render_output, input_signal_flat, v)
# Normalize state before input it to the update module
if self.cfg.input_normalize_state:
state_norm = state.norm(dim=1, keepdim=True) / math.sqrt(state.shape[-1]) # [BG, 1]
state = state / (state_norm + 1e-8) # [BG, C]
- normalized_input_signal = self.update_input_norm(tmp_input_signal)
+ normalized_input = self.input_norm(input_signal_flat)
if self.cfg.input_normalize_gaussians:
- tmp_gaussian_mean = tmp_gaussian.mean()
- tmp_gaussian_std = tmp_gaussian.std()
- tmp_gaussian = (tmp_gaussian - tmp_gaussian_mean) / (tmp_gaussian_std + 1e-8)
+ gaussians_flat_mean = gaussians_flat.mean()
+ gaussians_flat_std = gaussians_flat.std()
+ gaussians_flat = (gaussians_flat - gaussians_flat_mean) / (gaussians_flat_std + 1e-8)
with torch.amp.autocast(device_type='cuda', enabled=self.cfg.pt_update_amp, dtype=torch.bfloat16):
- point_cloud, tmp_gaussian, state, update_input = self.prepare_update_input(b, i, init_state,
- normalized_input_signal,
- latent_h,
- latent_w,
- local_window_update,
- point_cloud,
- tmp_gaussian,
- # gradients/errors + additional pixel related quantities
- state, v, window_end,
- window_start)
-
- # if self.cfg.refine_with_mv_attn:
- # state = concat
- # for i in range(len(self.update_module)):
- # print(i, len(self.update_module), self.update_module[i])
- # state = self.update_module[i]([point_cloud, state, offset]) # [N, C]
- # else:
- updated_state = self.apply_update_module(b, latent_h, latent_w, offset,
- point_cloud, update_input, v, state, i)
+ point_cloud, gaussians_flat, state, update_input = self.prepare_update_input(init_state, normalized_input,
+ point_cloud, gaussians_flat,
+ state)
+
+ updated_state = self._apply_point_transformer(b, latent_h, latent_w, offset,
+ point_cloud, update_input, v, state, i)
# Hard coded extract normalized gradients
if self.cfg.input_gradient and self.cfg.input_gradient_normalize:
- normalized_grads = normalized_input_signal
+ normalized_grads = normalized_input
else:
normalized_grads = None
@@ -1227,14 +1172,14 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
state_scale = state_scale / (state_scale.norm(p=2, dim=1, keepdim=True) + 1e-8)
else:
state_scale = torch.tensor([1], device=state.device, dtype=state.dtype)
- updated_state_scaled = state_scale * updated_state
+ updated_scaled_state = state_scale * updated_state
# optionally append time encodiing to normalize input
with TimeEncodingWrapper(self.cfg.use_time_encoding,
self.time_encoder_fn,
optimizer_output.t,
self.cfg.time_encoding_max_steps,
- updated_state_scaled) as embedded_state:
+ updated_scaled_state) as embedded_state:
if self.cfg.use_time_encoding:
assert not self.cfg.concat_init_state
assert not self.cfg.replace_init_state
@@ -1242,18 +1187,15 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
# delta gaussian head
delta_gaussians = self.apply_delta_gaussian_head(b, context, init_state, embedded_state, v)
- visibility_scale = None # disable for now
-
delta_means, delta_opacities, delta_rotations, delta_scales, delta_shs, init_repeat, delta_gaussians = (
- self.postprocess_deltas(b, delta_gaussians, gaussian_grads, gaussians_concat, grad_sign, latent_h, latent_w,
- local_window_update, normalized_grads, state, test_window_size, v, window_end,
- window_start, optimizer_output.t, optimizer_output.T, visibility_scale)
+ self.postprocess_deltas(b, delta_gaussians, gaussian_grads, gaussians_concat, grad_sign,
+ normalized_grads, state, optimizer_output.t)
)
means, opacities_raw, rotations_unnorm, scales, shs = self.repeat_gaussians(means, opacities_raw,
rotations_unnorm, scales, shs)
- covariances, means, scales, rotations, rotations_unnorm, opacities_raw, shs = self.update_gaussians_params(
+ covariances, means, scales, rotations, rotations_unnorm, opacities_raw, shs = self._apply_gaussian_deltas(
delta_means, delta_scales, delta_rotations, delta_opacities, delta_shs,
means, scales, rotations_unnorm, opacities_raw, shs, init_repeat)
@@ -1298,37 +1240,34 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
grads_raw = gaussian_grads.detach() if gaussian_grads is not None else None
grads_adam = normalized_grads.detach() if normalized_grads is not None else None
- return gaussians, updated_state, meta_for_adc, updates, grads_raw, grads_adam, updated_state_scaled, sel
+ return gaussians, updated_state, meta_for_adc, updates, grads_raw, grads_adam, updated_scaled_state, sel
- def postprocess_deltas(self, b, delta_gaussians, gaussian_grads, gaussians_concat, grad_sign, latent_h, latent_w,
- local_window_update, normalized_grads, state, test_window_size, v, window_end, window_start,
- t, T, visibility_scale):
+ def postprocess_deltas(self, b, delta_gaussians, gaussian_grads, gaussians_concat, grad_sign,
+ normalized_grads, state, t):
# Updates for gradient input (scale, log scale, )
delta_gaussians_raw = delta_gaussians
- delta_gaussians = self.update_delta_for_gradients_input(delta_gaussians_raw, grad_sign, normalized_grads,
- visibility_scale)
+ delta_gaussians = self._postprocess_delta_for_gradient_input(delta_gaussians_raw, grad_sign, normalized_grads,
+ )
# Rearrange back to [B, N, C]
- delta_gaussians, delta_gaussians_raw = self.rearrange_delta_gaussians(b, delta_gaussians,
- delta_gaussians_raw, latent_h,
- latent_w, local_window_update,
- gaussians_concat,
- test_window_size, v, window_end,
- window_start)
-
- # TODO Naama: shouldn't it be before rearranging?
- # multiple gaussian heads to predict multiple gaussians
+ delta_gaussians, delta_gaussians_raw = self.rearrange_delta_gaussians(b, delta_gaussians, delta_gaussians_raw)
+
+ # Extra non-weight-sharing heads, each predicting another N Gaussians from the state.
+ # NOTE: these run after rearrange and on the raw state, so their deltas bypass
+ # _postprocess_delta_for_gradient_input (scaling/negation/scalar_scale) that the main
+ # head's delta went through — an asymmetry that only matters if this path is used
+ # (gaussian_head_multiple is always 1 today). See the TODO (release) on the cfg field.
with torch.amp.autocast(device_type='cuda', enabled=self.cfg.pt_update_amp, dtype=torch.bfloat16):
if self.cfg.gaussian_head_multiple > 1:
num_additional_heads = self.cfg.gaussian_head_multiple - 1
delta_gaussian_list = [delta_gaussians] # list of [B, N, C]
- for i in range(num_additional_heads):
- curr_delta = self.update_head_list[i](state)
+ for j in range(num_additional_heads):
+ curr_delta = self.extra_gaussian_heads[j](state)
curr_delta = rearrange(curr_delta, "(b n) c -> b n c", b=b)
delta_gaussian_list.append(curr_delta)
delta_gaussians = torch.cat(delta_gaussian_list, dim=1) # [B, K*N, C]
- # Experimental overide deltas
+ # TODO (release): remove experimental delta override
if self.cfg.experimental_run:
self.experimental_update_deltas(delta_gaussians, gaussian_grads, normalized_grads)
@@ -1342,128 +1281,51 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
t, delta_means, delta_scales, delta_rotations, delta_opacities, delta_shs
)
- # Linear combination with adam normalized gradients
- if self.cfg.delta_adam_combine_step > 0 and normalized_grads is not None:
- assert t <= T
- if t > self.cfg.delta_adam_combine_step:
- alpha = 0.0
- beta = 1 - ((t - self.cfg.delta_adam_combine_step) / (T - self.cfg.delta_adam_combine_step)) ** alpha
- # Linear combination with adam normalized gradients
- # Use the inverse of the normalized gradients
- # TODO Naama: hard coded lr
- # means
- delta_means = beta * delta_means + (1 - beta) * -normalized_grads[
- ..., self.param_slices["means"]] * 1.6e-4
- # scales
- delta_scales = beta * delta_scales + (1 - beta) * -normalized_grads[
- ..., self.param_slices["scales"]] * 5e-3
- # rotations
- delta_rotations = beta * delta_rotations + (1 - beta) * -normalized_grads[
- ..., self.param_slices["quats"]] * 1e-3
- # opacities
- delta_opacities = beta * delta_opacities + (1 - beta) * -normalized_grads[
- ..., self.param_slices["opacities"]] * 5e-2
- # sh0 - use view instead of rearrange for speed
- delta_shs_bgdc = delta_shs.view(delta_shs.shape[0], delta_shs.shape[1], 3, -1) # [b, g, 3, c]
- delta_sh0 = delta_shs_bgdc[..., 0] # [b, g, 3]
- delta_shN = delta_shs_bgdc[..., 1:] # [b, g, 3, d-1]
- delta_shN = delta_shN.flatten(-2) # [b, g, 3*(d-1)] - faster than rearrange
-
- new_delta_sh0 = beta * delta_sh0 + (1 - beta) * -normalized_grads[
- ..., self.param_slices["sh0"]] * 2.5e-3
- new_delta_shN = beta * delta_shN + (1 - beta) * -normalized_grads[
- ..., self.param_slices["shN"]] * 1.25e-4
- new_delta_shN = new_delta_shN.view(new_delta_shN.shape[0], new_delta_shN.shape[1], 3,
- -1) # [b, g, 3, d-1]
- delta_shs[..., ::self.cfg.sh_d] = new_delta_sh0
- # shN
- for i in range(1, self.cfg.sh_d):
- delta_shs[..., i::self.cfg.sh_d] = new_delta_shN[..., i - 1]
-
return delta_means, delta_opacities, delta_rotations, delta_scales, delta_shs, init_repeat, delta_gaussians
- def handle_zero_grad_gaussians(self, context, context_render_output, gaussian_grads, gaussian_grads_raw, gaussians,
- grad_sign, init_state, input_signal, means2d_grads, meta_for_adc, optimizer_input,
- state):
- # Compute a mask for gaussian that did not contribute to any pixel of context views
- # Their gradients are strictly zero.
- # We don't want to prune them, as they might be relevant in other views (in dense views).
- if self.cfg.prune_invisible_gaussians:
- gaussian_grads, gaussians, grad_sign, input_signal, state = self.prune_invisible_gaussians(context,
- context_render_output,
- gaussian_grads,
- gaussian_grads_raw,
- gaussians,
- grad_sign,
- input_signal,
- means2d_grads,
- meta_for_adc,
- optimizer_input,
- state)
- else:
- assert not self.cfg.local_prune_zero_radii
- assert not self.cfg.local_prune_low_weights
- assert gaussian_grads.shape[0] == 1, "Batch size > 1 not supported with mask"
-
- # radii_mask = (context_render_output.radii != 0).all(1).all(-1) # [B, G]
- # valid_mask = valid_mask & radii_mask # only consider gaussians with non-zero radius as valid
-
- # radii = context_render_output.radii # [B, V, G, 2]
- #
- # # XOR on radii last dimension to find gaussians that have zero radius in only one dimension
- # assert ((radii[..., 0] == 0) ^ (radii[..., 1] == 0)).sum() == 0 # [B, V, G]
- #
- # # Check that all zero radius gaussians are in the zero gradient mask (but not necessarily the opposite)
- # zero_radius_mask = (radii == 0).any(1).any(-1) # [B, G]
- # zero_grad_mask = ~valid_mask # [B, G]
- # zero_radius_cnt = zero_radius_mask.sum()
- # zero_grad_of_zero_radii_cnt = zero_grad_mask[zero_radius_mask].sum()
- # assert zero_grad_of_zero_radii_cnt == zero_radius_cnt, (f"All zero radius gaussians should have zero "
- # f"gradients. Found {zero_radius_cnt} zero radius gaussians, but only {zero_grad_of_zero_radii_cnt} of "
- # f"them have zero gradients.")
- # print(f"Found {zero_grad_of_zero_radii_cnt} / {zero_radius_cnt} zero radius gaussians with zero gradients.")
-
- # Contribution of zero gradient gaussians
- # gaussian_grads_zero_radii = gaussian_grads[zero_radius_mask] # [G_zero_radius, D]
- # assert gaussian_grads_zero_radii.abs().sum() == 0, "Gaussians with zero radius should have zero gradients."
-
- # radii of zero gradient gaussians
- # radii_zero_grad = radii[:, :, zero_grad_mask[0]] # [G_zero_grad, V, 2]
- # zero_grad_radii_cont = radii_zero_grad.sum()
-
- # Compute [G] mask without materializing [B,G,D] bool
- # any() on floats treats nonzero as True
- valid_g = gaussian_grads[0].any(dim=-1) # [G] bool
- sel = None
+ def filter_gaussians_for_optimizer_step(self, gaussian_grads, gaussian_grads_raw, gaussians, grad_sign,
+ init_state, input_signal, state):
+ # Compute a mask for gaussians that did not contribute to any pixel of context views.
+ # Their gradients are strictly zero; skip them for the update but don't prune (they
+ # may be relevant in other views).
+ assert gaussian_grads.shape[0] == 1, "Batch size > 1 not supported with mask"
+
+ # A gaussian is valid if it has a nonzero gradient in any channel
+ valid_g = gaussian_grads[0].any(dim=-1) # [G] bool
+ sel = None
- # if everything is valid, skip all slicing work
- if not valid_g.all():
- sel = valid_g.nonzero(as_tuple=True)[0] # [G_valid]
+ # if everything is valid, skip all slicing work
+ if not valid_g.all():
+ sel = valid_g.nonzero(as_tuple=True)[0] # [G_valid]
- input_signal = input_signal[:, sel, :] # [B, G_valid, C]
+ input_signal = input_signal[:, sel, :] # [B, G_valid, C]
- gaussian_grads = gaussian_grads[:, sel, :] # [B, G_valid, D]
- if gaussian_grads_raw is not None:
- gaussian_grads_raw = gaussian_grads_raw[:, sel, :]
- if grad_sign is not None:
- grad_sign = grad_sign[:, sel, :]
+ gaussian_grads = gaussian_grads[:, sel, :] # [B, G_valid, D]
+ if gaussian_grads_raw is not None:
+ gaussian_grads_raw = gaussian_grads_raw[:, sel, :]
+ if grad_sign is not None:
+ grad_sign = grad_sign[:, sel, :]
- state = state[sel, :] # [G_valid, C]
+ state = state[sel, :] # [G_valid, C]
+ if init_state is not None:
init_state = init_state[sel, :] # [G_valid, C]
- valid_mask = valid_g.unsqueeze(0) # [1, G]
- gaussians.sel = sel
+ gaussians.sel = sel
- if self.cfg.input_gradient_normalize_type == "adam":
- self.update_input_norm.sel = sel
+ if self.cfg.input_gradient_normalize_type == "adam":
+ self.input_norm.sel = sel
return gaussian_grads, gaussian_grads_raw, gaussians, grad_sign, init_state, input_signal, state
- def prune_invisible_gaussians(self, context, context_render_output, gaussian_grads, gaussian_grads_raw, gaussians,
- grad_sign, input_signal, means2d_grads, meta_for_adc, optimizer_input, state):
+ # TODO (release): remove in public code
+ def experimental_prune_invisible_gaussians(self, context, context_render_output, gaussian_grads, gaussian_grads_raw,
+ gaussians,
+ grad_sign, input_signal, means2d_grads, meta_for_adc, optimizer_input,
+ state):
# Get visible gaussians mask, based on the last rendering
with torch.no_grad():
- visible_mask = self.get_visible_gaussian_mask(gaussian_grads, gaussians,
- context_render_output.visibility_filter, context) # [B, N, 1]
+ visible_mask = self.experimental_get_visible_gaussian_mask(gaussian_grads, gaussians,
+ context_render_output.visibility_filter,
+ context) # [B, N, 1]
if visible_mask is None:
return gaussian_grads, gaussians, grad_sign, input_signal, state
assert visible_mask.shape[0] == 1
@@ -1474,8 +1336,6 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
input_signal = input_signal[:, visible_mask] # [B, N, C]
if gaussian_grads is not None:
gaussian_grads = gaussian_grads[:, visible_mask] # [B, N, C]
- if gaussian_grads_raw is not None:
- gaussian_grads_raw = gaussian_grads_raw[:, visible_mask] # [B, N, C]
if grad_sign is not None:
grad_sign = grad_sign[:, visible_mask] # [B, N, C]
meta_for_adc["visibility_filter"] = context_render_output.visibility_filter[:, :, visible_mask]
@@ -1483,20 +1343,22 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
if means2d_grads is not None:
meta_for_adc["means_2d_grads"] = means2d_grads[:, :, visible_mask]
if self.cfg.input_gradient_normalize and self.cfg.input_gradient_normalize_type == "adam":
- if not self.update_input_norm.is_reset():
+ if not self.input_norm.is_reset():
prune_mask = ~visible_mask
- self.update_input_norm.prune(prune_mask) # the prune fn will invert the mask again
+ self.input_norm.prune(prune_mask) # the prune fn will invert the mask again
if self.cfg.any_adc:
optimizer_input.prev_output.state.adc_state.external_pruning(visible_mask)
return gaussian_grads, gaussians, grad_sign, input_signal, state
- def deactivate_updates(self, subset, gaussians, radii_vis_mask, deltas, gaussian_grads):
+ # TODO (release): remove
+ def experimental_deactivate_updates(self, subset, gaussians, radii_vis_mask, deltas, gaussian_grads):
""" Deactivate updates for gaussians that are not visible in any view """
- visible_mask = self.get_visible_gaussian_mask(gaussian_grads, gaussians, radii_vis_mask, subset)
+ visible_mask = self.experimental_get_visible_gaussian_mask(gaussian_grads, gaussians, radii_vis_mask, subset)
deltas = deltas * visible_mask # [B, N, C]
return deltas
- def get_visible_gaussian_mask(self, gaussian_grads, gaussians, radii_vis_mask, subset):
+ # TODO (release): remove
+ def experimental_get_visible_gaussian_mask(self, gaussian_grads, gaussians, radii_vis_mask, subset):
"""
Get mask for gaussians that are visible in at least one view.
@@ -1546,11 +1408,10 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
# Update this parameter
use_norm_grad = getattr(self.cfg.experimental_use_norm_grads, cfg_attr)
use_grad = self.cfg.experimental_use_grads and not use_norm_grad
- use_resplat = not use_grad and not use_norm_grad
assert not (use_grad and use_norm_grad)
if use_grad:
- # Use the inverse of the gradients
- # TODO Naama: hard coded learning rate for SGD
+ # Use the inverse of the gradients (plain gradient-descent step, for comparison)
+ # TODO (release): hard-coded SGD step size (30) in this experimental gradient-baseline path
deltas[..., param_slices[cfg_attr]] = -(grads[..., param_slices[cfg_attr]]).to(deltas.dtype) * 30
elif use_norm_grad:
# Use the inverse of the normalized gradients
@@ -1601,19 +1462,34 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
delta_shs = delta_shs.flatten(-2) # [b, g, d*c] - faster than rearrange
return delta_means, delta_scales, delta_rotations, delta_opacities, delta_shs
- def append_to_input_signal(self, b, context, context_render, tmp_input_signal, v):
+ def _unshuffle_to_flat(self, x: torch.Tensor, b: int, v: int) -> torch.Tensor:
+ """Pixel-unshuffle a (b*v, c, h, w) tensor and flatten spatial dims: -> (b*v*h*w, c)."""
+ x = F.pixel_unshuffle(x, downscale_factor=self.cfg.latent_downsample)
+ return rearrange(x, "(b v) c h w -> (b v h w) c", b=b, v=v)
+
+ def _downsample_error(self, x: torch.Tensor, b: int, v: int) -> torch.Tensor:
+ """Downsample a (b*v, c, h, w) error tensor to latent resolution: -> (b, v*h*w, c)."""
+ if self.cfg.input_error_rgb_no_shuffle:
+ x = F.interpolate(x, scale_factor=1. / self.cfg.latent_downsample,
+ mode='bilinear', align_corners=True)
+ else:
+ # NOTE: pixel_unshuffle assumes per-pixel, grid-aligned Gaussians. A non-grid
+ # subsampling (e.g. farthest-point sampling) would need a different error->Gaussian
+ # mapping to stay aligned — not supported here.
+ x = F.pixel_unshuffle(x, downscale_factor=self.cfg.latent_downsample)
+ return rearrange(x, "(b v) c h w -> b (v h w) c", b=b, v=v)
+
+ def append_to_input_signal(self, b, context, context_render, input_signal_flat, v):
if self.cfg.input_alpha:
- render_alpha = rearrange(context_render.accumulated_alpha, "b v h w -> (b v) () h w")
- render_alpha = F.pixel_unshuffle(render_alpha, downscale_factor=self.cfg.latent_downsample)
- render_alpha = rearrange(render_alpha, "(b v) c h w -> (b v h w) c", b=b, v=v)
- tmp_input_signal = torch.cat((tmp_input_signal, render_alpha), dim=-1)
+ render_alpha = rearrange(context_render.accumulated_alpha.detach(), "b v h w -> (b v) () h w")
+ render_alpha = self._unshuffle_to_flat(render_alpha, b, v)
+ input_signal_flat = torch.cat((input_signal_flat, render_alpha), dim=-1)
if self.cfg.input_depth:
- render_depth = rearrange(context_render.depth, "b v h w -> (b v) () h w")
- render_depth = F.pixel_unshuffle(render_depth, downscale_factor=self.cfg.latent_downsample)
- render_depth = rearrange(render_depth, "(b v) c h w -> (b v h w) c", b=b, v=v)
- tmp_input_signal = torch.cat((tmp_input_signal, render_depth), dim=-1)
+ render_depth = rearrange(context_render.depth.detach(), "b v h w -> (b v) () h w")
+ render_depth = self._unshuffle_to_flat(render_depth, b, v)
+ input_signal_flat = torch.cat((input_signal_flat, render_depth), dim=-1)
if self.cfg.input_depth_smooth_error:
- disp = 1. / context_render.depth.clamp(min=1e-3, max=1000.) # [B, V, H, W]
+ disp = 1. / context_render.depth.detach().clamp(min=1e-3, max=1000.) # [B, V, H, W]
disp = rearrange(disp, "b v h w -> (b v) () h w")
mean_disp = disp.mean(2, True).mean(3, True)
@@ -1622,11 +1498,9 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
tmp_imgs = rearrange(context["image"], "b v c h w -> (b v) c h w")
depth_smooth_error = get_smooth_loss(norm_disp, tmp_imgs, no_mean=True)
-
- depth_smooth_error = F.pixel_unshuffle(depth_smooth_error, downscale_factor=self.cfg.latent_downsample)
- depth_smooth_error = rearrange(depth_smooth_error, "(b v) c h w -> (b v h w) c", b=b, v=v)
- tmp_input_signal = torch.cat((tmp_input_signal, depth_smooth_error), dim=-1)
- return tmp_input_signal
+ depth_smooth_error = self._unshuffle_to_flat(depth_smooth_error, b, v)
+ input_signal_flat = torch.cat((input_signal_flat, depth_smooth_error), dim=-1)
+ return input_signal_flat
def repeat_gaussians(self, prev_means, prev_opacities_raw, prev_rotations_unnorm, prev_scales, prev_shs):
if self.cfg.gaussian_head_multiple > 1:
@@ -1638,7 +1512,7 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
1) / self.cfg.gaussian_head_multiple # smaller opacities, important
prev_shs = prev_shs.repeat(1, self.cfg.gaussian_head_multiple, 1)
# NOTE: only repeat at the first iteration
- refine_repeat = self.cfg.refine_gaussian_multiple
+ refine_repeat = self.cfg.delta_gaussian_multiple
if refine_repeat > 1:
# predict multiple gaussians for each point
prev_means = prev_means.repeat(1, refine_repeat, 1)
@@ -1649,34 +1523,28 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
return prev_means, prev_opacities_raw, prev_rotations_unnorm, prev_scales, prev_shs
def split_delta_gaussians(self, delta_gaussians):
- delta_rotations = None
-
- if self.cfg.init_gaussian_multiple > 1 and not self.cfg.refine_same_num_points:
+ if self.cfg.init_gaussian_multiple > 1 and not self.cfg.same_num_points:
init_repeat = self.cfg.init_gaussian_multiple
else:
init_repeat = 1
p = get_gaussian_param_sizes(self.cfg.sh_d)
- if self.cfg.refine_sh_only:
- delta_shs = delta_gaussians
- delta_means = delta_scales = delta_opacities = 0.
- elif self.cfg.no_refine_rotation:
- delta_means, delta_scales, delta_opacities, delta_shs = delta_gaussians.split(
- (p["means"] * init_repeat, p["scales"] * init_repeat, p["opacities"] * init_repeat,
- p["shs"] * init_repeat), dim=-1
- )
- elif self.cfg.no_refine_mean:
- delta_scales, delta_rotations, delta_opacities, delta_shs = delta_gaussians.split(
- (p["scales"] * init_repeat, p["quats"] * init_repeat, p["opacities"] * init_repeat,
- p["shs"] * init_repeat), dim=-1
- )
- delta_means = torch.zeros_like(delta_scales)
- else:
- delta_means, delta_scales, delta_rotations, delta_opacities, delta_shs = delta_gaussians.split(
- (p["means"] * init_repeat, p["scales"] * init_repeat, p["quats"] * init_repeat,
- p["opacities"] * init_repeat, p["shs"] * init_repeat), dim=-1
- )
+ frozen = self._frozen_param_groups()
+
+ # Split the head output among the predicted (non-frozen) groups; frozen groups get
+ # an additive-identity zero delta so downstream apply/rearrange is uniform.
+ b, n = delta_gaussians.shape[0], delta_gaussians.shape[1]
+ active = [(name, key) for name, key in _GAUSSIAN_GROUPS if name not in frozen]
+ active_chunks = delta_gaussians.split([p[key] * init_repeat for _, key in active], dim=-1)
+ deltas = dict(zip((name for name, _ in active), active_chunks))
+ for name, key in _GAUSSIAN_GROUPS:
+ if name in frozen:
+ deltas[name] = delta_gaussians.new_zeros(b, n, p[key] * init_repeat)
+
+ delta_means, delta_scales = deltas["means"], deltas["scales"]
+ delta_rotations, delta_opacities, delta_shs = deltas["rotations"], deltas["opacities"], deltas["shs"]
+
if (
- self.cfg.refine_gaussian_multiple > 1 or self.cfg.init_gaussian_multiple > 1) and not self.cfg.refine_same_num_points:
+ self.cfg.delta_gaussian_multiple > 1 or self.cfg.init_gaussian_multiple > 1) and not self.cfg.same_num_points:
delta_means = rearrange(delta_means, "b n (c k) -> b (n k) c", k=init_repeat)
delta_scales = rearrange(delta_scales, "b n (c k) -> b (n k) c", k=init_repeat)
delta_rotations = rearrange(delta_rotations, "b n (c k) -> b (n k) c", k=init_repeat)
@@ -1684,63 +1552,36 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
delta_shs = rearrange(delta_shs, "b n (c k) -> b (n k) c", k=init_repeat)
return delta_means, delta_scales, delta_rotations, delta_opacities, delta_shs, init_repeat
- def rearrange_delta_gaussians(self, b, delta_gaussians, delta_gaussians_raw, latent_h, latent_w,
- local_window_update, prev_gaussians_concat, test_window_size, v, window_end,
- window_start):
+ def rearrange_delta_gaussians(self, b, delta_gaussians, delta_gaussians_raw):
# [BV, C]
# update gaussian parameters
delta_gaussians = rearrange(delta_gaussians, "(b n) c -> b n c", b=b)
delta_gaussians_raw = rearrange(delta_gaussians_raw, "(b n) c -> b n c", b=b)
- if local_window_update and not self.cfg.local_gaussian_render:
- # zero padding for non-updated gaussians
- # curr_v = self.cfg.update_window_size if self.training else test_window_size
- curr_v = test_window_size
- tmp_delta = rearrange(delta_gaussians, "b (v h w) c -> b v h w c", b=b, v=curr_v, h=latent_h,
- w=latent_w)
-
- all_delta = []
- # padding
- if window_start > 0:
- tmp_size = rearrange(prev_gaussians_concat, "b (v h w) c -> b v h w c", b=b, v=v, h=latent_h,
- w=latent_w)
- pad_left = torch.zeros_like(tmp_size[:, :window_start, :, :, :], requires_grad=False)
- all_delta.append(pad_left)
-
- all_delta.append(tmp_delta)
-
- if window_end < v:
- tmp_size = rearrange(prev_gaussians_concat, "b (v h w) c -> b v h w c", b=b, v=v, h=latent_h,
- w=latent_w)
- pad_right = torch.zeros_like(tmp_size[:, window_end:, :, :, :], requires_grad=False)
- all_delta.append(pad_right)
-
- tmp_delta = torch.cat(all_delta, dim=1) # [B, V, H, W, C]
- delta_gaussians = rearrange(tmp_delta, "b v h w c -> b (v h w) c") # [B, N, C]
return delta_gaussians, delta_gaussians_raw
- def update_gaussians_params(self, delta_means, delta_scales, delta_rotations, delta_opacities, delta_shs,
- means, scales, rotations_unnorm, opacities_raw, shs,
- repeat):
- means = self.update_means(delta_means, means)
+ def _apply_gaussian_deltas(self, delta_means, delta_scales, delta_rotations, delta_opacities, delta_shs,
+ means, scales, rotations_unnorm, opacities_raw, shs,
+ repeat):
+ means = self._apply_mean_delta(delta_means, means)
# clamp the scale
- scales = self.update_scales(delta_scales, scales, repeat)
+ scales = self._apply_scale_delta(delta_scales, scales, repeat)
if self.cfg.opt_scales_before_act:
scales = scales.exp()
- if not self.cfg.no_refine_rotation:
- rotations, rotations_unnorm = self.update_rotations(delta_rotations, rotations_unnorm)
+ if not self.cfg.freeze_rotation:
+ rotations, rotations_unnorm = self._apply_rotation_delta(delta_rotations, rotations_unnorm)
else:
rotations = F.normalize(rotations_unnorm, dim=-1)
# compute covariance
covariances = build_covariance(scales, rotations) # ([1, VHW, 3, 3])
- opacities_raw = self.update_opacities(delta_opacities, opacities_raw, repeat)
- shs = self.update_shs(delta_shs, shs)
+ opacities_raw = self._apply_opacity_delta(delta_opacities, opacities_raw, repeat)
+ shs = self._apply_shs_delta(delta_shs, shs)
return covariances, means, scales, rotations, rotations_unnorm, opacities_raw, shs
- def update_shs(self, delta_shs, prev_shs):
+ def _apply_shs_delta(self, delta_shs, prev_shs):
shs = prev_shs + delta_shs # [B, N, 3*sh_d]
if self.cfg.clamp_shs_soft:
@@ -1751,7 +1592,7 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
return shs
- def update_opacities(self, delta_opacities, prev_opacities_raw, repeat):
+ def _apply_opacity_delta(self, delta_opacities, prev_opacities_raw, repeat):
# update init opacities when predicting multiple gaussians
if repeat > 1 and not self.cfg.multi_gaussian_scale_smaller and (self.cfg.init_gaussian_multiple == 1):
# Given y = sigmoid(x), to get new x' such that sigmoid(x') = y / K:
@@ -1766,34 +1607,46 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
return prev_opacities_raw
@staticmethod
- def update_rotations(delta_rotations, prev_rotations_unnorm):
+ def _apply_rotation_delta(delta_rotations, prev_rotations_unnorm):
assert delta_rotations is not None
prev_rotations_unnorm = prev_rotations_unnorm + delta_rotations
# normazlie
prev_rotations = prev_rotations_unnorm / (prev_rotations_unnorm.norm(dim=-1, keepdim=True) + 1e-8)
return prev_rotations, prev_rotations_unnorm
- def update_scales(self, delta_scales, prev_scales, repeat):
+ def _apply_scale_delta(self, delta_scales, prev_scales, repeat):
+ """Add the predicted scale delta and clamp the result.
+
+ When one point is represented by `repeat` overlapping gaussians (repeat > 1), the
+ cluster is too dense unless we compensate. `multi_gaussian_scale_smaller` chooses to
+ compensate on size: each gaussian's initial scale is divided by `repeat`, so the
+ cluster covers about the same footprint as the single original gaussian (opacity is
+ left unchanged). When the flag is off, size is kept and the opacity is reduced instead
+ (see _apply_opacity_delta).
+ """
if repeat > 1 and self.cfg.multi_gaussian_scale_smaller:
- # smaller initial scales
- new_scales = (prev_scales / repeat + delta_scales).clamp(min=self.cfg.gaussian_adapter.clamp_min_scale)
+ # smaller initial scales (the floor is applied by the shared clamp below).
+ # The divide-by-repeat is in activation space, so this branch assumes scales
+ # are not in raw/log space.
+ assert not self.cfg.opt_scales_before_act, \
+ "multi_gaussian_scale_smaller assumes opt_scales_before_act=False (activation-space scales)"
+ new_scales = prev_scales / repeat + delta_scales
else:
- new_scales = (prev_scales + delta_scales)
+ new_scales = prev_scales + delta_scales
if self.cfg.opt_scales_before_act:
min_scale = self.cfg.clamp_min_raw_scales
max_scale = self.cfg.clamp_max_raw_scales
else:
min_scale = self.cfg.clamp_min_scale
- max_scale = self.cfg.clamp_refine_max_scale
+ max_scale = self.cfg.clamp_max_scale
- new_scales = new_scales.clamp(min=min_scale)
- new_scales = new_scales.clamp(max=max_scale)
+ new_scales = new_scales.clamp(min=min_scale, max=max_scale)
return new_scales
@staticmethod
- def update_means(delta_means, prev_means):
+ def _apply_mean_delta(delta_means, prev_means):
prev_means = (prev_means + delta_means)
return prev_means
@@ -1807,11 +1660,11 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
self.reset_logs()
if self.cfg.input_gradient_normalize_type == "adam":
- self.update_input_norm.reset()
+ self.input_norm.reset()
nr_gaussians = rearrange(optimizer_input.prev_output.gaussians.means, "b n c -> (b n) c").shape[0]
param_num = self.gaussian_param_num
- self.update_input_norm.initialize(shape=(nr_gaussians, param_num),
- device=optimizer_input.prev_output.gaussians.means.device)
+ self.input_norm.initialize(shape=(nr_gaussians, param_num),
+ device=optimizer_input.prev_output.gaussians.means.device)
# make sure xyz are contiguous
optimizer_input.prev_output.gaussians.means = optimizer_input.prev_output.gaussians.means.contiguous()
@@ -1819,9 +1672,9 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
from_init = False
if self.cfg.input_gradient_normalize_type == "adam":
# Continuing previous optimization from replay buffer
- self.update_input_norm.update_state(optimizer_input.prev_output.state.adam_state)
+ self.input_norm.update_state(optimizer_input.prev_output.state.adam_state)
- # TODO Naama: logs are not handled right now for continuing from replay buffer
+ # Note: logs are not handled right now for continuing from replay buffer
self.reset_logs()
else:
raise ValueError(f"Unknown prev_output type {type(optimizer_input.prev_output)}")
@@ -1834,42 +1687,32 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
if from_init and self.cfg.any_adc:
self.initialize_adc_state(self.cfg, optimizer_input)
- def reshape_gaussians_to_nc(self, latent_h, latent_w, prev_gaussians_concat, v):
- if self.cfg.init_gaussian_multiple == 4 and not self.cfg.refine_same_num_points:
+ def reshape_gaussians_to_nc(self, latent_h, latent_w, gaussians_concat, v):
+ if self.cfg.init_gaussian_multiple > 1 and not self.cfg.same_num_points:
# gaussians are with more points, reshape
- tmp_gaussian = rearrange(prev_gaussians_concat, "b (v h x w y) c -> (b v h w) (c x y)",
- v=v, h=latent_h, w=latent_w, x=2, y=2)
- elif self.cfg.init_gaussian_multiple == 16 and not self.cfg.refine_same_num_points:
- tmp_gaussian = rearrange(prev_gaussians_concat, "b (v h x w y) c -> (b v h w) (c x y)",
- v=v, h=latent_h, w=latent_w, x=4, y=4)
+ factor = self._block_factor()
+ gaussians_flat = rearrange(gaussians_concat, "b (v h x w y) c -> (b v h w) (c x y)",
+ v=v, h=latent_h, w=latent_w, x=factor, y=factor)
else:
- tmp_gaussian = rearrange(prev_gaussians_concat, "b n c -> (b n) c")
- return tmp_gaussian
-
- def get_point_cloud(self, latent_h, latent_w, local_window_update, prev_means, test_window_size, v):
- # TODO: when the initial model predicts multiple gaussians, the number of points also increases
- if self.cfg.init_gaussian_multiple == 4 and not self.cfg.refine_same_num_points:
+ gaussians_flat = rearrange(gaussians_concat, "b n c -> (b n) c")
+ return gaussians_flat
+
+ def get_point_cloud(self, latent_h, latent_w, prev_means, v):
+ if self.cfg.init_gaussian_multiple > 1 and not self.cfg.same_num_points:
+ # The initializer predicts a (factor x factor) block of gaussians per latent pixel,
+ # so there are more points than the base grid. Subsample back to the base grid so the
+ # point cloud (the KNN structure for the transformer) stays at the base resolution.
+ factor = self._block_factor()
point_cloud = rearrange(prev_means, "b (v h w) c -> b v h w c",
- v=v, h=latent_h * 2, w=latent_w * 2,
+ v=v, h=latent_h * factor, w=latent_w * factor,
)
- tmp_batch_size = v * latent_h * latent_w
# simply use uniform grid subsample of point cloud to reduce points
- point_cloud = point_cloud[:, :, ::2, ::2]
+ point_cloud = point_cloud[:, :, ::factor, ::factor]
point_cloud = rearrange(point_cloud, "b v h w c -> (b v h w) c")
- elif self.cfg.init_gaussian_multiple == 16 and not self.cfg.refine_same_num_points:
- point_cloud = rearrange(prev_means, "b (v h w) c -> b v h w c",
- v=v, h=latent_h * 4, w=latent_w * 4,
- )
tmp_batch_size = v * latent_h * latent_w
- # simply use uniform grid subsample of point cloud to reduce points
- point_cloud = point_cloud[:, :, ::4, ::4]
- point_cloud = rearrange(point_cloud, "b v h w c -> (b v h w) c")
else:
point_cloud = rearrange(prev_means, "b n c -> (b n) c")
- if local_window_update:
- tmp_batch_size = test_window_size * latent_h * latent_w
- else:
- tmp_batch_size = prev_means.shape[1]
+ tmp_batch_size = prev_means.shape[1]
return point_cloud, tmp_batch_size
def get_vector_state(self, b, v, n, optimizer_input, from_init):
@@ -1895,25 +1738,11 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
state = self.get_state_from_condition_features(b, optimizer_input.prev_output.features,
v) # [B, N, C]
+ # combine gaussians of all scenes in the batch [B, N, C] -> [B*N, C]
+ state = rearrange(state, "b n c -> (b n) c")
else:
- # Restarting optimizing a scene from a replay buffer
+ # Restarting a scene from a replay buffer: the state is already stored flat as [B*N, C].
state = optimizer_input.prev_output.state.state
- # TODO Naama: need to understand why rearrange here, perhaps something with pruning
- state = rearrange(state, "(b n) c -> b n c", b=b)
-
- # combine gaussians of all scnes in the batch [B*N, C]
- state = rearrange(state, "b n c -> (b n) c") # [B*N, C]
-
- # Do something with window size
- _, _, _, h, w = optimizer_input.context["image"].shape # [B, V, C, H, W]
- local_window_update, test_window_size, window_end, window_start = optimizer_input.additional_info
- # select initial state
- if local_window_update and self.cfg.local_gaussian_render:
- state = rearrange(state, "(b v h w) c -> b v h w c", b=b, v=v,
- h=h // self.cfg.latent_downsample,
- w=w // self.cfg.latent_downsample)
- state = state[:, window_start:window_end, :, :, :]
- state = rearrange(state, "b v h w c -> (b v h w) c")
return state
@@ -1928,53 +1757,39 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
out.append(feat)
return out
+ def _block_factor(self) -> int:
+ """Per-pixel block side length when the initializer predicts multiple gaussians per pixel.
+
+ init_gaussian_multiple gaussians form a (factor x factor) block, so factor = sqrt(multiple).
+ """
+ multiple = self.cfg.init_gaussian_multiple
+ factor = math.isqrt(multiple)
+ assert factor * factor == multiple, f"init_gaussian_multiple must be a perfect square, got {multiple}"
+ return factor
+
def _get_latent_size(self, h: int, w: int) -> tuple[int, int]:
"""Compute latent (H, W) from image (H, W), accounting for init_gaussian_multiple upsampling."""
latent_h = h // self.cfg.latent_downsample
latent_w = w // self.cfg.latent_downsample
- if self.cfg.init_gaussian_multiple == 4 and self.cfg.refine_same_num_points:
- latent_h *= 2
- latent_w *= 2
- elif self.cfg.init_gaussian_multiple == 16 and self.cfg.refine_same_num_points:
- latent_h *= 4
- latent_w *= 4
+ if self.cfg.init_gaussian_multiple > 1 and self.cfg.same_num_points:
+ factor = self._block_factor()
+ latent_h *= factor
+ latent_w *= factor
return latent_h, latent_w
- def render_input_views_for_error_calc(self, context,
- local_window_update,
- prev_gaussians,
- renderer,
- window_end,
- window_start,
- num_refine,
- i):
+ def render_input_views_for_error_calc(self, context, prev_gaussians, renderer):
_, _, _, h, w = context["image"].shape # [B, V, C, H, W]
render_res = (h, w)
- # Default rendering parameters
- input_info = context
- start = None
- end = None
- cfg = self.cfg
-
# Use only first N views
- if cfg.input_error_num_views > 0:
- end = cfg.input_error_num_views
+ end = self.cfg.input_error_num_views if self.cfg.input_error_num_views > 0 else None
- # Local window update logic
- elif local_window_update:
- if i >= num_refine - 1:
- return None # Skip rendering on the last iteration
- start = window_start
- end = window_end
-
- # Final unified rendering call
return renderer.forward_batch_subset(
prev_gaussians,
- input_info,
+ context,
render_res,
- start=start,
+ start=None,
end=end,
return_radii=False
)
@@ -1983,87 +1798,31 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
with torch.amp.autocast(device_type='cuda', enabled=self.cfg.pt_update_amp, dtype=torch.bfloat16):
if not self.cfg.pt_update_amp and condition_features.dtype == torch.bfloat16:
condition_features = condition_features.float()
- state = self.update_proj(condition_features.detach()) # [B, C, H, W]
- if self.cfg.init_gaussian_multiple == 4 and self.cfg.refine_same_num_points:
- state = F.interpolate(state, scale_factor=2, mode='bilinear', align_corners=True)
- elif self.cfg.init_gaussian_multiple == 16 and self.cfg.refine_same_num_points:
- state = F.interpolate(state, scale_factor=4, mode='bilinear', align_corners=True)
- else:
- pass
+ state = self.state_proj(condition_features.detach()) # [B, C, H, W]
+ if self.cfg.init_gaussian_multiple > 1 and self.cfg.same_num_points:
+ factor = self._block_factor()
+ state = F.interpolate(state, scale_factor=factor, mode='bilinear', align_corners=True)
# Convert to vector of Gaussians per batch [B, N, C]
state = rearrange(state, "(b v) c h w -> b (v h w) c", b=b, v=v) # N = v * h * w
return state
- def get_window_size(self, v):
- test_window_size = None
- if self.cfg.update_window_size > 0:
-
- local_window_update = True
- # if self.training:
- # window_start = random.randint(0, v - self.cfg.update_window_size)
- # window_end = window_start + self.cfg.update_window_size
- # else:
- # fixed window at test time, uniformly move from left to right
- # TODO: loop closure, connect left and right
- if self.training:
- test_window_size = self.cfg.update_window_size
- window_start = random.randint(0, test_window_size)
- window_end = window_start + test_window_size
-
- if window_end == v:
- # restart
- window_start = random.randint(0, test_window_size)
- window_end = window_start + test_window_size
- else:
- # at least do a full pass of all input views
- # test_window_size = int(np.ceil(v / self.cfg.num_refine))
- test_window_size = self.cfg.update_window_size
- window_start = 0
- window_end = window_start + test_window_size
-
- else:
- local_window_update = False
- window_start = 0
- window_end = v
- return local_window_update, test_window_size, window_end, window_start
-
- def prepare_update_input(self, b, i, init_state, input_signal, latent_h, latent_w, local_window_update, point_cloud,
- tmp_gaussian, state, v, window_end, window_start):
+ def prepare_update_input(self, init_state, input_signal, point_cloud, gaussians_flat, state):
if self.cfg.replace_init_state:
state = init_state
if self.cfg.no_render_error:
- update_input = torch.cat((tmp_gaussian, state), dim=-1)
+ update_input = torch.cat((gaussians_flat, state), dim=-1)
else:
- if local_window_update and not self.cfg.local_gaussian_render:
- # select local window
- tmp_gaussian = rearrange(tmp_gaussian, "(b v h w) c -> b v h w c", b=b, v=v, h=latent_h,
- w=latent_w)
- tmp_gaussian = tmp_gaussian[:, window_start:window_end, :, :, :]
- tmp_gaussian = rearrange(tmp_gaussian, "b v h w c -> (b v h w) c")
-
- if i == 0:
- state = rearrange(state, "(b v h w) c -> b v h w c", b=b, v=v, h=latent_h,
- w=latent_w)
- state = state[:, window_start:window_end, :, :, :]
- state = rearrange(state, "b v h w c -> (b v h w) c")
-
- # local point cloud
- point_cloud = rearrange(point_cloud, "(b v h w) c -> b v h w c", b=b, v=v, h=latent_h,
- w=latent_w)
- point_cloud = point_cloud[:, window_start:window_end, :, :, :]
- point_cloud = rearrange(point_cloud, "b v h w c -> (b v h w) c")
-
- update_input = torch.cat((tmp_gaussian, state, input_signal), dim=-1)
+ update_input = torch.cat((gaussians_flat, state, input_signal), dim=-1)
if self.cfg.concat_init_state:
update_input = torch.cat((update_input, init_state), dim=-1)
- return point_cloud, tmp_gaussian, state, update_input
+ return point_cloud, gaussians_flat, state, update_input
- def apply_update_module(self, b, latent_h, latent_w, offset, point_cloud, update_input, v, state, iter):
+ def _apply_point_transformer(self, b, latent_h, latent_w, offset, point_cloud, update_input, v, state, iter):
def recurrent_chunk(update_input, point_cloud, offset):
- pxo = self.update_module[0]([point_cloud, update_input, offset])
- state = self.update_module[1](pxo, iter=iter, b=b, v=v, h=latent_h, w=latent_w)
+ pxo = self.point_transformer[0]([point_cloud, update_input, offset])
+ state = self.point_transformer[1](pxo, iter=iter, b=b, v=v, h=latent_h, w=latent_w)
return state
if self.cfg.use_checkpointing or self.cfg.recurrent_use_checkpointing:
@@ -2080,23 +1839,21 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
return new_state
def apply_delta_gaussian_head(self, b, context, init_state, state, v):
- if self.cfg.update_head_concat_img:
- # pixel unshuffle image
+ if self.cfg.delta_head_concat_img:
img_unshuffle = rearrange(context["image"], "b v c h w -> (b v) c h w")
- img_unshuffle = F.pixel_unshuffle(img_unshuffle, downscale_factor=self.cfg.latent_downsample)
- img_unshuffle = rearrange(img_unshuffle, "(b v) c h w -> (b v h w) c", b=b, v=v)
+ img_unshuffle = self._unshuffle_to_flat(img_unshuffle, b, v)
head_input = torch.cat((state, img_unshuffle), dim=-1)
else:
- if self.cfg.refine_residual_init_state:
+ if self.cfg.residual_init_state:
head_input = state + init_state
else:
head_input = state
- if self.cfg.update_head_per_param_heads:
+ if self.cfg.delta_head_per_param_heads:
delta_gaussians = self._apply_per_param_heads(head_input)
else:
- delta_gaussians = self.update_head(head_input)
+ delta_gaussians = self.delta_head(head_input)
return delta_gaussians
@@ -2108,7 +1865,7 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
"""
deltas = []
for name, dim in self._per_param_group_dims.items():
- raw = self.update_head[name](head_input) # [N, dim+1]
+ raw = self.delta_head[name](head_input) # [N, dim+1]
scale = self.scale_act(raw[:, -1:]) # [N, 1]
delta = raw[:, :-1] # [N, dim]
if dim > 1:
@@ -2119,9 +1876,7 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
deltas.append(delta)
return torch.cat(deltas, dim=-1)
- def apply_global_attn(self, b, h, input_signal, latent_h, latent_w,
- local_window_update, test_window_size, v, w):
- # TODO Naama: do we need local_window?
+ def apply_global_attn(self, h, input_signal, latent_h, latent_w, v, w):
assert self.cfg.input_error_resnet_feature
assert self.cfg.input_error
@@ -2131,13 +1886,12 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
input_render_error = input_signal
with torch.amp.autocast(device_type='cuda', enabled=self.cfg.use_amp, dtype=torch.bfloat16):
- for blk in self.update_error_attn:
- if self.cfg.refine_same_num_points:
+ for blk in self.error_mv_attn:
+ if self.cfg.same_num_points:
# no downsample, for re10k 256
input_render_error = blk(input_render_error, v=v, h=h, w=w)
else:
- curr_v = test_window_size if local_window_update else v
- input_render_error = blk(input_render_error, v=curr_v, h=latent_h, w=latent_w)
+ input_render_error = blk(input_render_error, v=v, h=latent_h, w=latent_w)
if self.cfg.input_gradient and self.cfg.input_error:
input_signal[..., :self.error_features_channels] = input_render_error
@@ -2146,16 +1900,16 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
return input_signal
- def prepare_input_signal(self, context, i, gaussians,
- local_window_update, renderer,
- window_end, window_start, num_refine):
- # TODO Naama: review
+ def prepare_input_signal(self, context, gaussians, renderer):
+ """Build the per-Gaussian signal that drives the update: render gradients and/or feature error.
+
+ Renders the context views once and turns the errors with the inputs into the network's input
+ signal (gradients from `_calc_input_gradients`, render error, or both concatenated).
+ """
# make sure at least one of the following is True
assert self.cfg.input_gradient or self.cfg.input_error
- input_view_features = None
input_signal = None
input_render_error = None
- context_render_output = None
gaussian_grads_raw = None
gaussian_grads = None
grad_sign = None
@@ -2168,25 +1922,16 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
)
input_signal = gaussian_grads_raw
-
- # When using gradients, context_render_output cannot be used for the meta-training,
- # because there was already one backward pass.
- # So we render again if in training.
- if context_render_output is None or self.training:
- context_render_output = self.render_input_views_for_error_calc(context, local_window_update,
- gaussians, renderer, window_end,
- window_start, num_refine, i)
+ else:
+ # Render if input_gradient=False. Otherwise, the gradients render can be reused for the error
+ # calculation, all further usage detach before use.
+ context_render_output = self.render_input_views_for_error_calc(context, gaussians, renderer)
# calculate input rendering errors
if self.cfg.input_error:
if means2d_grads is None and self.cfg.need_2d_grads:
raise NotImplementedError("Calculating 2dgrad for ADC is not implemented for error input alone")
- input_render_error = self._calc_input_errors(context, i, context_render_output,
- input_view_features,
- local_window_update,
- gaussians.means.detach(),
- window_end,
- window_start)
+ input_render_error = self._calc_input_errors(context, context_render_output)
input_signal = input_render_error
if self.cfg.input_gradient and self.cfg.input_error:
@@ -2195,22 +1940,6 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
return input_signal, gaussian_grads_raw, gaussian_grads, grad_sign, context_render_output, means2d_grads
- 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
-
- @property
- def sampler(self):
- return None
-
def debug_reprojection_error(self, means, debug_dict, context, i, latent_h, latent_w):
# Prepare means (remove singleton dim)
means = rearrange(means, "b (v h w) c -> b v (h w) c", h=latent_h, w=latent_w) # [B, V, H*W, 3]
@@ -2254,9 +1983,7 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
# plt.ylabel("Frequency")
# plt.show()
- def _calc_input_errors(self, context, i, input_render, input_view_features,
- local_window_update, prev_means,
- window_end, window_start):
+ def _calc_input_errors(self, context, input_render):
b, v, _, h, w = context["image"].shape
# Detach the last rendered object
input_rgb = input_render.color.detach()
@@ -2265,123 +1992,80 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
input0 = rearrange(input_rgb, "b v c h w -> (b v) c h w")
if self.cfg.input_error_num_views > 0:
gt_input = context["image"][:, :self.cfg.input_error_num_views, :, :, :]
- elif local_window_update:
- gt_input = context["image"][:, window_start:window_end, :, :, :]
else:
gt_input = context["image"]
input1 = rearrange(gt_input, "b v c h w -> (b v) c h w")
transform = _IMAGENET_NORM
+ concat = torch.cat((input0, input1), dim=0)
- if input_view_features is None:
- assert i == 0
- # first time: extract all features
- concat = torch.cat((input0, input1), dim=0)
+ input_tensor = transform(concat)
+ with torch.amp.autocast(device_type='cuda', enabled=self.cfg.pt_update_amp,
+ dtype=torch.bfloat16):
+ with torch.no_grad():
+ features = self.error_feature_extractor(input_tensor)
- input_tensor = transform(concat)
- with torch.amp.autocast(device_type='cuda', enabled=self.cfg.pt_update_amp,
- dtype=torch.bfloat16):
- # Extract features
- with torch.no_grad():
- features = self.update_feature(input_tensor)
-
- # align to the latent resolution
- latent_h, latent_w = self._get_latent_size(h, w)
-
- all_features = torch.cat(self._align_features(features, latent_h, latent_w), dim=1)
-
- render_view_features = all_features[:input0.shape[0]]
- input_view_features = all_features[input0.shape[0]:]
-
- else:
- # only extract render view features
- with torch.amp.autocast(device_type='cuda', enabled=self.cfg.pt_update_amp,
- dtype=torch.bfloat16):
- # Extract features
- with torch.no_grad():
- features = self.update_feature(transform(input0))
+ # align to the latent resolution
+ latent_h, latent_w = self._get_latent_size(h, w)
- # align to the latent resolution
- latent_h, latent_w = self._get_latent_size(h, w)
+ all_features = torch.cat(self._align_features(features, latent_h, latent_w), dim=1)
- render_view_features = torch.cat(self._align_features(features, latent_h, latent_w), dim=1)
+ render_view_features = all_features[:input0.shape[0]]
+ gt_view_features = all_features[input0.shape[0]:]
- corr = render_view_features - input_view_features
+ feature_error = render_view_features - gt_view_features
if self.cfg.input_error_num_views > 0:
# pad to V views
curr_v = self.cfg.input_error_num_views
indices = torch.arange(v) * curr_v // v
- corr = rearrange(corr, "(b v) c h w -> b v c h w", b=b)
- corr = corr[torch.arange(b).unsqueeze(1), indices, :, :, :]
- input_render_error = rearrange(corr, "b v c h w -> b (v h w) c")
+ feature_error = rearrange(feature_error, "(b v) c h w -> b v c h w", b=b)
+ feature_error = feature_error[torch.arange(b).unsqueeze(1), indices, :, :, :]
+ input_render_error = rearrange(feature_error, "b v c h w -> b (v h w) c")
else:
- input_render_error = rearrange(corr, "(b v) c h w -> b (v h w) c", b=b)
+ input_render_error = rearrange(feature_error, "(b v) c h w -> b (v h w) c", b=b)
else:
input_render_error = (input_render.color - context["image"]).abs() # [B, V, 3, H, W]
input_render_error = rearrange(input_render_error, "b v c h w -> (b v) c h w")
- if self.cfg.input_error_rgb_no_shuffle:
- # bilinear
- input_render_error = F.interpolate(input_render_error,
- scale_factor=1. / self.cfg.latent_downsample,
- mode='bilinear', align_corners=True)
- else:
- # pixel unshuffle
- # TODO: when fps is used, how to reshape the render error to make sure its somehow pixel aligned to the gaussians
- input_render_error = F.pixel_unshuffle(input_render_error,
- downscale_factor=self.cfg.latent_downsample)
-
- input_render_error = rearrange(input_render_error, "(b v) c h w -> b (v h w) c", b=b,
- v=v) # [B, N, C]
+ input_render_error = self._downsample_error(input_render_error, b, v) # [B, N, C]
# include both feature error and image error
if self.cfg.input_error_add_rgb_feature:
rgb_render_error = input_render.color - context["image"]
rgb_render_error = rearrange(rgb_render_error, "b v c h w -> (b v) c h w")
- if self.cfg.input_error_rgb_no_shuffle:
- # bilinear
- rgb_render_error = F.interpolate(rgb_render_error, scale_factor=1. / self.cfg.latent_downsample,
- mode='bilinear', align_corners=True)
- else:
- # pixel unshuffle
- # TODO: when fps is used, how to reshape the render error to make sure its somehow pixel aligned to the gaussians
- rgb_render_error = F.pixel_unshuffle(rgb_render_error,
- downscale_factor=self.cfg.latent_downsample)
+ rgb_render_error = self._downsample_error(rgb_render_error, b, v) # [B, N, C]
- rgb_render_error = rearrange(rgb_render_error, "(b v) c h w -> b (v h w) c", b=b, v=v) # [B, N, C]
-
- rgb_render_error = self.update_rgb_error_proj(rgb_render_error)
+ rgb_render_error = self.rgb_error_proj(rgb_render_error)
input_render_error = input_render_error + rgb_render_error
return input_render_error
def get_input_error_feature_extractor(self):
- update_feature = None
+ error_feature_extractor = None
# resnet feature
if self.cfg.input_error_resnet_feature:
- update_feature = ResNetFeatureWarpper(
+ error_feature_extractor = ResNetFeatureWarpper(
shallow_resnet_feature=self.cfg.input_error_shallow_resnet_feature)
if self.cfg.input_error_no_freeze_resnet_feature:
# remove unused layers
# NOTE: layer 3 is also not used
- update_feature.layer3 = nn.Identity()
- update_feature.train()
- for params in update_feature.parameters():
+ error_feature_extractor.layer3 = nn.Identity()
+ error_feature_extractor.train()
+ for params in error_feature_extractor.parameters():
params.requires_grad = True
else:
- update_feature.eval()
+ error_feature_extractor.eval()
- for params in update_feature.parameters():
+ for params in error_feature_extractor.parameters():
params.requires_grad = False
- return update_feature
+ return error_feature_extractor
- def update_delta_for_gradients_input(self, delta_gaussians, grad_sign, normalized_grad,
- visibility_scale: Tensor | None = None):
+ def _postprocess_delta_for_gradient_input(self, delta_gaussians, grad_sign, normalized_grad):
if self.cfg.input_gradient:
delta_gaussians = delta_gaussians / self.cfg.input_gradient_scale
if self.cfg.input_gradient_log:
@@ -2397,21 +2081,13 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
clip_mask = delta_gaussians.abs() > clip_value
delta_gaussians[clip_mask] = delta_gaussians[clip_mask].sign() * clip_value
- # TODO Naama: move these two, as they are not related to gradients
- if self.cfg.update_head_scale_mag:
- out_channels = delta_gaussians.shape[-1]
- param_num = out_channels / 2
- assert param_num.is_integer()
- param_num = int(param_num)
- scale = delta_gaussians[:, :param_num]
- mag = delta_gaussians[:, param_num:]
- delta_gaussians = scale * 0.01 * torch.exp(mag * 0.01)
-
- if self.cfg.update_head_scalar_scale:
- if self.cfg.update_head_per_param_heads:
+ # NOTE: scalar_scale is a delta-head concern, not a gradient one, but it lives
+ # here because it only runs in the gradient path today.
+ if self.cfg.delta_head_scalar_scale:
+ if self.cfg.delta_head_per_param_heads:
# Already handled in _apply_per_param_heads — nothing to do here
pass
- elif self.cfg.update_head_per_param_scales:
+ elif self.cfg.delta_head_per_param_scales:
# Feature B: per-group scalar scales
num_groups = len(self._per_param_group_dims)
scales = delta_gaussians[:, -num_groups:] # [G, num_groups]
@@ -2438,11 +2114,8 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
deltas_norm = deltas_unnorm / (deltas_unnorm.norm(p=2, dim=1, keepdim=True) + 1e-8) # [G, D]
delta_gaussians = deltas_norm * scale
- if visibility_scale is not None:
- delta_gaussians = delta_gaussians * visibility_scale
-
if self.cfg.scale_residual_grads:
- delta_gaussians = delta_gaussians * normalized_grad * self.cfg.gradient_update_scale # 1.0
+ delta_gaussians = delta_gaussians * normalized_grad * self.cfg.residual_grad_scale # 1.0
# To match the default behavior of SGD, Adam, and other optimizers, deltas are negated.
# SGD applies the gradients as `x = x - lr * grad`, while resaplt applies them as `x = x + lr * deltas`.
@@ -2451,6 +2124,12 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
return delta_gaussians
def _calc_input_gradients(self, context, gaussians, renderer):
+ """Render the context views and backprop the reconstruction loss to per-Gaussian parameter gradients.
+
+ These gradients are a feature fed to the network (the "what is wrong" signal), not used to update
+ weights — the graph is detached. Chunked over views to bound memory; optionally also returns the
+ 2D means gradients ADC consumes.
+ """
assert not self.cfg.input_gradient_same_loss, "input_gradient_same_loss is not implemented"
_, v, _, h, w = context["image"].shape
@@ -2465,7 +2144,7 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
detach=True,
clone=False,
requires_grad=True,
- scales_lims=(self.cfg.clamp_min_scale, self.cfg.clamp_refine_max_scale),
+ scales_lims=self.cfg.scales_clamp_lims,
raw_opacities_lims=(self.cfg.clamp_min_raw_opacities, self.cfg.clamp_max_raw_opacities)
)
@@ -2570,58 +2249,3 @@ class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]):
return gaussian_grads_raw, gaussian_grads, grads_sign, context_render_output, means2d_grads
-def select_gaussian_subset(gaussians, window_start, window_end, v, h, w):
- """Select a subset of gaussians based on view window. Optimized to avoid rearrange overhead."""
- b = gaussians.means.shape[0]
- hw = h * w
- window_v = window_end - window_start
- new_n = window_v * hw
-
- # Helper to slice view dimension efficiently using view+slice+reshape instead of rearrange
- def slice_tensor(t, extra_dims):
- # t shape: [b, v*h*w, *extra_dims] -> [b, window_v*h*w, *extra_dims]
- shape = (b, v, hw) + extra_dims
- new_shape = (b, new_n) + extra_dims
- return t.view(shape)[:, window_start:window_end, :].reshape(new_shape)
-
- means = slice_tensor(gaussians.means, (3,))
- covariances = slice_tensor(gaussians.covariances, (3, 3)) if gaussians.covariances is not None else None
- shs = slice_tensor(gaussians.harmonics, gaussians.harmonics.shape[2:])
- opacities = slice_tensor(gaussians.opacities.unsqueeze(-1), ()).squeeze(-1)
- scales = slice_tensor(gaussians.scales, (3,))
- rotations = slice_tensor(gaussians.rotations, (4,)) if gaussians.rotations is not None else None
- rotations_unnorm = slice_tensor(gaussians.rotations_unnorm, (4,))
-
- return Gaussians(
- means=means,
- covariances=covariances,
- harmonics=shs,
- opacities=opacities,
- scales=scales,
- rotations=rotations,
- rotations_unnorm=rotations_unnorm,
- )
-
-
-def replace_window(original, window, window_start, window_end, dim=1):
- slices = []
- if window_start > 0:
- # TODO: detach or not
- # slices.append(original[:, :window_start].detach())
- slices.append(original[:, :window_start])
- slices.append(window)
- if window_end < original.shape[dim]:
- # TODO: detach or not
- # slices.append(original[:, window_end:].detach())
- slices.append(original[:, window_end:])
- return torch.cat(slices, dim=dim)
-
-
-def freeze_batchnorm_layers(model):
- import torch.nn as nn
- for module in model.modules():
- if isinstance(module, nn.BatchNorm2d) or isinstance(module, nn.BatchNorm1d) or isinstance(module,
- nn.BatchNorm3d):
- module.eval() # Set to evaluation mode
- for param in module.parameters():
- param.requires_grad = False # Freeze parameters
diff --git a/optgs/scene_trainer/optimizer/optimizer_resplat.py b/optgs/scene_trainer/optimizer/optimizer_resplat.py
index 1439c605eae63fabcc9503bdea13187cef6c0ade..e1b059f1e6919a1afbca7612eb344b3a62813af6 100644
--- a/optgs/scene_trainer/optimizer/optimizer_resplat.py
+++ b/optgs/scene_trainer/optimizer/optimizer_resplat.py
@@ -1,6 +1,4 @@
-from typing import Literal
-
-from optgs.scene_trainer.optimizer import KnnBasedOptimizer, KnnBasedOptimizerCfg
+from optgs.scene_trainer.optimizer import KnnBasedOptimizer
class ResplatOptimizerV1(KnnBasedOptimizer):
diff --git a/optgs/scene_trainer/optimizer/optimizer_utils.py b/optgs/scene_trainer/optimizer/optimizer_utils.py
index 1a4d04e72b00836b64c11658175b619037cbdf3f..7b07a32ce7f271b4c57b988b00b07ac17b14dcdd 100644
--- a/optgs/scene_trainer/optimizer/optimizer_utils.py
+++ b/optgs/scene_trainer/optimizer/optimizer_utils.py
@@ -13,9 +13,7 @@ from optgs.model.decoder.decoder import Decoder, DecoderOutput
from optgs.model.types import Gaussians
from einops import rearrange
from tqdm import tqdm
-import gc
-import torch.autograd.profiler as profiler
-from optgs.misc.memory_profiler import profile_gpu_memory, report_gpu_tensors
+from optgs.misc.memory_profiler import profile_gpu_memory
T = TypeVar("T")
GPU_MEM_PROFILING = False # set to True to enable GPU memory profiling
@@ -30,9 +28,12 @@ def split_grads(grads_tensor, cfg):
assert grads_tensor.shape[0] == 1, "Batch size > 1 not supported for grads_tensor with ndim 3"
grads_tensor = grads_tensor.squeeze(0) # [N, D]
- # Split the last dimension
+ # Split the last dimension into the packed Gaussian parameter groups
+ sizes = get_gaussian_param_sizes(cfg.sh_d)
means, scales, rotations, opacities, shs = torch.split(
- grads_tensor, (3, 3, 4, 1, 3 * cfg.sh_d), dim=-1
+ grads_tensor,
+ (sizes["means"], sizes["scales"], sizes["quats"], sizes["opacities"], sizes["shs"]),
+ dim=-1,
)
shs = rearrange(shs, "n (c x) -> n c x", c=3, x=cfg.sh_d) # [N, 3, sh_d]
@@ -122,14 +123,6 @@ def chunk_ranges(v: int, chunk_size: int) -> List[Tuple[int, int]]:
start = stop
return ranges
-def chunk_slices(v: int, chunk_size: int, dim: int = 1) -> List[slice]:
- """
- Return a list of slice objects that slice along axis `dim`.
- Use like: tensor[(slice(None), slice_start_stop, ...)] — easier: use helper below.
- NOTE: slice objects don't encode the axis; they only give start/stop; see usage.
- """
- return [slice(s, e) for s, e in chunk_ranges(v, chunk_size)]
-
def chunk_index_iter(v: int, chunk_size: int) -> Iterator[Tuple[int,int,int]]:
"""
Iterate chunk info as (chunk_idx, start, stop) for convenience.
@@ -336,28 +329,30 @@ def unpack_gaussians(
detach: bool = True,
clone: bool = False,
requires_grad: bool = False,
- scales_lims: tuple | None = None, # post activation (1e-6, 3)
- raw_opacities_lims: tuple | None = None, # pre activation (-7, 7)
+ scales_lims: tuple | None = None, # in scales' refine space: log if scales_log else activation
+ raw_opacities_lims: tuple | None = None, # logit (pre-sigmoid) space, e.g. (-7, 7)
):
""" Unpack Gaussian parameters and invert opacities and scales.
- # TODO Naama: fix this
- Clamp values for scales are in post-activation space, i.e., after exponentiation.
- Clamp values for opacities are in pre-activation space, i.e., before sigmoid
-
+ Clamp limits apply in the space each parameter is refined in: scales_lims in log space
+ when scales_log is set (else activation space), raw_opacities_lims in logit (pre-sigmoid)
+ space.
"""
+ # Restrict to the valid (non-padding) gaussians when a selection is set.
+ gaussians = gaussians.select_valid()
+
# Means
means = gaussians.means # [B, N, 3]
# Scales
scales = gaussians.scales # [B, N, 3]
- if scales_lims is not None:
- scales = torch.clamp(scales, scales_lims[0], scales_lims[1])
- # if self.cfg.opt_scales_before_act:
if scales_log:
# Invert also scales
scales = torch.log(scales + 1e-8)
+ # Clamp in the same space the scales now live in (log if scales_log, else activation).
+ if scales_lims is not None:
+ scales = torch.clamp(scales, scales_lims[0], scales_lims[1])
# Quaternions
# use unnormalized rotations since we are going to refine the unnormed rotations
@@ -379,15 +374,6 @@ def unpack_gaussians(
shs = gaussians.harmonics # [B, N, 3, 9]
shs = shs.flatten(-2) # [B, N, C] - faster than rearrange
- if gaussians.sel is not None:
- # TODO Naama: move method to Gaussians class
- sel = gaussians.sel # [B, N]
- means = means[:, sel]
- opacities_raw = opacities_raw[:, sel]
- rotations_unnorm = rotations_unnorm[:, sel]
- scales = scales[:, sel]
- shs = shs[:, sel]
-
if detach:
means = means.detach()
opacities_raw = opacities_raw.detach()
@@ -409,29 +395,6 @@ def unpack_gaussians(
scales.requires_grad_(True)
shs.requires_grad_(True)
- # # predicting multiple gaussians per point, init new gaussians by copying with scaled opacities
- # if self.cfg.reinit_gaussian_when_refine_multiple and self.cfg.refine_gaussian_multiple > 1:
- # raise NotImplementedError
- # # This should only be called at the first iteration
- # # TODO Naama: might be bug if we use replay buffer
- # repeat = self.cfg.refine_gaussian_multiple
- # prev_means = prev_means.repeat(1, repeat, 1)
- # prev_scales = prev_scales.repeat(1, repeat, 1)
- # prev_rotations_unnorm = prev_rotations_unnorm.repeat(1, repeat, 1)
- #
- # # scale down opacities
- # prev_opacities_raw = prev_opacities_raw.repeat(1, repeat, 1) # smaller opacities, important
- # # Given y = sigmoid(x), to get new x' such that sigmoid(x') = y / K:
- # # The formula is: x' = x + log((1 - y) / (K - y))
- # # This adjusts x so that the sigmoid output is scaled down by a factor of K
- # tmp_sigmoid = prev_opacities_raw.sigmoid()
- # # print(tmp_sigmoid.mean().item())
- # prev_opacities_raw = prev_opacities_raw + torch.log((1 - tmp_sigmoid) / (repeat - tmp_sigmoid))
- #
- # prev_shs = prev_shs.repeat(1, repeat, 1)
- #
- # # TODO: this part not ready
-
return means, scales, rotations_unnorm, opacities_raw, shs
@@ -484,6 +447,8 @@ def pack_gaussians(
def get_visibility_contribution_from_gaussian_obj(views_info, gaussians, image_shape=None, render_image=False) -> tuple[Tensor, dict]:
"""
+ # TODO (release): experimental code
+
Args:
views_info: dict containing:
"extrinsics": Tensor of shape [B, V, 4, 4]
@@ -501,8 +466,6 @@ def get_visibility_contribution_from_gaussian_obj(views_info, gaussians, image_s
out[k] = sum_{c,i,j}^{C, H, W} w_{k,c,i,j}
️
"""
- # Context can be either context or target
- # TODO Naama: check visibility for both context and target views
b = gaussians.means.shape[0]
assert b == 1
# Data preparation
diff --git a/optgs/scene_trainer/optimizer/time_embed.py b/optgs/scene_trainer/optimizer/time_embed.py
index 51cce544b0e448cd31e07308ecb8cdc7847cb0c0..fa9f6f4b2286b69b9f678ce7ea0f81a51fee1f86 100644
--- a/optgs/scene_trainer/optimizer/time_embed.py
+++ b/optgs/scene_trainer/optimizer/time_embed.py
@@ -1,7 +1,6 @@
# From: https://github.com/ingra14m/Deformable-3D-Gaussians/blob/main/utils/time_utils.py
import torch
-import torch.nn as nn
def get_embedder(multires):
diff --git a/optgs/scene_trainer/postprocessing.py b/optgs/scene_trainer/postprocessing.py
index 99c78755a13db00bf2c6ac84265767fb18bdc6fd..da1c95303bc3c7b961bd1355784d7480a35f13a6 100644
--- a/optgs/scene_trainer/postprocessing.py
+++ b/optgs/scene_trainer/postprocessing.py
@@ -1,4 +1,12 @@
-from dataclasses import dataclass, field
+"""Test-time 3DGS post-processing applied after the learned optimizer.
+
+`PostProcessing3DGS` runs extra vanilla-3DGS refinement on the optimized Gaussians at test time: a torch
+Adam/SGD optimizer descends the parameters directly (built as a `GaussiansModule`), optionally with
+adaptive density control (clone/split/prune). This is the `meta_trainer.test.postprocessing` step and
+also provides the standalone 3DGS baseline numbers; it is off by default.
+"""
+
+from dataclasses import dataclass
from typing import List
import tqdm as tqdm
import numpy as np
@@ -7,9 +15,9 @@ from torch import Tensor
import math
from pytorch_optimizer import load_optimizer
from torch.optim.lr_scheduler import LambdaLR
-import torch.nn.functional as F
from einops import rearrange
from optgs.evaluation.metrics import compute_rgb_metrics
+from optgs.misc.benchmarker import Benchmarker
from optgs.misc.io import FrequencyScheduler
from optgs.scene_trainer.gaussian_module import GaussiansModule, gaussians2module, module2gaussians
from optgs.model.types import Gaussians
@@ -19,7 +27,7 @@ from optgs.misc.detaching_cpu_list import DetachingCPUList
from optgs.dataset.camera_datasets.camera import get_scene_scale
from optgs.misc.general_utils import get_expon_lr_func
from fused_ssim import fused_ssim
-from optgs.model.decoder.decoder import Decoder, DecoderOutput
+from optgs.model.decoder.decoder import DecoderOutput
@dataclass
@@ -134,21 +142,42 @@ class PostProcessing3DGS:
self.cfg = cfg
self.save_every = save_every
- # Timing
- self.iter_start = torch.cuda.Event(enable_timing=True)
- self.iter_end = torch.cuda.Event(enable_timing=True)
+ # Per-iteration timing + stats. benchmarker.time("iter") records CUDA events with a deferred
+ # sync, so the loop is not stalled every step; the *_log properties read back from it.
+ self.benchmarker = Benchmarker()
self.reset_logs()
def reset_logs(self):
- self.radii_max_log = []
- self.grads_max_log = []
- self.nr_cloned_log = []
- self.nr_splitted_log = []
- self.nr_pruned_log = []
- self.nr_gaussians_log = []
- self.nr_nonzero_grad_log = []
- self.iter_time_log = []
+ # All per-iteration logs (timings + counts) live in self.benchmarker; one clear resets them.
+ self.benchmarker.clear_history()
+
+ @property
+ def iter_time_log(self) -> list[float]:
+ """Total ms per post-processing iteration (from benchmarker.time("iter"))."""
+ return self.benchmarker.execution_times["iter"]
+
+ # Per-iteration stats, stored in self.benchmarker next to the timings (one entry per step).
+ # Counts are ints, so the element type is float | int (see Optimizer for the same pattern).
+ @property
+ def nr_gaussians_log(self) -> list[float | int]:
+ return self.benchmarker.execution_times["gaussians"]
+
+ @property
+ def nr_nonzero_grad_log(self) -> list[float | int]:
+ return self.benchmarker.execution_times["nonzero_grads"]
+
+ @property
+ def nr_cloned_log(self) -> list[float | int]:
+ return self.benchmarker.execution_times["cloned"]
+
+ @property
+ def nr_splitted_log(self) -> list[float | int]:
+ return self.benchmarker.execution_times["splitted"]
+
+ @property
+ def nr_pruned_log(self) -> list[float | int]:
+ return self.benchmarker.execution_times["pruned"]
def _calc_loss(
self, context, output_renderer: DecoderOutput
@@ -316,9 +345,9 @@ class PostProcessing3DGS:
changed = True
print(f"Opacity reset @ iter {step}")
- self.nr_cloned_log.append(nr_cloned)
- self.nr_splitted_log.append(nr_splitted)
- self.nr_pruned_log.append(nr_pruned)
+ self.benchmarker.record("cloned", nr_cloned)
+ self.benchmarker.record("splitted", nr_splitted)
+ self.benchmarker.record("pruned", nr_pruned)
if changed:
# Rebuild GaussiansModule from modified Gaussians
@@ -338,6 +367,10 @@ class PostProcessing3DGS:
visualization_dump=None
) -> OptimizerOutput | None:
+ # This instance is reused across scenes (created once in scene_trainer), so clear the
+ # per-iteration logs (benchmarker: iter_time_log, nr_gaussians_log, ...) for a per-scene reset.
+ self.reset_logs()
+
target_render_list = DetachingCPUList()
context_render_list = DetachingCPUList()
@@ -400,63 +433,60 @@ class PostProcessing3DGS:
pbar_postfix = {}
for i in pbar:
- self.iter_start.record()
-
- with torch.enable_grad():
+ with self.benchmarker.time("iter"):
+ with torch.enable_grad():
- # Log number of gaussians
- self.nr_gaussians_log.append(gaussian_module.means.shape[0])
+ # Log number of gaussians
+ self.benchmarker.record("gaussians", gaussian_module.means.shape[0])
- # reset gradients
- optimizer.zero_grad()
+ # reset gradients
+ optimizer.zero_grad()
- # Sample context views using the same strategy as the optimizer
- iter_context, _ = batchify_fn(batch, "context")
+ # Sample context views using the same strategy as the optimizer
+ iter_context, _ = batchify_fn(batch, "context")
- # [Improvement 4] Render in chunks, accumulate gradients, collect ADC metadata
- meta_for_adc = self._chunked_forward_backward(
- gaussian_module, iter_context, decoder, render_res, adc_state
- )
-
- # step
- optimizer.step()
-
- # update scheduler
- if scheduler is not None:
- scheduler.step()
-
- # [Improvement 3] ADC: update state and apply densification/pruning
- if adc_state is not None and meta_for_adc is not None:
- from optgs.scene_trainer.adc.vanilla import update_vanilla_strategy_state
-
- v_chunk = iter_context["image"].shape[1]
- update_vanilla_strategy_state(
- adc_state,
- radii_2d=meta_for_adc["radii"],
- means2d_grads=meta_for_adc["means_2d_grads"],
- visibility_mask=meta_for_adc["visibility_filter"],
- v=v_chunk,
- w=w,
- h=h,
- )
-
- gaussian_module, adc_changed = self._apply_adc(i, gaussian_module, adc_state, device)
- if adc_changed:
- # Rebuild optimizer and scheduler after ADC changed Gaussian count
- optimizer = self.get_optimizer(gaussian_module, scene_scale)
- scheduler = self.get_scheduler(
- optimizer, scene_scale=scene_scale, prior_steps=self.cfg.prior_steps
+ # [Improvement 4] Render in chunks, accumulate gradients, collect ADC metadata
+ meta_for_adc = self._chunked_forward_backward(
+ gaussian_module, iter_context, decoder, render_res, adc_state
)
- # Fast-forward scheduler to current step
- for _ in range(i + 1):
- scheduler.step() if scheduler is not None else None
- # Timing
- self.iter_end.record()
- torch.cuda.synchronize()
+ # step
+ optimizer.step()
+
+ # update scheduler
+ if scheduler is not None:
+ scheduler.step()
+
+ # [Improvement 3] ADC: update state and apply densification/pruning
+ if adc_state is not None and meta_for_adc is not None:
+ from optgs.scene_trainer.adc.vanilla import update_vanilla_strategy_state
+
+ v_chunk = iter_context["image"].shape[1]
+ update_vanilla_strategy_state(
+ adc_state,
+ radii_2d=meta_for_adc["radii"],
+ means2d_grads=meta_for_adc["means_2d_grads"],
+ visibility_mask=meta_for_adc["visibility_filter"],
+ v=v_chunk,
+ w=w,
+ h=h,
+ )
- elapsed_time = self.iter_start.elapsed_time(self.iter_end)
- self.iter_time_log.append(elapsed_time)
+ gaussian_module, adc_changed = self._apply_adc(i, gaussian_module, adc_state, device)
+ if adc_changed:
+ # Rebuild optimizer and scheduler after ADC changed the Gaussian count.
+ # Caveat: rebuilding drops the Adam moments (exp_avg/exp_avg_sq) for every
+ # Gaussian, so they restart from zero after each densification. Reference 3DGS
+ # instead resizes the optimizer state in place to keep the moments across
+ # clone/split/prune (as the in-loop optimizers do via AdamInputSmoothing); the
+ # nn.Parameter + torch.optim path here cannot do that without rebuilding.
+ optimizer = self.get_optimizer(gaussian_module, scene_scale)
+ scheduler = self.get_scheduler(
+ optimizer, scene_scale=scene_scale, prior_steps=self.cfg.prior_steps
+ )
+ # Fast-forward scheduler to current step
+ for _ in range(i + 1):
+ scheduler.step() if scheduler is not None else None
if self.save_every(i + 1, tag="context"):
with torch.no_grad():
@@ -524,7 +554,7 @@ class PostProcessing3DGS:
def get_optimizer(self, gaussians: GaussiansModule, scene_scale: float):
- # TODO Naama: support different batch sizes
+ # NOTE: post-processing runs one scene at a time (batch size 1)
batch_size: int = 1
# Build params list (name, parameter, lr)
diff --git a/optgs/scene_trainer/scene_trainer.py b/optgs/scene_trainer/scene_trainer.py
index 83821a01b5a8104ad274d376e6519c5fc4cf2703..c4dd2a11ef0ef4dd05674f745b8f9a022291b6e9 100644
--- a/optgs/scene_trainer/scene_trainer.py
+++ b/optgs/scene_trainer/scene_trainer.py
@@ -1,7 +1,16 @@
+"""Inner per-scene loop of the two-level pipeline.
+
+`SceneTrainer` composes the three pipeline stages for a single scene: an Initializer produces the
+starting Gaussians from the input views, an Optimizer iteratively refines them, and a Decoder renders
+novel views. Each stage is chosen by name from its registry (see the package `__init__.py` files). The
+outer meta-loop that drives many scenes and computes losses lives in `meta_trainer.py`.
+"""
+
import math
import random
from pathlib import Path
from typing import Optional, Mapping, Any
+from warnings import warn
import torch
from einops import rearrange
@@ -14,9 +23,10 @@ from optgs.dataset.data_types import BatchedExample
from optgs.dataset.view_sampler.view_sampler_bounded_v2 import farthest_point_sample
from optgs.loss.loss_monodepth import get_monodepth_model
from optgs.misc.benchmarker import Benchmarker
+from optgs.misc.general_utils import SkipBatchException
from optgs.misc.io import FrequencyScheduler
from optgs.misc.step_tracker import StepTracker
-from optgs.model.decoder.decoder import Decoder
+from optgs.model.decoder.decoder import Decoder, DepthRenderingMode
from optgs.model.types import Gaussians
from optgs.paths import DEBUG
from optgs.scene_trainer.initializer import get_scene_initializer
@@ -24,8 +34,8 @@ from optgs.scene_trainer.initializer.initializer import InitializerOutput
from optgs.scene_trainer.optimizer import get_scene_optimizer
from optgs.scene_trainer.optimizer.optimizer import OptimizerInput, Optimizer, OptimizerOutput, OptimizerPreviousOutput
from optgs.scene_trainer.postprocessing import PostProcessing3DGS
-from optgs.scene_trainer.scene_trainer_cfg import SceneTrainerCfg, TestCfg, TrainCfg
-from optgs.scripts.dev.debugging_optimizer import debugging_convergence, debugging_invisible_gaussians
+from optgs.scene_trainer.scene_trainer_cfg import SceneTrainerCfg
+from optgs.meta_trainer.meta_trainer_cfg import TestCfg, TrainCfg
class SceneTrainer(nn.Module):
@@ -43,6 +53,7 @@ class SceneTrainer(nn.Module):
scene_trainer_cfg: SceneTrainerCfg,
decoder: Decoder,
step_tracker: StepTracker | None,
+ benchmarker: Benchmarker,
eval_data_cfg: Optional[DatasetCfg | None] = None,
) -> None:
super().__init__()
@@ -71,7 +82,8 @@ class SceneTrainer(nn.Module):
self.decoder = decoder
- self.benchmarker = Benchmarker()
+ # Shared with the MetaTrainer so init/decoder timings land in one place.
+ self.benchmarker = benchmarker
if self.train_cfg.monodepth_loss_weight > 0:
self.pretrained_monodepth = get_monodepth_model()
@@ -135,6 +147,7 @@ class SceneTrainer(nn.Module):
debug_dict=None,
num_update_steps=None,
disable_tqdm=False,
+ allow_partial_on_skip=False,
**kwargs,
) -> OptimizerOutput:
"""
@@ -155,6 +168,10 @@ class SceneTrainer(nn.Module):
Should be 0 when starting a new scene.
debug_dict: Optional[dict], a dictionary to store debug information.
num_update_steps: Optional[int], number of update steps to perform. If None, use default from config.
+ allow_partial_on_skip: if True, a SkipBatchException or CUDA OOM raised during an update step
+ stops the loop early and returns the steps completed so far (with any partial appends from
+ the failed step rolled back) instead of propagating. Used by the test path so a scene that
+ fails mid-optimization still reports metrics for its successful steps.
Returns:
OptimizerOutput: The output from the optimizer containing the optimized Gaussians and renderings for
intermediate optimization steps.
@@ -170,7 +187,6 @@ class SceneTrainer(nn.Module):
target=batch["target"],
prev_output=prev_output,
renderer=self.decoder,
- num_refine=num_update_steps,
iter_batch_size=self.scene_trainer_cfg.iter_batch_size, # For rendering in batches
debug_dict=debug_dict,
)
@@ -189,10 +205,9 @@ class SceneTrainer(nn.Module):
# Insert the initialization into position 0 of the output lists so downstream
# consumers (evaluation, plotting, replay buffer) can treat init uniformly with
- # the optimizer steps. No-op when there's no init render to attach (train path
- # before init renders are wired through; replay buffer continuations).
- if isinstance(prev_output, InitializerOutput):
- self._insert_init_into_output(optimizer_output, prev_output)
+ # the optimizer steps. Handles both InitializerOutput and OptimizerPreviousOutput,
+ # and skips insertion during the rollout that saves no per-step outputs.
+ self._insert_init_into_output(optimizer_output, prev_output)
# SH degree scheduling (inspired by gsplat simple_trainer):
# sh_degree_to_use = min(step // sh_degree_interval, max_sh_degree)
@@ -200,11 +215,27 @@ class SceneTrainer(nn.Module):
if sh_degree_interval > 0:
max_sh_degree = int(math.sqrt(optimizer_input.prev_output.gaussians.harmonics.shape[-1])) - 1
+ # Test path only (allow_partial_on_skip): references to the output lists that must stay
+ # mutually aligned for the save/score paths. Built once -- the list objects don't change
+ # across steps (the optimizer mutates them in place); only their lengths grow. Off
+ # (training/validation): skipped entirely and the step loop runs exactly as before.
+ if allow_partial_on_skip:
+ aligned_lists = [
+ optimizer_output.target_render_list, optimizer_output.context_render_list,
+ optimizer_output.gaussian_list,
+ optimizer_output.target_index_list, optimizer_output.context_index_list,
+ ]
+
# Loop over update steps
for step in tqdm(range(num_update_steps),
disable=(self.training or num_update_steps < 20 or DEBUG) or disable_tqdm,
total=num_update_steps):
+ # Snapshot the aligned lists' lengths so a step that fails partway can be rolled back
+ # to the steps that fully completed (lengths grow per step, so this is per-step).
+ if allow_partial_on_skip:
+ pre_step_lens = [len(lst) for lst in aligned_lists]
+
# Sample minibatch of context/target views and move to device
optimizer_input.context, batch_idx = self.batchify_views(batch, "context", self.device)
if batch_idx is not None:
@@ -220,20 +251,42 @@ class SceneTrainer(nn.Module):
# Single optimization step
# Optimizer output is updated in place, but we return it for clarity
- optimizer_output = self.optimizer(
- step,
- optimizer_input,
- optimizer_output,
- full_context=batch["context"],
- full_target=batch["target"],
- **step_kwargs
- )
+ try:
+ optimizer_output = self.optimizer(
+ step,
+ optimizer_input,
+ optimizer_output,
+ full_context=batch["context"],
+ full_target=batch["target"],
+ **step_kwargs
+ )
+ # torch.cuda.OutOfMemoryError is an alias of torch.OutOfMemoryError; one covers both.
+ # Kept specific to the two expected failures so a generic RuntimeError (a real bug)
+ # still propagates and drops the scene rather than being half-recovered.
+ except (SkipBatchException, torch.OutOfMemoryError) as e:
+ if not allow_partial_on_skip:
+ raise
+ torch.cuda.empty_cache()
+ # No step completed -> nothing to score; re-raise so the caller drops the scene.
+ if optimizer_output.t == curr_iter:
+ raise
+ # >=1 step completed: roll back the failed step's partial appends so the output lists
+ # stay aligned to the fully-completed steps, then stop early and score those. (At test
+ # time SkipBatch fires before any append, so the rollback only matters for an OOM
+ # mid-step.) The per-iteration module logs may keep one extra (unread) entry; downstream
+ # indexes them by completed-step number, so that is harmless.
+ warn(f"Optimization stopped early at step {step}/{num_update_steps} ({type(e).__name__}: {e}); "
+ f"scoring the {step} completed step(s).")
+ for lst, n in zip(aligned_lists, pre_step_lens):
+ del lst[n:]
+ break
optimizer_output.t += 1
# Sync GPU before reading scene_start elapsed time (events were recorded before the loop).
torch.cuda.synchronize()
- self.optimizer.scene_start_ms = self.optimizer.scene_start_event_start.elapsed_time(
- self.optimizer.scene_start_event_end
+ self.optimizer.benchmarker.record(
+ "scene_start",
+ self.optimizer.scene_start_event_start.elapsed_time(self.optimizer.scene_start_event_end),
)
self.optimizer.on_scene_end()
@@ -485,6 +538,7 @@ class SceneTrainer(nn.Module):
depths=all_pred_depths)
def debugging(self, optimizer_output, output_path: Path, scene_name: str): # TODO (release): remove in public code
+ from optgs.scripts.dev.debugging_optimizer import debugging_convergence, debugging_invisible_gaussians
# Debugging reprojection errors
# if 'reprojection_error' in visualization_dump:
@@ -544,68 +598,142 @@ class SceneTrainer(nn.Module):
def _insert_init_into_output(
self,
optimizer_output: OptimizerOutput,
- init_output: InitializerOutput,
+ prev_output: InitializerOutput | OptimizerPreviousOutput,
) -> None:
- """Prepend the init render/gaussians at position 0 of the optimizer_output lists.
-
- No-op if no init render exists yet (train path before init renders are wired
- through). Otherwise inserts gaussians plus whichever of context/target renders
- are populated. detach_and_cpu mirrors the per-step append policy.
+ """Insert the init/resumed render + gaussians at position 0 of the optimizer_output lists.
+
+ Accepts either InitializerOutput (a fresh scene) or OptimizerPreviousOutput (a ckpt-buffer
+ resume); both expose `gaussians`, `context_render`, `target_render`. Position 0 holds the
+ init, so the per-step lists line up by position (gaussian_list[k] and the render lists at
+ index k refer to the same step). The init/resume views are rendered before optimizing, so a
+ missing render here means they were not rendered and we raise.
+
+ The exception is the rollout that promotes a sample to the ckpt buffer: there the optimizer
+ runs in eval mode with all save tags off and saves no per-step renders or gaussians, so
+ position 0 has nothing to align with. In that case there is nothing to insert (and
+ prev_output carries no render), so this returns early.
+
+ Inserts the gaussians together with whichever of context/target renders are present.
+ detach_and_cpu matches the per-step append policy.
+ View indices (set when rendering a subset of views in training) are inserted into the index
+ lists too, so index_list[0] = init and index_list[1..N] = steps, aligned with the render lists.
"""
- if init_output.context_render is None and init_output.target_render is None:
+ # Per-step renders are saved every step in training, and at the save_every iterations at test
+ # time; the rollout turns both off, so no per-step entry exists for position 0 to sit before.
+ saves_per_step = self.training or any(
+ self.optimizer.save_every.enabled_tags[tag] for tag in ("target", "context")
+ )
+ if not saves_per_step:
return
- optimizer_output.gaussian_list.insert(0, init_output.gaussians)
- if init_output.context_render is not None:
- optimizer_output.context_render_list.insert(
- 0, init_output.context_render, detach_and_cpu=not self.training
- )
- if init_output.target_render is not None:
- optimizer_output.target_render_list.insert(
- 0, init_output.target_render, detach_and_cpu=not self.training
+ if prev_output.context_render is None and prev_output.target_render is None:
+ raise ValueError(
+ "prev_output has no render attached, so the init cannot be placed at position 0 of "
+ "the optimizer output. Render the init/resume views before optimizing."
)
+ # Insert the init gaussians together with the init render so that position k of gaussian_list
+ # and of the render lists always refer to the same step.
+ optimizer_output.gaussian_list.insert(0, prev_output.gaussians)
+ for tag in ("context", "target"):
+ render = prev_output.get_render(tag)
+ if render is None:
+ continue
+ optimizer_output.get_render_list(tag).insert(0, render, detach_and_cpu=not self.training)
+ index = prev_output.get_render_index(tag)
+ if index is not None:
+ optimizer_output.get_index_list(tag).insert(0, index)
+
def init_gaussians_and_render(
self, batch, visualization_dump,
render_context: bool, render_target: bool, grad_enabled: bool,
+ depth_mode: DepthRenderingMode | None = None,
+ benchmark_decoder: bool = False,
+ to_cpu: bool | None = None,
**kwargs,
) -> InitializerOutput:
"""Run the initializer and optionally render its output to context/target views.
- Used in both training (grad_enabled=True, outputs stay on GPU so the init-loss term
- can backward through them) and test (grad_enabled=False, outputs moved to CPU to save
- memory since they're only consumed for evaluation/saving).
- """
+ Single entry point for "initialize then render" across training and the test/validation
+ eval pipelines.
+
+ By default renders go off the GPU when grads are disabled (test/eval, to save memory) and
+ stay on the GPU with grads enabled (training, so the init-loss term can backward through
+ them). Pass `to_cpu` to override -- the eval pipeline keeps the render on the GPU even with
+ grads disabled because it feeds straight into the optimizer.
- # run initializer
+ The "initializer" timer always records (the init model is timed in every path). The init
+ render's "decoder" timer is gated by `benchmark_decoder`: only the eval pipeline reports it,
+ so the main test path's "decoder" metric stays the optimizer-phase decode time.
+ """
with self.benchmarker.time("initializer"):
init_output: InitializerOutput = self.get_init_gaussians(batch, is_training=grad_enabled, **kwargs)
+ self.render_init_views(batch, init_output, render_context, render_target, grad_enabled,
+ depth_mode, benchmark_decoder=benchmark_decoder, to_cpu=to_cpu)
+ return init_output
+
+ def render_init_views(
+ self,
+ batch,
+ init_output: InitializerOutput | OptimizerPreviousOutput,
+ render_context: bool,
+ render_target: bool,
+ grad_enabled: bool,
+ depth_mode: DepthRenderingMode | None = None,
+ benchmark_decoder: bool = False,
+ to_cpu: bool | None = None,
+ ) -> None:
+ """Render context/target views into any output that exposes get_render/set_render.
+
+ Mutates `init_output` in place; returns None. Skips a view if it's already rendered.
+ Accepts either InitializerOutput (dataloader path) or OptimizerPreviousOutput
+ (replay-buffer path, where renders are needed to populate the splice but no
+ initializer model runs).
+ """
# to_cpu freezes the render off the GPU for evaluation/saving; with grads enabled we
- # keep it on GPU so the init-loss term can backward through it.
- to_cpu = not grad_enabled
+ # keep it on GPU so the init-loss term can backward through it. The eval pipeline overrides
+ # it (to_cpu=False) since the render feeds straight into the optimizer.
+ if to_cpu is None:
+ to_cpu = not grad_enabled
+
+ # In optimizer-training mode (self.training=True, grad_enabled=False), render only a
+ # random subset of views — matching what the optimizer does per step — to avoid the
+ # cost of rendering all views (which may not fit in one batch).
+ use_subset = self.training and not grad_enabled
with torch.set_grad_enabled(grad_enabled):
for input_str, should_render in (
("context", render_context),
("target", render_target),
):
- attr = f"{input_str}_render"
- if not should_render or getattr(init_output, attr) is not None:
+ if not should_render or init_output.get_render(input_str) is not None:
continue
- views = batch[input_str]
- h, w = views["image"].shape[-2:]
- rendered = self.decoder.forward_batch(
- init_output.gaussians.to(batch["target"]["image"].device),
- batch, (h, w),
- input_str=input_str,
- to_cpu=to_cpu,
- iter_batch_size=self.scene_trainer_cfg.iter_batch_size,
- )
- # TODO Naama: should we make it a render list as in OptimizerOutput and then the flow will be more unified?
- setattr(init_output, attr, rendered)
-
- return init_output
+ if use_subset:
+ subset_views, index = self.batchify_views(batch, input_str, self.device)
+ rendered = self.decoder.forward_batch_subset(
+ init_output.gaussians.to(self.device),
+ subset_views,
+ iter_batch_size=self.scene_trainer_cfg.iter_batch_size,
+ )
+ init_output.set_render_index(input_str, index)
+ else:
+ views = batch[input_str]
+ h, w = views["image"].shape[-2:]
+ n_views = views["image"].shape[1]
+ # The initializer produces a single render per view set, so it is stored as one
+ # DecoderOutput per tag; the optimizer keeps one render per step in a list. The
+ # init render becomes position 0 of that list (see _insert_init_into_output).
+ with self.benchmarker.time("decoder", num_calls=n_views, disable=not benchmark_decoder):
+ rendered = self.decoder.forward_batch(
+ init_output.gaussians.to(batch["target"]["image"].device),
+ batch, (h, w),
+ input_str=input_str,
+ to_cpu=to_cpu,
+ iter_batch_size=self.scene_trainer_cfg.iter_batch_size,
+ depth_mode=depth_mode,
+ )
+ init_output.set_render(input_str, rendered)
def test_postprocess_gaussians(self, batch, gaussians, visualization_dump) -> OptimizerOutput | None:
"""Run optional post-processing on the final Gaussians. Returns None if disabled."""
diff --git a/optgs/scene_trainer/scene_trainer_cfg.py b/optgs/scene_trainer/scene_trainer_cfg.py
index f333152a3fe795fb17c51eba8edc5d3fc4b6c772..2069c8665f4e4ce7d9983fa281a72e5adef866a4 100644
--- a/optgs/scene_trainer/scene_trainer_cfg.py
+++ b/optgs/scene_trainer/scene_trainer_cfg.py
@@ -1,13 +1,9 @@
from dataclasses import dataclass
-from pathlib import Path
from typing import Literal
-from optgs.meta_trainer.replay_buffer import ReplayBufferCfg
from .initializer import InitializerCfg
from .optimizer import SceneOptimizerCfg
-from .postprocessing import PostProcessCfg
from ..model.decoder import DecoderCfg
-from ..model.decoder.decoder import DepthRenderingMode
@dataclass
@@ -33,125 +29,3 @@ class SceneTrainerCfg:
def __post_init__(self):
if self.scene_optimizer is not None:
self.scene_optimizer.update(self.scene_initializer)
-
-
-# TODO Naama, probably need to move into meta_trainer cfg file
-
-@dataclass
-class MetaOptimizerCfg:
- lr: float
- warm_up_steps: int
- lr_monodepth: float
- lr_depth: float
- weight_decay: float
- warm_up_ratio: float
- adamw_8bit: bool
-
-
-@dataclass
-class TestCfg:
- output_path: Path | None
- compute_scores: bool
- compute_scores_metrics: list[str] | None
- metrics_batch_size: int
- eval_initialization: bool
- save_render_image: bool
- save_render_image_last_only: bool
- save_gt_image: bool
- save_render_depth: bool
- save_gt_depth: bool
- save_error_image: bool
- save_video: bool
- save_video_fixed_view: bool
- save_video_fixed_view_index: int
- save_video_fixed_view_duplicate: int
- save_video_fixed_iteration: bool
- save_video_fixed_iteration_indices: list | None
- save_video_fixed_iteration_render_fixed_view: bool
- save_video_combined: bool
- save_video_combined_iterations: list | None
- save_video_combined_fixed_iteration_length: int
- eval_time_skip_steps: int
- save_gaussian: bool
- save_poses: bool
- save_cameras_json: bool
- save_cameras_npz: bool
- save_point_cloud: bool
- render_chunk_size: int | None
- stablize_camera: bool
- stab_camera_kernel: int
- eval_context_views: bool
- inference_window_size: int | None
- profile_model: bool
- save_colmap_train_test_views: bool
- ori_colmap_data_path: str | None
- postprocessing: PostProcessCfg | None
- save_at_iters: list[int] | None
- save_every_freq: list[int] | None
- save_every_steps: list[int] | None
- skip_if_outputs_exist: bool
- scenes_filter: list[str] | None
-
- experimental_add_noise_to_images: bool
- experimental_add_noise_to_images_std: float | int | None
-
-
-# TODO Naama split into scene and meta trainer cfgs
-@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_deterministic: bool
- eval_time_skip_steps: int
- eval_save_model: bool
- l1_loss: bool
- intermediate_loss_weight: float
- no_viz_video: bool
- eval_depth: bool
- train_ignore_large_loss: float
- no_log_projections: bool
-
- depth_loss_weight: float
- log_depth_loss: bool
- depth_smooth_loss_weight: float
- depth_teacher_loss_weight: float
- viz_depth_teacher: bool
- depth_smooth_loss_nonorm: bool
- depth_smooth_loss_weight_nvs: float # for novel views
- monodepth_loss_weight: float # for monocular depth loss
-
- eval_render_depth: bool
- render_depth_loss_weight: float
- viz_render_depth: bool
- viz_depth_separate: bool
-
- use_gt_depth_range: bool
- depth_range_from_disparity: bool
- max_disparity: float
- min_disparity: float
-
- no_log_video: bool
-
- # when doing refinement, supervise input view or not since we also render input views
- loss_on_target_views: bool
- loss_on_target_views_num: int
- loss_on_input_views: bool
- loss_on_input_views_num: int
-
- # half res lpips loss to save memory
- half_res_lpips_loss: bool
-
- # local window training
- train_window_size: int | None
-
- # Replay buffer
- use_replay_buffer: bool
- replay_buffer_cfg: ReplayBufferCfg | None
-
- # L2 weight decay regularization on Gaussian properties (meta-loss)
- scale_l2_loss_weight: float
- sh_l2_loss_weight: float
- opacity_l2_loss_weight: float
diff --git a/optgs/scripts/analyze_metrics.sh b/optgs/scripts/analyze_metrics.sh
old mode 100644
new mode 100755
diff --git a/optgs/scripts/convert_dl3dv_test.py b/optgs/scripts/convert_dl3dv_test.py
index 54071abbb2310eae250f5591302d118f67b41ded..1de3d4cabc270f668bc973ceb8c40c298c41748b 100644
--- a/optgs/scripts/convert_dl3dv_test.py
+++ b/optgs/scripts/convert_dl3dv_test.py
@@ -42,9 +42,10 @@ def legal_check_for_all_scenes(root_dir, target_shape, img_dir):
# images_4 folder has resolution 270x480
for sub_folder in tqdm(sub_folders, desc="checking scenes..."):
- # img_dir = os.path.join(sub_folder, "images_8") # 270x480
- img_dir = os.path.join(sub_folder, img_dir) # 540x960
- if not is_image_shape_matched(Path(img_dir), target_shape):
+ # Use a per-scene variable; reusing img_dir here would clobber the parameter and
+ # build a wrong path for every scene after the first.
+ scene_img_dir = os.path.join(sub_folder, img_dir) # e.g. /nerfstudio/images_4
+ if not is_image_shape_matched(Path(scene_img_dir), target_shape):
print(f"image shape does not match for {sub_folder}")
continue
pose_file = os.path.join(sub_folder, "transforms.json")
diff --git a/optgs/scripts/convert_dl3dv_utils.py b/optgs/scripts/convert_dl3dv_utils.py
index 7186fbf7b1f807ee1345535787f9bfe622ddb58c..f4a37e82c050f370fbb5c3add684a59708943083 100644
--- a/optgs/scripts/convert_dl3dv_utils.py
+++ b/optgs/scripts/convert_dl3dv_utils.py
@@ -7,7 +7,7 @@ from typing import TypedDict
import numpy as np
import torch
-from PIL.Image import Image
+from PIL import Image
from jaxtyping import UInt8, Int, Float
from torch import Tensor
diff --git a/optgs/scripts/dev/debugging_optimizer.py b/optgs/scripts/dev/debugging_optimizer.py
index 10b79f5737addc62dbafd278b40a2423b6a605e3..fdd2eb4efe02211bc8c397b1935a6172340f4e6b 100644
--- a/optgs/scripts/dev/debugging_optimizer.py
+++ b/optgs/scripts/dev/debugging_optimizer.py
@@ -416,9 +416,7 @@ def make_gaussians_dashboard_video_with_ani(data_groups, psnrs, iters, out_path=
import numpy as np
import torch
-from torch import Tensor
from pathlib import Path
-import matplotlib.pyplot as plt
def calc_hist(data, max_percentile=99.9, min_percentile=0.1, density=False):
diff --git a/optgs/scripts/dev/export_resplat_init.py b/optgs/scripts/dev/export_resplat_init.py
new file mode 100644
index 0000000000000000000000000000000000000000..6979abf5f3a68e7acd943f7b088ac3a5fe9f6498
--- /dev/null
+++ b/optgs/scripts/dev/export_resplat_init.py
@@ -0,0 +1,78 @@
+"""Export the trained ReSplat initializer for public release.
+
+Takes the in-repo training run (Lightning checkpoint + its v0 config.yaml) and
+produces a clean, release-ready pair under a new directory:
+
+ /config.yaml migrated to the current config version
+ /checkpoints/ slimmed to just the initializer weights
+
+The config is run through ``config_migrate.migrate`` (v0 -> current) and its
+``output_dir`` is repointed at the release directory. The checkpoint keeps only
+``state_dict`` (the ``initializer.*`` weights), dropping the Lightning training
+state (optimizer moments, loops, callbacks, lr schedulers) that inference never
+reads. The resulting layout matches what ``_find_config_for_checkpoint`` expects,
+so it loads directly as ``pretrained_initializer`` (locally or via ``hf://``).
+
+Usage:
+ python -m optgs.scripts.dev.export_resplat_init \
+ --src checkpoints/optgs/unified-dl3dv-8views/init \
+ --out checkpoints/learn2splat/resplat_init
+"""
+
+import argparse
+from pathlib import Path
+
+import torch
+from omegaconf import OmegaConf
+
+from optgs.config_migrate import migrate
+from optgs.misc.checkpointing import _rename_optimizer_attrs, find_latest_ckpt
+from optgs.misc.io import cyan, read_omega_cfg
+
+
+def export_config(src_cfg: Path, out_cfg: Path, out_dir: Path) -> None:
+ cfg = read_omega_cfg(src_cfg)
+ cfg = migrate(cfg)
+ # Repoint output_dir at the release directory as a plain string (the source
+ # carried a pickled PosixPath, which serializes as an ugly !!python tag).
+ cfg.output_dir = str(out_dir)
+ out_cfg.parent.mkdir(parents=True, exist_ok=True)
+ OmegaConf.save(cfg, out_cfg)
+ print(cyan(f"Wrote migrated config -> {out_cfg}"))
+
+
+def export_checkpoint(src_ckpt: Path, out_ckpt: Path) -> None:
+ ckpt = torch.load(src_ckpt, map_location="cpu")
+ state_dict = ckpt["state_dict"] if "state_dict" in ckpt else ckpt
+ # Keep only the initializer weights and apply the optimizer-attr renames
+ # (a no-op here, since an init-only checkpoint has no optimizer keys).
+ state_dict = {k: v for k, v in state_dict.items() if k.startswith("initializer.")}
+ state_dict = _rename_optimizer_attrs(state_dict)
+ out_ckpt.parent.mkdir(parents=True, exist_ok=True)
+ torch.save({"state_dict": state_dict}, out_ckpt)
+ src_mb = src_ckpt.stat().st_size / 1e6
+ out_mb = out_ckpt.stat().st_size / 1e6
+ print(cyan(f"Wrote slimmed checkpoint -> {out_ckpt} "
+ f"({len(state_dict)} keys, {src_mb:.0f}MB -> {out_mb:.0f}MB)"))
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--src", type=Path, required=True,
+ help="Source run dir containing config.yaml and checkpoints/")
+ parser.add_argument("--out", type=Path, required=True,
+ help="Release dir to write config.yaml and checkpoints/ into")
+ parser.add_argument("--ckpt-name", type=str, default=None,
+ help="Checkpoint filename under src/checkpoints (default: latest by step)")
+ args = parser.parse_args()
+
+ src_ckpt = (args.src / "checkpoints" / args.ckpt_name) if args.ckpt_name \
+ else find_latest_ckpt(args.src / "checkpoints")
+ out_ckpt = args.out / "checkpoints" / src_ckpt.name
+
+ export_config(args.src / "config.yaml", args.out / "config.yaml", args.out)
+ export_checkpoint(src_ckpt, out_ckpt)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/optgs/scripts/dev/upload_resplat_init.sh b/optgs/scripts/dev/upload_resplat_init.sh
new file mode 100644
index 0000000000000000000000000000000000000000..e46810d48a5aed2012a6a661269b9f50d418b815
--- /dev/null
+++ b/optgs/scripts/dev/upload_resplat_init.sh
@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+# Upload the released ReSplat initializer (config + slimmed checkpoint) to the
+# Hugging Face model repo, into the resplat_init/ subfolder.
+#
+# Pairs with export_resplat_init.py: that script stages the release-ready pair
+# under $LOCAL_DIR; this script pushes that folder to the Hub at $REMOTE_DIR so
+# it loads as pretrained_initializer=hf://$REPO/$REMOTE_DIR/checkpoints/.
+#
+# Rerun after retraining: re-stage with export_resplat_init.py, then run this
+# again. Uploads are git commits that overwrite the same paths, so no deletion
+# is needed between versions.
+#
+# Requires a WRITE token with access to the autonomousvision org:
+# conda run -n resplat hf auth login # paste a Write token
+#
+# Usage:
+# bash optgs/scripts/dev/upload_resplat_init.sh
+set -euo pipefail
+
+REPO="autonomousvision/learn2splat"
+LOCAL_DIR="checkpoints/learn2splat/resplat_init"
+REMOTE_DIR="resplat_init"
+
+# 1. Upload the staged folder. Mirrors $LOCAL_DIR onto $REPO at $REMOTE_DIR:
+# $REMOTE_DIR/config.yaml
+# $REMOTE_DIR/checkpoints/
+conda run -n resplat hf upload "$REPO" "$LOCAL_DIR" "$REMOTE_DIR" \
+ --repo-type model \
+ --commit-message "Add ReSplat initializer (resplat_init): ckpt + migrated config"
+
+# 2. List the resulting repo files to confirm the layout.
+conda run -n resplat python -c "from huggingface_hub import HfApi; print('\n'.join(HfApi().list_repo_files('$REPO')))"
+
+# --- One-time rename note -----------------------------------------------------
+# The initializer originally lived under init/. To drop that stale folder
+# (including subdirs) after the first resplat_init/ upload, run once:
+# conda run -n resplat hf repo-files "$REPO" delete "init/**" --repo-type model \
+# --commit-message "Remove misnamed init/ folder (renamed to resplat_init)"
\ No newline at end of file
diff --git a/optgs/scripts/dl3dv_benchmark_hf_download.py b/optgs/scripts/dl3dv_benchmark_hf_download.py
deleted file mode 100644
index 865f18ba96f53d208fa7f70a9a759148b01b3022..0000000000000000000000000000000000000000
--- a/optgs/scripts/dl3dv_benchmark_hf_download.py
+++ /dev/null
@@ -1,220 +0,0 @@
-""" This script is used to download the DL3DV benchmark from the huggingface repo.
-
- The benchmark is composed of 140 different scenes covering different scene complexities (reflection, transparency, indoor/outdoor, etc.)
-
- The whole benchmark is very large: 2.1 TB. So we provide this script to download the subset of the dataset based on common needs.
-
-
- - [x] Full benchmark downloading
- Full download can directly be done by git clone (w. lfs installed).
-
- - [x] scene downloading based on scene hash code
-
- Option:
- - [x] images_4 (960 x 540 resolution) level dataset (approx 50G)
-
-"""
-
-import argparse
-import os
-import pickle
-import shutil
-import traceback
-from os.path import join
-
-import pandas as pd
-from huggingface_hub import HfApi
-from tqdm import tqdm
-
-from optgs.misc.io import CustomPath
-
-api = HfApi()
-repo_root = 'DL3DV/DL3DV-10K-Benchmark'
-
-
-def hf_download_path(repo_path: str, odir: str, max_try: int = 5):
- """ hf api is not reliable, retry when failed with max tries
-
- :param repo_path: The path of the repo to download
- :param odir: output path
- """
- rel_path = os.path.relpath(repo_path, repo_root)
-
- counter = 0
- while True:
- if counter >= max_try:
- print("ERROR: Download {} failed.".format(repo_path))
- return False
-
- try:
- api.hf_hub_download(repo_id=repo_root, filename=rel_path, repo_type='dataset', local_dir=odir,
- cache_dir=join(odir, '.cache'))
- return True
-
- except BaseException as e:
- traceback.print_exc()
- counter += 1
- print(f'Retry {counter}')
-
-
-def clean_huggingface_cache(cache_dir: str):
- """ Huggingface cache may take too much space, we clean the cache to save space if necessary
-
- :param cache_dir: the current cache directory
- """
- # Current huggingface hub does not provide good practice to clean the space.
- # We mannually clean the cache directory if necessary.
- try:
- shutil.rmtree(join(cache_dir, 'datasets--DL3DV--DL3DV-10K-Benchmark'))
- except Exception as e:
- pass
-
-
-def download_by_hash(filepath_dict: dict, odir: str, hash: str, only_level4: bool, only_sfm: bool):
- """ Given a hash, download the relevant data from the huggingface repo
-
- :param filepath_dict: the cache dict that stores all the file relative paths
- :param odir: the download directory
- :param hash: the hash code for the scene
- :param only_level4: the images_4 resolution level, if true, only the images_4 resolution level will be downloaded
- """
- all_files = filepath_dict[hash]
- download_files = [join(repo_root, f) for f in all_files]
-
- if only_level4: # only download images_4 level data
- download_files = []
- for f in all_files:
- subdirname = os.path.basename(os.path.dirname(f))
-
- if 'images' in f and subdirname != 'images_4' or 'input' in f:
- continue
-
- download_files.append(join(repo_root, f))
-
- if only_sfm: # only download nerfstudio colmap data
- download_files = list(filter(lambda x:
- 'nerfstudio' in x and
- ('.json' in x or '.bin' in x),
- all_files))
- download_files = [join(repo_root, f) for f in download_files]
-
- for f in download_files:
- if hf_download_path(f, odir) == False:
- return False
-
- if only_sfm:
- # Move files to the scene root directory
- # /nerfstudio/transforms.json --> /transforms.json
- # /nerfstudio/colmap/sparse --> /sparse
-
- # transforms.json
- src_transforms_path = CustomPath(odir) / hash / 'nerfstudio' / 'transforms.json'
- dst_transforms_path = CustomPath(odir) / hash / 'transforms.json'
- shutil.move(src_transforms_path, dst_transforms_path)
-
- # sparse
- src_sparse_path = CustomPath(odir) / hash / 'nerfstudio' / 'colmap' / 'sparse'
- dst_sparse_path = CustomPath(odir) / hash / 'sparse'
- try:
- shutil.move(src_sparse_path, dst_sparse_path)
- except Exception as e:
- print(f'Warning: {hash} sparse already exists. Overwriting.')
- shutil.rmtree(dst_sparse_path)
- shutil.move(src_sparse_path, dst_sparse_path)
-
- # remove empty nerfstudio directory
- nerfstudio_dir = CustomPath(odir) / hash / 'nerfstudio'
- # check if the colmap directory is empty
- if len(list(nerfstudio_dir.iterdir())) == 1 and len(list((nerfstudio_dir / 'colmap').iterdir())) == 0:
- shutil.rmtree(nerfstudio_dir)
-
- return True
-
-
-def download_benchmark(args):
- """ Download the benchmark based on the user inputs.
-
- 1. download the benchmark-meta.csv
- 2. based on the args, download the specific subset
- a. full benchmark
- b. full benchmark in images_4 resolution level
- c. full benchmark only with nerfstudio colmaps (w.o. gaussian splatting colmaps)
- d. specific scene based on the index in [0, 140)
-
- :param args: argparse args. Used to decide the subset.
- :return: download success or not
- """
- output_dir = args.odir
- subset_opt = args.subset
- level4_opt = args.only_level4
- hash_name = args.hash
- is_clean_cache = args.clean_cache
- only_sfm = args.only_sfm
-
- # import pdb; pdb.set_trace()
- os.makedirs(output_dir, exist_ok=True)
-
- # STEP 1: download the benchmark-meta.csv and .cache/filelist.bin
- meta_repo_path = join(repo_root, 'benchmark-meta.csv')
- cache_file_path = join(repo_root, '.cache/filelist.bin')
- if hf_download_path(meta_repo_path, output_dir) == False:
- print('ERROR: Download benchmark-meta.csv failed.')
- return False
-
- if hf_download_path(cache_file_path, output_dir) == False:
- print('ERROR: Download .cache/filelist.bin failed.')
- return False
-
- # STEP 2: download the specific subset
- df = pd.read_csv(join(output_dir, 'benchmark-meta.csv'))
- filepath_dict = pickle.load(open(join(output_dir, '.cache/filelist.bin'), 'rb'))
- hashlist = df['hash'].tolist()
- download_list = hashlist
-
- # sanity check
- if subset_opt == 'hash':
- if hash_name not in hashlist:
- print(f'ERROR: hash {hash_name} not in the benchmark-meta.csv')
- return False
-
- # if subset is hash, only download the specific hash
- download_list = [hash_name]
-
- # download the dataset
- for cur_hash in tqdm(download_list):
- if download_by_hash(filepath_dict, output_dir, cur_hash, level4_opt, only_sfm) == False:
- return False
-
- if is_clean_cache:
- clean_huggingface_cache(join(output_dir, '.cache'))
-
- return True
-
-
-if __name__ == '__main__':
- parser = argparse.ArgumentParser()
- parser.add_argument('--odir', type=str, help='output directory', default='DL3DV-10K-Benchmark')
- parser.add_argument('--subset', choices=['full', 'hash'], help='The subset of the benchmark to download',
- required=True)
- parser.add_argument('--only_level4', action='store_true',
- help='If set, only the images_4 resolution level will be downloaded to save space')
- parser.add_argument('--only_sfm', action='store_true',
- help='If set, only the nerfstudio colmap data will be downloaded to save space')
- parser.add_argument('--clean_cache', action='store_true',
- help='If set, will clean the huggingface cache to save space')
- parser.add_argument('--hash', type=str, help='If set subset=hash, this is the hash code of the scene to download',
- default='')
- params = parser.parse_args()
-
- # Check huggingface login
- try:
- user = api.whoami()
- print(f'Logged in as {user["name"]}')
- except Exception as e:
- print('ERROR: Huggingface login failed. Please check your internet connection and huggingface token.')
- exit(1)
-
- if download_benchmark(params):
- print('Download Done. Refer to', params.odir)
- else:
- print(f'Download to {params.odir} Failed. See error messsage.')
diff --git a/optgs/scripts/dl3dv_download.py b/optgs/scripts/dl3dv_download.py
new file mode 100644
index 0000000000000000000000000000000000000000..c6670d5f3dc1c84e8c6a74be6da62ea72c90acf6
--- /dev/null
+++ b/optgs/scripts/dl3dv_download.py
@@ -0,0 +1,452 @@
+# ---------------------------------------------------------------------------
+# Adapted (with modifications) from the DL3DV-10K download scripts:
+# https://github.com/DL3DV-10K/Dataset/blob/main/scripts/download.py
+# https://huggingface.co/datasets/DL3DV/DL3DV-Benchmark/blob/main/download.py
+#
+# Original work © the DL3DV-10K authors, licensed under the Creative Commons
+# Attribution-NonCommercial 4.0 International License (CC BY-NC 4.0):
+# https://creativecommons.org/licenses/by-nc/4.0/
+#
+# This file is a modified version of that work and is distributed under the
+# same CC BY-NC 4.0 license. It is provided "as is", and may be used for NonCommercial purposes only.
+#
+# NOTE: This per-file license differs from the repository's main license;
+# ---------------------------------------------------------------------------
+"""Unified DL3DV downloader.
+
+Downloads DL3DV data from Hugging Face for two sources:
+
+ --source all the 10K corpus repos (DL3DV-ALL-{480P,960P,2K,4K},
+ DL3DV-ALL-ColmapCache), selected by batch subset
+ (1K..11K) or a single --hash.
+ --source benchmark the 140-scene DL3DV-10K-Benchmark (the evaluation set).
+
+For either source, --content chooses what to fetch:
+
+ images image frames + poses (the data WITHOUT the SfM).
+ sfm only the COLMAP sparse SfM (cameras/images/points3D.bin),
+ rearranged to /sparse for the colmap initializer.
+ images+sfm both.
+
+Both DL3DV repos are gated: you must request access on Hugging Face and be
+logged in (`huggingface-cli login` or HF_TOKEN). This script downloads nothing
+itself — it requires each user to obtain their own access — so it does not
+redistribute the dataset.
+
+Example:
+ python -m optgs.scripts.dl3dv_download --source benchmark \
+ --odir datasets/dl3dv-benchmark --content images+sfm --clean_cache
+ python -m optgs.scripts.dl3dv_download --source all \
+ --odir datasets/dl3dv-colmap-sfm --content sfm --subset 1K --clean_cache
+"""
+
+import argparse
+import os
+import pathlib
+import shutil
+import traceback
+import urllib.request
+import zipfile
+from os.path import join
+
+import pandas as pd
+from huggingface_hub import HfApi, HfFileSystem, snapshot_download
+from tqdm import tqdm
+
+api = HfApi()
+
+# --- source: all (10K corpus) -------------------------------------------------
+RESOLUTION2REPO = {
+ '480P': 'DL3DV/DL3DV-ALL-480P',
+ '960P': 'DL3DV/DL3DV-ALL-960P',
+ '2K': 'DL3DV/DL3DV-ALL-2K',
+ '4K': 'DL3DV/DL3DV-ALL-4K',
+}
+COLMAP_CACHE_REPO = 'DL3DV/DL3DV-ALL-ColmapCache'
+ALL_META_LINK = 'https://raw.githubusercontent.com/DL3DV-10K/Dataset/main/cache/DL3DV-valid.csv'
+
+# --- source: benchmark (140-scene eval set) -----------------------------------
+BENCHMARK_REPO = 'DL3DV/DL3DV-10K-Benchmark'
+
+SFM_BIN_FILES = {"cameras.bin", "images.bin", "points3D.bin"}
+
+
+# =============================================================================
+# Shared helpers
+# =============================================================================
+def hf_download_path(repo: str, rel_path: str, odir: str, max_try: int = 5):
+ """Download a single file from a HF dataset repo, retrying on failure."""
+ counter = 0
+ while True:
+ if counter >= max_try:
+ print(f"ERROR: Download {repo}/{rel_path} failed.")
+ return False
+ try:
+ api.hf_hub_download(repo_id=repo, filename=rel_path, repo_type='dataset',
+ local_dir=odir, cache_dir=join(odir, '.cache'))
+ return True
+ except KeyboardInterrupt:
+ print('Keyboard Interrupt. Exit.')
+ raise SystemExit(1)
+ except BaseException:
+ traceback.print_exc()
+ counter += 1
+
+
+def download_from_url(url: str, ofile: str):
+ """Download a file from a plain URL to ofile."""
+ try:
+ urllib.request.urlretrieve(url, ofile)
+ return True
+ except Exception as e:
+ print(f"An error occurred while downloading the file: {e}")
+ return False
+
+
+def clean_huggingface_cache(output_dir: str):
+ """Remove the local HF cache dir to save space (no good hub API for this)."""
+ cur_cache_dir = join(output_dir, '.cache')
+ if os.path.exists(cur_cache_dir):
+ shutil.rmtree(cur_cache_dir)
+
+
+def verify_access(repo: str) -> bool:
+ """Return True if the logged-in user can list the (gated) repo."""
+ fs = HfFileSystem()
+ try:
+ fs.ls(f'datasets/{repo}')
+ return True
+ except BaseException:
+ return False
+
+
+# =============================================================================
+# SfM cleanup / validation (shared by both sources; layouts differ per source)
+# =============================================================================
+def sfm_cleanup_scene(scene_dir: pathlib.Path):
+ """Keep only the COLMAP sparse SfM files, then move them to /sparse.
+
+ Used by the 'all' source, whose ColmapCache zips contain
+ /colmap/sparse/* alongside many other files.
+ """
+ print(f"Cleaning up SfM scene at {scene_dir.resolve()}")
+ scene_dir = scene_dir.resolve()
+ if not scene_dir.exists():
+ print(f"[WARN] {scene_dir} does not exist")
+ return
+
+ # First pass: delete everything except the sparse .bin files and transforms.json.
+ for path in scene_dir.rglob("*"):
+ if path.is_file():
+ is_bin_file = (path.name in SFM_BIN_FILES and
+ path.parent.name.isdigit() and
+ path.parent.parent.name == "sparse")
+ is_transforms_file = (path.name == "transforms.json" and path.parent == scene_dir)
+ if is_bin_file or is_transforms_file:
+ continue
+ path.unlink()
+
+ # Second pass: remove empty directories bottom-up.
+ for path in sorted(scene_dir.rglob("*"), reverse=True):
+ if path.is_dir() and not any(path.iterdir()):
+ path.rmdir()
+
+ # Third pass: lift the SfM up to // (out of the scratch/subset scene dir):
+ # /colmap/sparse -> //sparse
+ # /transforms.json -> //transforms.json (next to the SfM, for future use)
+ dataset_dir = scene_dir.parent.parent
+ colmap_dir = scene_dir / "colmap"
+ sparse_dir = colmap_dir / "sparse"
+ if sparse_dir.exists():
+ target_scene_dir = dataset_dir / scene_dir.name
+ target_scene_dir.mkdir(parents=True, exist_ok=True)
+ shutil.move(str(sparse_dir), str(target_scene_dir / "sparse"))
+ transforms = scene_dir / "transforms.json"
+ if transforms.is_file():
+ shutil.move(str(transforms), str(target_scene_dir / "transforms.json"))
+ if not any(colmap_dir.iterdir()):
+ colmap_dir.rmdir()
+ if not any(scene_dir.iterdir()):
+ scene_dir.rmdir()
+
+
+def validate_sfm_structure(scene_dir: pathlib.Path, unsucc_count: int = 0) -> bool:
+ """Validate that a cleaned scene dir holds sparse/0/{cameras,images,points3D}.bin."""
+ scene_dir = scene_dir.resolve()
+ if not scene_dir.exists():
+ print(f"[WARN: {unsucc_count}] {scene_dir} does not exist")
+ return False
+
+ sparse_0_dir = scene_dir / "sparse" / "0"
+ for bin_file in SFM_BIN_FILES:
+ if not (sparse_0_dir / bin_file).is_file():
+ print(f"[ERROR: {unsucc_count}] {bin_file} is missing in {sparse_0_dir}")
+ return False
+ return True
+
+
+# =============================================================================
+# Source: all (10K corpus)
+# =============================================================================
+def all_to_download_item(hash_name: str, reso: str, batch: str, file_type: str) -> dict:
+ if file_type == 'images':
+ repo = RESOLUTION2REPO[reso]
+ rel_path = f'{batch}/{hash_name}.zip'
+ elif file_type == 'sfm':
+ repo = COLMAP_CACHE_REPO
+ rel_path = f'{batch}/{hash_name}.zip'
+ else:
+ raise ValueError(f'Unknown file_type {file_type!r} for source=all.')
+ return {'repo': repo, 'rel_path': rel_path, 'file_type': file_type}
+
+
+def all_get_download_list(subset_opt, hash_name, reso_opt, file_types, output_dir, limit=None):
+ """Build the download list for the 'all' source from the DL3DV-valid.csv meta."""
+ cache_folder = join(output_dir, '.cache')
+ meta_file = join(cache_folder, 'DL3DV-valid.csv')
+ os.makedirs(cache_folder, exist_ok=True)
+ if not os.path.exists(meta_file):
+ assert download_from_url(ALL_META_LINK, meta_file), 'Download meta file failed.'
+
+ df = pd.read_csv(meta_file)
+
+ if hash_name:
+ assert hash_name in df['hash'].values, f'Hash {hash_name} not found in the meta file.'
+ rows = [(hash_name, df[df['hash'] == hash_name]['batch'].values[0])]
+ else:
+ subdf = df[df['batch'] == subset_opt]
+ rows = [(r['hash'], subset_opt) for _, r in subdf.iterrows()]
+
+ if limit is not None:
+ rows = rows[:limit]
+
+ ret = []
+ for h, batch in rows:
+ for file_type in file_types:
+ ret.append(all_to_download_item(h, reso_opt, batch, file_type))
+ return ret
+
+
+def _zip_common_prefix(zip_ref, hash_name):
+ """Return the zip's single top-level dir if it equals hash_name, else None."""
+ zip_contents = zip_ref.namelist()
+ if zip_contents and '/' in zip_contents[0]:
+ potential_prefix = zip_contents[0].split('/')[0] + '/'
+ if all(p.startswith(potential_prefix) for p in zip_contents if not p.endswith('/')):
+ if potential_prefix.rstrip('/') == hash_name:
+ return hash_name
+ return None
+
+
+def all_download(download_list, output_dir, is_clean_cache):
+ """Download and unzip the 'all'-source items.
+
+ Images and SfM go to separate directories so the (destructive) SfM cleanup never
+ touches the images:
+ images → /// (image frames + transforms.json)
+ sfm → //sparse/ (cleaned; extracted via a scratch dir)
+ """
+ output_dir = pathlib.Path(output_dir)
+ succ_count = 0
+ for item in tqdm(download_list, desc='Downloading'):
+ repo, rel_path, file_type = item['repo'], item['rel_path'], item['file_type']
+ hash_name = os.path.splitext(os.path.basename(rel_path))[0]
+ subset_name = os.path.dirname(rel_path)
+
+ # Skip if the (per-content) final target already exists.
+ if file_type == 'sfm':
+ if validate_sfm_structure(output_dir / hash_name):
+ succ_count += 1
+ continue
+ else: # images
+ if (output_dir / subset_name / hash_name).exists():
+ succ_count += 1
+ continue
+
+ if not hf_download_path(repo, rel_path, str(output_dir)):
+ print(f'Download {rel_path} failed')
+ continue
+ succ_count += 1
+ if is_clean_cache:
+ clean_huggingface_cache(str(output_dir))
+
+ if not rel_path.endswith('.zip'):
+ continue
+ zip_file = str(output_dir / rel_path)
+
+ if file_type == 'sfm':
+ # Extract into a scratch dir; sfm_cleanup_scene keeps only the sparse SfM and
+ # moves it to //sparse, away from the images at //.
+ scratch_base = output_dir / '.sfm_scratch'
+ scene_dir = scratch_base / hash_name
+ scene_dir.mkdir(parents=True, exist_ok=True)
+ with zipfile.ZipFile(zip_file, 'r') as zip_ref:
+ dest = scratch_base if _zip_common_prefix(zip_ref, hash_name) else scene_dir
+ zip_ref.extractall(str(dest))
+ sfm_cleanup_scene(scene_dir)
+ shutil.rmtree(scratch_base, ignore_errors=True)
+ else: # images
+ target_dir = output_dir / subset_name / hash_name
+ target_dir.mkdir(parents=True, exist_ok=True)
+ with zipfile.ZipFile(zip_file, 'r') as zip_ref:
+ dest = output_dir / subset_name if _zip_common_prefix(zip_ref, hash_name) else target_dir
+ zip_ref.extractall(str(dest))
+ os.remove(zip_file)
+
+ print(f'Summary: {succ_count}/{len(download_list)} files downloaded successfully')
+ return succ_count == len(download_list)
+
+
+# =============================================================================
+# Source: benchmark (140-scene eval set)
+# =============================================================================
+def _benchmark_surface_sfm(odir, scene_hash):
+ """Surface a benchmark scene's COLMAP SfM for the colmap initializer.
+
+ /nerfstudio/colmap/sparse -> /sparse
+ transforms.json is copied to the scene root (next to the SfM, for future colmap/test use)
+ while the original stays under nerfstudio/ where the image-conversion step reads it.
+ """
+ scene_root = pathlib.Path(odir) / scene_hash
+ nerf_dir = scene_root / 'nerfstudio'
+ nerf_transforms = nerf_dir / 'transforms.json'
+ if nerf_transforms.is_file():
+ shutil.copy2(str(nerf_transforms), str(scene_root / 'transforms.json'))
+ dst_sparse = scene_root / 'sparse'
+ if dst_sparse.exists():
+ shutil.rmtree(dst_sparse)
+ shutil.move(str(nerf_dir / 'colmap' / 'sparse'), str(dst_sparse))
+ colmap_dir = nerf_dir / 'colmap'
+ if colmap_dir.exists() and not any(colmap_dir.iterdir()):
+ colmap_dir.rmdir()
+
+
+def benchmark_download(subset_opt, hash_name, contents, output_dir, is_clean_cache, limit=None):
+ """Download benchmark scenes (full set or a single --hash) for the requested contents.
+
+ The benchmark stores each scene's frames as individual files, so we fetch them with a single
+ parallel snapshot_download using per-scene glob patterns. Only the nerfstudio assets are
+ pulled, skipping the heavy gaussian_splat/ MVS+dense cache:
+ images → nerfstudio/images_4/* (+ transforms.json)
+ sfm → nerfstudio/colmap/sparse/* (+ transforms.json), then surfaced to /sparse.
+ """
+ # benchmark-meta.csv gives the scene list (for --hash validation and --limit).
+ if not hf_download_path(BENCHMARK_REPO, 'benchmark-meta.csv', output_dir):
+ print('ERROR: Download benchmark-meta.csv failed.')
+ return False
+ hashlist = pd.read_csv(join(output_dir, 'benchmark-meta.csv'))['hash'].tolist()
+
+ if subset_opt == 'hash':
+ if hash_name not in hashlist:
+ print(f'ERROR: hash {hash_name} not in benchmark-meta.csv')
+ return False
+ download_list = [hash_name]
+ else:
+ download_list = hashlist
+ if limit is not None:
+ download_list = download_list[:limit]
+
+ # Glob patterns for exactly the files we want (snapshot_download matches with fnmatch,
+ # where '*' also spans '/').
+ patterns = []
+ for h in download_list:
+ if 'images' in contents:
+ patterns += [f'{h}/nerfstudio/images_4/*', f'{h}/nerfstudio/transforms.json']
+ if 'sfm' in contents:
+ patterns += [f'{h}/nerfstudio/colmap/sparse/*', f'{h}/nerfstudio/transforms.json']
+
+ print(f'Downloading {len(download_list)} benchmark scene(s) in parallel...')
+ snapshot_download(repo_id=BENCHMARK_REPO, repo_type='dataset', local_dir=output_dir,
+ allow_patterns=patterns, cache_dir=join(output_dir, '.cache'))
+
+ if 'sfm' in contents:
+ for h in download_list:
+ _benchmark_surface_sfm(output_dir, h)
+ if is_clean_cache:
+ clean_huggingface_cache(output_dir)
+ return True
+
+
+# =============================================================================
+# CLI
+# =============================================================================
+def _contents_from_arg(content: str) -> list[str]:
+ return {'images': ['images'], 'sfm': ['sfm'], 'images+sfm': ['images', 'sfm']}[content]
+
+
+def main():
+ parser = argparse.ArgumentParser(description=__doc__,
+ formatter_class=argparse.RawDescriptionHelpFormatter)
+ parser.add_argument('--source', choices=['all', 'benchmark'], required=True,
+ help='all = 10K corpus repos; benchmark = 140-scene eval set')
+ parser.add_argument('--odir', type=str, default=None,
+ help='output directory (default: datasets/dl3dv-colmap-sfm for --source all, '
+ 'datasets/dl3dv-benchmark for --source benchmark)')
+ parser.add_argument('--content', choices=['images', 'sfm', 'images+sfm'], default='images+sfm',
+ help='what to download: image frames+poses, the COLMAP SfM, or both')
+ parser.add_argument('--hash', type=str, default='', help='download a single scene by hash')
+ parser.add_argument('--clean_cache', action='store_true',
+ help='remove the HF cache after each download to save space')
+ parser.add_argument('--limit', default=None,
+ help='download at most N scenes (a debug-sized set); empty or 0 means no limit')
+ # source=all options
+ parser.add_argument('--subset', choices=['1K', '2K', '3K', '4K', '5K', '6K', '7K', '8K',
+ '9K', '10K', '11K'],
+ help='[all] batch subset to download (ignored if --hash is set)')
+ parser.add_argument('--resolution', choices=['480P', '960P', '2K', '4K'], default='480P',
+ help='[all] image resolution repo')
+ # (benchmark images are always fetched at images_4 / 960x540, skipping the heavy
+ # gaussian_splat MVS+dense cache.)
+ args = parser.parse_args()
+
+ # Default output dir matches the existing per-source conventions:
+ # all → train SfM dir, benchmark → test SfM dir.
+ if args.odir is None:
+ args.odir = {'all': 'datasets/dl3dv-colmap-sfm',
+ 'benchmark': 'datasets/dl3dv-benchmark'}[args.source]
+
+ # Empty string or 0 means "no limit" (download the whole subset / benchmark).
+ limit = None if args.limit in (None, '', '0', 0) else int(args.limit)
+
+ os.makedirs(args.odir, exist_ok=True)
+ contents = _contents_from_arg(args.content)
+
+ if args.source == 'all':
+ if not args.hash and not args.subset:
+ parser.error('source=all requires --subset (or --hash).')
+ file_types = ['images' if c == 'images' else 'sfm' for c in contents]
+ # Verify access to every repo we will hit.
+ repos = set()
+ if 'images' in contents:
+ repos.add(RESOLUTION2REPO[args.resolution])
+ if 'sfm' in contents:
+ repos.add(COLMAP_CACHE_REPO)
+ for repo in repos:
+ if not verify_access(repo):
+ print(f'No access to https://huggingface.co/datasets/{repo} — request access and '
+ f'log in (huggingface-cli login).')
+ raise SystemExit(1)
+ download_list = all_get_download_list(args.subset, args.hash, args.resolution,
+ file_types, args.odir, limit=limit)
+ ok = all_download(download_list, args.odir, args.clean_cache)
+ else:
+ try:
+ user = api.whoami()
+ print(f'Logged in as {user["name"]}')
+ except Exception:
+ print('ERROR: Hugging Face login failed. Check your connection and token '
+ '(huggingface-cli login / HF_TOKEN).')
+ raise SystemExit(1)
+ subset_opt = 'hash' if args.hash else 'full'
+ ok = benchmark_download(subset_opt, args.hash, contents, args.odir, args.clean_cache,
+ limit=limit)
+
+ if ok:
+ print('Download Done. Refer to', args.odir)
+ else:
+ print(f'Download to {args.odir} failed. See error message above.')
+ raise SystemExit(1)
+
+
+if __name__ == '__main__':
+ main()
diff --git a/optgs/scripts/dl3dv_hf_download.py b/optgs/scripts/dl3dv_hf_download.py
deleted file mode 100644
index b59c3a905b3cf010146187c3ee6b148879e4c188..0000000000000000000000000000000000000000
--- a/optgs/scripts/dl3dv_hf_download.py
+++ /dev/null
@@ -1,421 +0,0 @@
-""" This script is used to download the DL3DV-10 dataset for all resolution levels from the huggingface repo.
- As the whole dataset is too large for most users, we provide this script so that you can download the dataset efficiently based on your needs.
- We provide several options to download the dataset (image frames with poses):
- - [X] Resolution level: 4K, 2K, 960P, 480P
- - [X] Subset of the 10K, e.g. 1K(0~1K), 2K(1K~2K), 3K(2K~3K), etc
- - [X] specific hash
- - [X] file_type: raw video | images+poses | colmap cache
-
- Notes:
- - file_type + resolution will decide which dataset repo to download the files
- - subset will decide which subdir will be used
- - if hash is set, only the specific hash will be downloaded
-
- example usage: python dl3dv_hf_download.py --odir ../../datasets/dl3dv-colmap-sfm --file_type colmap_sfm --hash e2cedefea8a0ed2d0ffbd5bdc08acbe7e1f85c96f72f7b790e9dfe1c98963047 --clean_cache --subset 1K --resolution 480P
-
-
-"""
-
-import os
-import pathlib
-from os.path import join
-import pandas as pd
-from tqdm import tqdm
-from huggingface_hub import HfApi
-import argparse
-import traceback
-import shutil
-import urllib.request
-import zipfile
-from huggingface_hub import HfFileSystem
-
-api = HfApi()
-resolution2repo = {
- '480P': 'DL3DV/DL3DV-ALL-480P',
- '960P': 'DL3DV/DL3DV-ALL-960P',
- '2K': 'DL3DV/DL3DV-ALL-2K',
- '4K': 'DL3DV/DL3DV-ALL-4K'
-}
-
-
-def verify_access(repo: str):
- """ This function can be used to verify if the user has access to the repo.
-
- :param repo: the repo name
- :return: True if the user has access, False otherwise
- """
- fs = HfFileSystem()
- try:
- fs.ls(f'datasets/{repo}')
- return True
- except BaseException as e:
- return False
-
-
-def hf_download_path(repo: str, rel_path: str, odir: str, max_try: int = 5):
- """ hf api is not reliable, retry when failed with max tries
-
- :param repo: The huggingface dataset repo
- :param rel_path: The relative path in the repo
- :param odir: output path
- :param max_try: As the downloading is not a reliable process, we will retry for max_try times
- """
- counter = 0
- while True:
- if counter >= max_try:
- print(f"ERROR: Download {repo}/{rel_path} failed.")
- return False
- try:
- api.hf_hub_download(repo_id=repo,
- filename=rel_path,
- repo_type='dataset',
- local_dir=odir,
- cache_dir=join(odir, '.cache'))
- return True
-
- except KeyboardInterrupt:
- print('Keyboard Interrupt. Exit.')
- exit()
- except BaseException as e:
- traceback.print_exc()
- counter += 1
-
-
-def download_from_url(url: str, ofile: str):
- """ Download a file from the url to ofile
-
- :param url: The url link
- :param ofile: The output path
- :return: True if download success, False otherwise
- """
- try:
- # Use urllib.request.urlretrieve to download the file from `url` and save it locally at `local_file_path`
- urllib.request.urlretrieve(url, ofile)
- return True
- except Exception as e:
- print(f"An error occurred while downloading the file: {e}")
- return False
-
-
-def clean_huggingface_cache(output_dir: str, repo: str):
- """ Huggingface cache may take too much space, we clean the cache to save space if necessary
-
- Current huggingface hub does not provide good practice to clean the space.
- We mannually clean the cache directory if necessary.
-
- :param output_dir: the current output directory
- :param output_dir: the huggingface repo
- """
- repo_cache_dir = repo.replace('/', '--')
- # cur_cache_dir = join(output_dir, '.cache', f'datasets--{repo_cache_dir}')
- cur_cache_dir = join(output_dir, '.cache')
-
- if os.path.exists(cur_cache_dir):
- shutil.rmtree(cur_cache_dir)
-
-
-def get_download_list(subset_opt: str, hash_name: str, reso_opt: str, file_type: str, output_dir: str):
- """ Get the download list based on the subset and hash name
-
- 1. Get the meta file
- 2. Select the subset. Based on reso_opt, get the downloading list prepared.
- 3. Return the download list.
-
- :param subset_opt: Subset of the 10K, e.g. 1K(0~1K), 2K(1K~2K), 3K(2K~3K), etc
- :param hash_name: If provided a non-empty string, ignore the subset_opt and only download the specific hash
- :param reso_opt: The resolution to download.
- :param file_type: The file type to download: video | images+poses | colmap_cache
- :param output_dir: The output directory.
- """
-
- def to_download_item(hash_name, reso, batch, file_type):
- if file_type == 'images+poses':
- repo = resolution2repo[reso]
- rel_path = f'{batch}/{hash_name}.zip'
- elif file_type == 'video':
- repo = 'DL3DV/DL3DV-ALL-video'
- rel_path = f'{batch}/{hash_name}/video.mp4'
- elif file_type in ['colmap_cache', 'colmap_sfm']:
- repo = 'DL3DV/DL3DV-ALL-ColmapCache'
- rel_path = f'{batch}/{hash_name}.zip'
- else:
- raise ValueError('Unknown file_type option.')
-
- # return f'{repo}/{batch}/{hash_name}'
- return {'repo': repo, 'rel_path': rel_path}
-
- ret = []
-
- meta_link = 'https://raw.githubusercontent.com/DL3DV-10K/Dataset/main/cache/DL3DV-valid.csv'
- cache_folder = join(output_dir, '.cache')
- meta_file = join(cache_folder, 'DL3DV-valid.csv')
- os.makedirs(cache_folder, exist_ok=True)
- if not os.path.exists(meta_file):
- assert download_from_url(meta_link, meta_file), 'Download meta file failed.'
-
- df = pd.read_csv(meta_file)
-
- # if hash is set, ignore the subset_opt
- if hash_name != '':
- assert hash_name in df['hash'].values, f'Hash {hash_name} not found in the meta file.'
-
- batch = df[df['hash'] == hash_name]['batch'].values[0]
- link = to_download_item(hash_name, reso_opt, batch, file_type)
- ret = [link]
- return ret
-
- # if hash not set, we download the whole subset
- subdf = df[df['batch'] == subset_opt]
- for i, r in subdf.iterrows():
- hash_name = r['hash']
- ret.append(to_download_item(hash_name, reso_opt, subset_opt, file_type))
-
- return ret
-
-
-SFM_BIN_FILES = {"cameras.bin", "images.bin", "points3D.bin"}
-
-
-def sfm_cleanup_scene(scene_dir: pathlib.Path):
- """
- Keep only COLMAP sparse SfM files:
- cameras.bin, images.bin, points3D.bin
- Delete everything else.
- """
- print(f"Cleaning up SfM scene at {scene_dir.resolve()}")
- scene_dir = scene_dir.resolve()
-
- if not scene_dir.exists():
- print(f"[WARN] {scene_dir} does not exist")
- return
-
- # First pass: delete unwanted files
- for path in scene_dir.rglob("*"):
- if path.is_file():
- # keep sparse/[0-9]+/{cameras,images,points3D}.bin and transforms.json
- is_bin_file = (path.name in SFM_BIN_FILES and
- path.parent.name.isdigit() and
- path.parent.parent.name == "sparse")
- is_transforms_file = (path.name == "transforms.json" and path.parent == scene_dir)
- if is_bin_file or is_transforms_file:
- continue
-
- path.unlink()
-
- # Second pass: remove empty directories bottom-up
- for path in sorted(scene_dir.rglob("*"), reverse=True):
- if path.is_dir() and not any(path.iterdir()):
- path.rmdir()
-
- # Third pass: rearrange files to fit clogs training
- # move //colmap/sparse/* to //*
- # and remove empty dirs
- subset_dir = scene_dir.parent
- dataset_dir = scene_dir.parent.parent
- colmap_dir = scene_dir / "colmap"
- sparse_dir = colmap_dir / "sparse"
- if sparse_dir.exists():
- # move the sparse dir to a scene_dir inside dataset dir
- target_sparse_dir = dataset_dir / scene_dir.name / "sparse"
- target_sparse_dir.parent.mkdir(parents=True, exist_ok=True)
- shutil.move(str(sparse_dir), str(target_sparse_dir))
-
- # remove empty dirs
- if not any(colmap_dir.iterdir()):
- colmap_dir.rmdir()
- if not any(scene_dir.iterdir()):
- scene_dir.rmdir()
-
-
-def validate_sfm_structure(scene_dir: pathlib.Path, unsucc_count: int):
- """
- Validate the SfM cleanup by checking if the scene dir only contains the sparse/0/cameras.bin, images.bin, points3D.bin files and transforms.json
- """
- scene_dir = scene_dir.resolve()
-
- if not scene_dir.exists():
- print(f"[WARN: {unsucc_count}] {scene_dir} does not exist")
- return False
-
- # Check if transforms.json exists in the root of the scene dir
- # transforms_file = scene_dir / "transforms.json"
- # TODO Naama: skipping transforms.json, but will need to redownload
- # if not transforms_file.is_file():
- # print(f"[ERROR] transforms.json is missing in {scene_dir}")
- # return False
-
- # Check if sparse/0/cameras.bin, images.bin, points3D.bin exist
- sparse_0_dir = scene_dir / "sparse" / "0"
- for bin_file in SFM_BIN_FILES:
- if not (sparse_0_dir / bin_file).is_file():
- print(f"[ERROR: {unsucc_count}] {bin_file} is missing in {sparse_0_dir}")
- return False
-
- # Check if there are any other files or directories in the scene dir
- for path in scene_dir.rglob("*"):
- if path.is_file():
- is_bin_file = (path.name in SFM_BIN_FILES and
- path.parent.name == "0" and
- path.parent.parent.name == "sparse")
- is_transforms_file = (path.name == "transforms.json" and path.parent == scene_dir)
- is_image_file = (path.suffix in ['.jpg', '.png'] and path.parent.name.startswith("images") and path.parent.parent == scene_dir)
- if not (is_bin_file or is_transforms_file or is_image_file):
- # print(f"*********** [WARN: {unsucc_count}] Unexpected file {path} found in {scene_dir}")
- # remove the unexpected file
- # path.unlink()
- pass
- elif path.is_dir():
- # if there is any dir other than sparse/0, it's unexpected
- is_sparse_dir = (path.name == "sparse" and path.parent == scene_dir)
- is_sparse_0_dir = (path.name == "0" and path.parent.name == "sparse")
- # For test scenes we might have images* dir
- is_images_dir = ("images" in path.name and path.parent == scene_dir)
- if not (is_sparse_0_dir or is_sparse_dir or is_images_dir):
- # print(f"*********** [WARN: {unsucc_count}] Unexpected directory {path} found in {scene_dir}")
- # remove the unexpected dir
- # shutil.rmtree(path)
- pass
- return True
-
-
-def download(download_list: list, output_dir: str, is_clean_cache: bool, only_sfm: bool = False):
- """ Download the dataset based on the download_list and user options.
-
- :param download_list: the list of files to download, [{'repo', 'rel_path'}]
- :param output_dir: the output directory
- :param reso_opt: the resolution option
- :param is_clean_cache: if set, will clean the huggingface cache to save space
- :param only_sfm: if set, only download the colmap sfm files (remove all other files)
- """
- succ_count = 0
- unsucc_count = 0
-
- for item in tqdm(download_list, desc='Downloading'):
- repo = item['repo']
- rel_path = item['rel_path']
-
- output_path = os.path.join(output_dir, rel_path)
- output_path = output_path.replace('.zip', '')
- # skip if already exists locally
- # scene dir can be moved from root/subset/hash/ to root/hash/ after sfm_cleanup, so we need to check both paths
- output_path_without_subset = pathlib.Path(output_path).parent.parent / pathlib.Path(output_path).name
- # if os.path.exists(output_path):
- # print(f"File {output_path} already exists, skip downloading.")
- # succ_count += 1
- # continue
- if output_path_without_subset.exists():
- # For sfm, verify donwload based on the hash dir after cleanup, which is moved to root/hash/
- if only_sfm:
- if validate_sfm_structure(output_path_without_subset, unsucc_count):
- succ_count += 1
- continue
- else:
- succ_count += 1
- continue
- unsucc_count += 1
- succ = hf_download_path(repo, rel_path, output_dir)
-
- if succ:
- succ_count += 1
- if is_clean_cache:
- clean_huggingface_cache(output_dir, repo)
-
- # unzip the file
- if rel_path.endswith('.zip'):
- zip_file = join(output_dir, rel_path)
- hash_name = os.path.splitext(os.path.basename(rel_path))[0]
- subset_name = os.path.dirname(rel_path)
- target_dir = join(output_dir, subset_name, hash_name)
-
- # Ensure target directory exists
- os.makedirs(target_dir, exist_ok=True)
-
- with zipfile.ZipFile(zip_file, 'r') as zip_ref:
- # Get list of files in the zip
- zip_contents = zip_ref.namelist()
-
- # Check if all files are under a single directory that matches the hash
- common_prefix = None
- if zip_contents:
- # Check if there's a common directory prefix
- first_path = zip_contents[0]
- if '/' in first_path:
- potential_prefix = first_path.split('/')[0] + '/'
- if all(path.startswith(potential_prefix) for path in zip_contents if
- not path.endswith('/')):
- common_prefix = potential_prefix.rstrip('/')
-
- # Extract files
- if common_prefix == hash_name:
- # Files are already under hash directory, extract normally
- zip_ref.extractall(join(output_dir, subset_name))
- else:
- # Extract directly to target hash directory
- zip_ref.extractall(target_dir)
-
- if only_sfm:
- scene_dir = pathlib.Path(target_dir)
- sfm_cleanup_scene(scene_dir)
- os.remove(zip_file)
-
-
- else:
- print(f'Download {rel_path} failed')
-
- print(f'Summary: {succ_count}/{len(download_list)} files downloaded successfully')
- return succ_count == len(download_list)
-
-
-def download_dataset(args):
- """ Download the dataset based on the user inputs.
-
- :param args: argparse args. Used to decide the subset.
- :return: download success or not
- """
- output_dir = args.odir
- subset_opt = args.subset
- reso_opt = args.resolution
- hash_name = args.hash
- file_type = args.file_type
- is_clean_cache = args.clean_cache
-
- os.makedirs(output_dir, exist_ok=True)
-
- download_list = get_download_list(subset_opt, hash_name, reso_opt, file_type, output_dir)
- return download(download_list, output_dir, is_clean_cache, only_sfm=file_type == 'colmap_sfm')
-
-
-if __name__ == '__main__':
- parser = argparse.ArgumentParser()
- parser.add_argument('--odir', type=str, help='output directory', required=True)
- parser.add_argument('--subset', choices=['1K', '2K', '3K', '4K', '5K', '6K', '7K', '8K', '9K', '10K', '11K'],
- help='The subset of the benchmark to download', required=True)
- parser.add_argument('--resolution', choices=['4K', '2K', '960P', '480P'], help='The resolution to donwnload',
- default='480P')
- parser.add_argument('--file_type', choices=['images+poses', 'video', 'colmap_cache', 'colmap_sfm'],
- help='The file type to download', required=True, default='images+poses')
- parser.add_argument('--hash', type=str, help='If set subset=hash, this is the hash code of the scene to download',
- default='')
- parser.add_argument('--clean_cache', action='store_true',
- help='If set, will clean the huggingface cache to save space')
- params = parser.parse_args()
-
- assert params.file_type in ['images+poses', 'video', 'colmap_cache', 'colmap_sfm',
- 'only_sfm'], 'Check the file_type input.'
-
- if params.file_type == 'images+poses':
- repo = resolution2repo[params.resolution]
- elif params.file_type == 'video':
- repo = 'DL3DV/DL3DV-ALL-video'
- elif params.file_type in ['colmap_cache', 'colmap_sfm']:
- repo = 'DL3DV/DL3DV-ALL-ColmapCache'
-
- if not verify_access(repo):
- print(
- f'You have not grant the access yet. Go to relevant huggingface repo (https://huggingface.co/datasets/{repo}) and apply for the access.')
- exit(1)
-
- if download_dataset(params):
- print('Download Done. Refer to', params.odir)
- else:
- print(f'Download to {params.odir} failed. See error messsage.')
diff --git a/optgs/scripts/dl3dv_verify_colmap_poses.py b/optgs/scripts/dl3dv_verify_colmap_poses.py
index ef819e324c1130ec1bbad14355efd376f163eec9..1e5d3786c32a78504b2246ea6f971a918c1d0fc9 100644
--- a/optgs/scripts/dl3dv_verify_colmap_poses.py
+++ b/optgs/scripts/dl3dv_verify_colmap_poses.py
@@ -5,7 +5,6 @@ import torch
from optgs.dataset.dataset_colmap import Parser
import json
-from optgs.scripts.convert_dl3dv_utils import load_metadata
from einops import rearrange, repeat
if __name__ == '__main__':
diff --git a/optgs/scripts/generate_dl3dv_index.py b/optgs/scripts/generate_dl3dv_index.py
index 655ae0f5a29fa5dc8d4a67dba0451165ba19e27d..580fdb9580b7b2885755f12ec46bf77efc26647e 100644
--- a/optgs/scripts/generate_dl3dv_index.py
+++ b/optgs/scripts/generate_dl3dv_index.py
@@ -1,12 +1,18 @@
import json
+import argparse
from pathlib import Path
import torch
from tqdm import tqdm
-DATASET_PATH = Path("/capstor/store/cscs/swissai/a03/hxu/datasets/dl3dv_2kres")
if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--data_dir", type=str, required=True, help="path to dataset directory")
+ args = parser.parse_args()
+
+ DATASET_PATH = Path(args.data_dir)
+
# "train" or "test"
for stage in ["test"]:
stage = DATASET_PATH / stage
@@ -20,4 +26,4 @@ if __name__ == "__main__":
for example in chunk:
index[example["key"]] = str(chunk_path.relative_to(stage))
with (stage / "index.json").open("w") as f:
- json.dump(index, f)
+ json.dump(index, f)
\ No newline at end of file
diff --git a/optgs/scripts/preextract_colmap_npz.py b/optgs/scripts/preextract_colmap_npz.py
index 41df909a0e3ddc8454133580398e2fb80e320690..0dae8663d83152140e52b6b0593b264ebc01e6bf 100644
--- a/optgs/scripts/preextract_colmap_npz.py
+++ b/optgs/scripts/preextract_colmap_npz.py
@@ -15,7 +15,6 @@ automatically and skip the full SceneManager parse.
import argparse
import concurrent.futures
-import os
import sys
import traceback
from pathlib import Path
diff --git a/optgs/scripts/verify_dl3dv_colmap_download.py b/optgs/scripts/verify_dl3dv_colmap_download.py
index 0202d4dca9b2fc1932cfd6687a834c55ae880922..3fd634fd375ba2bf909accad5a64b57b6450b713 100644
--- a/optgs/scripts/verify_dl3dv_colmap_download.py
+++ b/optgs/scripts/verify_dl3dv_colmap_download.py
@@ -7,7 +7,7 @@ import torch
from tqdm import tqdm
from optgs.misc.io import CustomPath
-from optgs.scripts.dl3dv_hf_download import validate_sfm_structure
+from optgs.scripts.dl3dv_download import validate_sfm_structure
if __name__ == '__main__':
chunk_dir = CustomPath("datasets/dl3dv-480p-chunks/train")
diff --git a/optgs/visualization/drawing/cameras.py b/optgs/visualization/drawing/cameras.py
index 6a60beb53a1a24931cd4b822e5ebe65b6425a6d6..cc66f398d3d08e66f847b5f9ae504eb0fca071f2 100644
--- a/optgs/visualization/drawing/cameras.py
+++ b/optgs/visualization/drawing/cameras.py
@@ -114,7 +114,6 @@ def draw_cameras(
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)
diff --git a/optgs/visualization/plots3d/matplotlib.py b/optgs/visualization/plots3d/matplotlib.py
index 5e2a8805e2bfb706e7aeb62c26f3c70bed3f1e84..80cf7879dbf5a57c2509a3cf8b34694684d6e911 100644
--- a/optgs/visualization/plots3d/matplotlib.py
+++ b/optgs/visualization/plots3d/matplotlib.py
@@ -1,9 +1,9 @@
import torch
import numpy as np
import matplotlib.pyplot as plt
-from typing import Optional, Literal, Tuple, List, Union
+from typing import Optional, Literal, Tuple, List
from pathlib import Path
-from itertools import product, combinations
+from itertools import combinations
from optgs.visualization.plots3d.utils import PointCloud, Camera
from optgs.dataset.camera_datasets.camera import get_scene_scale
diff --git a/optgs/visualization/vis_depth.py b/optgs/visualization/vis_depth.py
index 30f7d22fc9bb6e440b9755f5600831991a4b465c..2c019ad731e6eebd5fdf3af408887f82bb7832e9 100644
--- a/optgs/visualization/vis_depth.py
+++ b/optgs/visualization/vis_depth.py
@@ -1,9 +1,7 @@
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
diff --git a/requirements.txt b/requirements.txt
index b19d042603e2fcfd0f928df6130da272ac1b1b19..a5a6b0f418d586dddbf5bd15626becea13241869 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -15,6 +15,7 @@ Pillow==11.1.0
plyfile==1.1
pytorch_lightning==2.4.0
viser==1.0.29
+gradio==6.19.0
rich==14.3.4
scikit-image==0.24.0
sk-video==1.1.10
@@ -28,6 +29,7 @@ pytest
natsort
typeguard
scikit-learn
+scipy
pandas
albumentations
h5py