Add OF3GS ZeroGPU demo

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/teaser.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.py[cod]
+.space_outputs/
+weights/
+outputs/
+output-debug/

LICENSE

@@ -0,0 +1,25 @@
+OF3GS Research License
+
+Copyright (c) 2026 OF3GS authors.
+
+Permission is granted to use, copy, modify, and distribute this repository for
+non-commercial research and educational purposes, subject to the conditions
+below.
+
+1. This repository includes third-party code and derived code. Those components
+   remain governed by their original licenses and copyright notices.
+2. If a third-party component imposes stricter terms than this license, the
+   stricter terms apply to that component and to derivative work that depends on
+   it.
+3. Redistribution must preserve this license and all copyright and license
+   headers present in source files.
+4. Commercial use is not granted by this license. Contact the authors and the
+   relevant third-party rights holders for commercial licensing.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE, AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES, OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT, OR OTHERWISE, ARISING FROM,
+OUT OF, OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,13 +1,351 @@
 ---
-title: Of3gs Demo
-emoji: 📚
-colorFrom: indigo
+title: OF3GS
+emoji: "🧭"
+colorFrom: blue
 colorTo: green
 sdk: gradio
 sdk_version: 6.24.0
-python_version: '3.12'
+python_version: 3.10.13
 app_file: app.py
 pinned: false
+models:
+  - richardchencccc/OF3GS
+preload_from_hub:
+  - richardchencccc/OF3GS model.ckpt
+startup_duration_timeout: 1h
+short_description: Feed-forward 3D Gaussian Splatting from unposed images
+tags:
+  - 3d
+  - gaussian-splatting
+  - 3d-reconstruction
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+<h2 align="center">
+  OF<sup>3</sup>GS: <u>O</u>n-the-<u>F</u>ly <u>F</u>eed-<u>F</u>orward 3D <u>G</u>aussian <u>S</u>platting from Unposed Images
+</h2>
+
+<p align="center">
+  <a href="https://richardchen225.github.io/of3gs/">
+    <img alt="Website" src="https://img.shields.io/badge/OF3GS-Website-2ea44f?style=flat&logo=googlechrome&logoColor=white">
+  </a>
+  <a href="https://arxiv.org/abs/2606.03254">
+    <img alt="Paper" src="https://img.shields.io/badge/arXiv-Paper-b31b1b?style=flat&logo=arxiv&logoColor=white">
+  </a>
+  <a href="https://github.com/richardchen225/OF3GS_code">
+    <img alt="Code" src="https://img.shields.io/badge/GitHub-Code-181717?style=flat&logo=github&logoColor=white">
+  </a>
+  <a href="https://huggingface.co/richardchencccc/OF3GS/tree/main">
+    <img alt="Model" src="https://img.shields.io/badge/Hugging%20Face-Model-ffcc4d?style=flat&logo=huggingface&logoColor=black">
+  </a>
+</p>
+
+<p align="center">
+  Ruiyang Chen<sup>1</sup>,
+  Feiran Li<sup>2</sup>,
+  <a href="https://fourson.github.io/">Chu Zhou</a><sup>3</sup>,
+  Zonglin Li<sup>1</sup>,
+  <a href="https://zhanyuma.cn/">Zhanyu Ma</a><sup>1</sup>,
+  <a href="https://gh-home.github.io/">Heng Guo</a><sup>1,*</sup>
+</p>
+
+<p align="center">
+  <sup>1</sup>Beijing University of Posts and Telecommunications &nbsp;&nbsp;
+  <sup>2</sup>Independent Researcher &nbsp;&nbsp;
+  <sup>3</sup>National Institute of Informatics
+</p>
+
+<p align="center">
+  <sup>*</sup>Corresponding author
+</p>
+
+<div align="center">
+<img src="assets/teaser.png" alt="OF3GS teaser" width="100%">
+</div>
+
+Official implementation of OF3GS.
+
+## Status
+
+- [x] Release training and testing code.
+- [x] Release the OF3GS checkpoint.
+- [ ] Release the evaluation code and metric scripts.
+- [ ] Add more demo assets and usage examples.
+
+## Hugging Face Space
+
+The repository includes a ZeroGPU-compatible Gradio demo in `app.py`. It accepts
+1-64 ordered, overlapping views and returns a scale-filtered Gaussian PLY with an
+interactive `Model3D` viewer. New-view rendering is intentionally not run by the
+demo. Uploaded images follow the OF3GS data pipeline: they are rescaled to cover
+the trained `252 x 518` input rectangle and center-cropped, rather than forced into
+a square. The PLY export removes Gaussians whose largest axis scale exceeds the
+selected threshold.
+
+To deploy it, create a public Gradio Space with ZeroGPU, then push this repository
+to the Space. The README metadata preloads `richardchencccc/OF3GS/model.ckpt`
+during the Space build.
+
+For local GPU development with the Space dependency set:
+
+```bash
+pip install -r requirements.txt
+python app.py
+```
+
+The demo targets Python 3.10, PyTorch 2.8, and current ZeroGPU runtimes. The
+original training environment remains in `requirements-train.txt`.
+
+## Training Environment
+
+The original training dependency set targets:
+
+- Linux
+- Python 3.10
+- CUDA 12.1
+- PyTorch 2.2.x
+- `gsplat==1.4.0` built for `pt22cu121`
+
+Create an environment and install dependencies:
+
+```bash
+conda create -n of3gs python=3.10 -y
+conda activate of3gs
+
+cd OF3GS
+pip install torch==2.2.0 torchvision==0.17.0 --index-url https://download.pytorch.org/whl/cu121
+pip install git+https://github.com/facebookresearch/pytorch3d.git --no-build-isolation
+pip install -r requirements-train.txt
+```
+
+## Download Pretrained Weights
+
+Run:
+
+```bash
+bash scripts/download_weights.sh weights
+```
+
+The script writes files to `weights/` and prints the environment variables expected by `train.sh` and `test.sh`.
+
+| File | Used For | Default Source |
+| --- | --- | --- |
+| `weights/pre_wm.safetensors` | Training warm start / GS head initialization | `hf://tencent/HunyuanWorld-Mirror/model.safetensors` |
+| `weights/pre_vggt.safetensors` | VGGT initialization during training | `hf://facebook/VGGT-1B/model.safetensors` |
+| `weights/pre_svggt.safetensors` | StreamVGGT initialization during training | `hf://lch01/StreamVGGT/model.safetensors` |
+| `weights/pre_dav3.safetensors` | Depth Anything 3 / camera decoder initialization during training | `hf://depth-anything/DA3-LARGE/model.safetensors` |
+| `weights/of3gs.ckpt` | OF3GS checkpoint for testing | `hf://richardchencccc/OF3GS/model.ckpt` |
+
+The OF3GS testing checkpoint is available at [richardchencccc/OF3GS](https://huggingface.co/richardchencccc/OF3GS/tree/main). You can override any source:
+
+```bash
+PRE_OF3GS_URL=https://huggingface.co/richardchencccc/OF3GS/resolve/main/model.ckpt \
+bash scripts/download_weights.sh weights
+```
+
+Supported source formats:
+
+```text
+https://...                direct URL
+hf://namespace/repo/file   Hugging Face Hub file
+TODO:message               skip with a reminder
+```
+
+After downloading, the default scripts use these paths automatically:
+
+```bash
+export PRE_WM_PATH=weights/pre_wm.safetensors
+export PRE_VGGT_PATH=weights/pre_vggt.safetensors
+export PRE_SVGGT_PATH=weights/pre_svggt.safetensors
+export PRE_DAV3_PATH=weights/pre_dav3.safetensors
+export PRE_OF3GS_PATH=weights/of3gs.ckpt
+```
+
+## Download DL3DV
+
+The training code supports DL3DV. The helper script downloads `DL3DV/DL3DV-ALL-480P` `images+poses` subsets from `1K` through `11K`, then generates `train_index.json` and `test_index.json`.
+
+First request access to the dataset on Hugging Face if needed, then log in:
+
+```bash
+huggingface-cli login
+```
+
+Download:
+
+```bash
+bash scripts/download_dl3dv.sh datasets/dl3dv
+```
+
+You can download to any location:
+
+```bash
+bash scripts/download_dl3dv.sh /data/DL3DV-ALL-480P
+```
+
+The expected layout after download is:
+
+```text
+datasets/dl3dv/
+  train_index.json
+  test_index.json
+  1K/
+    <scene_name>/
+      ...
+        transforms.json
+        ...
+          images_8/
+  ...
+  11K/
+    <scene_name>/
+      ...
+        transforms.json
+        ...
+          images_8/
+```
+
+`DatasetDL3DV` reads `<root>/<split>_index.json`, then recursively resolves each indexed scene's `transforms.json` and `images_8` directory. Index entries may point to a parent scene directory such as `1K/<scene_name>` even when the actual files are nested deeper, for example `1K/<scene_name>/<subdir>/<subdir>/images_8`. The download helper generates parent-level entries when possible.
+
+## RE10K Data
+
+Prepare RE10K under one root directory:
+
+```text
+re10k/
+  train.pickle.gz
+  test.pickle.gz
+  train_index.json
+  test_index.json
+  train/
+    <scene_name>/
+      000000.png
+      000001.png
+  test/
+    <scene_name>/
+      000000.png
+      000001.png
+```
+
+The index files should be JSON arrays of scene folder names, for example `["scene_000", "scene_001"]`. `Datasetre10k` reads `<root>/<split>_index.json`, metadata from `<root>/<split>.pickle.gz`, and images from `<root>/<split>/<scene_name>/`.
+
+## Training
+
+`train.sh` runs distributed training with `mode=train` and `model.encoder.mode=train` explicitly set. By default it uses the files under `weights/`. For multi-dataset training, DL3DV and RE10K use separate root directories.
+
+```bash
+GPU_NUM=4 \
+DL3DV_ROOT=datasets/dl3dv \
+RE10K_ROOT=/path/to/re10k \
+bash train.sh
+```
+
+Equivalent explicit checkpoint paths:
+
+```bash
+GPU_NUM=4 \
+PRE_WM_PATH=weights/pre_wm.safetensors \
+PRE_VGGT_PATH=weights/pre_vggt.safetensors \
+PRE_SVGGT_PATH=weights/pre_svggt.safetensors \
+PRE_DAV3_PATH=weights/pre_dav3.safetensors \
+DL3DV_ROOT=datasets/dl3dv \
+RE10K_ROOT=/path/to/re10k \
+bash train.sh
+```
+
+Single-dataset training is also supported:
+
+```bash
+EXPERIMENT=dl3dv DL3DV_ROOT=datasets/dl3dv bash train.sh
+EXPERIMENT=re10k RE10K_ROOT=/path/to/re10k bash train.sh
+```
+
+You can still override the dataset roots directly with Hydra:
+
+```bash
+bash train.sh \
+  dataset.dl3dv.roots='[datasets/dl3dv]' \
+  dataset.re10k.roots='[/path/to/re10k]'
+```
+
+Expected output:
+
+- Hydra run directory under `exp_${wandb.name}/...`
+- Lightning checkpoints under the run directory's `checkpoints/`
+- Offline W&B logs unless `wandb.mode` is changed
+
+## Testing
+
+`test.sh` runs inference/evaluation with `mode=test` and `model.encoder.mode=test` explicitly set. Test mode only loads `weights/of3gs.ckpt`; the OF3GS checkpoint contains the full model state.
+
+The current code supports evaluation-style rendering on DL3DV and RE10K. Set `EXPERIMENT=dl3dv` or `EXPERIMENT=re10k` to choose the dataset configuration. NYUv2 is not wired as a built-in dataset; use the existing DL3DV and RE10K dataset classes and configs as references if you want to add a local NYUv2 adapter.
+
+DL3DV example:
+
+```bash
+EXPERIMENT=dl3dv \
+GPU_NUM=4 \
+bash test.sh \
+  dataset.dl3dv.roots='[datasets/dl3dv]' \
+  dataset.dl3dv.mode=test \
+  dataset.dl3dv.ctx_list='[0,8,16]' \
+  dataset.dl3dv.tgt_list='[4,12,20]'
+```
+
+RE10K example:
+
+```bash
+EXPERIMENT=re10k \
+GPU_NUM=4 \
+bash test.sh \
+  dataset.re10k.roots='[/path/to/re10k]'
+```
+
+Equivalent explicit checkpoint paths:
+
+```bash
+EXPERIMENT=dl3dv \
+GPU_NUM=4 \
+PRE_OF3GS_PATH=weights/of3gs.ckpt \
+TEST_OUTPUT_PATH=outputs/of3gs \
+bash test.sh \
+  dataset.dl3dv.roots='[datasets/dl3dv]' \
+  dataset.dl3dv.mode=test \
+  dataset.dl3dv.ctx_list='[0,8,16]' \
+  dataset.dl3dv.tgt_list='[4,12,20]'
+```
+
+Enable PSNR/SSIM/LPIPS during testing:
+
+```bash
+EXPERIMENT=dl3dv COMPUTE_SCORES=true GPU_NUM=4 bash test.sh \
+  dataset.dl3dv.roots='[datasets/dl3dv]' \
+  dataset.dl3dv.mode=test \
+  dataset.dl3dv.ctx_list='[0,8,16]' \
+  dataset.dl3dv.tgt_list='[4,12,20]'
+```
+
+Expected output:
+
+- Rendered target images under `${TEST_OUTPUT_PATH:-outputs/of3gs}/${wandb.name}/<scene>/color/`, or `<psnr>_<scene>/color/` when scores are enabled
+- Optional PSNR/SSIM/LPIPS logging when `test.compute_scores=true`
+- Benchmark timing summary printed at test end
+
+## License
+
+See `LICENSE`.
+
+## Citation
+
+If you find OF3GS useful, please cite:
+
+```bibtex
+@article{chen2026of3gs,
+  title={OF3GS: On-the-Fly Feed-Forward 3D Gaussian Splatting from Unposed Images},
+  author={Chen, Ruiyang and Li, Feiran and Zhou, Chu and Li, Zonglin and Ma, Zhanyu and Guo, Heng},
+  journal={arXiv preprint arXiv:2606.03254},
+  year={2026}
+}
+```
+
+## Acknowledgements
+
+We thank the authors of [VGGT](https://github.com/facebookresearch/vggt), [StreamVGGT](https://github.com/wzzheng/StreamVGGT), and [AnySplat](https://github.com/InternRobotics/AnySplat) for their excellent work and open-source contributions.

app.py

@@ -0,0 +1,328 @@
+import gc
+import os
+import re
+import shutil
+from pathlib import Path
+from typing import Any
+
+import spaces
+import torch
+
+import gradio as gr
+import numpy as np
+from huggingface_hub import hf_hub_download
+from PIL import Image
+from torchvision.transforms.functional import to_tensor
+
+from src.model.decoder.decoder_splatting_cuda import DecoderSplattingCUDACfg
+from src.model.encoder.common.gaussian_adapter import GaussianAdapterCfg
+from src.model.encoder.of3gs import EncoderOF3GSCfg, OpacityMappingCfg
+from src.model.encoder.visualization.encoder_visualizer_epipolar_cfg import (
+    EncoderVisualizerEpipolarCfg,
+)
+from src.model.model.of3gs import OF3GS
+from src.model.ply_export import export_ply
+
+
+MODEL_REPO = os.getenv("OF3GS_MODEL_REPO", "richardchencccc/OF3GS")
+MODEL_FILE = os.getenv("OF3GS_MODEL_FILE", "model.ckpt")
+# These are the dimensions used by the OF3GS training experiments
+# (config/experiment/{dl3dv,re10k}.yaml), not the standalone VGGT helpers.
+INPUT_HEIGHT = 252
+INPUT_WIDTH = 518
+MIN_VIEWS = 1
+MAX_VIEWS = 64
+OUTPUT_ROOT = (Path.cwd() / ".space_outputs").resolve()
+OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)
+
+
+def _build_model() -> OF3GS:
+    encoder_cfg = EncoderOF3GSCfg(
+        name="of3gs",
+        anchor_feat_dim=128,
+        voxel_size=0.002,
+        n_offsets=2,
+        d_feature=32,
+        add_view=False,
+        num_monocular_samples=32,
+        backbone=None,
+        visualizer=EncoderVisualizerEpipolarCfg(
+            num_samples=8,
+            min_resolution=256,
+            export_ply=False,
+        ),
+        gaussian_adapter=GaussianAdapterCfg(
+            gaussian_scale_min=0.5,
+            gaussian_scale_max=15.0,
+            sh_degree=4,
+        ),
+        apply_bounds_shim=True,
+        opacity_mapping=OpacityMappingCfg(initial=0.0, final=0.0, warm_up=1),
+        gaussians_per_pixel=1,
+        num_surfaces=1,
+        gs_params_head_type="dpt_gs",
+        pose_free=True,
+        pred_pose=True,
+        gs_prune=False,
+        pred_head_type="depth",
+        freeze_backbone=True,
+        freeze_module="patch_embed",
+        distill=True,
+        render_conf=True,
+        conf_threshold=0.05,
+        voxelize=True,
+        intermediate_layer_idx=[4, 11, 17, 23],
+        mode="test",
+    )
+    decoder_cfg = DecoderSplattingCUDACfg(
+        name="splatting_cuda",
+        background_color=[1.0, 1.0, 1.0],
+        make_scale_invariant=False,
+    )
+    return OF3GS(encoder_cfg, decoder_cfg)
+
+
+def _checkpoint_state(path: str) -> dict[str, torch.Tensor]:
+    try:
+        checkpoint = torch.load(
+            path,
+            map_location="cpu",
+            weights_only=True,
+            mmap=True,
+        )
+    except TypeError:
+        checkpoint = torch.load(path, map_location="cpu")
+
+    if "state_dict" in checkpoint:
+        checkpoint = checkpoint["state_dict"]
+    elif "model" in checkpoint:
+        checkpoint = checkpoint["model"]
+
+    if any(key.startswith("model.") for key in checkpoint):
+        checkpoint = {
+            key.removeprefix("model."): value
+            for key, value in checkpoint.items()
+            if key.startswith("model.")
+        }
+    return {
+        key.removeprefix("module."): value for key, value in checkpoint.items()
+    }
+
+
+def _load_model() -> tuple[OF3GS, torch.device]:
+    checkpoint_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
+    model = _build_model()
+    state = _checkpoint_state(checkpoint_path)
+    try:
+        incompatible = model.load_state_dict(state, strict=False, assign=True)
+    except TypeError:
+        incompatible = model.load_state_dict(state, strict=False)
+
+    if incompatible.missing_keys:
+        print(f"OF3GS checkpoint missing {len(incompatible.missing_keys)} keys")
+    if incompatible.unexpected_keys:
+        print(f"OF3GS checkpoint has {len(incompatible.unexpected_keys)} extra keys")
+
+    del state
+    gc.collect()
+
+    running_on_space = bool(os.getenv("SPACE_ID"))
+    device = torch.device(
+        "cuda" if torch.cuda.is_available() or running_on_space else "cpu"
+    )
+    model = model.to(device).eval()
+    model.requires_grad_(False)
+    return model, device
+
+
+MODEL, DEVICE = _load_model()
+
+
+def _file_path(file_data: Any) -> str:
+    if isinstance(file_data, str):
+        return file_data
+    if isinstance(file_data, dict):
+        return file_data.get("path") or file_data.get("name")
+    return str(getattr(file_data, "path", getattr(file_data, "name", file_data)))
+
+
+def _image_tensor(image_source: str | np.ndarray | Image.Image) -> torch.Tensor:
+    if isinstance(image_source, (str, Path)):
+        # Uploaded files are temporary; load pixels before closing the file handle.
+        with Image.open(image_source) as opened_image:
+            image = opened_image.copy()
+    elif isinstance(image_source, np.ndarray):
+        image = Image.fromarray(image_source)
+    else:
+        image = image_source.copy()
+
+    image = image.convert("RGB")
+    width, height = image.size
+    # Match dataset.shims.crop_shim.rescale_and_crop(intr_aug=False): scale
+    # until the requested rectangle is covered, then center-crop to 252x518.
+    scale_factor = max(INPUT_HEIGHT / height, INPUT_WIDTH / width)
+    scaled_height = round(height * scale_factor)
+    scaled_width = round(width * scale_factor)
+    image = image.resize((scaled_width, scaled_height), Image.Resampling.LANCZOS)
+    left = (scaled_width - INPUT_WIDTH) // 2
+    top = (scaled_height - INPUT_HEIGHT) // 2
+    image = image.crop((left, top, left + INPUT_WIDTH, top + INPUT_HEIGHT))
+    # This is equivalent to the dataset's torchvision ToTensor conversion.
+    return to_tensor(image).contiguous()
+
+
+def _prepare_images(files: list[Any] | None) -> torch.Tensor:
+    if not files:
+        raise gr.Error("Upload at least one image.")
+    if isinstance(files, (str, dict)):
+        files = [files]
+    if not MIN_VIEWS <= len(files) <= MAX_VIEWS:
+        raise gr.Error(f"Select between {MIN_VIEWS} and {MAX_VIEWS} images.")
+
+    tensors = [_image_tensor(_file_path(file_data)) for file_data in files]
+    return torch.stack(tensors).unsqueeze(0)
+
+
+def _session_directory(request: gr.Request | None) -> Path:
+    session_hash = request.session_hash if request and request.session_hash else "local"
+    safe_hash = re.sub(r"[^a-zA-Z0-9_-]", "_", session_hash)
+    directory = (OUTPUT_ROOT / safe_hash).resolve()
+    if directory.parent != OUTPUT_ROOT:
+        raise RuntimeError("Invalid output directory")
+    if directory.exists():
+        shutil.rmtree(directory)
+    directory.mkdir(parents=True)
+    return directory
+
+
+def _run_reconstruction(
+    images: torch.Tensor,
+    scale_threshold: float,
+    request: gr.Request | None,
+    progress: gr.Progress,
+):
+    if DEVICE.type != "cuda":
+        raise gr.Error("OF3GS inference requires a CUDA GPU Space.")
+
+    output_directory = _session_directory(request)
+    progress(0.1, desc="Preparing views")
+    images = images.to(DEVICE, non_blocking=True)
+    # The encoder uses target views only for optional pose/render evaluation.
+    # Keeping this empty makes every uploaded image a context view and stops at GS output.
+    target_images = images[:, :0]
+    gaussians = scale = keep_mask = harmonics = None
+
+    try:
+        progress(0.2, desc="Reconstructing Gaussian scene")
+        with torch.no_grad():
+            gaussians, _, _ = MODEL.encoder.inference(
+                images,
+                target_images,
+                global_step=0,
+            )
+
+        progress(0.8, desc="Filtering and writing Gaussians")
+        ply_path = output_directory / "of3gs_gaussians.ply"
+        scale = gaussians.scales[0]
+        keep_mask = scale.max(dim=-1).values <= scale_threshold
+        kept_count = int(keep_mask.sum().item())
+        if kept_count == 0:
+            raise gr.Error("The scale threshold removed every Gaussian.")
+        harmonics = gaussians.harmonics[0]
+        export_ply(
+            gaussians.means[0],
+            gaussians.scales[0],
+            gaussians.rotations[0],
+            harmonics,
+            gaussians.opacities[0],
+            ply_path,
+            scale_threshold=float(scale_threshold),
+            save_sh_dc_only=True,
+        )
+        progress(1.0, desc="Complete")
+
+        status = f"{images.shape[1]} views | {kept_count:,} Gaussians kept"
+        return str(ply_path), str(ply_path), status
+    finally:
+        del images, target_images, gaussians, scale, keep_mask, harmonics
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+@spaces.GPU(duration=300)
+def reconstruct_images(
+    files: list[Any] | None,
+    scale_threshold: float,
+    request: gr.Request,
+    progress: gr.Progress = gr.Progress(track_tqdm=False),
+) -> tuple[str, str, str]:
+    """Reconstruct an interactive Gaussian PLY from 1-64 ordered images."""
+    images = _prepare_images(files)
+    return _run_reconstruction(
+        images,
+        float(scale_threshold),
+        request,
+        progress,
+    )
+
+
+def cleanup_session(request: gr.Request) -> None:
+    """Remove temporary PLY files when a browser session is closed."""
+    if not request.session_hash:
+        return
+    safe_hash = re.sub(r"[^a-zA-Z0-9_-]", "_", request.session_hash)
+    directory = (OUTPUT_ROOT / safe_hash).resolve()
+    if directory.parent == OUTPUT_ROOT:
+        shutil.rmtree(directory, ignore_errors=True)
+
+
+CSS = """
+.gradio-container { max-width: 1180px !important; }
+#project-title { margin-bottom: 0.25rem; }
+#project-links { color: var(--body-text-color-subdued); margin-bottom: 1rem; }
+"""
+
+with gr.Blocks(title="OF3GS Demo", css=CSS) as demo:
+    gr.Markdown("# OF3GS", elem_id="project-title")
+    gr.Markdown(
+        "On-the-Fly Feed-Forward 3D Gaussian Splatting from Unposed Images | "
+        "[Project](https://richardchen225.github.io/of3gs/) | "
+        "[Paper](https://arxiv.org/abs/2606.03254) | "
+        "[Code](https://github.com/richardchen225/OF3GS_code)",
+        elem_id="project-links",
+    )
+
+    with gr.Row(equal_height=False):
+        with gr.Column(scale=5):
+            image_files = gr.File(
+                label="Ordered views (1-64 images)",
+                file_count="multiple",
+                file_types=["image"],
+                type="filepath",
+            )
+            scale_threshold = gr.Slider(
+                minimum=0.005,
+                maximum=0.30,
+                value=0.08,
+                step=0.005,
+                label="Maximum Gaussian scale",
+            )
+            image_submit = gr.Button("Reconstruct", variant="primary")
+
+        with gr.Column(scale=7):
+            reconstruction = gr.Model3D(label="Gaussian scene", height=520)
+            ply_download = gr.File(label="Download PLY")
+            status = gr.Textbox(label="Result", interactive=False)
+
+    outputs = [reconstruction, ply_download, status]
+    image_submit.click(
+        fn=reconstruct_images,
+        inputs=[image_files, scale_threshold],
+        outputs=outputs,
+    )
+    demo.unload(cleanup_session)
+
+demo.queue()
+
+if __name__ == "__main__":
+    demo.launch(show_error=True, mcp_server=True)

config/compute_metrics.yaml

@@ -0,0 +1,28 @@
+defaults:
+  - model/encoder: noposplat
+  - loss: []
+  - override dataset/view_sampler@dataset.re10k.view_sampler: evaluation
+
+dataset:
+  re10k:
+    view_sampler:
+      index_path: assets/evaluation_index_re10k.json
+
+data_loader:
+  train:
+    num_workers: 0
+    persistent_workers: true
+    batch_size: 1
+    seed: 1234
+  test:
+    num_workers: 4
+    persistent_workers: false
+    batch_size: 1
+    seed: 2345
+  val:
+    num_workers: 0
+    persistent_workers: true
+    batch_size: 1
+    seed: 3456
+
+seed: 111123

config/dataset/base_dataset.yaml

@@ -0,0 +1,7 @@
+make_baseline_1: true
+relative_pose: true
+augment: true
+background_color: [1.0, 1.0, 1.0]
+overfit_to_scene: null
+skip_bad_shape: true
+rescale_to_1cube: false

config/dataset/dl3dv.yaml

@@ -0,0 +1,19 @@
+defaults:
+  - base_dataset
+  - view_sampler: bounded
+
+name: dl3dv
+roots: []
+
+input_image_shape: [256, 256]
+original_image_shape: [270, 480]
+cameras_are_circular: false
+
+baseline_min: 1e-3
+baseline_max: 1e2
+max_fov: 100.0
+avg_pose: false
+
+rescale_to_1cube: true
+make_baseline_1: false
+intr_augment: true

config/dataset/re10k.yaml

@@ -0,0 +1,19 @@
+defaults:
+  - base_dataset
+  - view_sampler: bounded
+
+name: re10k
+roots: []
+
+input_image_shape: [256, 256]
+original_image_shape: [270, 480]
+cameras_are_circular: false
+
+baseline_min: 1e-3
+baseline_max: 1e2
+max_fov: 100.0
+avg_pose: false
+
+rescale_to_1cube: true
+make_baseline_1: false
+intr_augment: true

config/dataset/view_sampler/all.yaml

@@ -0,0 +1 @@
+name: all

config/dataset/view_sampler/arbitrary.yaml

@@ -0,0 +1,7 @@
+name: arbitrary
+
+num_target_views: 1
+num_context_views: 2
+
+# If you want to hard-code context views, do so here.
+context_views: null

config/dataset/view_sampler/bounded.yaml

@@ -0,0 +1,16 @@
+name: bounded
+
+num_target_views: 2
+num_context_views: 16
+
+min_distance_between_context_views: 2
+max_distance_between_context_views: 6
+min_distance_to_context_views: 0
+
+warm_up_steps: 0
+initial_min_distance_between_context_views: 2
+initial_max_distance_between_context_views: 6
+
+max_img_per_gpu: 16
+min_gap_multiplier: 3
+max_gap_multiplier: 5

config/dataset/view_sampler/evaluation.yaml

@@ -0,0 +1,4 @@
+name: evaluation
+
+index_path: assets/evaluation_index_re10k.json
+num_context_views: 2

config/dataset/view_sampler/rank.yaml

@@ -0,0 +1,14 @@
+name: rank
+
+num_target_views: 4
+num_context_views: 16
+
+min_distance_between_context_views: 4
+max_distance_between_context_views: 22
+min_distance_to_context_views: 0
+
+warm_up_steps: 0
+initial_min_distance_between_context_views: 5
+initial_max_distance_between_context_views: 7
+
+max_img_per_gpu: 16

config/experiment/dl3dv.yaml

@@ -0,0 +1,87 @@
+# @package _global_
+
+defaults:
+  - /dataset@_group_.dl3dv: dl3dv
+  - override /model/encoder: of3gs
+  - override /model/encoder/backbone: croco
+  - override /loss: [mse, lpips, ssim, depth]
+
+wandb:
+  name: dl3dv
+  tags: [dl3dv, 448x448]
+  
+model:
+  encoder:
+    mode: train
+    pre_vggt_path:  ""
+    pre_svggt_path: ""
+    pre_dav3_path: ""
+    gs_params_head_type: dpt_gs
+    pose_free: true
+    intrinsics_embed_loc: encoder
+    intrinsics_embed_type: token
+    pretrained_weights: ''
+    voxel_size: 0.002
+    pred_pose: true
+    anchor_feat_dim: 128
+    gs_prune: false 
+    pred_head_type: depth
+    freeze_backbone: true
+    distill: true
+    render_conf: true
+    conf_threshold: 0.05
+    freeze_module: patch_embed
+    voxelize: true
+    intermediate_layer_idx: [4, 11, 17, 23]
+    
+dataset:
+  dl3dv:
+    mode: train
+    ctx_list: []
+    tgt_list: []
+    input_image_shape: [252, 518]
+    view_sampler:
+      num_target_views: 9
+      min_distance_between_context_views: 32
+      max_distance_between_context_views: 256
+      min_gap_multiplier: 3
+      max_gap_multiplier: 5
+    avg_pose: false
+    intr_augment: true
+    normalize_by_pts3d: false
+    rescale_to_1cube: false
+
+optimizer:
+  lr: 2e-4
+  warm_up_steps: 1000
+  backbone_lr_multiplier: 0.1
+
+data_loader:
+  train:
+    batch_size: 1 
+    
+trainer:
+  max_steps: 10000
+  val_check_interval: 100
+  num_nodes: 1
+  accumulate_grad_batches: 1
+  precision: 16-mixed
+
+checkpointing:
+  load: null
+  every_n_train_steps: 500
+  save_weights_only: true
+  save_top_k: 10
+  
+train:
+  pose_loss_alpha: 1.0
+  pose_loss_delta: 1.0
+  cxt_depth_weight: 1.0
+  weight_pose: 10.0
+  weight_depth: 1.0
+  weight_normal: 0.1
+
+hydra:
+  run:
+    dir: ./exp_${wandb.name}/${now:%Y-%m-%d_%H-%M-%S}
+

config/experiment/multi-dataset.yaml

@@ -0,0 +1,102 @@
+# @package _global_
+
+defaults:
+  - /dataset@_group_.dl3dv: dl3dv
+  - /dataset@_group_.re10k: re10k
+  - override /model/encoder: of3gs
+  - override /model/encoder/backbone: croco
+  - override /loss: [mse, lpips, ssim, depth]
+
+wandb:
+  name: multidataset-16gpu
+  tags: [multidataset, 448x448]
+  
+model:
+  encoder:
+    mode: train
+    pre_vggt_path: "" 
+    pre_svggt_path: ""
+    pre_dav3_path: ""
+    gs_params_head_type: dpt_gs
+    pose_free: true
+    intrinsics_embed_loc: encoder
+    intrinsics_embed_type: token
+    pretrained_weights: ''
+    voxel_size: 0.002
+    pred_pose: true
+    anchor_feat_dim: 128
+    gs_prune: false 
+    pred_head_type: depth
+    freeze_backbone: true
+    distill: true
+    render_conf: true
+    conf_threshold: 0.05
+    freeze_module: patch_embed
+    voxelize: true
+    intermediate_layer_idx: [4, 11, 17, 23]
+    
+dataset:
+  dl3dv:
+    mode: train
+    ctx_list: []
+    tgt_list: []
+    input_image_shape: [252, 518]
+    view_sampler:
+      num_target_views: 9
+      min_distance_between_context_views: 32
+      max_distance_between_context_views: 256
+      min_gap_multiplier: 3
+      max_gap_multiplier: 5
+    avg_pose: false
+    intr_augment: true
+    normalize_by_pts3d: false
+    rescale_to_1cube: false
+
+  re10k:
+    input_image_shape: [252, 518]
+    view_sampler:
+      num_target_views: 9
+      min_distance_between_context_views: 16
+      max_distance_between_context_views: 128
+      min_gap_multiplier: 3
+      max_gap_multiplier: 5
+      max_img_per_gpu: 10
+    avg_pose: false
+    intr_augment: true
+    normalize_by_pts3d: false
+    rescale_to_1cube: true
+
+optimizer:
+  lr: 1e-4
+  warm_up_steps: 0
+  backbone_lr_multiplier: 0.1
+
+data_loader:
+  train:
+    batch_size: 1 
+    
+trainer:
+  max_steps: 10000
+  val_check_interval: 100
+  num_nodes: 1
+  accumulate_grad_batches: 1
+  precision: 16-mixed
+
+checkpointing:
+  load: null 
+  every_n_train_steps: 500
+  save_weights_only: true
+  save_top_k: 10
+  
+train:
+  pose_loss_alpha: 1.0
+  pose_loss_delta: 1.0
+  cxt_depth_weight: 1.0
+  weight_pose: 10.0
+  weight_depth: 1.0
+  weight_normal: 0.1
+
+hydra:
+  run:
+    dir: ./exp_${wandb.name}/${now:%Y-%m-%d_%H-%M-%S}
+

config/experiment/re10k.yaml

@@ -0,0 +1,84 @@
+# @package _global_
+
+defaults:
+  - /dataset@_group_.re10k: re10k
+  - override /model/encoder: of3gs
+  - override /model/encoder/backbone: croco
+  - override /loss: [mse, lpips, ssim, depth]
+
+wandb:
+  name: re10k
+  tags: [re10k, 448x448]
+
+model:
+  encoder:
+    mode: train
+    pre_vggt_path: ""
+    pre_svggt_path: ""
+    pre_dav3_path: ""
+    gs_params_head_type: dpt_gs
+    pose_free: true
+    intrinsics_embed_loc: encoder
+    intrinsics_embed_type: token
+    pretrained_weights: ''
+    voxel_size: 0.002
+    pred_pose: true
+    anchor_feat_dim: 128
+    gs_prune: false
+    pred_head_type: depth
+    freeze_backbone: true
+    distill: true
+    render_conf: true
+    conf_threshold: 0.05
+    freeze_module: patch_embed
+    voxelize: true
+    intermediate_layer_idx: [4, 11, 17, 23]
+
+dataset:
+  re10k:
+    input_image_shape: [252, 518]
+    view_sampler:
+      num_target_views: 9
+      min_distance_between_context_views: 16
+      max_distance_between_context_views: 128
+      min_gap_multiplier: 3
+      max_gap_multiplier: 5
+      max_img_per_gpu: 10
+    avg_pose: false
+    intr_augment: true
+    normalize_by_pts3d: false
+    rescale_to_1cube: true
+
+optimizer:
+  lr: 2e-4
+  warm_up_steps: 1000
+  backbone_lr_multiplier: 0.1
+
+data_loader:
+  train:
+    batch_size: 1
+
+trainer:
+  max_steps: 10000
+  val_check_interval: 100
+  num_nodes: 1
+  accumulate_grad_batches: 1
+  precision: 16-mixed
+
+checkpointing:
+  load: null
+  every_n_train_steps: 500
+  save_weights_only: true
+  save_top_k: 10
+
+train:
+  pose_loss_alpha: 1.0
+  pose_loss_delta: 1.0
+  cxt_depth_weight: 1.0
+  weight_pose: 10.0
+  weight_depth: 1.0
+  weight_normal: 0.1
+
+hydra:
+  run:
+    dir: ./exp_${wandb.name}/${now:%Y-%m-%d_%H-%M-%S}

config/generate_evaluation_index.yaml

@@ -0,0 +1,36 @@
+defaults:
+  - dataset: re10k
+  - optional dataset/view_sampler_dataset_specific_config: ${dataset/view_sampler}_${dataset}
+  - override dataset/view_sampler: all
+
+dataset:
+  overfit_to_scene: null
+
+data_loader:
+  train:
+    num_workers: 0
+    persistent_workers: true
+    batch_size: 1
+    seed: 1234
+  test:
+    num_workers: 8
+    persistent_workers: false
+    batch_size: 1
+    seed: 2345
+  val:
+    num_workers: 0
+    persistent_workers: true
+    batch_size: 1
+    seed: 3456
+
+index_generator:
+  num_target_views: 3
+  min_overlap: 0.6
+  max_overlap: 1.0
+  min_distance: 45
+  max_distance: 135
+  output_path: outputs/evaluation_index_re10k
+  save_previews: false
+  seed: 123
+
+seed: 456

config/loss/depth.yaml

@@ -0,0 +1,5 @@
+depth:
+  weight: 1.0
+  sigma_image: null
+  use_second_derivative: false
+  weights_path: ${model.encoder.pre_dav3_path}

config/loss/lpips.yaml

@@ -0,0 +1,3 @@
+lpips:
+  weight: 0.05
+  apply_after_step: 0

config/loss/mse.yaml

@@ -0,0 +1,3 @@
+mse:
+  weight_ctx: 1.0
+  weight_novel: 1.5

config/loss/ssim.yaml

@@ -0,0 +1,3 @@
+ssim:
+  weight: 0.05
+  apply_after_step: 0

config/main.yaml

@@ -0,0 +1,78 @@
+defaults:
+  - model/encoder: of3gs
+  - model/decoder: splatting_cuda
+  - loss: [mse]
+  
+wandb:
+  project: of3gs
+  entity: null
+  name: debug
+  mode: offline
+mode: test
+pre_wm_path: ""
+pre_of3gs_path: "" 
+
+data_loader:
+  train:
+    num_workers: 8 # 16
+    persistent_workers: false
+    batch_size: 1
+    seed: 1234
+  test:
+    num_workers: 1
+    persistent_workers: false
+    batch_size: 1
+    seed: 1234
+  val:
+    num_workers: 1
+    persistent_workers: true
+    batch_size: 1
+    seed: 1234
+
+optimizer:
+  lr: 1.5e-4
+  warm_up_steps: 2000
+  backbone_lr_multiplier: 0.1
+
+checkpointing:
+  load: null
+  every_n_train_steps: 5000
+  save_top_k: 1
+  save_weights_only: true
+
+train:
+  output_path: outputs/train
+  depth_mode: null
+  extended_visualization: false
+  print_log_every_n_steps: 10
+  distiller: ''
+  distill_max_steps: 1000000
+  random_context_views: false
+  
+test:
+  output_path: outputs/me_suppl_64
+  align_pose: false
+  pose_align_steps: 100
+  rot_opt_lr:  0.005
+  trans_opt_lr:  0.005
+  compute_scores: false
+  save_image: true
+  save_video: true
+  save_compare: true
+  generate_video: true
+  mode: inference
+  image_folder: examples/bungeenerf
+
+seed: 111123
+
+trainer:
+  max_steps: -1
+  val_check_interval: 250
+  gradient_clip_val: 0.5
+  num_nodes: 1
+  devices: 4
+  accumulate_grad_batches: 1
+
+hydra:
+  run:
+    dir: output-debug/exp_${wandb.name}/${now:%Y-%m-%d_%H-%M-%S}

config/model/decoder/splatting_cuda.yaml

@@ -0,0 +1,3 @@
+name: splatting_cuda
+background_color: [1.0, 1.0, 1.0]
+make_scale_invariant: false

config/model/encoder/backbone/croco.yaml

@@ -0,0 +1,9 @@
+name: croco
+
+model: ViTLarge_BaseDecoder
+patch_embed_cls: PatchEmbedDust3R
+asymmetry_decoder: true
+
+intrinsics_embed_loc: 'encoder'
+intrinsics_embed_degree: 4
+intrinsics_embed_type: 'token'

config/model/encoder/of3gs.yaml

@@ -0,0 +1,42 @@
+defaults:
+  - backbone: croco
+
+name: of3gs
+
+opacity_mapping:
+  initial: 0.0
+  final: 0.0
+  warm_up: 1
+
+num_monocular_samples: 32
+num_surfaces: 1
+predict_opacity: false
+
+gaussians_per_pixel: 1
+
+gaussian_adapter:
+  gaussian_scale_min: 0.5
+  gaussian_scale_max: 15.0
+  sh_degree: 4
+
+d_feature: 32
+
+visualizer:
+  num_samples: 8
+  min_resolution: 256
+  export_ply: false
+
+apply_bounds_shim: true
+
+gs_params_head_type: dpt_gs
+pose_free: true
+pretrained_weights: ""
+scale_align: false
+
+voxel_size: 0.001
+n_offsets: 2
+anchor_feat_dim: 83 # 32
+add_view: false
+color_attr: 3D # 3D or RGB
+mlp_type: unified
+scaffold: true

requirements-train.txt

@@ -0,0 +1,34 @@
+numpy==1.25.0
+wheel
+tqdm
+lightning
+black
+ruff
+hydra-core
+jaxtyping
+beartype
+wandb
+einops
+colorama
+scikit-image
+colorspacious
+matplotlib
+moviepy==1.0.3
+imageio
+timm
+dacite
+lpips
+e3nn
+plyfile
+tabulate
+svg.py
+scikit-video
+opencv-python
+Pillow
+huggingface_hub
+gradio
+xformers==0.0.24
+pydantic
+open3d
+safetensors
+https://github.com/nerfstudio-project/gsplat/releases/download/v1.4.0/gsplat-1.4.0%2Bpt22cu121-cp310-cp310-linux_x86_64.whl

requirements.txt

@@ -0,0 +1,16 @@
+--extra-index-url https://download.pytorch.org/whl/cu128
+
+torch==2.8.0
+torchvision==0.23.0
+transformers==4.57.6
+gsplat==1.5.3
+ninja
+numpy==1.26.4
+einops
+jaxtyping
+safetensors
+e3nn
+scipy
+plyfile
+Pillow
+rich>=12

scripts/download_dl3dv.sh

@@ -0,0 +1,146 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+usage() {
+  cat <<'USAGE'
+Usage:
+  bash scripts/download_dl3dv.sh <output_dir>
+
+Examples:
+  bash scripts/download_dl3dv.sh datasets/dl3dv
+  bash scripts/download_dl3dv.sh /data/DL3DV-ALL-480P
+
+This script downloads DL3DV/DL3DV-ALL-480P images+poses from 1K through 11K.
+You must request dataset access on Hugging Face first and run `huggingface-cli login`
+if the dataset is gated for your account.
+USAGE
+}
+
+if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
+  usage
+  exit 0
+fi
+
+if [[ $# -ne 1 ]]; then
+  echo "Missing output directory." >&2
+  usage
+  exit 1
+fi
+
+OUTPUT_DIR="$1"
+RESOLUTION="${DL3DV_RESOLUTION:-480P}"
+FILE_TYPE="${DL3DV_FILE_TYPE:-images+poses}"
+SUBSETS=(${DL3DV_SUBSETS:-1K 2K 3K 4K 5K 6K 7K 8K 9K 10K 11K})
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
+TEST_INDEX_SOURCE="${DL3DV_TEST_INDEX_PATH:-$REPO_ROOT/test_index.json}"
+
+mkdir -p "$OUTPUT_DIR"
+
+python -m pip install -U huggingface_hub pandas tqdm
+
+DOWNLOAD_SCRIPT="$OUTPUT_DIR/dl3dv_download.py"
+RAW_URL="https://raw.githubusercontent.com/DL3DV-10K/Dataset/main/scripts/download.py"
+
+if command -v wget >/dev/null 2>&1; then
+  wget -O "$DOWNLOAD_SCRIPT" "$RAW_URL"
+elif command -v curl >/dev/null 2>&1; then
+  curl -L "$RAW_URL" -o "$DOWNLOAD_SCRIPT"
+else
+  echo "wget or curl is required to fetch the official DL3DV download script." >&2
+  exit 1
+fi
+
+for subset in "${SUBSETS[@]}"; do
+  echo "Downloading DL3DV/DL3DV-ALL-480P subset $subset to $OUTPUT_DIR"
+  python "$DOWNLOAD_SCRIPT" \
+    --odir "$OUTPUT_DIR" \
+    --subset "$subset" \
+    --resolution "$RESOLUTION" \
+    --file_type "$FILE_TYPE" \
+    --clean_cache
+done
+
+python - "$OUTPUT_DIR" "$TEST_INDEX_SOURCE" <<'PY'
+import json
+import shutil
+import sys
+from pathlib import Path
+
+root = Path(sys.argv[1]).resolve()
+test_index_source = Path(sys.argv[2]).resolve()
+subsets = {f"{idx}K" for idx in range(1, 12)}
+
+test_items = set()
+if test_index_source.exists():
+    with test_index_source.open("r", encoding="utf-8") as f:
+        loaded = json.load(f)
+    if not isinstance(loaded, list):
+        raise TypeError(f"{test_index_source} must be a JSON array")
+    test_items = {str(item).strip("/") for item in loaded}
+
+scenes = []
+for transforms in sorted(root.rglob("transforms.json")):
+    scene_dir = transforms.parent
+    if not (scene_dir / "images_8").is_dir():
+        continue
+
+    rel = scene_dir.relative_to(root).as_posix()
+    rel_parts = rel.split("/")
+    subset_scene_key = (
+        "/".join(rel_parts[:2])
+        if len(rel_parts) >= 2 and rel_parts[0] in subsets
+        else rel
+    )
+    scenes.append((rel, subset_scene_key))
+
+seen = set()
+unique_scenes = []
+for rel, subset_scene_key in sorted(scenes):
+    if rel in seen:
+        continue
+    seen.add(rel)
+    unique_scenes.append((rel, subset_scene_key))
+
+test_index_items = sorted(
+    rel for rel, subset_scene_key in unique_scenes
+    if rel in test_items or subset_scene_key in test_items
+)
+test_index_set = set(test_index_items)
+train_items = sorted(rel for rel, _ in unique_scenes if rel not in test_index_set)
+index_path = root / "train_index.json"
+test_index_path = root / "test_index.json"
+
+if train_items:
+    index_path.write_text(json.dumps(train_items, indent=2) + "\n", encoding="utf-8")
+    print(f"Wrote {index_path} with {len(train_items)} scenes.")
+else:
+    print(
+        "No DL3DV scenes were found for train_index.json. "
+        "Check that downloaded scenes contain images_8 and transforms.json.",
+        file=sys.stderr,
+    )
+
+if test_items:
+    test_index_path.write_text(json.dumps(test_index_items, indent=2) + "\n", encoding="utf-8")
+    print(
+        f"Wrote {test_index_path} with {len(test_index_items)} scenes "
+        f"matched from {test_index_source}."
+    )
+    missing = sorted(test_items - {rel for rel, _ in unique_scenes} - {key for _, key in unique_scenes})
+    if missing:
+        print(
+            f"Warning: {len(missing)} test-index entries were not found under {root}.",
+            file=sys.stderr,
+        )
+elif test_index_source.exists():
+    shutil.copy2(test_index_source, test_index_path)
+    print(f"Copied empty test index {test_index_source} to {test_index_path}.")
+else:
+    print(
+        f"No test index source found at {test_index_source}; "
+        "only train_index.json was generated."
+    )
+PY
+
+echo "DL3DV/DL3DV-ALL-480P download finished under $OUTPUT_DIR"

scripts/download_weights.sh

@@ -0,0 +1,127 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+usage() {
+  cat <<'USAGE'
+Usage:
+  bash scripts/download_weights.sh [weights_dir]
+
+Examples:
+  bash scripts/download_weights.sh
+  bash scripts/download_weights.sh /data/of3gs_weights
+
+Environment overrides:
+  PRE_WM_URL       Source for pre_wm.safetensors
+  PRE_VGGT_URL     Source for pre_vggt.safetensors
+  PRE_SVGGT_URL    Source for pre_svggt.safetensors
+  PRE_DAV3_URL     Source for pre_dav3.safetensors
+  PRE_OF3GS_URL   Source for of3gs.ckpt
+
+Supported source formats:
+  https://...                         direct URL
+  hf://namespace/repo/path/in/repo    Hugging Face file
+  provide                             skip until a source is provided
+USAGE
+}
+
+if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
+  usage
+  exit 0
+fi
+
+WEIGHTS_DIR="${1:-weights}"
+
+PRE_WM_URL="${PRE_WM_URL:-provide}"
+PRE_VGGT_URL="${PRE_VGGT_URL:-provide}"
+PRE_SVGGT_URL="${PRE_SVGGT_URL:-provide}"
+PRE_DAV3_URL="${PRE_DAV3_URL:-provide}"
+PRE_OF3GS_URL="${PRE_OF3GS_URL:-provide}"
+
+mkdir -p "$WEIGHTS_DIR"
+
+download_hf() {
+  local output="$1"
+  local source="$2"
+  local spec="${source#hf://}"
+  IFS='/' read -r -a parts <<< "$spec"
+  if [[ "${#parts[@]}" -lt 3 ]]; then
+    echo "Invalid Hugging Face source: $source" >&2
+    echo "Expected hf://namespace/repo/path/in/repo" >&2
+    exit 1
+  fi
+  local repo_id="${parts[0]}/${parts[1]}"
+  local filename
+  filename="$(IFS=/; echo "${parts[*]:2}")"
+
+  python - "$repo_id" "$filename" "$WEIGHTS_DIR/$output" <<'PY'
+from pathlib import Path
+import shutil
+import sys
+
+from huggingface_hub import hf_hub_download
+
+repo_id, filename, output = sys.argv[1:4]
+output_path = Path(output)
+path = Path(hf_hub_download(repo_id=repo_id, filename=filename))
+shutil.copy2(path, output_path)
+print(f"Downloaded hf://{repo_id}/{filename} -> {output_path}")
+PY
+}
+
+download_url() {
+  local output="$1"
+  local url="$2"
+
+  if command -v wget >/dev/null 2>&1; then
+    wget -O "$WEIGHTS_DIR/$output" "$url"
+  elif command -v curl >/dev/null 2>&1; then
+    curl -L "$url" -o "$WEIGHTS_DIR/$output"
+  else
+    echo "wget or curl is required for direct URL downloads." >&2
+    exit 1
+  fi
+}
+
+download() {
+  local output="$1"
+  local source="$2"
+
+  if [[ -z "$source" || "$source" == "provide" ]]; then
+    echo "Skip $output: provide a source via the corresponding environment variable."
+    return 0
+  fi
+
+  if [[ -f "$WEIGHTS_DIR/$output" ]]; then
+    echo "Skip $output: already exists."
+    return 0
+  fi
+
+  case "$source" in
+    hf://*) download_hf "$output" "$source" ;;
+    http://*|https://*) download_url "$output" "$source" ;;
+    *)
+      echo "Unsupported source for $output: $source" >&2
+      exit 1
+      ;;
+  esac
+}
+
+python -m pip install -U huggingface_hub
+
+download "pre_wm.safetensors" "$PRE_WM_URL"
+download "pre_vggt.safetensors" "$PRE_VGGT_URL"
+download "pre_svggt.safetensors" "$PRE_SVGGT_URL"
+download "pre_dav3.safetensors" "$PRE_DAV3_URL"
+download "of3gs.ckpt" "$PRE_OF3GS_URL"
+
+cat <<EOF
+
+Weights are under: $WEIGHTS_DIR
+
+Suggested environment variables:
+  export PRE_WM_PATH="$WEIGHTS_DIR/pre_wm.safetensors"
+  export PRE_VGGT_PATH="$WEIGHTS_DIR/pre_vggt.safetensors"
+  export PRE_SVGGT_PATH="$WEIGHTS_DIR/pre_svggt.safetensors"
+  export PRE_DAV3_PATH="$WEIGHTS_DIR/pre_dav3.safetensors"
+  export PRE_OF3GS_PATH="$WEIGHTS_DIR/of3gs.ckpt"
+EOF

src/__init__.py

@@ -0,0 +1 @@
+"""OF3GS package."""

src/config.py

@@ -0,0 +1,111 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal, Optional, Type, TypeVar
+
+from dacite import Config, from_dict
+from omegaconf import DictConfig, OmegaConf
+
+from .dataset import DatasetCfgWrapper
+from .dataset.data_module import DataLoaderCfg
+from .loss import LossCfgWrapper
+from .model.decoder import DecoderCfg
+from .model.encoder import EncoderCfg
+from .model.model_wrapper import OptimizerCfg, TestCfg, TrainCfg
+
+
+@dataclass
+class CheckpointingCfg:
+    load: Optional[str]  # Not a path, since it could be something like wandb://...
+    every_n_train_steps: int
+    save_top_k: int
+    save_weights_only: bool
+
+
+@dataclass
+class ModelCfg:
+    decoder: DecoderCfg
+    encoder: EncoderCfg
+
+
+@dataclass
+class TrainerCfg:
+    max_steps: int
+    val_check_interval: int | float | None
+    gradient_clip_val: int | float | None
+    num_nodes: int = 1
+    devices: int | str = "auto"
+    accumulate_grad_batches: int = 1
+    precision: Literal["32-true", "16-mixed", "bf16-mixed"] = "32"
+
+
+@dataclass
+class RootCfg:
+    wandb: dict
+    mode: Literal["train", "test"]
+    pre_wm_path: str | None
+    pre_of3gs_path: str | None
+    dataset: list[DatasetCfgWrapper]
+    data_loader: DataLoaderCfg
+    model: ModelCfg
+    optimizer: OptimizerCfg
+    checkpointing: CheckpointingCfg
+    trainer: TrainerCfg
+    loss: list[LossCfgWrapper]
+    test: TestCfg
+    train: TrainCfg
+    seed: int
+
+
+TYPE_HOOKS = {
+    Path: Path,
+}
+
+
+T = TypeVar("T")
+
+
+def load_typed_config(
+    cfg: DictConfig,
+    data_class: Type[T],
+    extra_type_hooks: dict = {},
+) -> T:
+    return from_dict(
+        data_class,
+        OmegaConf.to_container(cfg, resolve=True),
+        config=Config(type_hooks={**TYPE_HOOKS, **extra_type_hooks}),
+    )
+
+
+def separate_loss_cfg_wrappers(joined: dict) -> list[LossCfgWrapper]:
+    # The dummy allows the union to be converted.
+    @dataclass
+    class Dummy:
+        dummy: LossCfgWrapper
+
+    return [
+        load_typed_config(DictConfig({"dummy": {k: v}}), Dummy).dummy
+        for k, v in joined.items()
+    ]
+
+
+def separate_dataset_cfg_wrappers(joined: dict) -> list[DatasetCfgWrapper]:
+    # The dummy allows the union to be converted.
+    @dataclass
+    class Dummy:
+        dummy: DatasetCfgWrapper
+
+    return [
+        load_typed_config(DictConfig({"dummy": {k: v}}), Dummy).dummy
+        for k, v in joined.items()
+    ]
+
+
+def load_typed_root_config(cfg: DictConfig) -> RootCfg:
+    return load_typed_config(
+        cfg,
+        RootCfg,
+        {
+            list[LossCfgWrapper]: separate_loss_cfg_wrappers,
+            list[DatasetCfgWrapper]: separate_dataset_cfg_wrappers,
+        },
+    )

src/dataset/__init__.py

@@ -0,0 +1,94 @@
+from dataclasses import fields
+from typing import Callable
+from torch.utils.data import Dataset, ConcatDataset
+import bisect
+
+from ..misc.step_tracker import StepTracker
+from .types import Stage
+from .view_sampler import get_view_sampler
+from .dataset_dl3dv import DatasetDL3DV, DatasetDL3DVCfgWrapper
+from .dataset_re10k import Datasetre10k, Datasetre10kCfgWrapper
+DATASETS: dict[str, Dataset] = {
+    "dl3dv": DatasetDL3DV,
+    "re10k": Datasetre10k,
+}
+
+DatasetCfgWrapper = DatasetDL3DVCfgWrapper | Datasetre10kCfgWrapper
+
+class TestDatasetWarpper(Dataset):
+    def __init__(self, dataset: Dataset):
+        self.dataset = dataset
+
+    def __getitem__(self, idx):
+
+        return self.dataset[(idx, self.dataset.view_sampler.num_context_views, self.dataset.cfg.input_image_shape[1] // 14)] 
+    
+    def __len__(self):
+        return len(self.dataset)
+
+        
+    
+class CustomConcatDataset(ConcatDataset):
+
+    def __getitem__(self, idx_tuple):
+
+        if isinstance(idx_tuple, list):
+            idx_tuple = idx_tuple[0]
+
+        idx = idx_tuple[0]
+        if idx < 0:
+            if -idx > len(self):
+                raise ValueError("absolute value of index should not exceed dataset length")
+            idx = len(self) + idx
+        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
+        if dataset_idx == 0:
+            sample_idx = idx
+        else:
+            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
+        return self.datasets[dataset_idx][(sample_idx, idx_tuple[1], idx_tuple[2])]
+
+
+def get_dataset(
+    cfgs: list[DatasetCfgWrapper],
+    stage: Stage,
+    step_tracker: StepTracker | None,
+    dataset_shim: Callable[[Dataset, str], Dataset]
+) -> list[Dataset]:
+    datasets = []
+    if stage != "test":
+        if stage == "val":
+            cfgs = [cfgs[0]]
+        for cfg in cfgs:
+            (field,) = fields(type(cfg))
+            cfg = getattr(cfg, field.name)
+            view_sampler = get_view_sampler(
+                cfg.view_sampler,
+                stage,
+                cfg.overfit_to_scene is not None,
+                cfg.cameras_are_circular,
+                step_tracker,
+            )
+            dataset = DATASETS[cfg.name](cfg, stage, view_sampler)
+            dataset = dataset_shim(dataset, stage)
+            datasets.append(dataset)
+
+        return CustomConcatDataset(datasets), datasets
+    elif stage == "test":
+        assert len(cfgs) == 1
+        cfg = cfgs[0]
+        (field,) = fields(type(cfg))
+        cfg = getattr(cfg, field.name)
+        
+        view_sampler = get_view_sampler(
+            cfg.view_sampler,
+            stage,
+            cfg.overfit_to_scene is not None,
+            cfg.cameras_are_circular,
+            step_tracker,
+        )
+        dataset = DATASETS[cfg.name](cfg, stage, view_sampler)
+        dataset = dataset_shim(dataset, stage)
+
+        return TestDatasetWarpper(dataset)
+    else:
+        NotImplementedError(f"Stage {stage} is not supported")

src/dataset/data_module.py

@@ -0,0 +1,196 @@
+import random
+from dataclasses import dataclass
+from typing import Callable
+
+import numpy as np
+import torch
+from lightning.pytorch import LightningDataModule
+from torch import Generator, nn
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+
+from src.dataset import *
+from src.global_cfg import get_cfg
+
+from ..misc.step_tracker import StepTracker
+from ..misc.utils import get_world_size, get_rank
+from . import DatasetCfgWrapper, get_dataset
+from .types import DataShim, Stage
+from .data_sampler import MixedBatchSampler
+
+def get_data_shim(encoder: nn.Module) -> DataShim:
+    """Get functions that modify the batch. It's sometimes necessary to modify batches
+    outside the data loader because GPU computations are required to modify the batch or
+    because the modification depends on something outside the data loader.
+    """
+
+    shims: list[DataShim] = []
+    if hasattr(encoder, "get_data_shim"):
+        shims.append(encoder.get_data_shim())
+
+    def combined_shim(batch):
+        for shim in shims:
+            batch = shim(batch)
+        return batch
+
+    return combined_shim
+
+# the training ratio of datasets (example)
+prob_mapping = {DatasetDL3DV: 0.5,
+               Datasetre10k: 0.5}
+
+@dataclass
+class DataLoaderStageCfg:
+    batch_size: int
+    num_workers: int
+    persistent_workers: bool
+    seed: int | None
+
+
+@dataclass
+class DataLoaderCfg:
+    train: DataLoaderStageCfg
+    test: DataLoaderStageCfg
+    val: DataLoaderStageCfg
+
+
+DatasetShim = Callable[[Dataset, Stage], Dataset]
+
+
+def worker_init_fn(worker_id: int) -> None:
+    random.seed(int(torch.utils.data.get_worker_info().seed) % (2**32 - 1))
+    np.random.seed(int(torch.utils.data.get_worker_info().seed) % (2**32 - 1))
+
+
+class DataModule(LightningDataModule):
+    dataset_cfgs: list[DatasetCfgWrapper]
+    data_loader_cfg: DataLoaderCfg
+    step_tracker: StepTracker | None
+    dataset_shim: DatasetShim
+    global_rank: int
+    
+    def __init__(
+        self,
+        dataset_cfgs: list[DatasetCfgWrapper],
+        data_loader_cfg: DataLoaderCfg,
+        step_tracker: StepTracker | None = None,
+        dataset_shim: DatasetShim = lambda dataset, _: dataset,
+        global_rank: int = 0,
+    ) -> None:
+        super().__init__()
+        self.dataset_cfgs = dataset_cfgs
+        self.data_loader_cfg = data_loader_cfg
+        self.step_tracker = step_tracker
+        self.dataset_shim = dataset_shim
+        self.global_rank = global_rank
+        
+    def get_persistent(self, loader_cfg: DataLoaderStageCfg) -> bool | None:
+        return None if loader_cfg.num_workers == 0 else loader_cfg.persistent_workers
+
+    def get_generator(self, loader_cfg: DataLoaderStageCfg) -> torch.Generator | None:
+        if loader_cfg.seed is None:
+            return None
+        generator = Generator()
+        generator.manual_seed(loader_cfg.seed + self.global_rank)
+        self.generator = generator
+        return self.generator
+        
+    def train_dataloader(self):
+        dataset, datasets_ls = get_dataset(self.dataset_cfgs, "train", self.step_tracker, self.dataset_shim)
+        world_size = get_world_size()
+        rank = get_rank()
+        prob_ls = [prob_mapping[type(dataset)] for dataset in datasets_ls]
+        # we assume all the dataset share the same num_context_views
+        
+        if len(datasets_ls) > 1:
+            prob = prob_ls
+            context_num_views = [dataset.cfg.view_sampler.num_context_views for dataset in datasets_ls]
+        else:
+            prob = None
+            dataset_key = next(iter(get_cfg()["dataset"]))
+            dataset_cfg = get_cfg()["dataset"][dataset_key]
+            context_num_views = dataset_cfg['view_sampler']['num_context_views']
+            
+        sampler = MixedBatchSampler(datasets_ls, 
+                                    batch_size=self.data_loader_cfg.train.batch_size, 
+                                    num_context_views=context_num_views, 
+                                    world_size=world_size, 
+                                    rank=rank,
+                                    prob=prob,
+                                    generator=self.get_generator(self.data_loader_cfg.train))
+        sampler.set_epoch(0)
+        self.train_loader = DataLoader(
+            dataset,
+            batch_sampler=sampler,
+            num_workers=self.data_loader_cfg.train.num_workers,
+            generator=self.generator,
+            worker_init_fn=worker_init_fn,
+            persistent_workers=self.get_persistent(self.data_loader_cfg.train),
+        )
+        # Set epoch for train and validation loaders (if applicable)
+        if hasattr(self.train_loader, "dataset") and hasattr(self.train_loader.dataset, "set_epoch"):
+            print("Training: Set Epoch in DataModule")
+            self.train_loader.dataset.set_epoch(0)
+        if hasattr(self.train_loader, "sampler") and hasattr(self.train_loader.sampler, "set_epoch"):
+            print("Training: Set Epoch in DataModule")
+            self.train_loader.sampler.set_epoch(0)
+        
+        return self.train_loader
+    
+    def val_dataloader(self):
+        dataset, datasets_ls = get_dataset(self.dataset_cfgs, "val", self.step_tracker, self.dataset_shim)
+        world_size = get_world_size()
+        rank = get_rank()
+        # here, we random select one dataset for val
+        dataset_key = next(iter(get_cfg()["dataset"]))
+        dataset_cfg = get_cfg()["dataset"][dataset_key]
+        if len(datasets_ls) > 1:
+             prob = [0.5] * len(datasets_ls)
+        else:
+            prob = None
+        sampler = MixedBatchSampler(datasets_ls, 
+                                    batch_size=self.data_loader_cfg.train.batch_size, 
+                                    num_context_views=dataset_cfg['view_sampler']['num_context_views'], 
+                                    world_size=1, 
+                                    rank=0,
+                                    prob=prob,
+                                    generator=self.get_generator(self.data_loader_cfg.train))
+        sampler.set_epoch(0)
+        self.val_loader = DataLoader(
+            dataset,
+            self.data_loader_cfg.val.batch_size,
+            num_workers=self.data_loader_cfg.val.num_workers,
+            sampler=sampler,
+            generator=self.get_generator(self.data_loader_cfg.val),
+            worker_init_fn=worker_init_fn,
+            persistent_workers=self.get_persistent(self.data_loader_cfg.val),
+            shuffle=False
+        )
+        if hasattr(self.val_loader, "dataset") and hasattr(self.val_loader.dataset, "set_epoch"):
+            print("Validation: Set Epoch in DataModule")
+            self.val_loader.dataset.set_epoch(0)
+        if hasattr(self.val_loader, "sampler") and hasattr(self.val_loader.sampler, "set_epoch"):
+            print("Validation: Set Epoch in DataModule")
+            self.val_loader.sampler.set_epoch(0)
+        return self.val_loader
+
+    def test_dataloader(self):
+        dataset = get_dataset(self.dataset_cfgs, "test", self.step_tracker, self.dataset_shim)
+        
+        sampler = DistributedSampler(
+            dataset,
+            num_replicas=torch.distributed.get_world_size(), 
+            rank=torch.distributed.get_rank(),  
+            shuffle=False  
+        )
+        sampler.set_epoch(0)
+        data_loader = DataLoader(
+            dataset,
+            self.data_loader_cfg.test.batch_size,
+            num_workers=self.data_loader_cfg.test.num_workers,
+            generator=self.get_generator(self.data_loader_cfg.test),
+            worker_init_fn=worker_init_fn,
+            persistent_workers=self.get_persistent(self.data_loader_cfg.test),
+            sampler=sampler,
+            shuffle=False
+        )        
+        return data_loader

src/dataset/data_sampler.py

@@ -0,0 +1,364 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Random sampling under a constraint
+# --------------------------------------------------------
+import numpy as np
+import torch
+from typing import Callable, Iterable, Optional
+from torch.utils.data import DistributedSampler, Sampler, BatchSampler
+import random
+
+def custom_collate_fn(batch):
+    """
+    Custom collate function to handle variable batch sizes
+    
+    Args:
+        batch: A list where each element could be either:
+            - A single tuple (idx, num_images, ...)
+            - A list of tuples [(idx1, num_images1, ...), (idx2, num_images2, ...)]
+    """
+    # If batch contains lists (variable batch size case)
+    if isinstance(batch[0], list):
+        # Flatten the batch
+        flattened = []
+        for item in batch:
+            flattened.extend(item)
+        batch = flattened
+    
+    # Now batch is a list of tuples, process normally
+    return torch.utils.data.default_collate(batch)
+
+class BatchedRandomSampler:
+    """Random sampling under a constraint: each sample in the batch has the same feature,
+    which is chosen randomly from a known pool of 'features' for each batch.
+
+    For instance, the 'feature' could be the image aspect-ratio.
+
+    The index returned is a tuple (sample_idx, feat_idx).
+    This sampler ensures that each series of `batch_size` indices has the same `feat_idx`.
+    """
+
+    def __init__(
+        self, dataset, batch_size, num_context_views, min_patch_num=20, max_patch_num=32, world_size=1, rank=0, drop_last=True
+    ):
+        self.batch_size = batch_size
+        self.num_context_views = num_context_views
+
+        self.len_dataset = N = len(dataset)
+        self.total_size = round_by(N, batch_size * world_size) if drop_last else N
+        self.min_patch_num = min_patch_num
+        self.max_patch_num = max_patch_num
+        assert (
+            world_size == 1 or drop_last
+        ), "must drop the last batch in distributed mode"
+
+        # distributed sampler
+        self.world_size = world_size
+        self.rank = rank
+        self.epoch = None
+
+    def __len__(self):
+
+
+        return self.total_size // self.world_size
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
+
+    def __iter__(self):
+        # prepare RNG
+        if self.epoch is None:
+            assert (
+                self.world_size == 1 and self.rank == 0
+            ), "use set_epoch() if distributed mode is used"
+            seed = int(torch.empty((), dtype=torch.int64).random_().item())
+        else:
+            seed = self.epoch + 777
+        rng = np.random.default_rng(seed=seed)
+        
+        # random indices (will restart from 0 if not drop_last)
+        sample_idxs = np.arange(self.total_size)
+        rng.shuffle(sample_idxs)
+        
+        # random feat_idxs (same across each batch)
+        n_batches = (self.total_size + self.batch_size - 1) // self.batch_size
+        num_imgs = rng.integers(low=2, high=self.num_context_views, size=n_batches)
+        # num_imgs = (np.ones(n_batches) * self.num_context_views).astype(np.int64) # same number of context views for each batch
+        num_imgs = np.broadcast_to(num_imgs[:, None], (n_batches, self.batch_size))
+        num_imgs = num_imgs.ravel()[: self.total_size]
+
+        # put them together
+        idxs = np.c_[sample_idxs, num_imgs]  # shape = (total_size, 2)
+
+        # Distributed sampler: we select a subset of batches
+        # make sure the slice for each node is aligned with batch_size
+        size_per_proc = self.batch_size * (
+            (self.total_size + self.world_size * self.batch_size - 1)
+            // (self.world_size * self.batch_size)
+        )
+        idxs = idxs[self.rank * size_per_proc : (self.rank + 1) * size_per_proc]
+
+        yield from (tuple(idx) for idx in idxs)
+
+class DynamicBatchSampler(Sampler):
+    """
+    A custom batch sampler that dynamically adjusts batch size, aspect ratio, and image number
+    for each sample. Batches within a sample share the same aspect ratio and image number.
+    """
+    def __init__(self,
+                 sampler,
+                 image_num_range,
+                 h_range,
+                 batch_size=1,
+                 epoch=0,
+                 seed=42,
+                 max_img_per_gpu=48):
+        """
+        Initializes the dynamic batch sampler.
+
+        Args:
+            sampler: Instance of DynamicDistributedSampler.
+            aspect_ratio_range: List containing [min_aspect_ratio, max_aspect_ratio].
+            image_num_range: List containing [min_images, max_images] per sample.
+            epoch: Current epoch number.
+            seed: Random seed for reproducibility.
+            max_img_per_gpu: Maximum number of images to fit in GPU memory.
+        """
+        self.sampler = sampler
+        self.image_num_range = image_num_range
+        self.h_range = h_range
+        self.rng = random.Random()
+        
+        # Uniformly sample from the range of possible image numbers
+        # For any image number, the weight is 1.0 (uniform sampling). You can set any different weights here.
+        self.image_num_weights = {num_images: float(num_images**2) for num_images in range(image_num_range[0], image_num_range[1]+1)}
+
+        # Possible image numbers, e.g., [2, 3, 4, ..., 24]
+        self.possible_nums = np.array([n for n in self.image_num_weights.keys()
+                                       if self.image_num_range[0] <= n <= self.image_num_range[1]])
+        
+        # Normalize weights for sampling
+        weights = [self.image_num_weights[n] for n in self.possible_nums]
+        self.normalized_weights = np.array(weights) / sum(weights)
+
+        # Maximum image number per GPU
+        self.max_img_per_gpu = max_img_per_gpu
+        self.batch_size = batch_size
+
+        # Set the epoch for the sampler
+        self.set_epoch(epoch + seed)
+
+    def set_epoch(self, epoch):
+        """
+        Sets the epoch for this sampler, affecting the random sequence.
+
+        Args:
+            epoch: The epoch number.
+        """
+        self.sampler.set_epoch(epoch)
+        self.epoch = epoch
+        self.rng.seed(epoch * 100)
+
+    def __iter__(self):
+        """
+        Yields batches of samples with synchronized dynamic parameters.
+
+        Returns:
+            Iterator yielding batches of indices with associated parameters.
+        """
+
+        sampler_iterator = iter(self.sampler)
+  
+        while True:
+            try:
+                # Sample random image number and aspect ratio
+                random_image_num = int(np.random.choice(self.possible_nums, p=self.normalized_weights))
+                random_ps_h = np.random.randint(low=(self.h_range[0] // 14), high=(self.h_range[1] // 14)+1)
+
+                # Update sampler parameters
+                self.sampler.update_parameters(
+                    image_num=random_image_num,
+                    ps_h=random_ps_h
+                )
+                
+                # Calculate batch size based on max images per GPU and current image number
+                dynamic_batch_size = self.max_img_per_gpu / random_image_num
+                dynamic_batch_size = max(1, int(np.floor(dynamic_batch_size)))
+                batch_size = min(self.batch_size, dynamic_batch_size)
+
+                # Collect samples for the current batch
+                current_batch = []
+                for _ in range(batch_size):
+                    try:
+                        item = next(sampler_iterator)  # item is (idx, aspect_ratio, image_num)
+                        current_batch.append(item)
+                    except StopIteration:
+                        break  # No more samples
+
+                if not current_batch:
+                    break  # No more data to yield
+
+                yield current_batch
+
+            except StopIteration:
+                break  # End of sampler's iterator
+       
+    def __len__(self):
+        # Return a large dummy length
+        return 1000000
+
+
+class DynamicDistributedSampler(DistributedSampler):
+    """
+    Extends PyTorch's DistributedSampler to include dynamic aspect_ratio and image_num
+    parameters, which can be passed into the dataset's __getitem__ method.
+    """
+    def __init__(
+        self,
+        dataset,
+        num_replicas: Optional[int] = None,
+        rank: Optional[int] = None,
+        shuffle: bool = False,
+        seed: int = 0,
+        drop_last: bool = False,
+    ):
+        super().__init__(
+            dataset,
+            num_replicas=num_replicas,
+            rank=rank,
+            shuffle=shuffle,
+            seed=seed,
+            drop_last=drop_last
+        )
+        self.image_num = None
+        self.ps_h = None
+
+    def __iter__(self):
+        """
+        Yields a sequence of (index, image_num, aspect_ratio).
+        Relies on the parent class's logic for shuffling/distributing
+        the indices across replicas, then attaches extra parameters.
+        """
+        indices_iter = super().__iter__()
+
+        for idx in indices_iter:
+            yield (idx, self.image_num, self.ps_h, )
+
+    def update_parameters(self, image_num, ps_h):
+        """
+        Updates dynamic parameters for each new epoch or iteration.
+
+        Args:
+            aspect_ratio: The aspect ratio to set.
+            image_num: The number of images to set.
+        """
+        self.image_num = image_num
+        self.ps_h = ps_h
+
+class MixedBatchSampler(BatchSampler):
+    """Sample one batch from a selected dataset with given probability.
+    Compatible with datasets at different resolution
+    """
+
+    def __init__(
+        self, src_dataset_ls, batch_size, num_context_views, world_size=1, rank=0, prob=None, sampler=None, generator=None
+    ):
+        self.base_sampler = None
+        self.batch_size = batch_size
+        self.num_context_views = num_context_views
+        self.world_size = world_size
+        self.rank = rank
+        self.drop_last = True
+        self.generator = generator
+
+        self.src_dataset_ls = src_dataset_ls
+        self.n_dataset = len(self.src_dataset_ls)
+        
+        # Dataset length
+        self.dataset_length = [len(ds) for ds in self.src_dataset_ls]
+        self.cum_dataset_length = [
+            sum(self.dataset_length[:i]) for i in range(self.n_dataset)
+        ]  # cumulative dataset length
+        
+        # BatchSamplers for each source dataset
+        self.src_batch_samplers = []
+        for ds in self.src_dataset_ls:
+            sampler = DynamicDistributedSampler(ds, num_replicas=self.world_size, rank=self.rank, seed=42, shuffle=True)
+            sampler.set_epoch(0)
+
+            if hasattr(ds, "epoch"):
+                ds.epoch = 0
+            if hasattr(ds, "set_epoch"):
+                ds.set_epoch(0)
+            batch_sampler = DynamicBatchSampler(
+                sampler, 
+                [2, ds.cfg.view_sampler.num_context_views], 
+                ds.cfg.input_image_shape,
+                batch_size=self.batch_size,
+                seed=42,
+                max_img_per_gpu=ds.cfg.view_sampler.max_img_per_gpu
+            )
+            self.src_batch_samplers.append(batch_sampler)
+        # set epoch here
+        print("Setting epoch for all underlying BatchedRandomSamplers")
+        self.raw_batches = [
+            list(bs) for bs in self.src_batch_samplers
+        ]  # index in original dataset
+  
+        self.n_batches = [len(b) for b in self.raw_batches]
+        self.n_total_batch = sum(self.n_batches)
+        print("Total batch num is ", self.n_total_batch)
+        # sampling probability
+        if prob is None:
+            # if not given, decide by dataset length
+            self.prob = torch.tensor(self.n_batches) / self.n_total_batch
+        else:
+            self.prob = torch.as_tensor(prob)
+    
+    def __iter__(self):
+        """Yields batches of indices in the format of (sample_idx, feat_idx) tuples,
+        where indices correspond to ConcatDataset of src_dataset_ls
+        """
+        for _ in range(self.n_total_batch):
+            idx_ds = torch.multinomial(
+                self.prob, 1, replacement=True, generator=self.generator
+            ).item()
+            
+            if 0 == len(self.raw_batches[idx_ds]):
+                self.raw_batches[idx_ds] = list(self.src_batch_samplers[idx_ds])
+            
+            # get a batch from list - this is already in (sample_idx, feat_idx) format
+            batch_raw = self.raw_batches[idx_ds].pop()
+
+            # shift only the sample_idx by cumulative dataset length, keep feat_idx unchanged
+            shift = self.cum_dataset_length[idx_ds]
+            processed_batch = []
+
+            for item in batch_raw:
+                # item[0] is the sample index, item[1] is the number of images
+                processed_item = (item[0] + shift, item[1], item[2])
+                processed_batch.append(processed_item)
+            yield processed_batch
+        
+    def set_epoch(self, epoch):
+        """Set epoch for all underlying BatchedRandomSamplers"""
+        for sampler in self.src_batch_samplers:
+            sampler.set_epoch(epoch)
+        # Reset raw_batches after setting new epoch
+        self.raw_batches = [list(bs) for bs in self.src_batch_samplers]
+
+    def __len__(self):
+        return self.n_total_batch
+
+def round_by(total, multiple, up=False):
+    if up:
+        total = total + multiple - 1
+    return (total // multiple) * multiple

src/dataset/dataset.py

@@ -0,0 +1,13 @@
+from dataclasses import dataclass
+
+from .view_sampler import ViewSamplerCfg
+
+
+@dataclass
+class DatasetCfgCommon:
+    original_image_shape: list[int]
+    input_image_shape: list[int]
+    background_color: list[float]
+    cameras_are_circular: bool
+    overfit_to_scene: str | None
+    view_sampler: ViewSamplerCfg

src/dataset/dataset_dl3dv.py

@@ -0,0 +1,452 @@
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import json
+from dataclasses import dataclass
+from functools import cached_property
+from pathlib import Path
+from typing import Literal
+import os
+import traceback
+import numpy as np
+import torch
+import torchvision.transforms as tf
+from einops import repeat
+from jaxtyping import Float
+from PIL import Image
+from torch import Tensor
+from torch.utils.data import Dataset
+from tqdm import tqdm
+from ..geometry.projection import get_fov
+from .dataset import DatasetCfgCommon
+from .shims.augmentation_shim import apply_augmentation_shim
+from .shims.crop_shim import apply_crop_shim
+from .types import Stage
+from .view_sampler import ViewSampler
+from ..misc.cam_utils import camera_normalization
+
+
+@dataclass
+class DL3DVScene:
+    index_item: str
+    scene_root: Path
+    transforms_path: Path
+    images_dir: Path
+    scene_id: str
+
+
+@dataclass
+class DatasetDl3dvCfg(DatasetCfgCommon):
+    mode: str
+    name: str
+    roots: list[Path]
+    baseline_min: float
+    baseline_max: float
+    max_fov: float
+    make_baseline_1: bool
+    augment: bool
+    relative_pose: bool
+    skip_bad_shape: bool
+    avg_pose: bool
+    rescale_to_1cube: bool
+    intr_augment: bool
+    normalize_by_pts3d: bool
+    ctx_list: list[int] | None
+    tgt_list: list[int] | None
+
+
+@dataclass
+class DatasetDL3DVCfgWrapper:
+    dl3dv: DatasetDl3dvCfg
+
+
+class DatasetDL3DV(Dataset):
+    cfg: DatasetDl3dvCfg
+    stage: Stage
+    view_sampler: ViewSampler
+
+    to_tensor: tf.ToTensor
+    chunks: list[Path]
+    near: float = 0.1
+    far: float = 100.0
+
+    def __init__(
+        self,
+        cfg: DatasetDl3dvCfg,
+        stage: Stage,
+        view_sampler: ViewSampler,
+    ) -> None:
+        super().__init__()
+        self.cfg = cfg
+        self.stage = stage
+        self.view_sampler = view_sampler
+        self.to_tensor = tf.ToTensor()
+        # load data
+        if not cfg.roots:
+            raise ValueError(
+                "DatasetDL3DV requires dataset.dl3dv.roots. Set "
+                "DL3DV_ROOT=/path/to/dl3dv when using train.sh or pass "
+                "dataset.dl3dv.roots='[/path/to/dl3dv]'."
+            )
+        self.data_root = Path(cfg.roots[0])
+        index_path = self.data_root / f"{self.data_stage}_index.json"
+        if not index_path.is_file():
+            raise FileNotFoundError(
+                f"DL3DV index file not found: {index_path}. Expected "
+                f"{self.data_stage}_index.json under dataset.dl3dv.roots[0]."
+            )
+        self.data_list = []
+        with index_path.open("r") as file:
+            data_index = json.load(file)
+
+        def filter_data_list(data_index):
+            data_list = []
+
+            for item in data_index:
+                scene = self.resolve_scene(str(item))
+                if scene is None:
+                    print(item)
+                else:
+                    data_list.append(scene)
+
+            return data_list
+
+        self.data_list = filter_data_list(data_index)
+        self.scene_ids = {}
+        self.scenes = {}
+        index = 0
+
+        if cfg.mode == "train":
+            with ThreadPoolExecutor(max_workers=64) as executor:
+                futures = [
+                    executor.submit(self.load_jsons, scene_path)
+                    for scene_path in self.data_list
+                ]
+                for future in tqdm(as_completed(futures), total=len(futures)):
+                    scene_frames, scene_id = future.result()
+                    self.scenes[scene_id] = scene_frames
+                    self.scene_ids[index] = scene_id
+                    index += 1
+        else:
+            futures = [self.load_jsons(scene_path) for scene_path in self.data_list]
+            for future in tqdm(futures, total=len(futures)):
+                scene_frames, scene_id = future
+                self.scenes[scene_id] = scene_frames
+                self.scene_ids[index] = scene_id
+                index += 1
+
+        print(f"DL3DV: {self.stage}: loaded {len(self.scene_ids)} scenes")
+
+    def resolve_scene(self, item: str) -> DL3DVScene | None:
+        scene_root = self.data_root / item
+        if not scene_root.exists():
+            return None
+
+        image_dirs = sorted(
+            (path for path in scene_root.rglob("images_8") if path.is_dir()),
+            key=lambda path: (-len(path.parts), path.as_posix()),
+        )
+        transforms_paths = sorted(
+            scene_root.rglob("transforms.json"),
+            key=lambda path: (-len(path.parts), path.as_posix()),
+        )
+        if not image_dirs or not transforms_paths:
+            return None
+
+        matches = []
+        for transforms_path in transforms_paths:
+            for images_dir in image_dirs:
+                try:
+                    transforms_path.relative_to(images_dir.parent)
+                    matches.append((transforms_path, images_dir))
+                    continue
+                except ValueError:
+                    pass
+
+                try:
+                    images_dir.relative_to(transforms_path.parent)
+                    matches.append((transforms_path, images_dir))
+                except ValueError:
+                    pass
+
+        if not matches:
+            matches = [(transforms_paths[0], image_dirs[0])]
+        if len(matches) > 1:
+            print(
+                f"Ambiguous DL3DV scene path for {item}: found {len(matches)} "
+                "matches; using the deepest transforms/images_8 pair."
+            )
+
+        transforms_path, images_dir = matches[0]
+        scene_rel = transforms_path.parent.relative_to(self.data_root).as_posix()
+        scene_id = scene_rel.replace("/", "_")
+        return DL3DVScene(
+            index_item=item,
+            scene_root=scene_root,
+            transforms_path=transforms_path,
+            images_dir=images_dir,
+            scene_id=scene_id,
+        )
+
+    def convert_intrinsics(self, meta_data):
+        store_h, store_w = meta_data["h"], meta_data["w"]
+        fx, fy, cx, cy = (
+            meta_data["fl_x"],
+            meta_data["fl_y"],
+            meta_data["cx"],
+            meta_data["cy"],
+        )
+        intrinsics = np.eye(3, dtype=np.float32)
+        intrinsics[0, 0] = float(fx) / float(store_w)
+        intrinsics[1, 1] = float(fy) / float(store_h)
+        intrinsics[0, 2] = float(cx) / float(store_w)
+        intrinsics[1, 2] = float(cy) / float(store_h)
+        return intrinsics
+
+    def blender2opencv_c2w(self, pose):
+        blender2opencv = np.array(
+            [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]]
+        )
+        opencv_c2w = np.array(pose) @ blender2opencv
+        return opencv_c2w.tolist()
+
+    def load_jsons(self, scene: DL3DVScene):
+        with scene.transforms_path.open("r") as f:
+            data = json.load(f)
+
+        scene_frames = []
+        for i, frame in enumerate(data["frames"]):
+            frame_tmp = {}
+            frame_tmp["file_path"] = self.resolve_image_path(
+                scene.images_dir,
+                frame["file_path"],
+            )
+            frame_tmp["intrinsics"] = self.convert_intrinsics(data).tolist()
+            frame_tmp["extrinsics"] = self.blender2opencv_c2w(frame["transform_matrix"])
+            scene_frames.append(frame_tmp)
+        return scene_frames, scene.scene_id
+
+    def load_frames(self, frames):
+        with ThreadPoolExecutor(max_workers=32) as executor:
+            # Create a list to store futures with their original indices
+            futures_with_idx = []
+            for idx, file_path in enumerate(frames):
+                file_path = file_path["file_path"]
+                futures_with_idx.append(
+                    (
+                        idx,
+                        executor.submit(
+                            lambda p: self.to_tensor(Image.open(p).convert("RGB")),
+                            file_path,
+                        ),
+                    )
+                )
+
+            # Pre-allocate list with correct size to maintain order
+            torch_images = [None] * len(frames)
+            for idx, future in futures_with_idx:
+                torch_images[idx] = future.result()
+            # Check if all images have the same size
+
+            sizes = set(img.shape for img in torch_images)
+            if len(sizes) == 1:
+                torch_images = torch.stack(torch_images)
+        return torch_images
+
+    def resolve_image_path(self, images_dir: Path, frame_path: str) -> str:
+        frame_rel = Path(frame_path)
+        filename = frame_rel.name
+        candidates = [
+            images_dir / filename,
+            images_dir / frame_rel,
+        ]
+        if "images" in frame_rel.parts:
+            image_idx = frame_rel.parts.index("images")
+            candidates.append(images_dir.joinpath(*frame_rel.parts[image_idx + 1 :]))
+
+        for candidate in candidates:
+            if candidate.is_file():
+                return str(candidate)
+
+        matches = sorted(images_dir.rglob(filename))
+        if matches:
+            return str(matches[0])
+        return str(images_dir / filename)
+
+    def getitem(self, index: int, num_context_views: int, patchsize: tuple) -> dict:
+
+        scene = self.scene_ids[index]
+
+        example = self.scenes[scene]
+        # load poses
+        extrinsics = []
+        intrinsics = []
+        for frame in example:
+            extrinsic = frame["extrinsics"]
+            intrinsic = frame["intrinsics"]
+            extrinsics.append(extrinsic)
+            intrinsics.append(intrinsic)
+
+        extrinsics = np.array(extrinsics)
+        intrinsics = np.array(intrinsics)
+        extrinsics = torch.tensor(extrinsics, dtype=torch.float32)
+        intrinsics = torch.tensor(intrinsics, dtype=torch.float32)
+
+        if self.cfg.mode == "train":
+            try:
+                context_indices, target_indices, overlap = self.view_sampler.sample(
+                    scene,
+                    num_context_views,
+                    extrinsics,
+                    intrinsics,
+                )
+            except ValueError:
+                # Skip because the example doesn't have enough frames.
+                raise Exception("Not enough frames")
+        else:
+            if not self.cfg.ctx_list or not self.cfg.tgt_list:
+                raise ValueError(
+                    "DL3DV test mode requires dataset.dl3dv.ctx_list and "
+                    "dataset.dl3dv.tgt_list to be non-empty lists."
+                )
+            context_indices = self.cfg.ctx_list
+            target_indices = self.cfg.tgt_list
+
+        if (get_fov(intrinsics).rad2deg() > self.cfg.max_fov).any():
+            raise Exception("Field of view too wide")
+
+        input_frames = [example[i] for i in context_indices]
+        target_frame = [example[i] for i in target_indices]
+
+        context_images = self.load_frames(input_frames)
+        target_images = self.load_frames(target_frame)
+        resize = tf.Resize((270, 480))
+        context_images = resize(context_images)
+        target_images = resize(target_images)
+
+        # Skip the example if the images don't have the right shape.
+        context_image_invalid = context_images.shape[1:] != (
+            3,
+            *self.cfg.original_image_shape,
+        )
+        target_image_invalid = target_images.shape[1:] != (
+            3,
+            *self.cfg.original_image_shape,
+        )
+        if self.cfg.skip_bad_shape and (context_image_invalid or target_image_invalid):
+            print(
+                f"Skipped bad example {scene}. Context shape was "
+                f"{context_images.shape} and target shape was "
+                f"{target_images.shape}."
+            )
+
+            raise Exception("Bad example image shape")
+
+        context_extrinsics = extrinsics[context_indices]
+        if self.cfg.make_baseline_1:
+            a, b = context_extrinsics[0, :3, 3], context_extrinsics[-1, :3, 3]
+            scale = (a - b).norm()
+            if scale < self.cfg.baseline_min or scale > self.cfg.baseline_max:
+                print(
+                    f"Skipped {scene} because of baseline out of range: " f"{scale:.6f}"
+                )
+                raise Exception("baseline out of range")
+            extrinsics[:, :3, 3] /= scale
+        else:
+            scale = 1
+
+        if self.cfg.relative_pose:
+            extrinsics = camera_normalization(
+                extrinsics[context_indices][0:1], extrinsics
+            )
+
+        if self.cfg.rescale_to_1cube:
+            scene_scale = torch.max(
+                torch.abs(extrinsics[context_indices][:, :3, 3])
+            )  
+            rescale_factor = 1 * scene_scale
+            extrinsics[:, :3, 3] /= rescale_factor
+
+        if torch.isnan(extrinsics).any() or torch.isinf(extrinsics).any():
+            raise Exception("encounter nan or inf in input poses")
+
+        example = {
+            "context": {
+                "extrinsics": extrinsics[context_indices],
+                "intrinsics": intrinsics[context_indices],
+                "image": context_images,
+                "near": self.get_bound("near", len(context_indices)) / scale,
+                "far": self.get_bound("far", len(context_indices)) / scale,
+                "index": context_indices,
+            },
+            "target": {
+                "extrinsics": extrinsics[target_indices],
+                "intrinsics": intrinsics[target_indices],
+                "image": target_images,
+                "near": self.get_bound("near", len(target_indices)) / scale,
+                "far": self.get_bound("far", len(target_indices)) / scale,
+                "index": target_indices,
+            },
+            "scene": "dl3dv_" + scene,
+        }
+        if self.stage == "train" and self.cfg.augment:
+            example = apply_augmentation_shim(example)
+
+        if self.stage == "train" and self.cfg.intr_augment:
+            intr_aug = True
+        else:
+            intr_aug = False
+
+        example = apply_crop_shim(
+            example, (patchsize[0] * 14, patchsize[1] * 14), intr_aug=intr_aug
+        )
+        return example
+
+    def __getitem__(self, index_tuple: tuple) -> dict:
+        index, num_context_views, patchsize_h = index_tuple
+        patchsize_w = self.cfg.input_image_shape[1] // 14
+
+        try:
+            return self.getitem(index, num_context_views, (patchsize_h, patchsize_w))
+        except Exception as e:
+            print(f"Error: {e}")
+            traceback.print_exc()
+            index = np.random.randint(len(self))
+            return self.__getitem__((index, num_context_views, patchsize_h))
+
+    def get_bound(
+        self,
+        bound: Literal["near", "far"],
+        num_views: int,
+    ) -> Float[Tensor, " view"]:
+        value = torch.tensor(getattr(self, bound), dtype=torch.float32)
+        return repeat(value, "-> v", v=num_views)
+
+    @property
+    def data_stage(self) -> Stage:
+        if self.cfg.overfit_to_scene is not None:
+            return "test"
+        if self.stage == "val":
+            return "test"
+        return self.stage
+
+    @cached_property
+    def index(self) -> dict[str, Path]:
+        merged_index = {}
+        data_stages = [self.data_stage]
+        if self.cfg.overfit_to_scene is not None:
+            data_stages = ("test", "train")
+        for data_stage in data_stages:
+            for root in self.cfg.roots:
+                # Load the root's index.
+                with (root / data_stage / "index.json").open("r") as f:
+                    index = json.load(f)
+                index = {k: Path(root / data_stage / v) for k, v in index.items()}
+
+                # The constituent datasets should have unique keys.
+                assert not (set(merged_index.keys()) & set(index.keys()))
+
+                # Merge the root's index into the main index.
+                merged_index = {**merged_index, **index}
+        return merged_index
+
+    def __len__(self) -> int:
+        return len(self.scene_ids)

src/dataset/dataset_re10k.py

@@ -0,0 +1,302 @@
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import json
+from dataclasses import dataclass
+from functools import cached_property
+from pathlib import Path
+from typing import Literal
+import os
+import traceback
+import numpy as np
+import torch
+import torchvision.transforms as tf
+from einops import repeat
+from jaxtyping import Float
+from PIL import Image
+from torch import Tensor
+from torch.utils.data import Dataset
+from tqdm import tqdm
+from .dataset import DatasetCfgCommon
+from .shims.augmentation_shim import apply_augmentation_shim
+from .shims.crop_shim import apply_crop_shim
+from .types import Stage
+from .view_sampler import ViewSampler
+import gzip
+import pickle
+
+
+@dataclass
+class Datasetre10kCfg(DatasetCfgCommon):
+    name: str
+    roots: list[Path]
+    baseline_min: float
+    baseline_max: float
+    max_fov: float
+    make_baseline_1: bool
+    augment: bool
+    relative_pose: bool
+    skip_bad_shape: bool
+    avg_pose: bool
+    rescale_to_1cube: bool
+    intr_augment: bool
+    normalize_by_pts3d: bool
+
+
+@dataclass
+class Datasetre10kCfgWrapper:
+    re10k: Datasetre10kCfg
+
+
+class Datasetre10k(Dataset):
+    cfg: Datasetre10kCfg
+    stage: Stage
+    view_sampler: ViewSampler
+
+    to_tensor: tf.ToTensor
+    chunks: list[Path]
+    near: float = 0.1
+    far: float = 100.0
+
+    def __init__(
+        self,
+        cfg: Datasetre10kCfg,
+        stage: Stage,
+        view_sampler: ViewSampler,
+    ) -> None:
+        super().__init__()
+        self.cfg = cfg
+        self.stage = stage
+        self.view_sampler = view_sampler
+        self.to_tensor = tf.ToTensor()
+        # load data
+        if not cfg.roots:
+            raise ValueError(
+                "Datasetre10k requires dataset.re10k.roots. Set "
+                "RE10K_ROOT=/path/to/re10k when using train.sh or pass "
+                "dataset.re10k.roots='[/path/to/re10k]'."
+            )
+        self.root = Path(cfg.roots[0])
+        self.split = self.data_stage
+        if len(cfg.roots) == 1:
+            self.data_root = self.root / self.split
+            self.index_root = self.root / f"{self.split}_index.json"
+            self.metadata_path = self.root / f"{self.split}.pickle.gz"
+        else:
+            self.data_root = self.root
+            self.index_root = Path(cfg.roots[1])
+            self.metadata_path = self.root / "train.pickle.gz"
+        if not self.index_root.is_file():
+            raise FileNotFoundError(
+                f"RE10K index file not found: {self.index_root}. Expected "
+                f"{self.split}_index.json under dataset.re10k.roots[0], or pass "
+                "a second roots entry pointing to the index file."
+            )
+        self.data_list = []
+        with self.index_root.open("r") as file:
+            self.data_index = json.load(file)
+
+        for item in self.data_index:
+            self.data_list.append(str(self.resolve_scene_path(str(item))))
+        self.scene_ids = {}
+        self.scenes = {}
+        index = 0
+
+        if not self.metadata_path.is_file() and self.split == "test":
+            fallback = self.root / "train.pickle.gz"
+            if fallback.is_file():
+                self.metadata_path = fallback
+        if not self.metadata_path.is_file():
+            raise FileNotFoundError(
+                f"RE10K metadata file not found: {self.metadata_path}. Expected "
+                f"{self.split}.pickle.gz under dataset.re10k.roots[0]."
+            )
+
+        with gzip.open(self.metadata_path, "rb") as f:
+            self.seq_data = pickle.load(f)
+
+        print("re10k loading data !!!")
+        with ThreadPoolExecutor(max_workers=64) as executor:
+            futures = [
+                executor.submit(self.load_jsons, scene_path)
+                for scene_path in self.data_list
+            ]
+            for future in tqdm(as_completed(futures), total=len(futures)):
+                scene_frames, scene_id = future.result()
+                self.scenes[scene_id] = scene_frames
+                self.scene_ids[index] = scene_id
+                index += 1
+        print(f"RE10k: {self.stage}: loaded {len(self.scene_ids)} scenes")
+
+    def resolve_scene_path(self, item: str) -> Path:
+        scene_path = self.data_root / item
+        if scene_path.is_dir():
+            return scene_path
+        return self.root / item
+
+    def load_jsons(self, scene_path):
+        files = os.listdir(scene_path)
+        files = [f for f in files if os.path.isfile(os.path.join(scene_path, f))]
+        files.sort()
+        scene_frames = []
+        scene_id = Path(scene_path).name
+        for i, frame in enumerate(files):
+            frame_tmp = {}
+
+            frame_tmp["file_path"] = os.path.join(scene_path, frame)
+            frame_tmp["extrinsics"] = np.array(
+                [[0.0, -0.5, 0, 0], [0.5, 0.866, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]
+            )
+            itmp = self.seq_data[scene_id]["intrinsics"][i]
+            intrinsics = np.eye(3, dtype=np.float32)
+            intrinsics[0, 0] = float(itmp[0])
+            intrinsics[1, 1] = float(itmp[1])
+            intrinsics[0, 2] = float(itmp[2])
+            intrinsics[1, 2] = float(itmp[3])
+            frame_tmp["intrinsics"] = intrinsics
+
+            scene_frames.append(frame_tmp)
+
+        return scene_frames, scene_id
+
+    def load_frames(self, frames):
+        with ThreadPoolExecutor(max_workers=32) as executor:
+            # Create a list to store futures with their original indices
+            futures_with_idx = []
+            for idx, file_path in enumerate(frames):
+                file_path = file_path["file_path"]
+                futures_with_idx.append(
+                    (
+                        idx,
+                        executor.submit(
+                            lambda p: self.to_tensor(Image.open(p).convert("RGB")),
+                            file_path,
+                        ),
+                    )
+                )
+
+            # Pre-allocate list with correct size to maintain order
+            torch_images = [None] * len(frames)
+            for idx, future in futures_with_idx:
+                torch_images[idx] = future.result()
+            # Check if all images have the same size
+
+            sizes = set(img.shape for img in torch_images)
+            if len(sizes) == 1:
+                torch_images = torch.stack(torch_images)
+        # Return as list if images have different sizes
+        return torch_images
+
+    def getitem(self, index: int, num_context_views: int, patchsize: tuple) -> dict:
+        scene = self.scene_ids[index]
+        example = self.scenes[scene]
+        # load poses
+        extrinsics = []
+        intrinsics = []
+        for frame in example:
+            extrinsic = frame["extrinsics"]
+            intrinsic = frame["intrinsics"]
+            extrinsics.append(extrinsic)
+            intrinsics.append(intrinsic)
+
+        extrinsics = np.array(extrinsics)
+        intrinsics = np.array(intrinsics)
+        extrinsics = torch.tensor(extrinsics, dtype=torch.float32)
+        intrinsics = torch.tensor(intrinsics, dtype=torch.float32)
+        try:
+            context_indices, target_indices, overlap = self.view_sampler.sample(
+                scene,
+                num_context_views,
+                extrinsics,
+                intrinsics,
+            )
+        except ValueError:
+            # Skip because the example doesn't have enough frames.
+            raise Exception("Not enough frames")
+
+        input_frames = [example[i] for i in context_indices]
+        target_frame = [example[i] for i in target_indices]
+
+        context_images = self.load_frames(input_frames)
+        target_images = self.load_frames(target_frame)
+
+        example = {
+            "context": {
+                "extrinsics": extrinsics[context_indices],
+                "intrinsics": intrinsics[context_indices],
+                "image": context_images,
+                "near": self.get_bound("near", len(context_indices)),
+                "far": self.get_bound("far", len(context_indices)),
+                "index": context_indices,
+            },
+            "target": {
+                "extrinsics": extrinsics[target_indices],
+                "intrinsics": intrinsics[target_indices],
+                "image": target_images,
+                "near": self.get_bound("near", len(target_indices)),
+                "far": self.get_bound("far", len(target_indices)),
+                "index": target_indices,
+            },
+            "scene": "re10k_" + scene,
+        }
+        if self.stage == "train" and self.cfg.augment:
+            example = apply_augmentation_shim(example)
+
+        if self.stage == "train" and self.cfg.intr_augment:
+            intr_aug = True
+        else:
+            intr_aug = False
+
+        example = apply_crop_shim(
+            example, (patchsize[0] * 14, patchsize[1] * 14), intr_aug=intr_aug
+        )
+
+        return example
+
+    def __getitem__(self, index_tuple: tuple) -> dict:
+        index, num_context_views, patchsize_h = index_tuple
+        patchsize_w = self.cfg.input_image_shape[1] // 14
+        try:
+            return self.getitem(index, num_context_views, (patchsize_h, patchsize_w))
+        except Exception as e:
+            print(f"Error: {e}")
+            traceback.print_exc()
+            index = np.random.randint(len(self))
+            return self.__getitem__((index, num_context_views, patchsize_h))
+
+    def get_bound(
+        self,
+        bound: Literal["near", "far"],
+        num_views: int,
+    ) -> Float[Tensor, " view"]:
+        value = torch.tensor(getattr(self, bound), dtype=torch.float32)
+        return repeat(value, "-> v", v=num_views)
+
+    @property
+    def data_stage(self) -> Stage:
+        if self.cfg.overfit_to_scene is not None:
+            return "test"
+        if self.stage == "val":
+            return "test"
+        return self.stage
+
+    @cached_property
+    def index(self) -> dict[str, Path]:
+        merged_index = {}
+        data_stages = [self.data_stage]
+        if self.cfg.overfit_to_scene is not None:
+            data_stages = ("test", "train")
+        for data_stage in data_stages:
+            for root in self.cfg.roots:
+                # Load the root's index.
+                with (root / data_stage / "index.json").open("r") as f:
+                    index = json.load(f)
+                index = {k: Path(root / data_stage / v) for k, v in index.items()}
+
+                # The constituent datasets should have unique keys.
+                assert not (set(merged_index.keys()) & set(index.keys()))
+
+                # Merge the root's index into the main index.
+                merged_index = {**merged_index, **index}
+        return merged_index
+
+    def __len__(self) -> int:
+        return len(self.scene_ids)

src/dataset/shims/augmentation_shim.py

@@ -0,0 +1,219 @@
+import copy
+import random
+import numpy as np
+import torch
+from jaxtyping import Float
+from torch import Tensor
+
+from ..types import AnyExample, AnyViews
+
+
+def reflect_extrinsics(
+    extrinsics: Float[Tensor, "*batch 4 4"],
+) -> Float[Tensor, "*batch 4 4"]:
+    reflect = torch.eye(4, dtype=torch.float32, device=extrinsics.device)
+    reflect[0, 0] = -1
+    return reflect @ extrinsics @ reflect
+
+
+def reflect_views(views: AnyViews) -> AnyViews:
+    if "depth" in views.keys():
+        return {
+            **views,
+            "image": views["image"].flip(-1),
+            "extrinsics": reflect_extrinsics(views["extrinsics"]),
+            "depth": views["depth"].flip(-1),
+        }
+    else:
+        return {
+            **views,
+            "image": views["image"].flip(-1),
+            "extrinsics": reflect_extrinsics(views["extrinsics"]),
+        }
+
+
+def apply_augmentation_shim(
+    example: AnyExample,
+    generator: torch.Generator | None = None,
+) -> AnyExample:
+    """Randomly augment the training images."""
+    # Do not augment with 50% chance.
+    if torch.rand(tuple(), generator=generator) < 0.5:
+        return example
+    
+    return {
+        **example,
+        "context": reflect_views(example["context"]),
+        "target": reflect_views(example["target"]),
+    }
+
+def rotate_90_degrees(
+    image: torch.Tensor, depth_map: torch.Tensor | None, extri_opencv: torch.Tensor, intri_opencv: torch.Tensor, clockwise=True
+):
+    """
+    Rotates the input image, depth map, and camera parameters by 90 degrees.
+
+    Applies one of two 90-degree rotations:
+    - Clockwise 
+    - Counterclockwise (if clockwise=False)
+
+    The extrinsic and intrinsic matrices are adjusted accordingly to maintain
+    correct camera geometry.
+
+    Args:
+        image (torch.Tensor):
+            Input image tensor of shape (C, H, W).
+        depth_map (torch.Tensor or None):
+            Depth map tensor of shape (H, W), or None if not available.
+        extri_opencv (torch.Tensor):
+            Extrinsic matrix (3x4) in OpenCV convention.
+        intri_opencv (torch.Tensor):
+            Intrinsic matrix (3x3).
+        clockwise (bool):
+            If True, rotates the image 90 degrees clockwise; else 90 degrees counterclockwise.
+
+    Returns:
+        tuple:
+            (
+                rotated_image,
+                rotated_depth_map,
+                new_extri_opencv,
+                new_intri_opencv
+            )
+
+            Where each is the updated version after the rotation.
+    """
+    image_height, image_width = image.shape[-2:]
+    
+    # Rotate the image and depth map
+    rotated_image, rotated_depth_map = rotate_image_and_depth_rot90(image, depth_map, clockwise)
+    # Adjust the intrinsic matrix
+    new_intri_opencv = adjust_intrinsic_matrix_rot90(intri_opencv, image_width, image_height, clockwise)
+    # Adjust the extrinsic matrix
+    new_extri_opencv = adjust_extrinsic_matrix_rot90(extri_opencv, clockwise)
+
+    return (
+        rotated_image,
+        rotated_depth_map,
+        new_extri_opencv,
+        new_intri_opencv,
+    )
+
+
+def rotate_image_and_depth_rot90(image: torch.Tensor, depth_map: torch.Tensor | None, clockwise: bool):
+    """
+    Rotates the given image and depth map by 90 degrees (clockwise or counterclockwise).
+
+    Args:
+        image (torch.Tensor):
+            Input image tensor of shape (C, H, W).
+        depth_map (torch.Tensor or None):
+            Depth map tensor of shape (H, W), or None if not available.
+        clockwise (bool):
+            If True, rotate 90 degrees clockwise; else 90 degrees counterclockwise.
+
+    Returns:
+        tuple:
+            (rotated_image, rotated_depth_map)
+    """
+    rotated_depth_map = None
+    if clockwise:
+        rotated_image = torch.rot90(image, k=-1, dims=[-2, -1])
+        if depth_map is not None:
+            rotated_depth_map = torch.rot90(depth_map, k=-1, dims=[-2, -1])
+    else:
+        rotated_image = torch.rot90(image, k=1, dims=[-2, -1])
+        if depth_map is not None:
+            rotated_depth_map = torch.rot90(depth_map, k=1, dims=[-2, -1])
+    return rotated_image, rotated_depth_map
+
+
+def adjust_extrinsic_matrix_rot90(extri_opencv: torch.Tensor, clockwise: bool):
+    """
+    Adjusts the extrinsic matrix (3x4) for a 90-degree rotation of the image.
+
+    The rotation is in the image plane. This modifies the camera orientation
+    accordingly. The function applies either a clockwise or counterclockwise
+    90-degree rotation.
+
+    Args:
+        extri_opencv (torch.Tensor):
+            Extrinsic matrix (3x4) in OpenCV convention.
+        clockwise (bool):
+            If True, rotate extrinsic for a 90-degree clockwise image rotation;
+            otherwise, counterclockwise.
+
+    Returns:
+        torch.Tensor:
+            A new 3x4 extrinsic matrix after the rotation.
+    """
+    R = extri_opencv[:3, :3]
+    t = extri_opencv[:3, 3]
+
+    if clockwise:
+        R_rotation = torch.tensor([
+            [0, -1, 0],
+            [1,  0, 0],
+            [0,  0, 1]
+        ], dtype=extri_opencv.dtype, device=extri_opencv.device)
+    else:
+        R_rotation = torch.tensor([
+            [0, 1, 0],
+            [-1, 0, 0],
+            [0, 0, 1]
+        ], dtype=extri_opencv.dtype, device=extri_opencv.device)
+
+    new_R = torch.matmul(R_rotation, R)
+    new_t = torch.matmul(R_rotation, t)
+    new_extri_opencv = torch.cat((new_R, new_t.reshape(-1, 1)), dim=1)
+    new_extri_opencv = torch.cat((new_extri_opencv, 
+                                  torch.tensor([[0, 0, 0, 1]], 
+                                dtype=extri_opencv.dtype, device=extri_opencv.device)), dim=0)
+    return new_extri_opencv
+
+
+def adjust_intrinsic_matrix_rot90(intri_opencv: torch.Tensor, image_width: int, image_height: int, clockwise: bool):
+    """
+    Adjusts the intrinsic matrix (3x3) for a 90-degree rotation of the image in the image plane.
+
+    Args:
+        intri_opencv (torch.Tensor):
+            Intrinsic matrix (3x3).
+        image_width (int):
+            Original width of the image.
+        image_height (int):
+            Original height of the image.
+        clockwise (bool):
+            If True, rotate 90 degrees clockwise; else 90 degrees counterclockwise.
+
+    Returns:
+        torch.Tensor:
+            A new 3x3 intrinsic matrix after the rotation.
+    """
+    intri_opencv = copy.deepcopy(intri_opencv)
+    intri_opencv[0, :] *= image_width
+    intri_opencv[1, :] *= image_height
+
+    fx, fy, cx, cy = (
+        intri_opencv[0, 0],
+        intri_opencv[1, 1],
+        intri_opencv[0, 2],
+        intri_opencv[1, 2],
+    )
+
+    new_intri_opencv = torch.eye(3, dtype=intri_opencv.dtype, device=intri_opencv.device)
+    if clockwise:
+        new_intri_opencv[0, 0] = fy
+        new_intri_opencv[1, 1] = fx
+        new_intri_opencv[0, 2] = image_height - cy
+        new_intri_opencv[1, 2] = cx
+    else:
+        new_intri_opencv[0, 0] = fy
+        new_intri_opencv[1, 1] = fx
+        new_intri_opencv[0, 2] = cy
+        new_intri_opencv[1, 2] = image_width - cx
+    
+    new_intri_opencv[0, :] /= image_height
+    new_intri_opencv[1, :] /= image_width
+
+    return new_intri_opencv

src/dataset/shims/bounds_shim.py

@@ -0,0 +1,80 @@
+import torch
+from einops import einsum, reduce, repeat
+from jaxtyping import Float
+from torch import Tensor
+
+from ..types import BatchedExample
+
+
+def compute_depth_for_disparity(
+    extrinsics: Float[Tensor, "batch view 4 4"],
+    intrinsics: Float[Tensor, "batch view 3 3"],
+    image_shape: tuple[int, int],
+    disparity: float,
+    delta_min: float = 1e-6,  # This prevents motionless scenes from lacking depth.
+) -> Float[Tensor, " batch"]:
+    """Compute the depth at which moving the maximum distance between cameras
+    corresponds to the specified disparity (in pixels).
+    """
+
+    # Use the furthest distance between cameras as the baseline.
+    origins = extrinsics[:, :, :3, 3]
+    deltas = (origins[:, None, :, :] - origins[:, :, None, :]).norm(dim=-1)
+    deltas = deltas.clip(min=delta_min)
+    baselines = reduce(deltas, "b v ov -> b", "max")
+
+    # Compute a single pixel's size at depth 1.
+    h, w = image_shape
+    pixel_size = 1 / torch.tensor((w, h), dtype=torch.float32, device=extrinsics.device)
+    pixel_size = einsum(
+        intrinsics[..., :2, :2].inverse(), pixel_size, "... i j, j -> ... i"
+    )
+
+    # This wouldn't make sense with non-square pixels, but then again, non-square pixels
+    # don't make much sense anyway.
+    mean_pixel_size = reduce(pixel_size, "b v xy -> b", "mean")
+
+    return baselines / (disparity * mean_pixel_size)
+
+
+def apply_bounds_shim(
+    batch: BatchedExample,
+    near_disparity: float,
+    far_disparity: float,
+) -> BatchedExample:
+    """Compute reasonable near and far planes (lower and upper bounds on depth). This
+    assumes that all of an example's views are of roughly the same thing.
+    """
+
+    context = batch["context"]
+    _, cv, _, h, w = context["image"].shape
+    
+    # Compute near and far planes using the context views.
+    near = compute_depth_for_disparity(
+        context["extrinsics"],
+        context["intrinsics"],
+        (h, w),
+        near_disparity,
+    )
+    far = compute_depth_for_disparity(
+        context["extrinsics"],
+        context["intrinsics"],
+        (h, w),
+        far_disparity,
+    )
+
+    target = batch["target"]
+    _, tv, _, _, _ = target["image"].shape
+    return {
+        **batch,
+        "context": {
+            **context,
+            "near": repeat(near, "b -> b v", v=cv),
+            "far": repeat(far, "b -> b v", v=cv),
+        },
+        "target": {
+            **target,
+            "near": repeat(near, "b -> b v", v=tv),
+            "far": repeat(far, "b -> b v", v=tv),
+        },
+    }

src/dataset/shims/crop_shim.py

@@ -0,0 +1,196 @@
+import random
+import numpy as np
+import torch
+from einops import rearrange
+from jaxtyping import Float
+from PIL import Image
+from torch import Tensor
+import torchvision.transforms.functional as F
+import cv2
+
+from ..types import AnyExample, AnyViews
+
+
+def rescale(
+    image: Float[Tensor, "3 h_in w_in"],
+    shape: tuple[int, int],
+) -> Float[Tensor, "3 h_out w_out"]:
+    h, w = shape
+    image_new = (image * 255).clip(min=0, max=255).type(torch.uint8)
+    image_new = rearrange(image_new, "c h w -> h w c").detach().cpu().numpy()
+    image_new = Image.fromarray(image_new)
+    image_new = image_new.resize((w, h), Image.LANCZOS)
+    image_new = np.array(image_new) / 255
+    image_new = torch.tensor(image_new, dtype=image.dtype, device=image.device)
+    return rearrange(image_new, "h w c -> c h w")
+
+def rescale_depth(
+    depth: Float[Tensor, "1 h w"],
+    shape: tuple[int, int],
+) -> Float[Tensor, "1 h_out w_out"]:
+    h, w = shape
+    depth_new = depth.detach().cpu().numpy()
+    depth_new = cv2.resize(depth_new, (w,h), interpolation=cv2.INTER_NEAREST)
+    depth_new = torch.from_numpy(depth_new).to(depth.device)
+    return depth_new
+    
+def center_crop(
+    images: Float[Tensor, "*#batch c h w"],
+    intrinsics: Float[Tensor, "*#batch 3 3"],
+    shape: tuple[int, int],
+    depths: Float[Tensor, "*#batch 1 h w"] | None = None,
+) -> tuple[
+    Float[Tensor, "*#batch c h_out w_out"],  # updated images
+    Float[Tensor, "*#batch 3 3"],  # updated intrinsics
+    Float[Tensor, "*#batch 1 h_out w_out"] | None,  # updated depths
+]:
+    *_, h_in, w_in = images.shape
+    h_out, w_out = shape
+
+    # Note that odd input dimensions induce half-pixel misalignments.
+    row = (h_in - h_out) // 2
+    col = (w_in - w_out) // 2
+
+    # Center-crop the image.
+    images = images[..., :, row : row + h_out, col : col + w_out]
+
+    if depths is not None:
+        depths = depths[..., row : row + h_out, col : col + w_out]
+
+    # Adjust the intrinsics to account for the cropping.
+    intrinsics = intrinsics.clone()
+    intrinsics[..., 0, 0] *= w_in / w_out  # fx
+    intrinsics[..., 1, 1] *= h_in / h_out  # fy
+
+    
+    if depths is not None:
+        return images, intrinsics, depths
+    else:
+        return images, intrinsics
+
+
+def rescale_and_crop(
+    images: Float[Tensor, "*#batch c h w"],
+    intrinsics: Float[Tensor, "*#batch 3 3"],
+    shape: tuple[int, int],
+    intr_aug: bool = False,
+    scale_range: tuple[float, float] = (0.77, 1.0),
+    depths: Float[Tensor, "*#batch 1 h w"] | None = None,
+) -> tuple[
+    Float[Tensor, "*#batch c h_out w_out"],  # updated images
+    Float[Tensor, "*#batch 3 3"],  # updated intrinsics
+    Float[Tensor, "*#batch 1 h_out w_out"] | None,  # updated depths
+]:
+    if type(images) == list:
+        images_new = []
+        intrinsics_new = []
+        for i in range(len(images)):
+            image = images[i]
+            intrinsic = intrinsics[i]
+            
+            *_, h_in, w_in = image.shape
+            h_out, w_out = shape
+
+            scale_factor = max(h_out / h_in, w_out / w_in)
+            h_scaled = round(h_in * scale_factor)
+            w_scaled = round(w_in * scale_factor)
+            image = F.resize(image, (h_scaled, w_scaled))
+            image = F.center_crop(image, (h_out, w_out))
+            images_new.append(image)
+            
+            intrinsic_new = intrinsic.clone()
+            intrinsic_new[..., 0, 0] *= w_scaled / w_in  # fx
+            intrinsic_new[..., 1, 1] *= h_scaled / h_in  # fy
+            intrinsics_new.append(intrinsic_new)
+        
+        if depths is not None:
+            depths_new = []
+            for i in range(len(depths)):
+                depth = depths[i]
+                depth = rescale_depth(depth, (h_out, w_out))
+                depth = F.center_crop(depth, (h_out, w_out))
+                depths_new.append(depth)
+            return torch.stack(images_new), torch.stack(intrinsics_new), torch.stack(depths_new)
+        else:
+            return torch.stack(images_new), torch.stack(intrinsics_new)
+    
+    else:
+        # we only support intr_aug for clean datasets
+        *_, h_in, w_in = images.shape
+        h_out, w_out = shape
+        # assert h_out <= h_in and w_out <= w_in # to avoid the case that the image is too small, like co3d
+        
+        if intr_aug:
+            scale = random.uniform(*scale_range)
+            h_scale = round(h_out * scale)
+            w_scale = round(w_out * scale)
+        else:
+            h_scale = h_out
+            w_scale = w_out
+
+        scale_factor = max(h_scale / h_in, w_scale / w_in)
+        h_scaled = round(h_in * scale_factor)
+        w_scaled = round(w_in * scale_factor)
+        assert h_scaled == h_scale or w_scaled == w_scale
+
+        # Reshape the images to the correct size. Assume we don't have to worry about
+        # changing the intrinsics based on how the images are rounded.
+        *batch, c, h, w = images.shape
+        images = images.reshape(-1, c, h, w)
+        images = torch.stack([rescale(image, (h_scaled, w_scaled)) for image in images])
+        images = images.reshape(*batch, c, h_scaled, w_scaled)
+
+        if depths is not None:
+            if type(depths) == list:
+                depths_new = []
+                for i in range(len(depths)):
+                    depth = depths[i]
+                    depth = rescale_depth(depth, (h_scaled, w_scaled)) 
+                    depths_new.append(depth)
+                depths = torch.stack(depths_new)
+            else:
+                depths = depths.reshape(-1, h, w)
+                depths = torch.stack([rescale_depth(depth, (h_scaled, w_scaled)) for depth in depths])
+                depths = depths.reshape(*batch, h_scaled, w_scaled)
+            
+            images, intrinsics, depths = center_crop(images, intrinsics, (h_scale, w_scale), depths)
+
+            if intr_aug:
+                images = F.resize(images, size=(h_out, w_out), interpolation=F.InterpolationMode.BILINEAR)
+                depths = F.resize(depths, size=(h_out, w_out), interpolation=F.InterpolationMode.NEAREST)
+                
+            return images, intrinsics, depths
+        else:
+            images, intrinsics = center_crop(images, intrinsics, (h_scale, w_scale))
+
+            if intr_aug:
+                images = F.resize(images, size=(h_out, w_out))
+
+            return images, intrinsics
+
+
+def apply_crop_shim_to_views(views: AnyViews, shape: tuple[int, int], intr_aug: bool = False) -> AnyViews:
+    if "depth" in views.keys():
+        images, intrinsics, depths = rescale_and_crop(views["image"], views["intrinsics"], shape, depths=views["depth"], intr_aug=intr_aug)
+        return {
+            **views,
+            "image": images,
+            "intrinsics": intrinsics,
+            "depth": depths,
+        }
+    else:
+        images, intrinsics = rescale_and_crop(views["image"], views["intrinsics"], shape, intr_aug)
+        return {
+            **views,
+            "image": images,
+            "intrinsics": intrinsics,
+        }
+        
+
+def apply_crop_shim(example: AnyExample, shape: tuple[int, int], intr_aug: bool = False) -> AnyExample:
+    """Crop images in the example."""
+    return {
+        **example,
+        "context": apply_crop_shim_to_views(example["context"], shape, intr_aug),
+        "target": apply_crop_shim_to_views(example["target"], shape, intr_aug),
+    }

src/dataset/shims/geometry_shim.py

@@ -0,0 +1,383 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# geometry utilitary functions
+# --------------------------------------------------------
+import torch
+import numpy as np
+from scipy.spatial import cKDTree as KDTree
+
+def invalid_to_nans(arr, valid_mask, ndim=999):
+    if valid_mask is not None:
+        arr = arr.clone()
+        arr[~valid_mask] = float('nan')
+    if arr.ndim > ndim:
+        arr = arr.flatten(-2 - (arr.ndim - ndim), -2)
+    return arr
+
+
+def invalid_to_zeros(arr, valid_mask, ndim=999):
+    if valid_mask is not None:
+        arr = arr.clone()
+        arr[~valid_mask] = 0
+        nnz = valid_mask.view(len(valid_mask), -1).sum(1)
+    else:
+        nnz = arr.numel() // len(arr) if len(arr) else 0  # number of point per image
+    if arr.ndim > ndim:
+        arr = arr.flatten(-2 - (arr.ndim - ndim), -2)
+    return arr, nnz
+
+
+def xy_grid(W, H, device=None, origin=(0, 0), unsqueeze=None, cat_dim=-1, homogeneous=False, **arange_kw):
+    """ Output a (H,W,2) array of int32 
+        with output[j,i,0] = i + origin[0]
+             output[j,i,1] = j + origin[1]
+    """
+    if device is None:
+        # numpy
+        arange, meshgrid, stack, ones = np.arange, np.meshgrid, np.stack, np.ones
+    else:
+        # torch
+        arange = lambda *a, **kw: torch.arange(*a, device=device, **kw)
+        meshgrid, stack = torch.meshgrid, torch.stack
+        ones = lambda *a: torch.ones(*a, device=device)
+
+    tw, th = [arange(o, o + s, **arange_kw) for s, o in zip((W, H), origin)]
+    grid = meshgrid(tw, th, indexing='xy')
+    if homogeneous:
+        grid = grid + (ones((H, W)),)
+    if unsqueeze is not None:
+        grid = (grid[0].unsqueeze(unsqueeze), grid[1].unsqueeze(unsqueeze))
+    if cat_dim is not None:
+        grid = stack(grid, cat_dim)
+    return grid
+
+
+def geotrf(Trf, pts, ncol=None, norm=False):
+    """ Apply a geometric transformation to a list of 3-D points.
+
+    H: 3x3 or 4x4 projection matrix (typically a Homography)
+    p: numpy/torch/tuple of coordinates. Shape must be (...,2) or (...,3)
+
+    ncol: int. number of columns of the result (2 or 3)
+    norm: float. if != 0, the resut is projected on the z=norm plane.
+
+    Returns an array of projected 2d points.
+    """
+    assert Trf.ndim >= 2
+    if isinstance(Trf, np.ndarray):
+        pts = np.asarray(pts)
+    elif isinstance(Trf, torch.Tensor):
+        pts = torch.as_tensor(pts, dtype=Trf.dtype)
+
+    # adapt shape if necessary
+    output_reshape = pts.shape[:-1]
+    ncol = ncol or pts.shape[-1]
+
+    # optimized code
+    if (isinstance(Trf, torch.Tensor) and isinstance(pts, torch.Tensor) and
+            Trf.ndim == 3 and pts.ndim == 4):
+        d = pts.shape[3]
+        if Trf.shape[-1] == d:
+            pts = torch.einsum("bij, bhwj -> bhwi", Trf, pts)
+        elif Trf.shape[-1] == d + 1:
+            pts = torch.einsum("bij, bhwj -> bhwi", Trf[:, :d, :d], pts) + Trf[:, None, None, :d, d]
+        else:
+            raise ValueError(f'bad shape, not ending with 3 or 4, for {pts.shape=}')
+    else:
+        if Trf.ndim >= 3:
+            n = Trf.ndim - 2
+            assert Trf.shape[:n] == pts.shape[:n], 'batch size does not match'
+            Trf = Trf.reshape(-1, Trf.shape[-2], Trf.shape[-1])
+
+            if pts.ndim > Trf.ndim:
+                # Trf == (B,d,d) & pts == (B,H,W,d) --> (B, H*W, d)
+                pts = pts.reshape(Trf.shape[0], -1, pts.shape[-1])
+            elif pts.ndim == 2:
+                # Trf == (B,d,d) & pts == (B,d) --> (B, 1, d)
+                pts = pts[:, None, :]
+
+        if pts.shape[-1] + 1 == Trf.shape[-1]:
+            Trf = Trf.swapaxes(-1, -2)  # transpose Trf
+            pts = pts @ Trf[..., :-1, :] + Trf[..., -1:, :]
+        elif pts.shape[-1] == Trf.shape[-1]:
+            Trf = Trf.swapaxes(-1, -2)  # transpose Trf
+            pts = pts @ Trf
+        else:
+            pts = Trf @ pts.T
+            if pts.ndim >= 2:
+                pts = pts.swapaxes(-1, -2)
+
+    if norm:
+        pts = pts / pts[..., -1:]  # DONT DO /= BECAUSE OF WEIRD PYTORCH BUG
+        if norm != 1:
+            pts *= norm
+
+    res = pts[..., :ncol].reshape(*output_reshape, ncol)
+    return res
+
+
+def inv(mat):
+    """ Invert a torch or numpy matrix
+    """
+    if isinstance(mat, torch.Tensor):
+        return torch.linalg.inv(mat)
+    if isinstance(mat, np.ndarray):
+        return np.linalg.inv(mat)
+    raise ValueError(f'bad matrix type = {type(mat)}')
+
+
+def depthmap_to_pts3d(depth, pseudo_focal, pp=None, **_):
+    """
+    Args:
+        - depthmap (BxHxW array):
+        - pseudo_focal: [B,H,W] ; [B,2,H,W] or [B,1,H,W]
+    Returns:
+        pointmap of absolute coordinates (BxHxWx3 array)
+    """
+
+    if len(depth.shape) == 4:
+        B, H, W, n = depth.shape
+    else:
+        B, H, W = depth.shape
+        n = None
+
+    if len(pseudo_focal.shape) == 3:  # [B,H,W]
+        pseudo_focalx = pseudo_focaly = pseudo_focal
+    elif len(pseudo_focal.shape) == 4:  # [B,2,H,W] or [B,1,H,W]
+        pseudo_focalx = pseudo_focal[:, 0]
+        if pseudo_focal.shape[1] == 2:
+            pseudo_focaly = pseudo_focal[:, 1]
+        else:
+            pseudo_focaly = pseudo_focalx
+    else:
+        raise NotImplementedError("Error, unknown input focal shape format.")
+
+    assert pseudo_focalx.shape == depth.shape[:3]
+    assert pseudo_focaly.shape == depth.shape[:3]
+    grid_x, grid_y = xy_grid(W, H, cat_dim=0, device=depth.device)[:, None]
+
+    # set principal point
+    if pp is None:
+        grid_x = grid_x - (W - 1) / 2
+        grid_y = grid_y - (H - 1) / 2
+    else:
+        grid_x = grid_x.expand(B, -1, -1) - pp[:, 0, None, None]
+        grid_y = grid_y.expand(B, -1, -1) - pp[:, 1, None, None]
+
+    if n is None:
+        pts3d = torch.empty((B, H, W, 3), device=depth.device)
+        pts3d[..., 0] = depth * grid_x / pseudo_focalx
+        pts3d[..., 1] = depth * grid_y / pseudo_focaly
+        pts3d[..., 2] = depth
+    else:
+        pts3d = torch.empty((B, H, W, 3, n), device=depth.device)
+        pts3d[..., 0, :] = depth * (grid_x / pseudo_focalx)[..., None]
+        pts3d[..., 1, :] = depth * (grid_y / pseudo_focaly)[..., None]
+        pts3d[..., 2, :] = depth
+    return pts3d
+
+
+def depthmap_to_camera_coordinates(depthmap, camera_intrinsics, pseudo_focal=None):
+    """
+    Args:
+        - depthmap (HxW array):
+        - camera_intrinsics: a 3x3 matrix
+    Returns:
+        pointmap of absolute coordinates (HxWx3 array), and a mask specifying valid pixels.
+    """
+    camera_intrinsics = np.float32(camera_intrinsics)
+    H, W = depthmap.shape
+
+    # Compute 3D ray associated with each pixel
+    # Strong assumption: there are no skew terms
+    assert camera_intrinsics[0, 1] == 0.0
+    assert camera_intrinsics[1, 0] == 0.0
+    if pseudo_focal is None:
+        fu = camera_intrinsics[0, 0]
+        fv = camera_intrinsics[1, 1]
+    else:
+        assert pseudo_focal.shape == (H, W)
+        fu = fv = pseudo_focal
+    cu = camera_intrinsics[0, 2]
+    cv = camera_intrinsics[1, 2]
+
+    u, v = np.meshgrid(np.arange(W), np.arange(H))
+    z_cam = depthmap
+    x_cam = (u - cu) * z_cam / fu
+    y_cam = (v - cv) * z_cam / fv
+    X_cam = np.stack((x_cam, y_cam, z_cam), axis=-1).astype(np.float32)
+    
+    # Mask for valid coordinates
+    valid_mask = (depthmap > 0.0)
+    return X_cam, valid_mask
+
+
+def depthmap_to_absolute_camera_coordinates(depthmap, camera_intrinsics, camera_pose=None, **kw):
+    """
+    Args:
+        - depthmap (HxW array):
+        - camera_intrinsics: a 3x3 matrix
+        - camera_pose: a 4x3 or 4x4 cam2world matrix
+    Returns:
+        pointmap of absolute coordinates (HxWx3 array), and a mask specifying valid pixels."""
+    X_cam, valid_mask = depthmap_to_camera_coordinates(depthmap, camera_intrinsics)
+    
+    X_world = X_cam # default
+    if camera_pose is not None:
+        # R_cam2world = np.float32(camera_params["R_cam2world"])
+        # t_cam2world = np.float32(camera_params["t_cam2world"]).squeeze()
+        R_cam2world = camera_pose[:3, :3]
+        t_cam2world = camera_pose[:3, 3]
+        
+        # Express in absolute coordinates (invalid depth values)
+        X_world = np.einsum("ik, vuk -> vui", R_cam2world, X_cam) + t_cam2world[None, None, :]
+
+    return X_world, valid_mask
+
+
+def colmap_to_opencv_intrinsics(K):
+    """
+    Modify camera intrinsics to follow a different convention.
+    Coordinates of the center of the top-left pixels are by default:
+    - (0.5, 0.5) in Colmap
+    - (0,0) in OpenCV
+    """
+    K = K.copy()
+    K[0, 2] -= 0.5
+    K[1, 2] -= 0.5
+    return K
+
+
+def opencv_to_colmap_intrinsics(K):
+    """
+    Modify camera intrinsics to follow a different convention.
+    Coordinates of the center of the top-left pixels are by default:
+    - (0.5, 0.5) in Colmap
+    - (0,0) in OpenCV
+    """
+    K = K.copy()
+    K[0, 2] += 0.5
+    K[1, 2] += 0.5
+    return K
+
+
+def normalize_pointcloud(pts1, pts2, norm_mode='avg_dis', valid1=None, valid2=None, ret_factor=False):
+    """ renorm pointmaps pts1, pts2 with norm_mode
+    """
+    assert pts1.ndim >= 3 and pts1.shape[-1] == 3
+    assert pts2 is None or (pts2.ndim >= 3 and pts2.shape[-1] == 3)
+    norm_mode, dis_mode = norm_mode.split('_')
+
+    if norm_mode == 'avg':
+        # gather all points together (joint normalization)
+        nan_pts1, nnz1 = invalid_to_zeros(pts1, valid1, ndim=3)
+        nan_pts2, nnz2 = invalid_to_zeros(pts2, valid2, ndim=3) if pts2 is not None else (None, 0)
+        all_pts = torch.cat((nan_pts1, nan_pts2), dim=1) if pts2 is not None else nan_pts1
+
+        # compute distance to origin
+        all_dis = all_pts.norm(dim=-1)
+        if dis_mode == 'dis':
+            pass  # do nothing
+        elif dis_mode == 'log1p':
+            all_dis = torch.log1p(all_dis)
+        elif dis_mode == 'warp-log1p':
+            # actually warp input points before normalizing them
+            log_dis = torch.log1p(all_dis)
+            warp_factor = log_dis / all_dis.clip(min=1e-8)
+            H1, W1 = pts1.shape[1:-1]
+            pts1 = pts1 * warp_factor[:, :W1 * H1].view(-1, H1, W1, 1)
+            if pts2 is not None:
+                H2, W2 = pts2.shape[1:-1]
+                pts2 = pts2 * warp_factor[:, W1 * H1:].view(-1, H2, W2, 1)
+            all_dis = log_dis  # this is their true distance afterwards
+        else:
+            raise ValueError(f'bad {dis_mode=}')
+
+        norm_factor = all_dis.sum(dim=1) / (nnz1 + nnz2 + 1e-8)
+    else:
+        # gather all points together (joint normalization)
+        nan_pts1 = invalid_to_nans(pts1, valid1, ndim=3)
+        nan_pts2 = invalid_to_nans(pts2, valid2, ndim=3) if pts2 is not None else None
+        all_pts = torch.cat((nan_pts1, nan_pts2), dim=1) if pts2 is not None else nan_pts1
+
+        # compute distance to origin
+        all_dis = all_pts.norm(dim=-1)
+
+        if norm_mode == 'avg':
+            norm_factor = all_dis.nanmean(dim=1)
+        elif norm_mode == 'median':
+            norm_factor = all_dis.nanmedian(dim=1).values.detach()
+        elif norm_mode == 'sqrt':
+            norm_factor = all_dis.sqrt().nanmean(dim=1)**2
+        else:
+            raise ValueError(f'bad {norm_mode=}')
+
+    norm_factor = norm_factor.clip(min=1e-8)
+    while norm_factor.ndim < pts1.ndim:
+        norm_factor.unsqueeze_(-1)
+
+    res = pts1 / norm_factor
+    if pts2 is not None:
+        res = (res, pts2 / norm_factor)
+    if ret_factor:
+        res = res + (norm_factor,)
+    return res
+
+
+@torch.no_grad()
+def get_joint_pointcloud_depth(z1, z2, valid_mask1, valid_mask2=None, quantile=0.5):
+    # set invalid points to NaN
+    _z1 = invalid_to_nans(z1, valid_mask1).reshape(len(z1), -1)
+    _z2 = invalid_to_nans(z2, valid_mask2).reshape(len(z2), -1) if z2 is not None else None
+    _z = torch.cat((_z1, _z2), dim=-1) if z2 is not None else _z1
+
+    # compute median depth overall (ignoring nans)
+    if quantile == 0.5:
+        shift_z = torch.nanmedian(_z, dim=-1).values
+    else:
+        shift_z = torch.nanquantile(_z, quantile, dim=-1)
+    return shift_z  # (B,)
+
+
+@torch.no_grad()
+def get_joint_pointcloud_center_scale(pts1, pts2, valid_mask1=None, valid_mask2=None, z_only=False, center=True):
+    # set invalid points to NaN
+    _pts1 = invalid_to_nans(pts1, valid_mask1).reshape(len(pts1), -1, 3)
+    _pts2 = invalid_to_nans(pts2, valid_mask2).reshape(len(pts2), -1, 3) if pts2 is not None else None
+    _pts = torch.cat((_pts1, _pts2), dim=1) if pts2 is not None else _pts1
+
+    # compute median center
+    _center = torch.nanmedian(_pts, dim=1, keepdim=True).values  # (B,1,3)
+    if z_only:
+        _center[..., :2] = 0  # do not center X and Y
+
+    # compute median norm
+    _norm = ((_pts - _center) if center else _pts).norm(dim=-1)
+    scale = torch.nanmedian(_norm, dim=1).values
+    return _center[:, None, :, :], scale[:, None, None, None]
+
+    
+def find_reciprocal_matches(P1, P2):
+    """
+    returns 3 values:
+    1 - reciprocal_in_P2: a boolean array of size P2.shape[0], a "True" value indicates a match
+    2 - nn2_in_P1: a int array of size P2.shape[0], it contains the indexes of the closest points in P1
+    3 - reciprocal_in_P2.sum(): the number of matches
+    """
+    tree1 = KDTree(P1)
+    tree2 = KDTree(P2)
+
+    _, nn1_in_P2 = tree2.query(P1, workers=8)
+    _, nn2_in_P1 = tree1.query(P2, workers=8)
+
+    reciprocal_in_P1 = (nn2_in_P1[nn1_in_P2] == np.arange(len(nn1_in_P2)))
+    reciprocal_in_P2 = (nn1_in_P2[nn2_in_P1] == np.arange(len(nn2_in_P1)))
+    assert reciprocal_in_P1.sum() == reciprocal_in_P2.sum()
+    return reciprocal_in_P2, nn2_in_P1, reciprocal_in_P2.sum()
+
+
+def get_med_dist_between_poses(poses):
+    from scipy.spatial.distance import pdist
+    return np.median(pdist([p[:3, 3].detach().cpu().numpy() for p in poses]))

src/dataset/shims/load_shim.py

@@ -0,0 +1,12 @@
+import cv2
+
+def imread_cv2(path, options=cv2.IMREAD_COLOR):
+    """Open an image or a depthmap with opencv-python."""
+    if path.endswith((".exr", "EXR")):
+        options = cv2.IMREAD_ANYDEPTH
+    img = cv2.imread(path, options)
+    if img is None:
+        raise IOError(f"Could not load image={path} with {options=}")
+    if img.ndim == 3:
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    return img

src/dataset/shims/normalize_shim.py

@@ -0,0 +1,27 @@
+import torch
+from einops import einsum, reduce, repeat
+from jaxtyping import Float
+from torch import Tensor
+
+from ..types import BatchedExample
+
+
+def inverse_normalize_image(tensor, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)):
+    mean = torch.as_tensor(mean, dtype=tensor.dtype, device=tensor.device).view(-1, 1, 1)
+    std = torch.as_tensor(std, dtype=tensor.dtype, device=tensor.device).view(-1, 1, 1)
+    return tensor * std + mean
+
+
+def normalize_image(tensor, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)):
+    mean = torch.as_tensor(mean, dtype=tensor.dtype, device=tensor.device).view(-1, 1, 1)
+    std = torch.as_tensor(std, dtype=tensor.dtype, device=tensor.device).view(-1, 1, 1)
+    return (tensor - mean) / std
+
+
+def apply_normalize_shim(
+    batch: BatchedExample,
+    mean: tuple[float, float, float] = (0.5, 0.5, 0.5),
+    std: tuple[float, float, float] = (0.5, 0.5, 0.5),
+) -> BatchedExample:
+    batch["context"]["image"] = normalize_image(batch["context"]["image"], mean, std)
+    return batch

src/dataset/shims/patch_shim.py

@@ -0,0 +1,38 @@
+from ..types import BatchedExample, BatchedViews
+
+
+def apply_patch_shim_to_views(views: BatchedViews, patch_size: int) -> BatchedViews:
+    _, _, _, h, w = views["image"].shape
+
+    # Image size must be even so that naive center-cropping does not cause misalignment.
+    assert h % 2 == 0 and w % 2 == 0
+
+    h_new = (h // patch_size) * patch_size
+    row = (h - h_new) // 2
+    w_new = (w // patch_size) * patch_size
+    col = (w - w_new) // 2
+
+    # Center-crop the image.
+    image = views["image"][:, :, :, row : row + h_new, col : col + w_new]
+
+    # Adjust the intrinsics to account for the cropping.
+    intrinsics = views["intrinsics"].clone()
+    intrinsics[:, :, 0, 0] *= w / w_new  # fx
+    intrinsics[:, :, 1, 1] *= h / h_new  # fy
+
+    return {
+        **views,
+        "image": image,
+        "intrinsics": intrinsics,
+    }
+
+
+def apply_patch_shim(batch: BatchedExample, patch_size: int) -> BatchedExample:
+    """Crop images in the batch so that their dimensions are cleanly divisible by the
+    specified patch size.
+    """
+    return {
+        **batch,
+        "context": apply_patch_shim_to_views(batch["context"], patch_size),
+        "target": apply_patch_shim_to_views(batch["target"], patch_size),
+    }

src/dataset/types.py

@@ -0,0 +1,51 @@
+from typing import Callable, Literal, TypedDict
+
+from jaxtyping import Float, Int64
+from torch import Tensor
+
+Stage = Literal["train", "val", "test"]
+
+
+# The following types mainly exist to make type-hinted keys show up in VS Code. Some
+# dimensions are annotated as "_" because either:
+# 1. They're expected to change as part of a function call (e.g., resizing the dataset).
+# 2. They're expected to vary within the same function call (e.g., the number of views,
+#    which differs between context and target BatchedViews).
+
+
+class BatchedViews(TypedDict, total=False):
+    extrinsics: Float[Tensor, "batch _ 4 4"]  # batch view 4 4
+    intrinsics: Float[Tensor, "batch _ 3 3"]  # batch view 3 3
+    image: Float[Tensor, "batch _ _ _ _"]  # batch view channel height width
+    near: Float[Tensor, "batch _"]  # batch view
+    far: Float[Tensor, "batch _"]  # batch view
+    index: Int64[Tensor, "batch _"]  # batch view
+    overlap: Float[Tensor, "batch _"]  # batch view
+
+
+class BatchedExample(TypedDict, total=False):
+    target: BatchedViews
+    context: BatchedViews
+    scene: list[str]
+
+
+class UnbatchedViews(TypedDict, total=False):
+    extrinsics: Float[Tensor, "_ 4 4"]
+    intrinsics: Float[Tensor, "_ 3 3"]
+    image: Float[Tensor, "_ 3 height width"]
+    near: Float[Tensor, " _"]
+    far: Float[Tensor, " _"]
+    index: Int64[Tensor, " _"]
+
+
+class UnbatchedExample(TypedDict, total=False):
+    target: UnbatchedViews
+    context: UnbatchedViews
+    scene: str
+
+
+# A data shim modifies the example after it's been returned from the data loader.
+DataShim = Callable[[BatchedExample], BatchedExample]
+
+AnyExample = BatchedExample | UnbatchedExample
+AnyViews = BatchedViews | UnbatchedViews

src/dataset/validation_wrapper.py

@@ -0,0 +1,33 @@
+from typing import Iterator, Optional
+
+import torch
+from torch.utils.data import Dataset, IterableDataset
+
+
+class ValidationWrapper(Dataset):
+    """Wraps a dataset so that PyTorch Lightning's validation step can be turned into a
+    visualization step.
+    """
+
+    dataset: Dataset
+    dataset_iterator: Optional[Iterator]
+    length: int
+
+    def __init__(self, dataset: Dataset, length: int) -> None:
+        super().__init__()
+        self.dataset = dataset
+        self.length = length
+        self.dataset_iterator = None
+
+    def __len__(self):
+        return self.length
+        
+    def __getitem__(self, index: tuple):
+        if isinstance(self.dataset, IterableDataset):
+            if self.dataset_iterator is None:
+                self.dataset_iterator = iter(self.dataset)
+            return next(self.dataset_iterator)
+        
+        random_index = torch.randint(0, len(self.dataset), tuple())
+        random_context_num = torch.randint(2, self.dataset.view_sampler.num_context_views + 1, tuple())
+        return self.dataset[random_index.item(), random_context_num.item()]

src/dataset/view_sampler/__init__.py

@@ -0,0 +1,40 @@
+from typing import Any
+
+from ...misc.step_tracker import StepTracker
+from ..types import Stage
+from .view_sampler import ViewSampler
+from .view_sampler_all import ViewSamplerAll, ViewSamplerAllCfg
+from .view_sampler_arbitrary import ViewSamplerArbitrary, ViewSamplerArbitraryCfg
+from .view_sampler_bounded import ViewSamplerBounded, ViewSamplerBoundedCfg
+from .view_sampler_evaluation import ViewSamplerEvaluation, ViewSamplerEvaluationCfg
+from .view_sampler_rank import ViewSamplerRank, ViewSamplerRankCfg
+VIEW_SAMPLERS: dict[str, ViewSampler[Any]] = {
+    "all": ViewSamplerAll,
+    "arbitrary": ViewSamplerArbitrary,
+    "bounded": ViewSamplerBounded,
+    "evaluation": ViewSamplerEvaluation,
+    "rank": ViewSamplerRank,
+}
+
+ViewSamplerCfg = (
+    ViewSamplerArbitraryCfg
+    | ViewSamplerBoundedCfg
+    | ViewSamplerEvaluationCfg
+    | ViewSamplerAllCfg
+    | ViewSamplerRankCfg
+)
+
+def get_view_sampler(
+    cfg: ViewSamplerCfg,
+    stage: Stage,
+    overfit: bool,
+    cameras_are_circular: bool,
+    step_tracker: StepTracker | None,
+) -> ViewSampler[Any]:
+    return VIEW_SAMPLERS[cfg.name](
+        cfg,
+        stage,
+        overfit,
+        cameras_are_circular,
+        step_tracker,
+    )

src/dataset/view_sampler/three_view_hack.py

@@ -0,0 +1,10 @@
+import torch
+from jaxtyping import Int
+from torch import Tensor
+
+
+def add_third_context_index(
+    indices: Int[Tensor, "*batch 2"]
+) -> Int[Tensor, "*batch 3"]:
+    left, right = indices.unbind(dim=-1)
+    return torch.stack((left, (left + right) // 2, right), dim=-1)