Upload 119 files

.gitattributes

@@ -3,3 +3,4 @@
 *.safetensors filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
 examples/omnirooms/*.jpg filter=lfs diff=lfs merge=lfs -text
+unisharp/cli/__pycache__/unified_trainer.cpython-313.pyc filter=lfs diff=lfs merge=lfs -text

unisharp/__init__.py

@@ -0,0 +1 @@
+DEFAULT_MAX_DEPTH_M: float = 100.0

unisharp/cli/__init__.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+import click
+
+from .train_feature import train_feature_cli
+
+
+@click.group()
+def main_cli():
+    pass
+
+main_cli.add_command(train_feature_cli, "train-feature")
+

unisharp/cli/__main__.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from unisharp.cli import main_cli
+
+
+def main() -> None:
+    main_cli()
+
+
+if __name__ == "__main__":
+    main()
+

unisharp/cli/__pycache__/unified_trainer.cpython-313.pyc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b2698667885fba54eef04bacbbee4bbf897c0dc3df57e6fe7d10ba185a76d2ed
+size 103553

unisharp/cli/mixed_sampler.py

@@ -0,0 +1,80 @@
+
+from __future__ import annotations
+
+import random
+from typing import Any, Iterator
+
+from torch.utils.data import Dataset, IterableDataset
+
+
+class LazyDataLoaderIterator:
+
+    def __init__(self, dataloader: Any):
+        self.dataloader = dataloader
+        self.iterator: Iterator[Any] | None = None
+
+    def __next__(self) -> Any:
+        if self.iterator is None:
+            self.iterator = iter(self.dataloader)
+        return next(self.iterator)
+
+
+class MixedDatasetSampler:
+
+    def __init__(
+        self,
+        datasets: dict[str, Dataset | IterableDataset],
+        weights: dict[str, float],
+        iterators: dict[str, Iterator[Any]],
+        seed: int | None = None,
+    ):
+        self.datasets = datasets
+        self.weights = weights
+        self.iterators = iterators
+        self._rng = random.Random(seed)
+
+        if len(weights) == 0:
+            raise ValueError("weights is empty")
+        for name, w in weights.items():
+            if float(w) <= 0.0:
+                raise ValueError(f"Dataset weight must be > 0, got {name}={float(w)}")
+            if name not in datasets:
+                raise ValueError(f"Unknown dataset in weights: {name}")
+            if name not in iterators:
+                raise ValueError(f"Missing iterator for dataset: {name}")
+
+        total_weight = float(sum(float(v) for v in weights.values()))
+        self.probs = {name: float(w) / total_weight for name, w in weights.items()}
+        self.dataset_names = list(datasets.keys())
+        self.prob_list = [self.probs[name] for name in self.dataset_names]
+
+    def sample(self) -> tuple[str, Any]:
+        dataset_name = self.choose_dataset_name()
+        batch = self.next_batch(dataset_name)
+        return dataset_name, batch
+
+    def choose_dataset_name(self, allowed_dataset_names: list[str] | None = None) -> str:
+        if allowed_dataset_names is None:
+            names = self.dataset_names
+            probs = self.prob_list
+        else:
+            names = [name for name in self.dataset_names if name in set(allowed_dataset_names)]
+            if len(names) == 0:
+                raise ValueError("No allowed dataset names available for sampling.")
+            probs = [self.probs[name] for name in names]
+        return self._rng.choices(names, weights=probs, k=1)[0]
+
+    def next_batch(self, dataset_name: str) -> Any:
+        if dataset_name not in self.iterators:
+            raise ValueError(f"Unknown dataset iterator: {dataset_name}")
+        try:
+            batch = next(self.iterators[dataset_name])
+        except StopIteration as exc:
+            raise StopIteration(f"Dataset {dataset_name} exhausted") from exc
+        return batch
+
+    def get_sampling_stats(self) -> dict[str, float]:
+        return {
+            "probabilities": self.probs.copy(),
+            "sampling": self.weights.copy(),
+        }

unisharp/cli/train_feature.py

@@ -0,0 +1,1410 @@
+from __future__ import annotations
+
+import csv
+import json
+import logging
+import os
+import random
+import sys
+import time
+from dataclasses import fields, is_dataclass, replace
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Any
+
+import click
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+from unisharp.datasets.re10k import Re10KDataset, re10k_collate, re10k_passthrough
+from unisharp.datasets.wildrgbd import WildRGBDDataset, wildrgbd_collate
+from unisharp.datasets.dl3dv import DL3DVDataset
+from unisharp.datasets.scannetpp_fisheye import ScannetppFisheyeDataset, scannetpp_fisheye_passthrough
+from unisharp.datasets.sim_panorama import SimPanoramaDataset
+from unisharp.datasets.panogs import PanOGSDataset, panogs_collate
+from unisharp.losses import UnisharpLoss, UnisharpLossWeights
+from unisharp.models.unisharp_feature import UnisharpFeatureModel, UnisharpFeatureConfig
+from unisharp.utils import logging as logging_utils
+from unisharp import DEFAULT_MAX_DEPTH_M
+from unisharp.utils.gsplat import GSplatRenderer
+from unisharp.utils.io import save_image
+from unisharp.utils.rayfit_camera import scale_pinhole_intrinsics
+from unisharp.utils.unified_vis import save_pair_visualization
+
+from .mixed_sampler import LazyDataLoaderIterator, MixedDatasetSampler  # type: ignore[import]
+from .train_utils import warmup_cosine_lr  # type: ignore[import]
+
+LOGGER = logging.getLogger(__name__)
+REPO_ROOT = Path(__file__).resolve().parents[2]
+
+
+def _default_dataset_manifest_file(name: str) -> Path:
+    parent_path = REPO_ROOT.parent / "dataset_manifests" / name
+    if parent_path.exists():
+        return parent_path
+    return REPO_ROOT / "dataset_manifests" / name
+
+
+DEFAULT_WILDRGBD_ROOTS_FILE = _default_dataset_manifest_file("wildrgbd_roots.txt")
+
+
+def _multiple_aligned_hw(hw: tuple[int, int], multiple: int) -> tuple[int, int]:
+    h, w = int(hw[0]), int(hw[1])
+    m = int(multiple)
+    if m <= 1:
+        return h, w
+    out_h = max(m, (h // m) * m)
+    out_w = max(m, (w // m) * m)
+    return min(out_h, h), min(out_w, w)
+
+
+def _erp_multiple_aligned_hw(hw: tuple[int, int], multiple: int) -> tuple[int, int]:
+    h, w = int(hw[0]), int(hw[1])
+    m = int(multiple)
+    if m <= 1:
+        return h, w
+    max_h_from_h = h // m
+    max_h_from_w = w // (2 * m)
+    h_units = min(max_h_from_h, max_h_from_w)
+    if h_units <= 0:
+        return h, w
+    out_h = h_units * m
+    return out_h, 2 * out_h
+
+
+def _resize_chw_tensor(x: torch.Tensor, dst_hw: tuple[int, int], *, kind: str) -> torch.Tensor:
+    if not torch.is_tensor(x) or x.ndim < 3:
+        return x
+    src_hw = (int(x.shape[-2]), int(x.shape[-1]))
+    if src_hw == tuple(int(v) for v in dst_hw):
+        return x
+    orig_dtype = x.dtype
+    flat = x.reshape(-1, int(x.shape[-3]), src_hw[0], src_hw[1]).to(dtype=torch.float32)
+    if kind == "image":
+        y = F.interpolate(flat, size=dst_hw, mode="bilinear", align_corners=False)
+        y = y.round().clamp(0.0, 255.0).to(dtype=orig_dtype) if orig_dtype == torch.uint8 else y.to(dtype=orig_dtype)
+    elif kind == "ray":
+        y = F.interpolate(flat, size=dst_hw, mode="bilinear", align_corners=False)
+        y = y / torch.linalg.vector_norm(y, dim=1, keepdim=True).clamp(min=1e-6)
+        y = y.to(dtype=orig_dtype)
+    else:
+        y = F.interpolate(flat, size=dst_hw, mode="nearest").to(dtype=orig_dtype)
+    return y.reshape(*x.shape[:-2], int(dst_hw[0]), int(dst_hw[1])).contiguous()
+
+
+def _resize_cube_tensor(x: torch.Tensor, dst_hw: tuple[int, int], *, kind: str) -> torch.Tensor:
+    if not torch.is_tensor(x) or x.ndim < 4:
+        return x
+    src_hw = (int(x.shape[-3]), int(x.shape[-2]))
+    if src_hw == tuple(int(v) for v in dst_hw):
+        return x
+    orig_dtype = x.dtype
+    channels = int(x.shape[-1])
+    flat = x.reshape(-1, src_hw[0], src_hw[1], channels).permute(0, 3, 1, 2).to(dtype=torch.float32)
+    if kind == "image":
+        y = F.interpolate(flat, size=dst_hw, mode="bilinear", align_corners=False)
+        y = y.round().clamp(0.0, 255.0).to(dtype=orig_dtype) if orig_dtype == torch.uint8 else y.to(dtype=orig_dtype)
+    else:
+        y = F.interpolate(flat, size=dst_hw, mode="nearest").to(dtype=orig_dtype)
+    y = y.permute(0, 2, 3, 1)
+    return y.reshape(*x.shape[:-3], int(dst_hw[0]), int(dst_hw[1]), channels).contiguous()
+
+
+def _training_batch_src_hw(batch: Any) -> tuple[int, int] | None:
+    for name in ("src_rgb_u8", "src_erp_rgb_u8"):
+        value = getattr(batch, name, None)
+        if torch.is_tensor(value) and value.ndim >= 3:
+            return int(value.shape[-2]), int(value.shape[-1])
+    return None
+
+
+def _scale_fisheye624_params_any(params: torch.Tensor, *, src_hw: tuple[int, int], dst_hw: tuple[int, int]) -> torch.Tensor:
+    if tuple(int(x) for x in src_hw) == tuple(int(x) for x in dst_hw):
+        return params
+    src_h, src_w = int(src_hw[0]), int(src_hw[1])
+    dst_h, dst_w = int(dst_hw[0]), int(dst_hw[1])
+    sx = float(dst_w) / float(max(src_w, 1))
+    sy = float(dst_h) / float(max(src_h, 1))
+    out = params.clone()
+    out[..., 0] *= sx
+    out[..., 1] *= sy
+    out[..., 2] = (out[..., 2] + 0.5) * sx - 0.5
+    out[..., 3] = (out[..., 3] + 0.5) * sy - 0.5
+    return out
+
+
+def _resize_training_batch_to_multiple(batch: Any, multiple: int) -> Any:
+    if int(multiple) <= 1 or not is_dataclass(batch):
+        return batch
+    src_hw = _training_batch_src_hw(batch)
+    if src_hw is None:
+        return batch
+
+    def _view_hw(prefix: str) -> tuple[int, int] | None:
+        for rgb_name in (f"{prefix}_rgb_u8", f"{prefix}_erp_rgb_u8"):
+            rgb = getattr(batch, rgb_name, None)
+            if torch.is_tensor(rgb) and rgb.ndim >= 3:
+                return int(rgb.shape[-2]), int(rgb.shape[-1])
+        return None
+
+    def _aligned_view_hw(prefix: str, hw: tuple[int, int]) -> tuple[int, int]:
+        is_view_erp = torch.is_tensor(getattr(batch, f"{prefix}_erp_rgb_u8", None))
+        return (
+            _erp_multiple_aligned_hw(hw, int(multiple))
+            if bool(is_view_erp)
+            else _multiple_aligned_hw(hw, int(multiple))
+        )
+
+    def _field_dst_hw(name: str, value: torch.Tensor) -> tuple[int, int]:
+        prefix = "tgt" if name.startswith("tgt_") else "src"
+        view_hw = _view_hw(prefix)
+        if view_hw is not None:
+            return _aligned_view_hw(prefix, view_hw)
+        hw = (int(value.shape[-2]), int(value.shape[-1]))
+        return _erp_multiple_aligned_hw(hw, int(multiple)) if "_erp_" in name else _multiple_aligned_hw(hw, int(multiple))
+
+    updates: dict[str, Any] = {}
+    for field in fields(batch):
+        name = field.name
+        value = getattr(batch, name)
+        if not torch.is_tensor(value):
+            continue
+        if name.endswith("_rgb_u8") and value.ndim >= 3:
+            if "_cube_" in name:
+                cube_hw = _multiple_aligned_hw((int(value.shape[-3]), int(value.shape[-2])), int(multiple))
+                updates[name] = _resize_cube_tensor(value, cube_hw, kind="image")
+            else:
+                updates[name] = _resize_chw_tensor(value, _field_dst_hw(name, value), kind="image")
+        elif name.endswith("_depth_m") and value.ndim >= 3:
+            if "_cube_" in name:
+                cube_hw = _multiple_aligned_hw((int(value.shape[-3]), int(value.shape[-2])), int(multiple))
+                updates[name] = _resize_cube_tensor(value, cube_hw, kind="depth")
+            else:
+                updates[name] = _resize_chw_tensor(value, _field_dst_hw(name, value), kind="depth")
+        elif name.endswith("_valid_mask") and value.ndim >= 3:
+            updates[name] = _resize_chw_tensor(value, _field_dst_hw(name, value), kind="depth")
+        elif name.endswith("_rays") and value.ndim >= 3:
+            updates[name] = _resize_chw_tensor(value, _field_dst_hw(name, value), kind="ray")
+
+    for intr_name in ("src_intrinsics", "tgt_intrinsics"):
+        intr = getattr(batch, intr_name, None)
+        if torch.is_tensor(intr):
+            prefix = "tgt" if intr_name.startswith("tgt_") else "src"
+            view_hw = _view_hw(prefix)
+            if view_hw is not None:
+                updates[intr_name] = scale_pinhole_intrinsics(
+                    intr,
+                    src_hw=view_hw,
+                    dst_hw=_aligned_view_hw(prefix, view_hw),
+                )
+    for params_name in ("src_camera_params", "tgt_camera_params"):
+        params = getattr(batch, params_name, None)
+        if torch.is_tensor(params):
+            prefix = "tgt" if params_name.startswith("tgt_") else "src"
+            view_hw = _view_hw(prefix)
+            if view_hw is not None:
+                updates[params_name] = _scale_fisheye624_params_any(
+                    params,
+                    src_hw=view_hw,
+                    dst_hw=_aligned_view_hw(prefix, view_hw),
+                )
+
+    return replace(batch, **updates) if updates else batch
+
+
+def _build_optimizer_param_groups(
+    raw_model: UnisharpFeatureModel,
+) -> tuple[list[torch.nn.Parameter], list[torch.nn.Parameter], list[torch.nn.Parameter]]:
+    base_params: list[torch.nn.Parameter] = []
+    unik3d_encoder_params: list[torch.nn.Parameter] = []
+    unik3d_decoder_params: list[torch.nn.Parameter] = []
+    for name, param in raw_model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if name.startswith("feature_extractor.unik3d.pixel_encoder."):
+            unik3d_encoder_params.append(param)
+        elif name.startswith("second_layer_depth_head."):
+            unik3d_decoder_params.append(param)
+        elif name.startswith("feature_extractor.unik3d."):
+            unik3d_decoder_params.append(param)
+        else:
+            base_params.append(param)
+    return base_params, unik3d_encoder_params, unik3d_decoder_params
+
+
+def _count_numel(params: list[torch.nn.Parameter]) -> int:
+    return int(sum(int(p.numel()) for p in params))
+
+
+def _configure_torchhub_cache() -> Path:
+    torchhub_dir = REPO_ROOT / "checkpoints" / "torchhub"
+    torchhub_dir.mkdir(parents=True, exist_ok=True)
+    os.environ["TORCH_HOME"] = str(torchhub_dir)
+    torch.hub.set_dir(str(torchhub_dir))
+    return torchhub_dir
+
+
+def _ddp_is_enabled() -> bool:
+    return int(os.environ.get("WORLD_SIZE", "1")) > 1
+
+
+def _ddp_setup(device: str, ddp_timeout_hours: float = 8.0) -> tuple[torch.device, int, int, bool]:
+    if not _ddp_is_enabled():
+        dev = torch.device(device)
+        return dev, 0, 1, True
+
+    if device != "cuda":
+        raise RuntimeError("DDP currently supports CUDA only.")
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA not available.")
+
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    rank = int(os.environ.get("RANK", "0"))
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    torch.cuda.set_device(local_rank)
+    timeout_hours = max(float(ddp_timeout_hours), 0.25)
+    if rank == 0:
+        print(
+            "[ddp_setup] init_process_group backend=nccl "
+            f"world_size={world_size} NCCL_NET={os.environ.get('NCCL_NET', '<unset>')} "
+            f"NCCL_IB_DISABLE={os.environ.get('NCCL_IB_DISABLE', '<unset>')}",
+            flush=True,
+        )
+    dist.init_process_group(backend="nccl", timeout=timedelta(hours=timeout_hours))
+    if rank == 0:
+        print("[ddp_setup] init_process_group done", flush=True)
+    dev = torch.device("cuda", local_rank)
+    return dev, rank, world_size, (rank == 0)
+
+
+def _ddp_broadcast_path(p: Path, is_main: bool) -> Path:
+    if not _ddp_is_enabled():
+        return p
+    obj_list: list[str] = [str(p) if is_main else ""]
+    dist.broadcast_object_list(obj_list, src=0)
+    return Path(obj_list[0])
+
+
+def _ddp_broadcast_str(value: str, is_main: bool) -> str:
+    if not _ddp_is_enabled():
+        return value
+    obj_list: list[str] = [str(value) if is_main else ""]
+    dist.broadcast_object_list(obj_list, src=0)
+    return str(obj_list[0])
+
+
+def _ddp_any_bool(flag: bool, device: torch.device) -> bool:
+    if not _ddp_is_enabled():
+        return bool(flag)
+    x = torch.tensor(1 if flag else 0, device=device, dtype=torch.int32)
+    dist.all_reduce(x, op=dist.ReduceOp.MAX)
+    return bool(int(x.item()) != 0)
+
+
+def _env_flag(name: str, default: bool = False) -> bool:
+    raw = os.environ.get(name)
+    if raw is None:
+        return bool(default)
+    return raw.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _is_oom_exception(exc: BaseException) -> bool:
+    if isinstance(exc, torch.cuda.OutOfMemoryError):
+        return True
+    msg = str(exc).lower()
+    oom_markers = (
+        "out of memory",
+        "cuda error: out of memory",
+        "cublas_status_alloc_failed",
+        "cudnn_status_alloc_failed",
+        "defaultcpuallocator",
+    )
+    return any(marker in msg for marker in oom_markers)
+
+
+def _ddp_barrier(device: torch.device) -> None:
+    if not _ddp_is_enabled():
+        return
+    if device.type == "cuda" and device.index is not None:
+        dist.barrier(device_ids=[device.index])
+    else:
+        dist.barrier()
+
+
+def _maybe_set_dataset_epoch(dataset: Any, epoch: int) -> None:
+    set_epoch = getattr(dataset, "set_epoch", None)
+    if callable(set_epoch):
+        set_epoch(int(epoch))
+
+
+def _ddp_mean(x: torch.Tensor) -> torch.Tensor:
+    if not _ddp_is_enabled():
+        return x
+    y = x.detach().clone()
+    dist.all_reduce(y, op=dist.ReduceOp.SUM)
+    y = y / float(dist.get_world_size())
+    return y
+
+
+def _save_train_vis(
+    out_dir: Path,
+    step: int,
+    src_gt: torch.Tensor,
+    src_pred: torch.Tensor,
+    src_alpha: torch.Tensor,
+    tgt_gt: torch.Tensor,
+    tgt_pred: torch.Tensor,
+    tgt_alpha: torch.Tensor,
+    src_gt_depth: torch.Tensor | None = None,
+    tgt_gt_depth: torch.Tensor | None = None,
+    src_pred_depth: torch.Tensor | None = None,
+    tgt_pred_depth: torch.Tensor | None = None,
+    src_unik3d_depth: torch.Tensor | None = None,
+    tgt_unik3d_depth: torch.Tensor | None = None,
+    dataset_name: str | None = None,
+    scene: str | None = None,
+    src_idx: int | None = None,
+    tgt_idx: int | None = None,
+    src_pose_w2c: torch.Tensor | None = None,
+    tgt_pose_w2c: torch.Tensor | None = None,
+    src_metric_mask: torch.Tensor | None = None,
+    tgt_metric_mask: torch.Tensor | None = None,
+    src_cube_gt_u8: torch.Tensor | None = None,
+    src_cube_pred_linear: torch.Tensor | None = None,
+    src_cube_alpha: torch.Tensor | None = None,
+    tgt_cube_gt_u8: torch.Tensor | None = None,
+    tgt_cube_pred_linear: torch.Tensor | None = None,
+    tgt_cube_alpha: torch.Tensor | None = None,
+) -> None:
+    vis_dir = out_dir / "vis"
+    vis_dir.mkdir(parents=True, exist_ok=True)
+    LOGGER.info("Saving train visualization: %s", str(vis_dir / f"step_{int(step):07d}.png"))
+    save_pair_visualization(
+        vis_dir / f"step_{int(step):07d}.png",
+        src_gt=src_gt,
+        src_pred=src_pred,
+        src_alpha=src_alpha,
+        tgt_gt=tgt_gt,
+        tgt_pred=tgt_pred,
+        tgt_alpha=tgt_alpha,
+        src_gt_depth=src_gt_depth,
+        tgt_gt_depth=tgt_gt_depth,
+        src_pred_depth=src_pred_depth,
+        tgt_pred_depth=tgt_pred_depth,
+        src_unik3d_depth=src_unik3d_depth,
+        tgt_unik3d_depth=tgt_unik3d_depth,
+        dataset_name=dataset_name,
+        scene=scene,
+        step=int(step),
+        src_idx=src_idx,
+        tgt_idx=tgt_idx,
+        src_pose_w2c=src_pose_w2c,
+        tgt_pose_w2c=tgt_pose_w2c,
+        src_cube_gt_u8=src_cube_gt_u8,
+        src_cube_pred_linear=src_cube_pred_linear,
+        src_cube_alpha=src_cube_alpha,
+        tgt_cube_gt_u8=tgt_cube_gt_u8,
+        tgt_cube_pred_linear=tgt_cube_pred_linear,
+        tgt_cube_alpha=tgt_cube_alpha,
+    )
+
+
+def _read_nonempty_lines(path: Path) -> list[str]:
+    return [line.strip() for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
+
+
+def _resolve_manifest_file(manifest_dir: Path | None, filename: str) -> Path | None:
+    if manifest_dir is None:
+        return None
+    path = Path(manifest_dir) / filename
+    return path if path.exists() else None
+
+
+@click.command()
+@click.option("--data-root-re10k", type=click.Path(path_type=Path, exists=True), default=None)
+@click.option("--data-root-hm3d", type=click.Path(path_type=Path, exists=True), default=Path("/media/team_data/ML4_team/datasets/panogs"))
+@click.option("--data-root-sim", type=click.Path(path_type=Path, exists=True), default=Path("/media/team_data/ML4_team/datasets/smx_sim"))
+@click.option("--sim-pose-root", type=click.Path(path_type=Path, exists=True), default=Path("/media/team_data/ML4_team/datasets/smx_sim/30cm"))
+@click.option("--data-root-wildrgbd", type=click.Path(path_type=Path, exists=True), default=None)
+@click.option("--wild-roots-file", type=click.Path(path_type=Path, exists=True, dir_okay=False), default=DEFAULT_WILDRGBD_ROOTS_FILE)
+@click.option("--data-root-dl3dv", type=click.Path(path_type=Path, exists=True), default=Path("/media/team_data/ML4_team/datasets/sharp/DL3DV-ALL-960P"))
+@click.option("--data-root-dl3dv-depth", type=click.Path(path_type=Path, exists=True), default=Path("/media/team_data/ML4_team/datasets/sharp/DL3DV-ALL-960P_da3_outputs"))
+@click.option("--data-root-scanetpp", type=click.Path(path_type=Path, exists=True), default=Path("/media/team_data/ML4_team/datasets/scan"))
+@click.option("--dataset-manifest-dir", type=click.Path(path_type=Path, file_okay=False), default=None)
+@click.option("--out-root", type=click.Path(path_type=Path, file_okay=False), required=True)
+@click.option("--run-name", type=str, default=None)
+@click.option("--steps", type=int, default=1000000)
+@click.option("--batch-size", type=int, default=2)
+@click.option("--num-workers", type=int, default=1)
+@click.option("--warmup", type=int, default=75000)
+@click.option("--lr0", type=float, default=1.2e-4)
+@click.option("--lr1", type=float, default=1.6e-5)
+@click.option("--unik3d-lr0", type=float, default=2.5e-5, help="UniK3D decoder/head peak LR.")
+@click.option("--unik3d-lr1", type=float, default=2.5e-6, help="UniK3D decoder/head final LR.")
+@click.option("--unik3d-encoder-lr0", type=float, default=1.5e-6, help="UniK3D pixel_encoder peak LR.")
+@click.option("--unik3d-encoder-lr1", type=float, default=1.5e-7, help="UniK3D pixel_encoder final LR.")
+@click.option("--grad-clip-norm", type=float, default=1.0, show_default=True)
+@click.option("--max-step-grad-norm", type=float, default=100000.0, show_default=True, help="Skip optimizer step when pre-clip grad norm exceeds this value. 0 disables.")
+@click.option("--max-depth-m", type=float, default=DEFAULT_MAX_DEPTH_M, show_default=True)
+@click.option("--sim-far-depth-invalid-m", type=float, default=30.0, show_default=True)
+@click.option("--sim-far-depth-invalid-max-frac", type=float, default=1.0, show_default=True)
+@click.option("--sim-max-long-edge", type=int, default=512, show_default=True, help="Resize SIM ERP frames before cubemap conversion. 0 keeps native resolution.")
+@click.option("--train-resize-multiple", type=int, default=256, show_default=True, help="Before model forward, downsize training inputs to the largest H/W divisible by this value. 0 disables.")
+@click.option("--pinhole-train-size", type=int, default=0, show_default=True, help="Resize pinhole training datasets to NxN before model forward. 0 keeps dataset native resolution.")
+@click.option("--scanetpp-fisheye-far-depth-invalid-m", type=float, default=30.0, show_default=True)
+@click.option("--max-index-gap", type=int, default=10)
+@click.option("--device", type=str, default="cuda")
+@click.option("--render-low-pass-filter-eps", type=float, default=1e-2, show_default=True)
+@click.option("--ddp-timeout-hours", type=float, default=8.0)
+@click.option("--save-every", type=int, default=5000)
+@click.option("--log-every", type=int, default=50)
+@click.option("--vis-every", type=int, default=500)
+@click.option("--unik3d-backbone", type=click.Choice(["vitb", "vitl"]), default="vitl")
+@click.option("--unik3d-resolution-level", type=click.IntRange(0, 9), default=0, show_default=True)
+@click.option("--initializer-stride", type=click.IntRange(1, 2), default=1)
+@click.option("--initializer-scale-factor", type=float, default=1.5, show_default=True)
+@click.option("--lambda-aux-ray", type=float, default=3.0)
+@click.option("--lambda-aux-depth-scale", type=float, default=3.0)
+@click.option("--lambda-aux-depth2-scale", type=float, default=1.0)
+@click.option("--lambda-color", type=float, default=1.0)
+@click.option("--lambda-alpha", type=float, default=1.5)
+@click.option("--alpha-tail-min", type=float, default=0.99, show_default=True, help="Alpha value below which local tail coverage loss is applied.")
+@click.option("--alpha-tail-weight", type=float, default=0.0, show_default=True, help="Extra normalized tail weight for local low-alpha holes.")
+@click.option("--lambda-percep", type=float, default=1.0)
+@click.option("--lambda-depth", type=float, default=0.5)
+@click.option("--lambda-tv", type=float, default=1.0)
+@click.option("--lambda-grad", type=float, default=1.0)
+@click.option("--lambda-grad-img", type=float, default=0.2)
+@click.option("--lambda-edge-rgb", type=float, default=0.0, show_default=True, help="Weight for GT RGB edge-band gradient matching.")
+@click.option("--lambda-delta", type=float, default=1.0)
+@click.option("--lambda-delta-rho", type=float, default=0.01, show_default=True)
+@click.option("--lambda-splat", type=float, default=1.0)
+@click.option("--lambda-edge-splat", type=float, default=0.0, show_default=True, help="Weight for stricter projected-sigma penalty on GT depth-edge bands.")
+@click.option("--lambda-grid", type=float, default=0.05, show_default=True, help="Weight for Gaussian-grid 2x2 checkerboard residual regularization.")
+@click.option("--delta-clip", type=float, default=10.0, show_default=True)
+@click.option("--raw-delta-clip", type=float, default=400.0, show_default=True)
+@click.option("--raw-delta-rho-clip", type=float, default=5.0, show_default=True)
+@click.option("--delta-rho-limit", type=float, default=2.0, show_default=True)
+@click.option("--splat-sigma-min", type=float, default=1e-1, show_default=True, help="Minimum projected screen-space variance for L_splat.")
+@click.option("--splat-sigma-max", type=float, default=1e2, show_default=True, help="Maximum projected screen-space variance for L_splat.")
+@click.option("--edge-splat-sigma-max", type=float, default=2.0, show_default=True, help="Maximum projected variance on depth-edge bands for L_edge_splat.")
+@click.option("--depth-edge-log-threshold", type=float, default=0.05, show_default=True, help="Log-depth jump threshold used to build L_edge_splat edge bands.")
+@click.option("--depth-edge-dilate-px", type=int, default=2, show_default=True, help="Dilation radius in pixels for L_edge_splat depth-edge bands.")
+@click.option("--target-mask-erode-px", type=int, default=0, show_default=True, help="Erode source-visible target masks by this many pixels before target supervision.")
+@click.option("--dataset-weight-re10k", type=float, default=1.0)
+@click.option("--dataset-weight-hm3d", type=float, default=1.0)
+@click.option("--dataset-weight-sim", type=float, default=1.0)
+@click.option("--dataset-weight-wildrgbd", type=float, default=1.0)
+@click.option("--dataset-weight-dl3dv", type=float, default=1.0)
+@click.option("--dataset-weight-scanetpp", type=float, default=0.0)
+@click.option(
+    "--re10k-pseudo-depth-root",
+    type=click.Path(path_type=Path, file_okay=False),
+    default=Path("/media/team_data/ML4_team/datasets/nopose/re10k_unik3d_pseudo_depth"),
+)
+@click.option("--re10k-pseudo-depth-autogen/--no-re10k-pseudo-depth-autogen", default=True)
+@click.option("--re10k-pseudo-depth-backbone", type=click.Choice(["vitb", "vitl"]), default="vitl")
+@click.option("--re10k-pseudo-depth-device", type=str, default="cpu")
+@click.option("--re10k-pseudo-lock-timeout-sec", type=float, default=120.0)
+@click.option("--re10k-pseudo-lock-stale-sec", type=float, default=1800.0)
+@click.option("--re10k-pseudo-far-depth-invalid-m", type=float, default=30.0)
+@click.option("--seed", type=int, default=None)
+@click.option("-v", "--verbose", is_flag=True)
+def train_feature_cli(
+    data_root_re10k: Path | None,
+    data_root_hm3d: Path | None,
+    data_root_sim: Path | None,
+    sim_pose_root: Path | None,
+    data_root_wildrgbd: Path | None,
+    wild_roots_file: Path,
+    data_root_dl3dv: Path | None,
+    data_root_dl3dv_depth: Path | None,
+    data_root_scanetpp: Path | None,
+    dataset_manifest_dir: Path | None,
+    out_root: Path,
+    run_name: str | None,
+    steps: int,
+    batch_size: int,
+    num_workers: int,
+    warmup: int,
+    lr0: float,
+    lr1: float,
+    unik3d_lr0: float,
+    unik3d_lr1: float,
+    unik3d_encoder_lr0: float,
+    unik3d_encoder_lr1: float,
+    grad_clip_norm: float,
+    max_step_grad_norm: float,
+    max_depth_m: float,
+    sim_far_depth_invalid_m: float,
+    sim_far_depth_invalid_max_frac: float,
+    sim_max_long_edge: int,
+    train_resize_multiple: int,
+    pinhole_train_size: int,
+    scanetpp_fisheye_far_depth_invalid_m: float,
+    max_index_gap: int,
+    device: str,
+    render_low_pass_filter_eps: float,
+    ddp_timeout_hours: float,
+    save_every: int,
+    log_every: int,
+    vis_every: int,
+    unik3d_backbone: str,
+    unik3d_resolution_level: int,
+    initializer_stride: int,
+    initializer_scale_factor: float,
+    lambda_aux_ray: float,
+    lambda_aux_depth_scale: float,
+    lambda_aux_depth2_scale: float,
+    lambda_color: float,
+    lambda_alpha: float,
+    alpha_tail_min: float,
+    alpha_tail_weight: float,
+    lambda_percep: float,
+    lambda_depth: float,
+    lambda_tv: float,
+    lambda_grad: float,
+    lambda_grad_img: float,
+    lambda_edge_rgb: float,
+    lambda_delta: float,
+    lambda_delta_rho: float,
+    lambda_splat: float,
+    lambda_edge_splat: float,
+    lambda_grid: float,
+    delta_clip: float,
+    raw_delta_clip: float,
+    raw_delta_rho_clip: float,
+    delta_rho_limit: float,
+    splat_sigma_min: float,
+    splat_sigma_max: float,
+    edge_splat_sigma_max: float,
+    depth_edge_log_threshold: float,
+    depth_edge_dilate_px: int,
+    target_mask_erode_px: int,
+    dataset_weight_re10k: float,
+    dataset_weight_hm3d: float,
+    dataset_weight_sim: float,
+    dataset_weight_wildrgbd: float,
+    dataset_weight_dl3dv: float,
+    dataset_weight_scanetpp: float,
+    re10k_pseudo_depth_root: Path,
+    re10k_pseudo_depth_autogen: bool,
+    re10k_pseudo_depth_backbone: str,
+    re10k_pseudo_depth_device: str,
+    re10k_pseudo_lock_timeout_sec: float,
+    re10k_pseudo_lock_stale_sec: float,
+    re10k_pseudo_far_depth_invalid_m: float,
+    seed: int | None,
+    verbose: bool,
+) -> None:
+    detach_init_layer0_distance = True
+
+    log_level = logging.DEBUG if verbose else logging.INFO
+    logging_utils.configure(log_level)
+    if float(max_depth_m) <= 0.0:
+        raise ValueError("--max-depth-m must be positive.")
+    if float(grad_clip_norm) <= 0.0:
+        raise ValueError("--grad-clip-norm must be positive.")
+    if float(max_step_grad_norm) < 0.0:
+        raise ValueError("--max-step-grad-norm must be non-negative.")
+    if float(render_low_pass_filter_eps) < 0.0:
+        raise ValueError("--render-low-pass-filter-eps must be non-negative.")
+    if not (0.0 <= float(sim_far_depth_invalid_max_frac) <= 1.0):
+        raise ValueError("--sim-far-depth-invalid-max-frac must be in [0, 1].")
+    if int(sim_max_long_edge) < 0:
+        raise ValueError("--sim-max-long-edge must be non-negative.")
+    if int(train_resize_multiple) < 0:
+        raise ValueError("--train-resize-multiple must be non-negative.")
+    if int(pinhole_train_size) < 0:
+        raise ValueError("--pinhole-train-size must be non-negative.")
+    if float(scanetpp_fisheye_far_depth_invalid_m) < 0.0:
+        raise ValueError("--scanetpp-fisheye-far-depth-invalid-m must be non-negative.")
+    if float(delta_clip) < 0.0:
+        raise ValueError("--delta-clip must be non-negative.")
+    if float(raw_delta_clip) < 0.0:
+        raise ValueError("--raw-delta-clip must be non-negative.")
+    if float(raw_delta_rho_clip) < 0.0:
+        raise ValueError("--raw-delta-rho-clip must be non-negative.")
+    if float(lambda_grid) < 0.0:
+        raise ValueError("--lambda-grid must be non-negative.")
+    if float(lambda_edge_rgb) < 0.0:
+        raise ValueError("--lambda-edge-rgb must be non-negative.")
+    if float(lambda_edge_splat) < 0.0:
+        raise ValueError("--lambda-edge-splat must be non-negative.")
+    if float(edge_splat_sigma_max) < 0.0:
+        raise ValueError("--edge-splat-sigma-max must be non-negative.")
+    if float(depth_edge_log_threshold) < 0.0:
+        raise ValueError("--depth-edge-log-threshold must be non-negative.")
+    if int(depth_edge_dilate_px) < 0:
+        raise ValueError("--depth-edge-dilate-px must be non-negative.")
+    if int(target_mask_erode_px) < 0:
+        raise ValueError("--target-mask-erode-px must be non-negative.")
+    if not (0.0 <= float(alpha_tail_min) <= 1.0):
+        raise ValueError("--alpha-tail-min must be in [0, 1].")
+    if float(alpha_tail_weight) < 0.0:
+        raise ValueError("--alpha-tail-weight must be non-negative.")
+    if float(delta_rho_limit) < 0.0:
+        raise ValueError("--delta-rho-limit must be non-negative.")
+    if float(splat_sigma_min) < 0.0:
+        raise ValueError("--splat-sigma-min must be non-negative.")
+    if float(splat_sigma_max) <= float(splat_sigma_min):
+        raise ValueError("--splat-sigma-max must be greater than --splat-sigma-min.")
+    dev, rank, world_size, is_main = _ddp_setup(device, ddp_timeout_hours=ddp_timeout_hours)
+
+    if seed is not None:
+        s = int(seed)
+        random.seed(s + rank)
+        np.random.seed(s + rank)
+        torch.manual_seed(s + rank)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(s + rank)
+
+    if is_main and (run_name is None or run_name.strip() == ""):
+        run_name = f"unified_feature_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+    if run_name is None:
+        run_name = "unified_feature_ddp"
+    out_dir = _ddp_broadcast_path(Path(out_root) / run_name, is_main=is_main)
+    logging_utils.configure(log_level)
+    if not is_main:
+        logging.getLogger().setLevel(logging.WARNING)
+        LOGGER.setLevel(logging.WARNING)
+    _configure_torchhub_cache()
+    re10k_enabled_for_train = bool(float(dataset_weight_re10k) > 0.0)
+    hm3d_enabled_for_train = bool(float(dataset_weight_hm3d) > 0.0)
+    sim_enabled_for_train = bool(float(dataset_weight_sim) > 0.0)
+    dl3dv_enabled_for_train = bool(float(dataset_weight_dl3dv) > 0.0)
+    scanetpp_enabled_for_train = bool(float(dataset_weight_scanetpp) > 0.0)
+    wild_roots = _read_nonempty_lines(wild_roots_file) if wild_roots_file.exists() else []
+    re10k_manifest = _resolve_manifest_file(dataset_manifest_dir, "re10k_train_chunks.txt")
+    hm3d_manifest = _resolve_manifest_file(dataset_manifest_dir, "hm3d_train_scenes.txt")
+    sim_manifest = _resolve_manifest_file(dataset_manifest_dir, "sim_train_scenes.txt")
+    wildrgbd_manifest = _resolve_manifest_file(dataset_manifest_dir, "wildrgbd_train_scenes.txt")
+    dl3dv_manifest = _resolve_manifest_file(dataset_manifest_dir, "dl3dv_train_scenes.txt")
+    scanetpp_manifest = _resolve_manifest_file(dataset_manifest_dir, "scanetpp_fisheye_train_scenes.txt")
+    wildrgbd_enabled_for_train = bool(
+        ((data_root_wildrgbd is not None) or bool(wild_roots)) and (float(dataset_weight_wildrgbd) > 0.0)
+    )
+    if re10k_enabled_for_train and data_root_re10k is None:
+        raise ValueError("dataset_weight_re10k>0 but --data-root-re10k is not provided.")
+    if hm3d_enabled_for_train and data_root_hm3d is None:
+        raise ValueError("dataset_weight_hm3d>0 but --data-root-hm3d is not provided.")
+    if sim_enabled_for_train and (data_root_sim is None or sim_pose_root is None):
+        raise ValueError("dataset_weight_sim>0 but --data-root-sim / --sim-pose-root is missing.")
+    if sim_enabled_for_train and sim_manifest is None:
+        raise ValueError("dataset_weight_sim>0 but sim_train_scenes.txt is missing from --dataset-manifest-dir.")
+    if float(dataset_weight_wildrgbd) > 0.0 and (data_root_wildrgbd is None) and (not wild_roots):
+        raise ValueError("dataset_weight_wildrgbd>0 but neither --data-root-wildrgbd nor --wild-roots-file is provided.")
+    if dl3dv_enabled_for_train and (data_root_dl3dv is None or data_root_dl3dv_depth is None):
+        raise ValueError("dataset_weight_dl3dv>0 but --data-root-dl3dv / --data-root-dl3dv-depth is missing.")
+    if scanetpp_enabled_for_train and data_root_scanetpp is None:
+        raise ValueError("dataset_weight_scanetpp>0 but --data-root-scanetpp is missing.")
+
+    if is_main:
+        out_dir.mkdir(parents=True, exist_ok=True)
+        LOGGER.info(
+            "Training start: out=%s branch=gt-override scratch_unik3d_pretrained backbone=%s steps=%d batch=%d",
+            str(out_dir),
+            str(unik3d_backbone),
+            int(steps),
+            int(batch_size),
+        )
+        LOGGER.info(
+            "Loss weights: color=%.3g alpha=%.3g depth=%.3g percep=%.3g aux_ray=%.3g aux_depth0=%.3g aux_depth1=%.3g",
+            float(lambda_color),
+            float(lambda_alpha),
+            float(lambda_depth),
+            float(lambda_percep),
+            float(lambda_aux_ray),
+            float(lambda_aux_depth_scale),
+            float(lambda_aux_depth2_scale),
+        )
+
+    dataset_seed = int(seed) if seed is not None else 12345
+    pinhole_output_h = int(pinhole_train_size) if int(pinhole_train_size) > 0 else None
+    pinhole_output_w = int(pinhole_train_size) if int(pinhole_train_size) > 0 else None
+
+    re10k_ds = None
+    if re10k_enabled_for_train:
+        re10k_ds = Re10KDataset(
+            root=data_root_re10k,
+            chunks_file=re10k_manifest,
+            split="train",
+            min_frame_gap=1,
+            max_frame_gap=int(max_index_gap),
+            pair_max_translation_m=0.5,
+            pair_min_overlap=0.6,
+            output_h=pinhole_output_h,
+            output_w=pinhole_output_w,
+            shuffle_chunk=True,
+            shuffle_example=True,
+            ddp_rank=rank,
+            ddp_world_size=world_size,
+            pseudo_depth_root=re10k_pseudo_depth_root,
+            pseudo_depth_autogen=bool(re10k_pseudo_depth_autogen),
+            pseudo_depth_backbone=str(re10k_pseudo_depth_backbone),
+            pseudo_depth_device=str(re10k_pseudo_depth_device),
+            pseudo_lock_timeout_sec=float(re10k_pseudo_lock_timeout_sec),
+            pseudo_lock_stale_sec=float(re10k_pseudo_lock_stale_sec),
+            batch_size_hint=int(batch_size),
+            depth_max_m=float(max_depth_m),
+            pseudo_far_depth_invalid_m=float(re10k_pseudo_far_depth_invalid_m),
+            seed=dataset_seed,
+        )
+    hm3d_train_root = None
+    if data_root_hm3d is not None:
+        hm3d_train_root = data_root_hm3d / "train" if (data_root_hm3d / "train").exists() else data_root_hm3d
+
+    hm3d_ds = None
+    if hm3d_enabled_for_train:
+        hm3d_ds = PanOGSDataset(
+            root=hm3d_train_root,
+            index_manifest_path=hm3d_manifest,
+            src_tgt_max_index_gap=int(max_index_gap),
+            use_cubemap_supervision=True,
+            pair_sampling=True,
+            pair_max_translation_m=0.5,
+            pair_min_depth_overlap=0.6,
+            pair_overlap_face_w=64,
+            pair_overlap_margin=1.05,
+            pair_max_tries=48,
+            depth_max_m=float(max_depth_m),
+        )
+    sim_ds = None
+    if sim_enabled_for_train:
+        sim_ds = SimPanoramaDataset(
+            root=data_root_sim,
+            pose_root=sim_pose_root,
+            scene_list_file=sim_manifest,
+            max_index_gap=int(max_index_gap),
+            pair_max_translation_m=0.5,
+            pair_min_depth_overlap=0.6,
+            pairs_per_chunk=15,
+            chunk_size=30,
+            shuffle_scene=True,
+            ddp_rank=rank,
+            ddp_world_size=world_size,
+            depth_max_m=float(max_depth_m),
+            far_depth_invalid_m=float(sim_far_depth_invalid_m),
+            far_depth_invalid_max_frac=float(sim_far_depth_invalid_max_frac),
+            max_long_edge=int(sim_max_long_edge),
+            seed=dataset_seed,
+        )
+    wildrgbd_ds = None
+    if wildrgbd_enabled_for_train:
+        wild_dataset_roots = [Path(p) for p in wild_roots]
+        if data_root_wildrgbd is not None:
+            wild_dataset_roots.append(data_root_wildrgbd)
+        wildrgbd_ds = WildRGBDDataset(
+            root=None,
+            scene_list_file=wildrgbd_manifest,
+            split="scenes",
+            min_frame_gap=1,
+            max_frame_gap=int(max_index_gap),
+            pair_max_translation_m=0.5,
+            pair_min_overlap=0.6,
+            output_h=pinhole_output_h,
+            output_w=pinhole_output_w,
+            shuffle_scene=True,
+            shuffle_frame=False,
+            ddp_rank=rank,
+            ddp_world_size=world_size,
+            roots=wild_dataset_roots,
+            depth_max_m=float(max_depth_m),
+            seed=dataset_seed,
+        )
+    dl3dv_ds = None
+    if dl3dv_enabled_for_train:
+        dl3dv_ds = DL3DVDataset(
+            root=data_root_dl3dv,
+            depth_root=data_root_dl3dv_depth,
+            scene_specs_file=dl3dv_manifest,
+            min_frame_gap=1,
+            max_frame_gap=int(max_index_gap),
+            pair_max_translation_m=0.5,
+            pair_min_overlap=0.6,
+            output_h=pinhole_output_h,
+            output_w=pinhole_output_w,
+            shuffle_scene=True,
+            shuffle_frame=False,
+            ddp_rank=rank,
+            ddp_world_size=world_size,
+            batch_size_hint=int(batch_size),
+            depth_max_m=float(max_depth_m),
+            seed=dataset_seed,
+        )
+
+    scanetpp_ds = None
+    if scanetpp_enabled_for_train:
+        scanetpp_ds = ScannetppFisheyeDataset(
+            root=data_root_scanetpp,
+            scene_list_file=scanetpp_manifest,
+            min_frame_gap=1,
+            max_frame_gap=int(max_index_gap),
+            pair_max_translation_m=0.5,
+            shuffle_scene=True,
+            shuffle_frame=False,
+            ddp_rank=rank,
+            ddp_world_size=world_size,
+            batch_size_hint=int(batch_size),
+            depth_max_m=float(max_depth_m),
+            far_depth_invalid_m=float(scanetpp_fisheye_far_depth_invalid_m),
+            seed=dataset_seed,
+        )
+
+    hm3d_sampler = None
+    if hm3d_ds is not None and _ddp_is_enabled():
+        hm3d_sampler = DistributedSampler(hm3d_ds, num_replicas=world_size, rank=rank, shuffle=True, drop_last=False)
+
+    re10k_num_workers = int(num_workers)
+    if re10k_ds is not None and bool(re10k_pseudo_depth_autogen) and re10k_num_workers > 0:
+        re10k_num_workers = 0
+        if is_main:
+            LOGGER.warning(
+                "RE10K pseudo-depth auto-generate enabled: force re10k dataloader num_workers=%d (requested=%d).",
+                int(re10k_num_workers),
+                int(num_workers),
+            )
+    if re10k_ds is not None and batch_size > 1 and re10k_num_workers > 0:
+        re10k_num_workers = 0
+        if is_main:
+            LOGGER.warning(
+                "Dynamic-resolution RE10K batching requires ordered same-resolution samples: force re10k dataloader num_workers=%d (requested=%d).",
+                int(re10k_num_workers),
+                int(num_workers),
+            )
+
+    highres_pin_memory = os.environ.get("HIGHRES_TRAIN_PIN_MEMORY", "0").strip().lower() in {"1", "true", "yes", "on"}
+    standard_pin_memory = os.environ.get("TRAIN_PIN_MEMORY", "1").strip().lower() in {"1", "true", "yes", "on"}
+    try:
+        train_prefetch_factor = max(1, int(os.environ.get("TRAIN_PREFETCH_FACTOR", "1").strip()))
+    except Exception:
+        train_prefetch_factor = 1
+    def _loader_worker_kwargs(worker_count: int, *, pin_memory: bool) -> dict[str, Any]:
+        kwargs: dict[str, Any] = {
+            "num_workers": int(worker_count),
+            "pin_memory": bool(pin_memory),
+        }
+        if int(worker_count) > 0:
+            kwargs["prefetch_factor"] = int(train_prefetch_factor)
+        return kwargs
+
+    re10k_dl = None
+    if re10k_ds is not None:
+        re10k_dl = DataLoader(
+            re10k_ds,
+            batch_size=None,
+            **_loader_worker_kwargs(re10k_num_workers, pin_memory=standard_pin_memory),
+            collate_fn=re10k_passthrough,
+        )
+
+    hm3d_dl = None
+    if hm3d_ds is not None:
+        hm3d_dl = DataLoader(
+            hm3d_ds,
+            batch_size=batch_size,
+            shuffle=(hm3d_sampler is None),
+            sampler=hm3d_sampler,
+            **_loader_worker_kwargs(num_workers, pin_memory=highres_pin_memory),
+            collate_fn=panogs_collate,
+        )
+
+    sim_dl = None
+    if sim_ds is not None:
+        sim_dl = DataLoader(
+            sim_ds,
+            batch_size=batch_size,
+            **_loader_worker_kwargs(num_workers, pin_memory=highres_pin_memory),
+            collate_fn=panogs_collate,
+        )
+
+    wildrgbd_dl = None
+    if wildrgbd_ds is not None:
+        wildrgbd_dl = DataLoader(
+            wildrgbd_ds,
+            batch_size=batch_size,
+            **_loader_worker_kwargs(num_workers, pin_memory=standard_pin_memory),
+            collate_fn=wildrgbd_collate,
+        )
+
+    dl3dv_dl = None
+    if dl3dv_ds is not None:
+        dl3dv_dl = DataLoader(
+            dl3dv_ds,
+            batch_size=None,
+            **_loader_worker_kwargs(num_workers, pin_memory=standard_pin_memory),
+            collate_fn=re10k_passthrough,
+        )
+
+    scanetpp_dl = None
+    if scanetpp_ds is not None:
+        scanetpp_dl = DataLoader(
+            scanetpp_ds,
+            batch_size=None,
+            **_loader_worker_kwargs(num_workers, pin_memory=highres_pin_memory),
+            collate_fn=scannetpp_fisheye_passthrough,
+        )
+
+    candidate_datasets: dict[str, Any] = {}
+    candidate_dataloaders: dict[str, DataLoader] = {}
+    candidate_weights: dict[str, float] = {}
+    if re10k_ds is not None and re10k_dl is not None:
+        candidate_datasets["re10k"] = re10k_ds
+        candidate_dataloaders["re10k"] = re10k_dl
+        candidate_weights["re10k"] = float(dataset_weight_re10k)
+    if hm3d_ds is not None and hm3d_dl is not None:
+        candidate_datasets["hm3d"] = hm3d_ds
+        candidate_dataloaders["hm3d"] = hm3d_dl
+        candidate_weights["hm3d"] = float(dataset_weight_hm3d)
+    if sim_ds is not None and sim_dl is not None:
+        candidate_datasets["sim"] = sim_ds
+        candidate_dataloaders["sim"] = sim_dl
+        candidate_weights["sim"] = float(dataset_weight_sim)
+    if wildrgbd_ds is not None and wildrgbd_dl is not None:
+        candidate_datasets["wildrgbd"] = wildrgbd_ds
+        candidate_dataloaders["wildrgbd"] = wildrgbd_dl
+        candidate_weights["wildrgbd"] = float(dataset_weight_wildrgbd)
+    if dl3dv_ds is not None and dl3dv_dl is not None:
+        candidate_datasets["dl3dv"] = dl3dv_ds
+        candidate_dataloaders["dl3dv"] = dl3dv_dl
+        candidate_weights["dl3dv"] = float(dataset_weight_dl3dv)
+    if scanetpp_ds is not None and scanetpp_dl is not None:
+        candidate_datasets["scanetpp_fisheye"] = scanetpp_ds
+        candidate_dataloaders["scanetpp_fisheye"] = scanetpp_dl
+        candidate_weights["scanetpp_fisheye"] = float(dataset_weight_scanetpp)
+
+    datasets: dict[str, Any] = {}
+    dataloaders: dict[str, DataLoader] = {}
+    sampling: dict[str, float] = {}
+    for name, w in candidate_weights.items():
+        if float(w) > 0.0:
+            datasets[name] = candidate_datasets[name]
+            dataloaders[name] = candidate_dataloaders[name]
+            sampling[name] = float(w)
+        elif is_main:
+            LOGGER.warning("Skip dataset in mixed sampler: %s (weight=%.4f <= 0)", name, float(w))
+
+    if len(datasets) == 0:
+        raise ValueError("No dataset selected for mixed sampler (all dataset weights <= 0).")
+    for name, dataset in datasets.items():
+        _maybe_set_dataset_epoch(dataset, 0)
+    iterators = {name: LazyDataLoaderIterator(dl) for name, dl in dataloaders.items()}
+    sampler_seed = int(seed + rank) if seed is not None else int(12345 + rank)
+    sampler = MixedDatasetSampler(
+        datasets=datasets,
+        weights=sampling,
+        iterators=iterators,
+        seed=sampler_seed,
+    )
+
+    config = UnisharpFeatureConfig(
+        unik3d_backbone=unik3d_backbone,
+        unik3d_resolution_level=int(unik3d_resolution_level),
+        initializer_stride=int(initializer_stride),
+        initializer_scale_factor=float(initializer_scale_factor),
+        detach_init_layer0_distance=bool(detach_init_layer0_distance),
+        delta_rho_limit=float(delta_rho_limit),
+    )
+    setattr(config, "max_distance_m", float(max_depth_m))
+    
+    model = UnisharpFeatureModel(config).to(dev).train()
+
+    if _ddp_is_enabled():
+        model = DDP(
+            model,
+            device_ids=[dev.index],
+            output_device=dev.index,
+            find_unused_parameters=True,
+            gradient_as_bucket_view=True,
+        )
+
+    raw_model = model.module if isinstance(model, DDP) else model
+    base_params, unik3d_encoder_params, unik3d_decoder_params = _build_optimizer_param_groups(raw_model)
+    unik3d_params = unik3d_encoder_params + unik3d_decoder_params
+    trainable_params = base_params + unik3d_params
+    if len(trainable_params) == 0:
+        raise RuntimeError("No trainable parameters found.")
+    if len(unik3d_params) == 0:
+        raise RuntimeError(
+            "No UniK3D parameters were collected for the default unfreeze training path. "
+            "Please check parameter naming."
+        )
+    depth_head_params = [p for p in raw_model.second_layer_depth_head.parameters() if p.requires_grad]
+    if len(depth_head_params) == 0:
+        raise RuntimeError("Depth heads have no trainable parameters; depth branch would not train.")
+
+    opt_groups: list[dict[str, Any]] = [{"params": base_params, "lr": float(lr0), "group_name": "base"}]
+    if len(unik3d_encoder_params) > 0:
+        opt_groups.append(
+            {
+                "params": unik3d_encoder_params,
+                "lr": float(unik3d_encoder_lr0),
+                "group_name": "unik3d_encoder",
+            }
+        )
+    if len(unik3d_decoder_params) > 0:
+        opt_groups.append(
+            {
+                "params": unik3d_decoder_params,
+                "lr": float(unik3d_lr0),
+                "group_name": "unik3d_decoder",
+            }
+        )
+    opt = torch.optim.Adam(opt_groups)
+    if is_main:
+        LOGGER.info(
+            "Model ready: scratch heads, pretrained UniK3D, trainable_params=%d",
+            _count_numel(trainable_params),
+        )
+    if dev.type == "cuda":
+        scaler = torch.amp.GradScaler("cuda", enabled=True)
+    else:
+        scaler = torch.amp.GradScaler("cpu", enabled=False)
+
+    renderer = GSplatRenderer(
+        color_space="sRGB",
+        background_color="black",
+        low_pass_filter_eps=float(render_low_pass_filter_eps),
+    ).to(dev)
+
+    loss_w = UnisharpLossWeights(
+        lambda_color=float(lambda_color),
+        lambda_alpha=float(lambda_alpha),
+        lambda_percep=float(lambda_percep),
+        lambda_depth=float(lambda_depth),
+        lambda_tv=float(lambda_tv),
+        lambda_grad=float(lambda_grad),
+        lambda_grad_img=float(lambda_grad_img),
+        lambda_edge_rgb=float(lambda_edge_rgb),
+        lambda_delta=float(lambda_delta),
+        lambda_delta_rho=float(lambda_delta_rho),
+        lambda_splat=float(lambda_splat),
+        lambda_edge_splat=float(lambda_edge_splat),
+        lambda_grid=float(lambda_grid),
+    )
+    loss_fn = UnisharpLoss(
+        weights=loss_w,
+        delta_clip=float(delta_clip),
+        raw_delta_clip=float(raw_delta_clip),
+        raw_delta_rho_clip=float(raw_delta_rho_clip),
+        alpha_tail_min=float(alpha_tail_min),
+        alpha_tail_weight=float(alpha_tail_weight),
+        splat_sigma_min=float(splat_sigma_min),
+        splat_sigma_max=float(splat_sigma_max),
+        edge_splat_sigma_max=float(edge_splat_sigma_max),
+        depth_edge_log_threshold=float(depth_edge_log_threshold),
+        depth_edge_dilate_px=int(depth_edge_dilate_px),
+    ).to(dev)
+    loss_fn.SUPERVISION_MAX_DEPTH_M = float(max_depth_m)
+
+    if is_main:
+        config_dict = {
+            "max_depth_m": float(max_depth_m),
+            "sim_far_depth_invalid_m": float(sim_far_depth_invalid_m),
+            "sim_far_depth_invalid_max_frac": float(sim_far_depth_invalid_max_frac),
+            "re10k_pseudo_far_depth_invalid_m": float(re10k_pseudo_far_depth_invalid_m),
+            "scanetpp_fisheye_far_depth_invalid_m": float(scanetpp_fisheye_far_depth_invalid_m),
+            "render_low_pass_filter_eps": float(render_low_pass_filter_eps),
+        }
+        (out_dir / "config.json").write_text(
+            json.dumps(config_dict, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+
+    loss_csv = out_dir / "losses.csv"
+    loss_csv_fields = [
+        "loss",
+        "src_loss",
+        "tgt_loss",
+        "dataset",
+    ]
+    if is_main:
+        with loss_csv.open("w", newline="") as f:
+            csv.DictWriter(f, fieldnames=loss_csv_fields).writeheader()
+
+    if is_main:
+        LOGGER.info("Training loop started.")
+
+    from unisharp.cli.unified_trainer import UnifiedTrainer
+    
+    trainer = UnifiedTrainer(
+        model=model,
+        renderer=renderer,
+        loss_fn=loss_fn,
+        device=dev,
+        max_depth_m=float(max_depth_m),
+        sim_far_depth_invalid_m=float(sim_far_depth_invalid_m),
+        re10k_pseudo_far_depth_invalid_m=float(re10k_pseudo_far_depth_invalid_m),
+        scanetpp_fisheye_far_depth_invalid_m=float(scanetpp_fisheye_far_depth_invalid_m),
+        aux_ray_loss_weight=float(lambda_aux_ray),
+        aux_depth_scale_loss_weight=float(lambda_aux_depth_scale),
+        aux_depth2_scale_loss_weight=float(lambda_aux_depth2_scale),
+        target_mask_erode_px=int(target_mask_erode_px),
+    )
+    skip_forward_oom = _env_flag("TRAIN_SKIP_FORWARD_OOM", default=True)
+
+    dataset_epochs: dict[str, int] = {name: 0 for name in dataloaders.keys()}
+    dataset_samplers: dict[str, DistributedSampler | None] = {"hm3d": hm3d_sampler}
+
+    for step in range(1, steps + 1):
+        lr = warmup_cosine_lr(step, warmup, steps, lr0, lr1)
+        lr_unik3d_encoder = warmup_cosine_lr(step, warmup, steps, unik3d_encoder_lr0, unik3d_encoder_lr1)
+        lr_unik3d_decoder = warmup_cosine_lr(step, warmup, steps, unik3d_lr0, unik3d_lr1)
+        for g in opt.param_groups:
+            if g.get("group_name") == "unik3d_encoder":
+                g["lr"] = lr_unik3d_encoder
+            elif g.get("group_name") == "unik3d_decoder":
+                g["lr"] = lr_unik3d_decoder
+            else:
+                g["lr"] = lr
+
+        if _ddp_is_enabled():
+            batch = None
+            available_dataset_names = list(dataloaders.keys())
+            dataset_name = ""
+            for _dataset_attempt in range(max(1, len(dataloaders))):
+                dataset_name = _ddp_broadcast_str(
+                    sampler.choose_dataset_name(available_dataset_names) if is_main else "",
+                    is_main=is_main,
+                )
+
+                local_exhausted = False
+                try:
+                    batch = sampler.next_batch(dataset_name)
+                except StopIteration:
+                    local_exhausted = True
+
+                exhausted_any = _ddp_any_bool(local_exhausted, device=dev)
+                if exhausted_any:
+                    dataset_epochs[dataset_name] = dataset_epochs.get(dataset_name, 0) + 1
+                    ds_sampler = dataset_samplers.get(dataset_name, None)
+                    if ds_sampler is not None:
+                        ds_sampler.set_epoch(dataset_epochs[dataset_name])
+                    _maybe_set_dataset_epoch(datasets[dataset_name], dataset_epochs[dataset_name])
+                    iterators[dataset_name] = iter(dataloaders[dataset_name])
+                    sampler.iterators = iterators
+                    batch = None
+
+                    local_exhausted = False
+                    try:
+                        batch = sampler.next_batch(dataset_name)
+                    except StopIteration:
+                        local_exhausted = True
+                    exhausted_any = _ddp_any_bool(local_exhausted, device=dev)
+
+                if not exhausted_any:
+                    break
+
+                batch = None
+                available_dataset_names = [name for name in available_dataset_names if name != dataset_name]
+                if len(available_dataset_names) == 0:
+                    break
+            if batch is None:
+                raise RuntimeError(f"Failed to fetch synchronized DDP batch for dataset={dataset_name}")
+        else:
+            try:
+                dataset_name, batch = sampler.sample()
+            except StopIteration as e:
+                msg = str(e)
+                exhausted_name = None
+                if msg.startswith("Dataset ") and msg.endswith(" exhausted"):
+                    exhausted_name = msg[len("Dataset ") : -len(" exhausted")]
+                if exhausted_name is None or exhausted_name not in dataloaders:
+                    raise
+                dataset_epochs[exhausted_name] = dataset_epochs.get(exhausted_name, 0) + 1
+                ds_sampler = dataset_samplers.get(exhausted_name, None)
+                if ds_sampler is not None:
+                    ds_sampler.set_epoch(dataset_epochs[exhausted_name])
+                _maybe_set_dataset_epoch(datasets[exhausted_name], dataset_epochs[exhausted_name])
+                iterators[exhausted_name] = iter(dataloaders[exhausted_name])
+                sampler.iterators = iterators
+                dataset_name, batch = sampler.sample()
+
+        batch = _resize_training_batch_to_multiple(batch, int(train_resize_multiple))
+
+        opt.zero_grad(set_to_none=True)
+
+        autocast_enabled = dev.type == "cuda"
+        if autocast_enabled and torch.cuda.is_bf16_supported():
+            autocast_dtype = torch.bfloat16
+        else:
+            autocast_dtype = torch.float16 if autocast_enabled else torch.bfloat16
+
+        need_vis = bool(is_main and vis_every > 0 and (step % vis_every == 0))
+        result: dict[str, Any] | None = None
+        forward_oom_local = False
+        forward_oom_error = ""
+        try:
+            with torch.autocast(device_type=dev.type, enabled=autocast_enabled, dtype=autocast_dtype):
+                result = trainer.process_batch(
+                    batch,
+                    dataset_name,
+                    step,
+                    need_vis=need_vis,
+                )
+        except Exception as e:
+            if skip_forward_oom and _is_oom_exception(e):
+                forward_oom_local = True
+                forward_oom_error = str(e)
+                opt.zero_grad(set_to_none=True)
+                if dev.type == "cuda":
+                    torch.cuda.empty_cache()
+            else:
+                raise
+
+        forward_oom_any = _ddp_any_bool(forward_oom_local, device=dev)
+        if forward_oom_any:
+            opt.zero_grad(set_to_none=True)
+            if result is not None:
+                del result
+                result = None
+            if dev.type == "cuda":
+                torch.cuda.empty_cache()
+            if is_main:
+                LOGGER.error(
+                    "Skipping optimizer step=%d because forward OOM occurred on at least one rank | dataset=%s",
+                    int(step),
+                    str(dataset_name),
+                )
+            continue
+
+        if result is None:
+            raise RuntimeError(f"Forward returned no result for dataset={dataset_name} step={step}")
+        total_loss = result["total"]
+        local_nonfinite_loss = not bool(torch.isfinite(total_loss.detach()).item())
+        nonfinite_loss_any = _ddp_any_bool(local_nonfinite_loss, device=dev)
+        if nonfinite_loss_any:
+            opt.zero_grad(set_to_none=True)
+            if is_main:
+                LOGGER.error(
+                    "Skipping optimizer step=%d because loss is non-finite on at least one rank | dataset=%s",
+                    int(step),
+                    str(dataset_name),
+                )
+            continue
+
+        try:
+            scaler.scale(total_loss).backward()
+        except Exception as e:
+            raise
+        try:
+            scaler.unscale_(opt)
+            grad_norm = torch.nn.utils.clip_grad_norm_(trainable_params, max_norm=float(grad_clip_norm))
+        except Exception as e:
+            LOGGER.error("Gradient unscale/clip failed at step=%d: %s", int(step), str(e))
+            raise
+        grad_norm_value = float(grad_norm.detach().to(dtype=torch.float32).cpu().item()) if torch.is_tensor(grad_norm) else float(grad_norm)
+        local_nonfinite_grad = not np.isfinite(grad_norm_value)
+        nonfinite_grad_any = _ddp_any_bool(local_nonfinite_grad, device=dev)
+        if nonfinite_grad_any:
+            opt.zero_grad(set_to_none=True)
+            scaler.update()
+            if is_main:
+                LOGGER.error(
+                    "Skipping optimizer step=%d because grad norm is non-finite on at least one rank | dataset=%s | local_grad_norm=%s",
+                    int(step),
+                    str(dataset_name),
+                    str(grad_norm_value),
+                )
+            continue
+        local_huge_grad = bool(float(max_step_grad_norm) > 0.0 and grad_norm_value > float(max_step_grad_norm))
+        huge_grad_any = _ddp_any_bool(local_huge_grad, device=dev)
+        if huge_grad_any:
+            opt.zero_grad(set_to_none=True)
+            scaler.update()
+            if is_main:
+                LOGGER.error(
+                    "Skipping optimizer step=%d because grad norm exceeded max-step-grad-norm on at least one rank | dataset=%s | local_grad_norm=%.6g | threshold=%.6g",
+                    int(step),
+                    str(dataset_name),
+                    float(grad_norm_value),
+                    float(max_step_grad_norm),
+                )
+            continue
+        scaler.step(opt)
+        scaler.update()
+
+        if log_every > 0 and step % log_every == 0:
+            loss_v = float(_ddp_mean(total_loss.detach()).item())
+            src_v = float(_ddp_mean(result["src"].detach()).item())
+            tgt_v = float(_ddp_mean(result["tgt"].detach()).item())
+            row = {
+                "loss": loss_v,
+                "src_loss": src_v,
+                "tgt_loss": tgt_v,
+                "dataset": str(dataset_name),
+            }
+            if is_main:
+                LOGGER.info(
+                    "step=%d dataset=%s loss=%.6f src_loss=%.6f tgt_loss=%.6f",
+                    step,
+                    dataset_name,
+                    loss_v,
+                    src_v,
+                    tgt_v,
+                )
+                row_csv = dict(row)
+                for k in ("loss", "src_loss", "tgt_loss"):
+                    v = float(row_csv.get(k, float("nan")))
+                    row_csv[k] = "" if not np.isfinite(v) else f"{v:.4f}"
+                with loss_csv.open("a", newline="") as f:
+                    csv.DictWriter(f, fieldnames=loss_csv_fields).writerow(row_csv)
+
+        if need_vis and result.get("vis_payload"):
+            vis = result["vis_payload"]
+            _save_train_vis(
+                out_dir,
+                step,
+                vis["src_gt"],
+                vis["src_pred"],
+                vis["src_alpha"],
+                vis["tgt_gt"],
+                vis["tgt_pred"],
+                vis["tgt_alpha"],
+                src_gt_depth=vis.get("src_gt_depth"),
+                tgt_gt_depth=vis.get("tgt_gt_depth"),
+                src_pred_depth=vis.get("src_pred_depth"),
+                tgt_pred_depth=vis.get("tgt_pred_depth"),
+                src_unik3d_depth=vis.get("src_unik3d_depth"),
+                tgt_unik3d_depth=vis.get("tgt_unik3d_depth"),
+                dataset_name=vis.get("dataset_name"),
+                scene=vis.get("scene"),
+                src_idx=vis.get("src_idx"),
+                tgt_idx=vis.get("tgt_idx"),
+                src_pose_w2c=vis.get("src_pose_w2c"),
+                tgt_pose_w2c=vis.get("tgt_pose_w2c"),
+                src_metric_mask=vis.get("src_metric_mask"),
+                tgt_metric_mask=vis.get("tgt_metric_mask"),
+                src_cube_gt_u8=vis.get("src_cube_gt_u8"),
+                src_cube_pred_linear=vis.get("src_cube_pred_linear"),
+                src_cube_alpha=vis.get("src_cube_alpha"),
+                tgt_cube_gt_u8=vis.get("tgt_cube_gt_u8"),
+                tgt_cube_pred_linear=vis.get("tgt_cube_pred_linear"),
+                tgt_cube_alpha=vis.get("tgt_cube_alpha"),
+            )
+
+        if need_vis:
+            if "vis" in locals():
+                del vis
+            if dev.type == "cuda":
+                torch.cuda.empty_cache()
+        del result
+        del total_loss
+        batch = None
+
+        if is_main and (save_every > 0) and (step % save_every == 0):
+            path = out_dir / f"step_{step:07d}.pt"
+            raw_model.save_checkpoint(str(path), step, opt)
+            LOGGER.info("💾 Saved checkpoint: %s", str(path))
+
+
+    if _ddp_is_enabled():
+        _ddp_barrier(dev)
+        dist.destroy_process_group()
+
+    if is_main:
+        LOGGER.info("✅ Training completed!")

unisharp/cli/train_utils.py

@@ -0,0 +1,130 @@
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+
+def quat_mul_wxyz(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor:
+    w1, x1, y1, z1 = q1.unbind(dim=-1)
+    w2, x2, y2, z2 = q2.unbind(dim=-1)
+    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
+    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
+    y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
+    z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
+    return torch.stack([w, x, y, z], dim=-1)
+
+
+def rotmat_to_quat_wxyz(Rm: torch.Tensor) -> torch.Tensor:
+    m00, m01, m02 = Rm[0, 0], Rm[0, 1], Rm[0, 2]
+    m10, m11, m12 = Rm[1, 0], Rm[1, 1], Rm[1, 2]
+    m20, m21, m22 = Rm[2, 0], Rm[2, 1], Rm[2, 2]
+    tr = m00 + m11 + m22
+    if tr > 0.0:
+        s = torch.sqrt(tr + 1.0) * 2.0
+        w = 0.25 * s
+        x = (m21 - m12) / s
+        y = (m02 - m20) / s
+        z = (m10 - m01) / s
+    elif (m00 > m11) and (m00 > m22):
+        s = torch.sqrt(1.0 + m00 - m11 - m22) * 2.0
+        w = (m21 - m12) / s
+        x = 0.25 * s
+        y = (m01 + m10) / s
+        z = (m02 + m20) / s
+    elif m11 > m22:
+        s = torch.sqrt(1.0 + m11 - m00 - m22) * 2.0
+        w = (m02 - m20) / s
+        x = (m01 + m10) / s
+        y = 0.25 * s
+        z = (m12 + m21) / s
+    else:
+        s = torch.sqrt(1.0 + m22 - m00 - m11) * 2.0
+        w = (m10 - m01) / s
+        x = (m02 + m20) / s
+        y = (m12 + m21) / s
+        z = 0.25 * s
+    q = torch.stack([w, x, y, z])
+    return q / q.norm().clamp(min=1e-8)
+
+
+def to_k4(k3: torch.Tensor) -> torch.Tensor:
+    b = k3.shape[0]
+    out = torch.eye(4, dtype=k3.dtype, device=k3.device).unsqueeze(0).repeat(b, 1, 1)
+    out[:, :3, :3] = k3
+    return out
+
+
+def warmup_cosine_lr(step: int, warmup: int, total: int, lr0: float, lr1: float) -> float:
+    if step <= warmup:
+        return lr0 * float(step) / float(max(1, warmup))
+    t = (step - warmup) / float(max(1, total - warmup))
+    cos = 0.5 * (1 + np.cos(np.pi * t))
+    return lr1 + (lr0 - lr1) * cos
+
+
+@torch.no_grad()
+def compute_frustum_mask(
+    depth: torch.Tensor,
+    tgt_w2c: torch.Tensor,
+    src_w2c: torch.Tensor,
+    src_k3: torch.Tensor,
+    tgt_k3: torch.Tensor,
+    img_h: int,
+    img_w: int,
+    source_img_h: int | None = None,
+    source_img_w: int | None = None,
+    depth_min: float = 0.05,
+    margin: float = 0.05,
+) -> torch.Tensor:
+    dev = depth.device
+    f32 = torch.float32
+    src_h = int(img_h if source_img_h is None else source_img_h)
+    src_w = int(img_w if source_img_w is None else source_img_w)
+
+    d = depth[0, 0].to(f32)
+    valid = d > depth_min
+
+    vy, vx = torch.meshgrid(
+        torch.arange(img_h, device=dev, dtype=f32),
+        torch.arange(img_w, device=dev, dtype=f32),
+        indexing="ij",
+    )
+
+    fx_t = tgt_k3[0, 0, 0].to(f32)
+    fy_t = tgt_k3[0, 1, 1].to(f32)
+    cx_t = tgt_k3[0, 0, 2].to(f32)
+    cy_t = tgt_k3[0, 1, 2].to(f32)
+    X_t = (vx - cx_t) / fx_t * d
+    Y_t = (vy - cy_t) / fy_t * d
+    Z_t = d
+    pts_t = torch.stack([X_t, Y_t, Z_t], dim=-1).reshape(-1, 3)
+
+    c2w_t = torch.linalg.inv(tgt_w2c[0].to(f32))
+    pts_w = pts_t @ c2w_t[:3, :3].T + c2w_t[:3, 3][None, :]
+
+    w2c_s = src_w2c[0].to(f32)
+    pts_s = pts_w @ w2c_s[:3, :3].T + w2c_s[:3, 3][None, :]
+
+    Z_s = pts_s[:, 2].clamp(min=1e-4)
+    fx_s = src_k3[0, 0, 0].to(f32)
+    fy_s = src_k3[0, 1, 1].to(f32)
+    cx_s = src_k3[0, 0, 2].to(f32)
+    cy_s = src_k3[0, 1, 2].to(f32)
+    u_s = pts_s[:, 0] / Z_s * fx_s + cx_s
+    v_s = pts_s[:, 1] / Z_s * fy_s + cy_s
+
+    half_w = (src_w - 1) * 0.5
+    half_h = (src_h - 1) * 0.5
+    x_ndc = (u_s - half_w) / half_w
+    y_ndc = (v_s - half_h) / half_h
+
+    in_frust = (
+        (x_ndc.abs() <= 1.0 + margin)
+        & (y_ndc.abs() <= 1.0 + margin)
+        & (pts_s[:, 2] > 0)
+    )
+
+    mask = in_frust.reshape(img_h, img_w).float()
+    mask = mask * valid.float()
+    return mask[None, None]

unisharp/cli/unified_trainer.py

@@ -0,0 +1,1966 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import os
+from typing import Any, Callable
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from unisharp.utils.gsplat import GSplatRenderer
+from unisharp.losses import UnisharpLoss
+from unisharp.utils.camera_utils import (
+    transform_gaussians_to_world,
+    to_k4,
+    compute_frustum_mask,
+)
+from unisharp.utils.fisheye_geer import (
+    compute_fisheye624_frustum_mask,
+    render_gaussians_fisheye624,
+)
+from unisharp.utils.camera_projection import cubemap_face_cameras, build_extrinsics_w2c, view_frustum_mask_cubemap_union
+from unisharp.utils.pano import Cube2Equirec, get_pinhole_intrinsics_4x4
+from unisharp import DEFAULT_MAX_DEPTH_M
+from unisharp.utils.pixel_convention import integer_pixel_center_grid
+
+
+@dataclass
+class _ModeStrategy:
+
+    batch_size: int
+    gaussians: Any
+    make_world_gaussians: Callable[[int, Any], Any]
+    make_sample: Callable[[int, Any, bool], dict[str, Any]]
+    collect_all_vis: bool = False
+
+
+class UnifiedTrainer:
+    
+    def __init__(
+        self,
+        model: nn.Module,
+        renderer: GSplatRenderer,
+        loss_fn: UnisharpLoss,
+        device: torch.device,
+        enable_tgt_unik3d_vis: bool = True,
+        max_depth_m: float = DEFAULT_MAX_DEPTH_M,
+        sim_far_depth_invalid_m: float = 30.0,
+        re10k_pseudo_far_depth_invalid_m: float = 30.0,
+        scanetpp_fisheye_far_depth_invalid_m: float = 30.0,
+        aux_ray_loss_weight: float = 3.0,
+        aux_depth_scale_loss_weight: float = 3.0,
+        aux_depth2_scale_loss_weight: float = 1.0,
+        target_mask_erode_px: int = 0,
+    ):
+        self.model = model
+        self.renderer = renderer
+        self.loss_fn = loss_fn
+        self.device = device
+        self.enable_tgt_unik3d_vis = bool(enable_tgt_unik3d_vis)
+        self.max_depth_m = float(max_depth_m)
+        self.sim_far_depth_invalid_m = float(sim_far_depth_invalid_m)
+        self.re10k_pseudo_far_depth_invalid_m = float(re10k_pseudo_far_depth_invalid_m)
+        self.scanetpp_fisheye_far_depth_invalid_m = float(scanetpp_fisheye_far_depth_invalid_m)
+        self.aux_ray_loss_weight = float(aux_ray_loss_weight)
+        self.aux_depth_scale_loss_weight = float(aux_depth_scale_loss_weight)
+        self.aux_depth2_scale_loss_weight = float(aux_depth2_scale_loss_weight)
+        self.target_mask_erode_px = max(int(target_mask_erode_px), 0)
+
+    @staticmethod
+    def _erode_supervision_mask(mask: torch.Tensor, radius_px: int, *, circular_h: bool = False) -> torch.Tensor:
+        radius = max(int(radius_px), 0)
+        if radius <= 0:
+            return mask
+        if not torch.is_tensor(mask):
+            return mask
+        m = mask.to(dtype=torch.float32).clamp(0.0, 1.0)
+        if m.ndim == 3:
+            m = m.unsqueeze(1)
+        invalid = 1.0 - m
+        kernel = 2 * radius + 1
+        if bool(circular_h):
+            invalid = F.pad(invalid, (radius, radius, 0, 0), mode="circular")
+            invalid = F.pad(invalid, (0, 0, radius, radius), mode="constant", value=0.0)
+            dilated_invalid = F.max_pool2d(invalid, kernel_size=kernel, stride=1)
+        else:
+            dilated_invalid = F.max_pool2d(invalid, kernel_size=kernel, stride=1, padding=radius)
+        return (m * (1.0 - dilated_invalid)).to(device=mask.device, dtype=mask.dtype)
+
+    def _aux_ray_losses(
+        self,
+        *,
+        pred_rays: torch.Tensor | None,
+        gt_rays: torch.Tensor | None,
+        mask: torch.Tensor | None,
+        pred_distance: torch.Tensor | None = None,
+        pred_distance2: torch.Tensor | None = None,
+        gt_distance: torch.Tensor | None = None,
+        gt_distance2: torch.Tensor | None = None,
+        depth_mask: torch.Tensor | None = None,
+        depth_mask2: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        out: dict[str, torch.Tensor] = {}
+        if torch.is_tensor(pred_rays) and torch.is_tensor(gt_rays) and self.aux_ray_loss_weight > 0.0:
+            out["unik3d_ray"] = self.aux_ray_loss_weight * self._unik3d_polar_ray_loss(
+                pred_rays,
+                gt_rays,
+                mask,
+            )
+        if torch.is_tensor(pred_distance) and torch.is_tensor(gt_distance):
+            out["unik3d_depth_scale"] = self.aux_depth_scale_loss_weight * self._unik3d_scale_depth_loss(
+                pred_distance,
+                gt_distance,
+                depth_mask if torch.is_tensor(depth_mask) else mask,
+            )
+        depth2_target = gt_distance2 if torch.is_tensor(gt_distance2) else gt_distance
+        if torch.is_tensor(pred_distance2) and torch.is_tensor(depth2_target):
+            depth2_mask = depth_mask2 if torch.is_tensor(depth_mask2) else depth_mask
+            out["unik3d_depth2_scale"] = self.aux_depth2_scale_loss_weight * self._unik3d_scale_depth_loss(
+                pred_distance2,
+                depth2_target,
+                depth2_mask if torch.is_tensor(depth2_mask) else mask,
+            )
+        return out
+
+    DEPTH_SUPERVISION_MAX_M: float = DEFAULT_MAX_DEPTH_M
+
+    def _distance_init_cap_for_dataset(self, dataset_name: str) -> float | None:
+        name = str(dataset_name).lower()
+        if name == "re10k" and self.re10k_pseudo_far_depth_invalid_m > 0.0:
+            return self.re10k_pseudo_far_depth_invalid_m
+        if name == "sim" and self.sim_far_depth_invalid_m > 0.0:
+            return self.sim_far_depth_invalid_m
+        if name in {"scanetpp_fisheye", "scannetpp_fisheye"} and self.scanetpp_fisheye_far_depth_invalid_m > 0.0:
+            return self.scanetpp_fisheye_far_depth_invalid_m
+        return None
+
+    @staticmethod
+    def _unik3d_polar_ray_loss(
+        pred_rays: torch.Tensor | None,
+        gt_rays: torch.Tensor | None,
+        mask: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if not torch.is_tensor(pred_rays) or not torch.is_tensor(gt_rays):
+            device = pred_rays.device if torch.is_tensor(pred_rays) else torch.device("cpu")
+            return torch.zeros((), device=device, dtype=torch.float32)
+        pred = pred_rays.to(dtype=torch.float32)
+        gt = gt_rays.to(device=pred.device, dtype=torch.float32)
+        if pred.ndim == 3:
+            pred = pred.unsqueeze(0)
+        if gt.ndim == 3:
+            gt = gt.unsqueeze(0)
+        if tuple(pred.shape) != tuple(gt.shape):
+            gt = F.interpolate(gt, size=pred.shape[-2:], mode="bilinear", align_corners=False)
+            gt = gt / torch.norm(gt, dim=1, keepdim=True).clamp(min=1e-5)
+        pred = pred / torch.norm(pred, dim=1, keepdim=True).clamp(min=1e-5)
+        gt = gt / torch.norm(gt, dim=1, keepdim=True).clamp(min=1e-5)
+
+        px, py, pz = pred.unbind(dim=1)
+        gx, gy, gz = gt.unbind(dim=1)
+        polar_pred = torch.acos(pz.clamp(min=-0.99999, max=0.99999))
+        polar_gt = torch.acos(gz.clamp(min=-0.99999, max=0.99999))
+        az_pred = torch.atan2(py, px.abs().clamp(min=1e-5) * (2.0 * (px > 0).to(px.dtype) - 1.0))
+        az_gt = torch.atan2(gy, gx.abs().clamp(min=1e-5) * (2.0 * (gx > 0).to(gx.dtype) - 1.0))
+        polar_error = (polar_pred - polar_gt).abs()
+        az_delta = az_pred - az_gt
+        az_error = torch.atan2(torch.sin(az_delta), torch.cos(az_delta)).abs()
+        quantile_weight = torch.ones_like(polar_error)
+        quantile_weight[(polar_gt > polar_pred) & (polar_gt > torch.pi / 2)] = 1.4
+        quantile_weight[(polar_gt <= polar_pred) & (polar_gt > torch.pi / 2)] = 0.6
+
+        if torch.is_tensor(mask):
+            m = mask.to(device=pred.device, dtype=torch.float32)
+            if m.ndim == 3:
+                m = m.unsqueeze(1)
+            if tuple(m.shape[-2:]) != tuple(pred.shape[-2:]):
+                m = F.interpolate(m, size=pred.shape[-2:], mode="nearest")
+            m = m[:, 0].clamp(0.0, 1.0)
+        else:
+            m = torch.ones_like(polar_error)
+        denom = m.sum(dim=(-1, -2), keepdim=False).clamp(min=1.0)
+        mean_polar = (polar_error * quantile_weight * m).sum(dim=(-1, -2)) / denom
+        mean_azimuth = (az_error * m).sum(dim=(-1, -2)) / denom
+        mean_error = (3.0 * mean_polar + mean_azimuth) / 4.0
+        return torch.sqrt(mean_error + 1e-4).mean()
+
+    @staticmethod
+    def _unik3d_scale_depth_loss(
+        pred_distance: torch.Tensor,
+        gt_distance: torch.Tensor,
+        mask: torch.Tensor | None,
+    ) -> torch.Tensor:
+        pred = UnifiedTrainer._as_b1hw_depth(pred_distance).to(dtype=torch.float32)
+        gt = UnifiedTrainer._as_b1hw_depth(gt_distance).to(device=pred.device, dtype=torch.float32)
+        if tuple(gt.shape[-2:]) != tuple(pred.shape[-2:]):
+            gt = F.interpolate(gt, size=pred.shape[-2:], mode="nearest")
+        valid = torch.isfinite(pred) & torch.isfinite(gt) & (pred > 0.0) & (gt > 0.0)
+        if torch.is_tensor(mask):
+            m = mask.to(device=pred.device)
+            if m.ndim == 3:
+                m = m.unsqueeze(1)
+            if tuple(m.shape[-2:]) != tuple(pred.shape[-2:]):
+                m = F.interpolate(m.to(dtype=torch.float32), size=pred.shape[-2:], mode="nearest")
+            valid = valid & (m[:, :1] > 0.5)
+        err = (gt.clamp(min=1e-4).log() - pred.clamp(min=1e-4).log()).abs()
+        err = torch.where(valid, err, torch.zeros_like(err))
+        denom = valid.to(dtype=err.dtype).sum(dim=(-2, -1)).clamp(min=1.0)
+        per_image = err.sum(dim=(-2, -1)) / denom
+        return torch.sqrt(per_image.clamp(min=0.0)).mean()
+
+    def _base_model(self) -> nn.Module:
+        return self.model.module if hasattr(self.model, "module") else self.model
+        
+    def process_batch(
+        self,
+        batch: Any,
+        dataset_name: str,
+        step: int,
+        need_vis: bool = False,
+    ) -> dict[str, Any]:
+        if hasattr(batch, "src_rgb_u8") and hasattr(batch, "src_intrinsics"):
+            strategy = self._build_pinhole_strategy(
+                batch,
+                step,
+                need_vis=need_vis,
+                dataset_name=str(dataset_name),
+            )
+        elif hasattr(batch, "src_rgb_u8") and hasattr(batch, "src_camera_params"):
+            strategy = self._build_fisheye_strategy(
+                batch,
+                step,
+                need_vis=need_vis,
+                dataset_name=str(dataset_name),
+            )
+        elif hasattr(batch, "src_erp_rgb_u8") and hasattr(batch, "src_cube_depth_m"):
+            strategy = self._build_spherical_strategy(
+                batch,
+                step,
+                need_vis=need_vis,
+                dataset_name=str(dataset_name),
+            )
+        else:
+            raise ValueError(f"Unknown batch schema for dataset={dataset_name}")
+        return self._run_strategy_loop(
+            strategy,
+            need_vis=need_vis,
+        )
+
+    def _run_strategy_loop(
+        self,
+        strategy: _ModeStrategy,
+        need_vis: bool = False,
+    ) -> dict[str, Any]:
+        total_loss = torch.zeros((), device=self.device)
+        src_sum = torch.zeros((), device=self.device)
+        tgt_sum = torch.zeros((), device=self.device)
+        src_log_sum: dict[str, torch.Tensor] = {}
+        tgt_log_sum: dict[str, torch.Tensor] = {}
+        aux_log_sum: dict[str, torch.Tensor] = {}
+        vis_payload = None
+        vis_payloads: list[dict[str, Any]] = []
+
+        def _accumulate_loss_terms(term_specs: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
+            merged: dict[str, torch.Tensor] = {}
+            for spec in term_specs:
+                term_losses = self._compute_view_loss(**spec)
+                for k, v in term_losses.items():
+                    merged[k] = merged.get(k, torch.zeros((), device=self.device)) + v
+            return merged
+
+        collect_all_vis = bool(getattr(strategy, "collect_all_vis", False))
+        for b in range(int(strategy.batch_size)):
+            g = strategy.gaussians
+            g_b = type(g)(
+                mean_vectors=g.mean_vectors[b : b + 1],
+                singular_values=g.singular_values[b : b + 1],
+                quaternions=g.quaternions[b : b + 1],
+                colors=g.colors[b : b + 1],
+                opacities=g.opacities[b : b + 1],
+            )
+            g_world = strategy.make_world_gaussians(b, g_b)
+            sample = strategy.make_sample(
+                b,
+                g_world,
+                bool(need_vis and (collect_all_vis or b == 0)),
+            )
+
+            if isinstance(sample.get("src_loss_terms", None), list):
+                src_losses = _accumulate_loss_terms(sample["src_loss_terms"])
+            else:
+                src_losses = self._compute_view_loss(
+                    pred_rgb_linear=sample["src_pred_rgb_linear"],
+                    pred_alpha=sample["src_pred_alpha"],
+                    pred_depth_m=sample["src_pred_depth_m"],
+                    pred_depth2_m=sample.get("src_pred_depth2_m", None),
+                    gt_rgb_u8=sample["src_gt_rgb_u8"],
+                    gt_depth_m=sample["src_gt_depth_m"],
+                    mask=sample["src_mask"],
+                    apply_color=bool(sample.get("src_apply_color", True)),
+                    apply_alpha=bool(sample.get("src_apply_alpha", True)),
+                    apply_depth=bool(sample.get("src_apply_depth", True)),
+                    apply_percep=False,
+                    apply_tv=True,
+                    apply_grad=bool(sample.get("src_apply_grad", True)),
+                    apply_grad_img=bool(sample.get("src_apply_grad_img", True)),
+                    apply_splat=bool(sample.get("src_apply_splat", True)),
+                    grad_img_circular_h=sample.get("src_grad_img_circular_h", None),
+                    gaussian_scales=sample.get("gaussian_scales", None),
+                    gaussian_quaternions=sample.get("gaussian_quaternions", None),
+                    gaussian_angular_cell=sample.get("gaussian_angular_cell", None),
+                    delta_xy=sample.get("delta_xy", None),
+                    delta_rho=sample.get("delta_rho", None),
+                    delta_grid=sample.get("delta_grid", None),
+                    gaussian_mean_vectors=sample.get("gaussian_mean_vectors", None),
+                    gaussian_base_mean_vectors=sample.get("gaussian_base_mean_vectors", None),
+                    gaussian_opacities=sample.get("gaussian_opacities", None),
+                    gauss_grid_shape=sample.get("gauss_grid_shape", None),
+                    projected_scale_factor=sample.get("projected_scale_factor", None),
+                    projection_model=sample.get("projection_model", None),
+                    projection_intrinsics=sample.get("projection_intrinsics", None),
+                    projection_camera_params=sample.get("projection_camera_params", None),
+                    depth_mask=sample.get("src_depth_mask", None),
+                )
+            if isinstance(sample.get("src_extra_loss_terms", None), list):
+                extra_src_losses = _accumulate_loss_terms(sample["src_extra_loss_terms"])
+                for k, v in extra_src_losses.items():
+                    src_losses[k] = src_losses.get(k, torch.zeros((), device=self.device)) + v
+            if isinstance(sample.get("tgt_loss_terms", None), list):
+                tgt_losses = _accumulate_loss_terms(sample["tgt_loss_terms"])
+            else:
+                tgt_losses = self._compute_view_loss(
+                    pred_rgb_linear=sample["tgt_pred_rgb_linear"],
+                    pred_alpha=sample["tgt_pred_alpha"],
+                    pred_depth_m=sample["tgt_pred_depth_m"],
+                    pred_depth2_m=sample.get("tgt_pred_depth2_m", None),
+                    gt_rgb_u8=sample["tgt_gt_rgb_u8"],
+                    gt_depth_m=sample["tgt_gt_depth_m"],
+                    mask=sample["tgt_mask"],
+                    apply_color=bool(sample.get("tgt_apply_color", True)),
+                    apply_alpha=bool(sample.get("tgt_apply_alpha", True)),
+                    apply_depth=bool(sample.get("tgt_apply_depth", True)),
+                    apply_percep=bool(sample.get("tgt_apply_percep", False)),
+                    apply_tv=False,
+                    apply_grad=False,
+                    apply_grad_img=bool(sample.get("tgt_apply_grad_img", True)),
+                    apply_splat=bool(sample.get("tgt_apply_splat", False)),
+                    grad_img_circular_h=sample.get("tgt_grad_img_circular_h", None),
+                    gaussian_scales=None,
+                    gaussian_quaternions=None,
+                    delta_xy=None,
+                    delta_rho=None,
+                    gaussian_mean_vectors=None,
+                    gaussian_base_mean_vectors=None,
+                    gaussian_opacities=None,
+                    gauss_grid_shape=None,
+                    projected_scale_factor=sample.get("projected_scale_factor", None),
+                    projection_model=sample.get("projection_model", None),
+                    projection_intrinsics=sample.get("projection_intrinsics", None),
+                    projection_camera_params=sample.get("projection_camera_params", None),
+                    depth_mask=sample.get("tgt_depth_mask", None),
+                )
+            if isinstance(sample.get("tgt_extra_loss_terms", None), list):
+                extra_tgt_losses = _accumulate_loss_terms(sample["tgt_extra_loss_terms"])
+                for k, v in extra_tgt_losses.items():
+                    tgt_losses[k] = tgt_losses.get(k, torch.zeros((), device=self.device)) + v
+
+            aux_total = torch.zeros((), device=self.device)
+            raw_aux = sample.get("aux_losses", None)
+            if isinstance(raw_aux, dict):
+                for k, v in raw_aux.items():
+                    if torch.is_tensor(v):
+                        vv = v.to(device=self.device)
+                    else:
+                        vv = torch.tensor(float(v), device=self.device, dtype=torch.float32)
+                    aux_total = aux_total + vv
+                    aux_log_sum[str(k)] = aux_log_sum.get(str(k), torch.zeros((), device=self.device)) + vv.detach()
+
+            src_sum = src_sum + src_losses["total"]
+            tgt_sum = tgt_sum + tgt_losses["total"]
+            total_loss = total_loss + src_losses["total"] + tgt_losses["total"] + aux_total
+            for k, v in src_losses.items():
+                src_log_sum[k] = src_log_sum.get(k, torch.zeros((), device=self.device)) + v.detach()
+            for k, v in tgt_losses.items():
+                tgt_log_sum[k] = tgt_log_sum.get(k, torch.zeros((), device=self.device)) + v.detach()
+
+            if need_vis and isinstance(sample.get("vis_payload", None), dict):
+                vis_payloads.append(sample["vis_payload"])
+                if b == 0:
+                    vis_payload = sample["vis_payload"]
+
+        bs = float(strategy.batch_size)
+        total_loss = total_loss / bs
+        src_sum = src_sum / bs
+        tgt_sum = tgt_sum / bs
+        loss_breakdown: dict[str, torch.Tensor] = {}
+        for k, v in src_log_sum.items():
+            loss_breakdown[f"src_{k}"] = v / bs
+        for k, v in tgt_log_sum.items():
+            loss_breakdown[f"tgt_{k}"] = v / bs
+        for k, v in aux_log_sum.items():
+            loss_breakdown[f"aux_{k}"] = v / bs
+        batch_stats = {
+            "batch_size": int(strategy.batch_size),
+            "gaussian_count": int(strategy.gaussians.mean_vectors.shape[1]),
+        }
+
+        return {
+            "total": total_loss,
+            "src": src_sum,
+            "tgt": tgt_sum,
+            "loss_breakdown": loss_breakdown,
+            "batch_stats": batch_stats,
+            "vis_payload": vis_payload,
+            "vis_payloads": vis_payloads,
+        }
+
+    @staticmethod
+    def _first_item(x: Any, default: Any = None) -> Any:
+        if x is None:
+            return default
+        if isinstance(x, (list, tuple)):
+            return x[0] if len(x) > 0 else default
+        if torch.is_tensor(x):
+            if x.numel() == 0:
+                return default
+            return x.flatten()[0].item()
+        return x
+
+    @staticmethod
+    def _item_at(x: Any, index: int, default: Any = None) -> Any:
+        if x is None:
+            return default
+        if isinstance(x, (list, tuple)):
+            return x[index] if 0 <= int(index) < len(x) else default
+        if torch.is_tensor(x):
+            if x.numel() == 0:
+                return default
+            if x.ndim == 0:
+                return x.item()
+            if 0 <= int(index) < int(x.shape[0]):
+                item = x[int(index)]
+                return item.item() if item.numel() == 1 else item
+            return default
+        return x
+
+    @staticmethod
+    def _finite_quantile(x: torch.Tensor, q: float, default: float = float("nan")) -> torch.Tensor:
+        vals = x[torch.isfinite(x)]
+        if int(vals.numel()) <= 0:
+            return torch.tensor(float(default), device=x.device, dtype=torch.float32)
+        vals = vals.to(torch.float32).flatten()
+        if int(vals.numel()) > 262144:
+            step = max(1, int(vals.numel()) // 262144)
+            vals = vals[::step]
+        return torch.quantile(vals, float(q))
+
+    def _clamp_distance_for_supervision(
+        self,
+        depth_m: torch.Tensor | None,
+        *,
+        max_depth_m: float | None = None,
+        clamp_max: bool = True,
+    ) -> torch.Tensor | None:
+        if not torch.is_tensor(depth_m):
+            return None
+        cap = float(self.max_depth_m if max_depth_m is None else max_depth_m)
+        out = depth_m.to(dtype=torch.float32)
+        valid = torch.isfinite(out) & (out > 0.0)
+        if bool(clamp_max):
+            sanitized = out.clamp(min=1e-4, max=cap)
+        else:
+            sanitized = out.clamp(min=1e-4)
+        return torch.where(valid, sanitized, torch.zeros_like(out))
+
+    @staticmethod
+    def _rendered_depth_valid_for_inv_loss(
+        depth_m: torch.Tensor,
+        alpha: torch.Tensor,
+        *,
+        alpha_min: float | None = None,
+        depth_min_m: float = 1e-3,
+    ) -> torch.Tensor:
+        depth = depth_m.detach()
+        valid = torch.isfinite(depth) & (depth > float(depth_min_m))
+        if alpha_min is not None:
+            a = alpha.detach().to(device=depth.device)
+            valid = valid & (a[:, :1] > float(alpha_min))
+        return valid.to(dtype=depth.dtype)
+
+    def _pinhole_z_to_supervision_distance(
+        self,
+        z_depth_b1hw: torch.Tensor | None,
+        k3_b33: torch.Tensor | None,
+        *,
+        clamp_max: bool = True,
+    ) -> torch.Tensor | None:
+        if not torch.is_tensor(z_depth_b1hw) or not torch.is_tensor(k3_b33):
+            return None
+        dist = self._z_depth_to_distance_pinhole(z_depth_b1hw, k3_b33)
+        return self._clamp_distance_for_supervision(dist, clamp_max=bool(clamp_max))
+
+    @staticmethod
+    def _sanitize_positive_depth(depth_m: torch.Tensor | None) -> torch.Tensor | None:
+        if not torch.is_tensor(depth_m):
+            return None
+        out = depth_m.to(dtype=torch.float32)
+        valid = torch.isfinite(out) & (out > 0.0)
+        return torch.where(valid, out, torch.zeros_like(out))
+
+    @staticmethod
+    def _as_b1hw_depth(depth: torch.Tensor) -> torch.Tensor:
+        if depth.ndim == 3:
+            return depth.unsqueeze(1)
+        if depth.ndim == 4 and depth.shape[1] == 1:
+            return depth
+        raise ValueError(f"Expected depth shape (B,H,W) or (B,1,H,W), got {tuple(depth.shape)}")
+
+    @staticmethod
+    def _as_bchw_rgb_u8(image: torch.Tensor) -> torch.Tensor:
+        if image.ndim == 3 and image.shape[0] == 3:
+            return image.unsqueeze(0)
+        if image.ndim == 4 and image.shape[1] == 3:
+            return image
+        raise ValueError(f"Expected image shape (3,H,W) or (B,3,H,W), got {tuple(image.shape)}")
+
+    @staticmethod
+    def _as_b33_intrinsics(intrinsics: torch.Tensor) -> torch.Tensor:
+        if intrinsics.ndim == 2 and tuple(intrinsics.shape) == (3, 3):
+            return intrinsics.unsqueeze(0)
+        if intrinsics.ndim == 3 and tuple(intrinsics.shape[1:]) == (3, 3):
+            return intrinsics
+        raise ValueError(
+            f"Expected intrinsics shape (3,3) or (B,3,3), got {tuple(intrinsics.shape)}"
+        )
+
+    @staticmethod
+    def _as_b9_camera_params(camera_params: torch.Tensor) -> torch.Tensor:
+        if camera_params.ndim == 1 and int(camera_params.shape[0]) == 9:
+            return camera_params.unsqueeze(0)
+        if camera_params.ndim == 2 and int(camera_params.shape[1]) == 9:
+            return camera_params
+        raise ValueError(f"Expected camera_params shape (9,) or (B,9), got {tuple(camera_params.shape)}")
+
+    @staticmethod
+    def _as_b16_camera_params(camera_params: torch.Tensor) -> torch.Tensor:
+        if camera_params.ndim == 1 and int(camera_params.shape[0]) == 16:
+            return camera_params.unsqueeze(0)
+        if camera_params.ndim == 2 and int(camera_params.shape[1]) == 16:
+            return camera_params
+        raise ValueError(f"Expected camera_params shape (16,) or (B,16), got {tuple(camera_params.shape)}")
+
+    @staticmethod
+    def _as_b44_pose(extrinsics: torch.Tensor) -> torch.Tensor:
+        if extrinsics.ndim == 2 and tuple(extrinsics.shape) == (4, 4):
+            return extrinsics.unsqueeze(0)
+        if extrinsics.ndim == 3 and tuple(extrinsics.shape[1:]) == (4, 4):
+            return extrinsics
+        raise ValueError(
+            f"Expected extrinsics shape (4,4) or (B,4,4), got {tuple(extrinsics.shape)}"
+        )
+
+    @staticmethod
+    def _pick_depth_for_pinhole_frustum_mask(
+        gt_depth: torch.Tensor | None,
+        pred_depth: torch.Tensor,
+        min_valid_px: int = 8,
+    ) -> torch.Tensor:
+        if torch.is_tensor(gt_depth):
+            gt_depth = UnifiedTrainer._as_b1hw_depth(gt_depth)
+            valid = torch.isfinite(gt_depth) & (gt_depth > 0.0)
+            if int(valid.sum().item()) >= int(min_valid_px):
+                return gt_depth
+        return pred_depth
+
+    @staticmethod
+    def _pick_depth_for_fisheye_frustum_mask(
+        gt_depth: torch.Tensor | None,
+        pred_depth: torch.Tensor,
+        gt_valid_mask: torch.Tensor | None = None,
+        min_valid_px: int = 8,
+    ) -> torch.Tensor:
+        if torch.is_tensor(gt_depth):
+            gt_depth = UnifiedTrainer._as_b1hw_depth(gt_depth)
+            if torch.is_tensor(gt_valid_mask):
+                gt_valid = gt_depth > 0.0
+                gt_valid = gt_valid & (gt_valid_mask > 0.5)
+            else:
+                gt_valid = torch.isfinite(gt_depth) & (gt_depth > 0.0)
+            if int(gt_valid.sum().item()) >= int(min_valid_px):
+                return gt_depth
+        return pred_depth
+
+    @staticmethod
+    def _as_cubemap_depth_hw1(depth: torch.Tensor) -> torch.Tensor:
+        if depth.ndim != 4:
+            raise ValueError(f"Expected 4D cubemap depth, got shape={tuple(depth.shape)}")
+        if depth.shape[-1] == 1:
+            return depth
+        if depth.shape[1] == 1:
+            return depth.permute(0, 2, 3, 1).contiguous()
+        raise ValueError(f"Unsupported cubemap depth shape={tuple(depth.shape)}")
+
+    def _pick_depth_for_cubemap_frustum_mask(
+        self,
+        gt_depth_cube: torch.Tensor | None,
+        pred_depth_cube: torch.Tensor,
+        face_w: int,
+        min_valid_px: int = 8,
+    ) -> torch.Tensor:
+        pred_hw1 = self._as_cubemap_depth_hw1(pred_depth_cube)
+        if torch.is_tensor(gt_depth_cube):
+            gt_hw1 = self._as_cubemap_depth_hw1(gt_depth_cube)
+            gt_dist = self._cubemap_z_depth_to_distance(gt_hw1)
+            gt_hw1 = self._as_cubemap_depth_hw1(gt_dist)
+            if gt_hw1.shape[1] != int(face_w) or gt_hw1.shape[2] != int(face_w):
+                gt_hw1 = F.interpolate(
+                    gt_hw1.permute(0, 3, 1, 2),
+                    size=(int(face_w), int(face_w)),
+                    mode="nearest",
+                ).permute(0, 2, 3, 1).contiguous()
+            valid = torch.isfinite(gt_hw1[..., 0]) & (gt_hw1[..., 0] > 0.0)
+            if int(valid.sum().item()) >= int(min_valid_px):
+                return gt_hw1
+        return pred_hw1
+
+    @staticmethod
+    def _distance_to_z_depth_pinhole(
+        distance_b1hw: torch.Tensor,
+        intrinsics_b33: torch.Tensor,
+    ) -> torch.Tensor:
+        distance_b1hw = UnifiedTrainer._as_b1hw_depth(distance_b1hw)
+        intrinsics_b33 = UnifiedTrainer._as_b33_intrinsics(intrinsics_b33)
+        b, _, h, w = distance_b1hw.shape
+        dev = distance_b1hw.device
+        dtype = distance_b1hw.dtype
+        uu, vv = integer_pixel_center_grid(h, w, device=dev, dtype=dtype)
+        uu = uu.unsqueeze(0).expand(b, -1, -1)
+        vv = vv.unsqueeze(0).expand(b, -1, -1)
+        fx = intrinsics_b33[:, 0, 0].view(b, 1, 1).to(dtype=dtype, device=dev)
+        fy = intrinsics_b33[:, 1, 1].view(b, 1, 1).to(dtype=dtype, device=dev)
+        cx = intrinsics_b33[:, 0, 2].view(b, 1, 1).to(dtype=dtype, device=dev)
+        cy = intrinsics_b33[:, 1, 2].view(b, 1, 1).to(dtype=dtype, device=dev)
+        x = (uu - cx) / fx
+        y = (vv - cy) / fy
+        ray_z = 1.0 / torch.sqrt(x * x + y * y + 1.0).clamp(min=1e-8)
+        return distance_b1hw * ray_z.unsqueeze(1)
+
+    @staticmethod
+    def _z_depth_to_distance_pinhole(
+        z_depth_b1hw: torch.Tensor,
+        intrinsics_b33: torch.Tensor,
+    ) -> torch.Tensor:
+        z_depth_b1hw = UnifiedTrainer._as_b1hw_depth(z_depth_b1hw)
+        intrinsics_b33 = UnifiedTrainer._as_b33_intrinsics(intrinsics_b33)
+        b, _, h, w = z_depth_b1hw.shape
+        dev = z_depth_b1hw.device
+        dtype = z_depth_b1hw.dtype
+        uu, vv = integer_pixel_center_grid(h, w, device=dev, dtype=dtype)
+        uu = uu.unsqueeze(0).expand(b, -1, -1)
+        vv = vv.unsqueeze(0).expand(b, -1, -1)
+        fx = intrinsics_b33[:, 0, 0].view(b, 1, 1).to(dtype=dtype, device=dev)
+        fy = intrinsics_b33[:, 1, 1].view(b, 1, 1).to(dtype=dtype, device=dev)
+        cx = intrinsics_b33[:, 0, 2].view(b, 1, 1).to(dtype=dtype, device=dev)
+        cy = intrinsics_b33[:, 1, 2].view(b, 1, 1).to(dtype=dtype, device=dev)
+        x = (uu - cx) / fx
+        y = (vv - cy) / fy
+        ray_z = 1.0 / torch.sqrt(x * x + y * y + 1.0).clamp(min=1e-8)
+        return z_depth_b1hw / ray_z.unsqueeze(1).clamp(min=1e-8)
+
+    def _cubemap_z_depth_to_distance(
+        self,
+        depth_cube: torch.Tensor,
+    ) -> torch.Tensor:
+        if depth_cube.ndim != 4:
+            raise ValueError(f"Expected 4D cubemap depth, got {tuple(depth_cube.shape)}")
+        if depth_cube.shape[-1] == 1:
+            depth_61hw = depth_cube.permute(0, 3, 1, 2).contiguous()
+        elif depth_cube.shape[1] == 1:
+            depth_61hw = depth_cube
+        else:
+            raise ValueError(f"Unsupported cubemap depth shape={tuple(depth_cube.shape)}")
+
+        _, _, h, w = depth_61hw.shape
+        intr = get_pinhole_intrinsics_4x4(int(w)).to(
+            device=depth_61hw.device,
+            dtype=depth_61hw.dtype,
+        )
+        fx = intr[0, 0]
+        fy = intr[1, 1]
+        cx = intr[0, 2]
+        cy = intr[1, 2]
+        uu, vv = integer_pixel_center_grid(h, w, device=depth_61hw.device, dtype=depth_61hw.dtype)
+        x = (uu - cx) / fx
+        y = (vv - cy) / fy
+        ray_z = 1.0 / torch.sqrt(x * x + y * y + 1.0).clamp(min=1e-8)
+        dist = depth_61hw / ray_z.view(1, 1, h, w).clamp(min=1e-8)
+        valid = torch.isfinite(dist) & (depth_61hw > 0.0)
+        dist = torch.where(valid, dist.clamp(min=1e-4), torch.zeros_like(dist))
+        return dist
+
+    def _collect_regularization_inputs(
+        self,
+        out: dict[str, Any],
+        gaussians: Any,
+        b: int,
+        projected_scale_factor: float | None,
+    ) -> dict[str, Any]:
+        delta_b = out.get("delta", None)
+        delta_xy_raw = None
+        if torch.is_tensor(delta_b):
+            delta_xy_raw = delta_b[b : b + 1, 0:2]
+            delta_rho_raw = delta_b[b : b + 1, 2:3]
+            delta_grid_raw = delta_b[b : b + 1]
+        else:
+            delta_rho_raw = None
+            delta_grid_raw = None
+        delta_rho_applied_all = out.get("delta_rho_applied", None)
+        delta_rho_applied = (
+            delta_rho_applied_all[b : b + 1]
+            if torch.is_tensor(delta_rho_applied_all)
+            else None
+        )
+        scale_factor_applied_all = out.get("scale_factor_applied", None)
+        scale_factor_applied = (
+            scale_factor_applied_all[b : b + 1]
+            if torch.is_tensor(scale_factor_applied_all)
+            else None
+        )
+
+        scales_b = gaussians.singular_values[b : b + 1]
+        means_b = gaussians.mean_vectors[b : b + 1]
+        quats_b = gaussians.quaternions[b : b + 1]
+        opac_b = gaussians.opacities[b : b + 1]
+
+        base_values = out.get("gaussian_base_values", None)
+        gauss_grid_shape = None
+        base_means_b = None
+        base_scales_b = None
+        angular_cell_b = None
+        if base_values is not None and hasattr(base_values, "rays"):
+            _, _, l, hb, wb = base_values.rays.shape
+            gauss_grid_shape = (int(l), int(hb), int(wb))
+            inv_dist_b = base_values.inv_distance[b : b + 1].clamp(min=1e-6)
+            base_rays_b = F.normalize(base_values.rays[b : b + 1], dim=1, eps=1e-6)
+            base_means_grid = base_rays_b / inv_dist_b
+            base_scales_b = base_values.scales[b : b + 1]
+            init_output = out.get("initializer_output", None)
+            global_scale = (
+                init_output.global_scale[b : b + 1]
+                if init_output is not None
+                and getattr(init_output, "global_scale", None) is not None
+                else None
+            )
+            if torch.is_tensor(global_scale):
+                base_means_grid = base_means_grid * global_scale.view(-1, 1, 1, 1, 1)
+                base_scales_b = base_scales_b * global_scale.view(-1, 1, 1, 1, 1)
+            base_means_b = base_means_grid.permute(0, 2, 3, 4, 1).flatten(1, 3)
+            angular_cell = getattr(base_values, "angular_cell", None)
+            angular_cell_b = angular_cell[b : b + 1] if torch.is_tensor(angular_cell) else None
+
+        return {
+            "delta_xy_eff": delta_xy_raw,
+            "delta_rho_raw": delta_rho_raw,
+            "delta_grid": delta_grid_raw,
+            "delta_rho_applied": delta_rho_applied,
+            "scale_factor_applied": scale_factor_applied,
+            "gaussian_scales": scales_b,
+            "gaussian_quaternions": quats_b,
+            "gaussian_angular_cell": angular_cell_b,
+            "gaussian_mean_vectors": means_b,
+            "gaussian_base_mean_vectors": base_means_b,
+            "gaussian_base_scales": base_scales_b,
+            "gaussian_opacities": opac_b,
+            "gauss_grid_shape": gauss_grid_shape,
+            "projected_scale_factor": projected_scale_factor,
+        }
+
+    def _compute_view_loss(
+        self,
+        *,
+        pred_rgb_linear: torch.Tensor,
+        pred_alpha: torch.Tensor,
+        pred_depth_m: torch.Tensor,
+        pred_depth2_m: torch.Tensor | None,
+        gt_rgb_u8: torch.Tensor,
+        gt_depth_m: torch.Tensor,
+        mask: torch.Tensor,
+        apply_color: bool,
+        apply_alpha: bool,
+        apply_depth: bool,
+        apply_percep: bool,
+        apply_tv: bool,
+        apply_grad: bool,
+        apply_grad_img: bool,
+        grad_img_circular_h: bool | None = None,
+        gaussian_scales: torch.Tensor | None = None,
+        gaussian_quaternions: torch.Tensor | None = None,
+        gaussian_angular_cell: torch.Tensor | None = None,
+        delta_xy: torch.Tensor | None = None,
+        delta_rho: torch.Tensor | None = None,
+        delta_grid: torch.Tensor | None = None,
+        gaussian_mean_vectors: torch.Tensor | None = None,
+        gaussian_base_mean_vectors: torch.Tensor | None = None,
+        gaussian_opacities: torch.Tensor | None = None,
+        gauss_grid_shape: tuple[int, int, int] | None = None,
+        projected_scale_factor: float | torch.Tensor | None = None,
+        projection_model: str | None = None,
+        projection_intrinsics: torch.Tensor | None = None,
+        projection_camera_params: torch.Tensor | None = None,
+        loss_scale: float = 1.0,
+        apply_splat: bool | None = None,
+        depth_mask: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        losses = self.loss_fn(
+            pred_rgb_linear=pred_rgb_linear,
+            pred_alpha=pred_alpha,
+            pred_depth_m=pred_depth_m,
+            pred_depth2_m=pred_depth2_m,
+            gt_rgb_u8=gt_rgb_u8,
+            gt_depth_m=gt_depth_m,
+            mask=mask,
+            depth_mask=depth_mask,
+            gaussian_scales=gaussian_scales,
+            gaussian_quaternions=gaussian_quaternions,
+            gaussian_angular_cell=gaussian_angular_cell,
+            delta_xy=delta_xy,
+            delta_rho=delta_rho,
+            delta_grid=delta_grid,
+            apply_color=bool(apply_color),
+            apply_alpha=bool(apply_alpha),
+            apply_depth=bool(apply_depth),
+            apply_percep=bool(apply_percep),
+            apply_tv=bool(apply_tv),
+            apply_grad=bool(apply_grad),
+            apply_grad_img=bool(apply_grad_img),
+            grad_img_circular_h=grad_img_circular_h,
+            apply_delta=bool(torch.is_tensor(delta_xy) or torch.is_tensor(delta_rho)),
+            apply_splat=bool(torch.is_tensor(gaussian_scales)) if apply_splat is None else bool(apply_splat),
+            gaussian_mean_vectors=gaussian_mean_vectors,
+            gaussian_base_mean_vectors=gaussian_base_mean_vectors,
+            gaussian_opacities=gaussian_opacities,
+            gauss_grid_shape=gauss_grid_shape,
+            projected_scale_factor=projected_scale_factor,
+            projection_model=projection_model,
+            projection_intrinsics=projection_intrinsics,
+            projection_camera_params=projection_camera_params,
+        )
+        scale = float(loss_scale)
+        if abs(scale - 1.0) > 1e-8:
+            losses = {k: (v * scale) for k, v in losses.items()}
+        return losses
+    
+    def _build_pinhole_strategy(
+        self,
+        batch: Any,
+        step: int,
+        need_vis: bool = False,
+        dataset_name: str = "re10k",
+    ) -> _ModeStrategy:
+        src_u8 = self._as_bchw_rgb_u8(batch.src_rgb_u8.to(self.device, non_blocking=True))
+        tgt_u8 = self._as_bchw_rgb_u8(batch.tgt_rgb_u8.to(self.device, non_blocking=True))
+        src_u8_orig = getattr(batch, "src_rgb_u8_orig", None)
+        tgt_u8_orig = getattr(batch, "tgt_rgb_u8_orig", None)
+        src_depth_gt = getattr(batch, "src_depth_m", None)
+        tgt_depth_gt = getattr(batch, "tgt_depth_m", None)
+        src_depth_gt_orig = getattr(batch, "src_depth_m_orig", None)
+        tgt_depth_gt_orig = getattr(batch, "tgt_depth_m_orig", None)
+        has_depth_gt = torch.is_tensor(src_depth_gt) and torch.is_tensor(tgt_depth_gt)
+        if has_depth_gt:
+            src_depth_gt = self._as_b1hw_depth(
+                src_depth_gt.to(self.device, non_blocking=True).to(torch.float32)
+            )
+            tgt_depth_gt = self._as_b1hw_depth(
+                tgt_depth_gt.to(self.device, non_blocking=True).to(torch.float32)
+            )
+        has_depth_gt_orig = torch.is_tensor(src_depth_gt_orig) and torch.is_tensor(tgt_depth_gt_orig)
+        if has_depth_gt_orig:
+            src_depth_gt_orig = self._as_b1hw_depth(
+                src_depth_gt_orig.to(self.device, non_blocking=True).to(torch.float32)
+            )
+            tgt_depth_gt_orig = self._as_b1hw_depth(
+                tgt_depth_gt_orig.to(self.device, non_blocking=True).to(torch.float32)
+            )
+        src_w2c = self._as_b44_pose(batch.src_w2c.to(self.device, non_blocking=True).to(torch.float32))
+        tgt_w2c = self._as_b44_pose(batch.tgt_w2c.to(self.device, non_blocking=True).to(torch.float32))
+        src_k3 = self._as_b33_intrinsics(batch.src_intrinsics.to(self.device, non_blocking=True).to(torch.float32))
+        tgt_k3 = self._as_b33_intrinsics(batch.tgt_intrinsics.to(self.device, non_blocking=True).to(torch.float32))
+        src_k3_orig = getattr(batch, "src_intrinsics_orig", None)
+        tgt_k3_orig = getattr(batch, "tgt_intrinsics_orig", None)
+        has_orig_vis = (
+            torch.is_tensor(src_u8_orig)
+            and torch.is_tensor(tgt_u8_orig)
+            and torch.is_tensor(src_k3_orig)
+            and torch.is_tensor(tgt_k3_orig)
+        )
+        if has_orig_vis:
+            src_u8_orig = self._as_bchw_rgb_u8(src_u8_orig.to(self.device, non_blocking=True))
+            tgt_u8_orig = self._as_bchw_rgb_u8(tgt_u8_orig.to(self.device, non_blocking=True))
+            src_k3_orig = self._as_b33_intrinsics(
+                src_k3_orig.to(self.device, non_blocking=True).to(torch.float32)
+            )
+            tgt_k3_orig = self._as_b33_intrinsics(
+                tgt_k3_orig.to(self.device, non_blocking=True).to(torch.float32)
+            )
+        src_depth_gt_dist = None
+        tgt_depth_gt_dist = None
+        src_unik3d_gt_dist = None
+        if has_depth_gt:
+            src_unik3d_gt_dist = self._pinhole_z_to_supervision_distance(src_depth_gt, src_k3)
+            src_depth_gt_dist = src_unik3d_gt_dist
+            tgt_depth_gt_dist = self._pinhole_z_to_supervision_distance(tgt_depth_gt, tgt_k3)
+        
+        src = src_u8.float().clamp(0, 255) / 255.0
+        tgt = tgt_u8.float().clamp(0, 255) / 255.0
+        distance_init_cap_m = self._distance_init_cap_for_dataset(dataset_name)
+        
+        share_src_forward = bool(getattr(batch, "share_src_forward", False)) and int(src.shape[0]) > 1
+
+        def _repeat_first_dim(value: Any, batch_size: int) -> Any:
+            if torch.is_tensor(value):
+                if value.ndim > 0 and int(value.shape[0]) == 1:
+                    return value.repeat(batch_size, *([1] * (value.ndim - 1)))
+                return value
+            if hasattr(value, "_fields"):
+                return type(value)(*[_repeat_first_dim(getattr(value, field), batch_size) for field in value._fields])
+            return value
+
+        if share_src_forward:
+            out_single = self.model(
+                image=src[0:1],
+                image_u8=src_u8[0:1],
+                camera_intrinsics=src_k3[0:1],
+                camera_model="pinhole",
+                depth_gt=(src_depth_gt_dist[0:1] if torch.is_tensor(src_depth_gt_dist) else None),
+                distance_init_cap_m=distance_init_cap_m,
+                return_aux=True,
+            )
+            out = {k: _repeat_first_dim(v, int(src.shape[0])) for k, v in out_single.items()}
+        else:
+            out = self.model(
+                image=src,
+                image_u8=src_u8,
+                camera_intrinsics=src_k3,
+                camera_model="pinhole",
+                depth_gt=src_depth_gt_dist,
+                distance_init_cap_m=distance_init_cap_m,
+                return_aux=True,
+            )
+        gaussians = out["gaussians"]
+        src_render_k3 = src_k3
+        tgt_render_k3 = tgt_k3
+        src_depth_gt_z_render = src_depth_gt if has_depth_gt else None
+        tgt_depth_gt_z_render = tgt_depth_gt if has_depth_gt else None
+        src_depth_gt_render_valid = (torch.isfinite(src_depth_gt) & (src_depth_gt > 0.0)) if has_depth_gt else None
+        tgt_depth_gt_render_valid = (torch.isfinite(tgt_depth_gt) & (tgt_depth_gt > 0.0)) if has_depth_gt else None
+        aux_ray_target_all = out.get("unik3d_gt_rays", None)
+        def make_world_gaussians(b: int, g_b: Any) -> Any:
+            return g_b
+
+        def make_sample(b: int, g_world: Any, enable_vis: bool) -> dict[str, Any]:
+            src_h = int(src_u8.shape[-2])
+            src_w = int(src_u8.shape[-1])
+            tgt_h = int(tgt_u8.shape[-2])
+            tgt_w = int(tgt_u8.shape[-1])
+            ident = torch.eye(4, dtype=src_w2c.dtype, device=self.device).unsqueeze(0)
+            rel_tgt_w2c = tgt_w2c[b : b + 1] @ torch.linalg.inv(src_w2c[b : b + 1])
+            src_k_render_b = src_render_k3[b : b + 1]
+            tgt_k_render_b = tgt_render_k3[b : b + 1]
+            src_out = self.renderer(
+                g_world,
+                extrinsics=ident,
+                intrinsics=to_k4(src_k_render_b),
+                image_width=src_w,
+                image_height=src_h,
+            )
+            tgt_out = self.renderer(
+                g_world,
+                extrinsics=rel_tgt_w2c,
+                intrinsics=to_k4(tgt_k_render_b),
+                image_width=tgt_w,
+                image_height=tgt_h,
+            )
+
+            zeros_src_depth = torch.zeros((1, 1, src_h, src_w), dtype=torch.float32, device=self.device)
+            zeros_tgt_depth = torch.zeros((1, 1, tgt_h, tgt_w), dtype=torch.float32, device=self.device)
+            ones_mask = torch.ones_like(zeros_src_depth)
+            fx_b = float(src_k_render_b[0, 0, 0].item())
+            fy_b = float(src_k_render_b[0, 1, 1].item())
+            proj_scale_pinhole = 0.5 * (fx_b + fy_b)
+            reg_inputs = self._collect_regularization_inputs(
+                out=out,
+                gaussians=gaussians,
+                b=b,
+                projected_scale_factor=proj_scale_pinhole,
+            )
+
+            src_depth_for_visibility = None
+            tgt_gt_depth_for_mask = (
+                tgt_depth_gt_z_render[b : b + 1]
+                if has_depth_gt and torch.is_tensor(tgt_depth_gt_z_render)
+                else None
+            )
+            if has_depth_gt:
+                src_depth_for_visibility = (
+                    src_depth_gt_z_render[b : b + 1]
+                    if torch.is_tensor(src_depth_gt_z_render)
+                    else src_depth_gt[b : b + 1]
+                )
+
+            tgt_depth_for_mask = self._pick_depth_for_pinhole_frustum_mask(
+                gt_depth=tgt_gt_depth_for_mask,
+                pred_depth=tgt_out.depth,
+            )
+            tgt_frustum_mask = compute_frustum_mask(
+                depth=tgt_depth_for_mask,
+                tgt_w2c=tgt_w2c[b : b + 1],
+                src_w2c=src_w2c[b : b + 1],
+                src_k3=src_k_render_b,
+                tgt_k3=tgt_k_render_b,
+                img_h=tgt_h,
+                img_w=tgt_w,
+                source_img_h=src_h,
+                source_img_w=src_w,
+                source_depth=src_depth_for_visibility,
+            )
+            tgt_frustum_mask_raw = tgt_frustum_mask
+            tgt_frustum_mask = self._erode_supervision_mask(
+                tgt_frustum_mask,
+                self.target_mask_erode_px,
+                circular_h=False,
+            )
+            src_depth_pred = self._clamp_distance_for_supervision(
+                out["distance_layers"][b : b + 1, 0:1],
+                clamp_max=False,
+            )
+            src_depth2_pred = (
+                self._clamp_distance_for_supervision(out["distance_layers"][b : b + 1, 1:2], clamp_max=False)
+                if out["distance_layers"] is not None and out["distance_layers"].shape[1] > 1
+                else None
+            )
+            src_depth2_gt_for_aux = (
+                src_unik3d_gt_dist[b : b + 1]
+                if torch.is_tensor(src_unik3d_gt_dist)
+                else None
+            )
+            src_depth2_mask_for_aux = src_depth_gt[b : b + 1] > 0.0 if has_depth_gt else None
+            tgt_depth_pred = self._pinhole_z_to_supervision_distance(
+                tgt_out.depth,
+                tgt_k_render_b,
+                clamp_max=False,
+            )
+            tgt_depth_loss_mask = self._rendered_depth_valid_for_inv_loss(tgt_depth_pred, tgt_out.alpha)
+            if torch.is_tensor(tgt_depth_gt_render_valid):
+                tgt_depth_loss_mask = tgt_depth_loss_mask * tgt_depth_gt_render_valid[b : b + 1].to(
+                    device=tgt_depth_loss_mask.device,
+                    dtype=tgt_depth_loss_mask.dtype,
+                )
+            tgt_extra_loss_terms: list[dict[str, Any]] = []
+            vis_payload = None
+            if enable_vis:
+                vis_src_u8 = src_u8[b : b + 1]
+                vis_tgt_u8 = tgt_u8[b : b + 1]
+                vis_src_depth_gt = (src_depth_gt[b : b + 1] if has_depth_gt else None)
+                vis_tgt_depth_gt = (tgt_depth_gt[b : b + 1] if has_depth_gt else None)
+                vis_src_out = src_out
+                vis_tgt_out = tgt_out
+                if has_orig_vis:
+                    vis_src_u8 = src_u8_orig[b : b + 1]
+                    vis_tgt_u8 = tgt_u8_orig[b : b + 1]
+                    vis_src_depth_gt = (src_depth_gt_orig[b : b + 1] if has_depth_gt_orig else None)
+                    vis_tgt_depth_gt = (tgt_depth_gt_orig[b : b + 1] if has_depth_gt_orig else None)
+                    vis_src_render_k3 = src_k3_orig[b : b + 1]
+                    vis_tgt_render_k3 = tgt_k3_orig[b : b + 1]
+                    vis_src_out = self.renderer(
+                        g_world,
+                        extrinsics=ident,
+                        intrinsics=to_k4(vis_src_render_k3),
+                        image_width=int(vis_src_u8.shape[-1]),
+                        image_height=int(vis_src_u8.shape[-2]),
+                    )
+                    vis_tgt_out = self.renderer(
+                        g_world,
+                        extrinsics=rel_tgt_w2c,
+                        intrinsics=to_k4(vis_tgt_render_k3),
+                        image_width=int(vis_tgt_u8.shape[-1]),
+                        image_height=int(vis_tgt_u8.shape[-2]),
+                    )
+                src_unik3d_depth = None
+                tgt_unik3d_depth = None
+                raw_dist = out.get("unik3d_distance", None)
+                if torch.is_tensor(raw_dist):
+                    try:
+                        conditioning_rays = out.get("unik3d_ray_conditioning_rays", None)
+                        if not torch.is_tensor(conditioning_rays):
+                            conditioning_rays = out.get("unik3d_rays", None)
+                        ray_z = (
+                            conditioning_rays[b : b + 1, 2:3].detach()
+                            if torch.is_tensor(conditioning_rays)
+                            else None
+                        )
+                        if torch.is_tensor(ray_z):
+                            if tuple(ray_z.shape[-2:]) != tuple(raw_dist.shape[-2:]):
+                                ray_z = F.interpolate(ray_z, size=raw_dist.shape[-2:], mode="bilinear", align_corners=False)
+                            src_unik3d_depth = raw_dist[b : b + 1, 0:1].detach() * ray_z
+                        else:
+                            src_unik3d_depth = self._distance_to_z_depth_pinhole(
+                                raw_dist[b : b + 1, 0:1].detach(),
+                                src_k_render_b,
+                            )
+                    except Exception:
+                        src_unik3d_depth = raw_dist[b : b + 1, 0:1].detach()
+                if self.enable_tgt_unik3d_vis:
+                    try:
+                        with torch.no_grad():
+                            from unisharp.utils.unik3d_adapter import forward_unik3d_pinhole
+
+                            unik_tgt = forward_unik3d_pinhole(
+                                self._base_model().feature_extractor.unik3d,
+                                rgb_u8=tgt_u8[b : b + 1],
+                                intrinsics=tgt_k3[b : b + 1],
+                                normalize=True,
+                            )
+                            dist_tgt = unik_tgt.get("distance", None) if isinstance(unik_tgt, dict) else None
+                            if torch.is_tensor(dist_tgt):
+                                try:
+                                    tgt_unik3d_depth = self._distance_to_z_depth_pinhole(
+                                        dist_tgt[:, 0:1].detach(),
+                                        tgt_k_render_b,
+                                    )
+                                except Exception:
+                                    tgt_unik3d_depth = dist_tgt[:, 0:1].detach()
+                    except Exception:
+                        tgt_unik3d_depth = None
+
+                vis_payload = {
+                    "src_gt": (vis_src_u8.float() / 255.0).detach(),
+                    "src_pred": vis_src_out.color.clamp(0, 1).detach(),
+                    "src_alpha": vis_src_out.alpha.detach(),
+                    "src_gt_depth": (vis_src_depth_gt.detach() if torch.is_tensor(vis_src_depth_gt) else None),
+                    "src_pred_depth": vis_src_out.depth.detach(),
+                    "src_unik3d_depth": src_unik3d_depth,
+                    "tgt_gt": (vis_tgt_u8.float() / 255.0).detach(),
+                    "tgt_pred": vis_tgt_out.color.clamp(0, 1).detach(),
+                    "tgt_alpha": vis_tgt_out.alpha.detach(),
+                    "tgt_gt_depth": (vis_tgt_depth_gt.detach() if torch.is_tensor(vis_tgt_depth_gt) else None),
+                    "tgt_pred_depth": vis_tgt_out.depth.detach(),
+                    "tgt_unik3d_depth": tgt_unik3d_depth,
+                    "dataset_name": str(dataset_name),
+                    "scene": str(self._item_at(getattr(batch, "scene", None), b, "unknown")),
+                    "src_idx": int(self._item_at(getattr(batch, "src_idx", None), b, -1)),
+                    "tgt_idx": int(self._item_at(getattr(batch, "tgt_idx", None), b, -1)),
+                    "src_pose_w2c": src_w2c[b : b + 1].detach(),
+                    "tgt_pose_w2c": tgt_w2c[b : b + 1].detach(),
+                    "tgt_metric_mask_raw": tgt_frustum_mask_raw.detach(),
+                    "tgt_metric_mask": tgt_frustum_mask.detach(),
+                }
+
+            return {
+                "src_pred_rgb_linear": src_out.color,
+                "src_pred_alpha": src_out.alpha,
+                "src_pred_depth_m": src_depth_pred,
+                "src_pred_depth2_m": src_depth2_pred,
+                "src_gt_rgb_u8": src_u8[b : b + 1],
+                "src_gt_depth_m": (src_depth_gt_dist[b : b + 1] if has_depth_gt and src_depth_gt_dist is not None else zeros_src_depth),
+                "src_mask": ones_mask,
+                "src_apply_depth": False,
+                "src_apply_grad": bool(has_depth_gt),
+                "src_apply_grad_img": bool(has_depth_gt),
+                "src_grad_img_circular_h": False,
+                "tgt_pred_rgb_linear": tgt_out.color,
+                "tgt_pred_alpha": tgt_out.alpha,
+                "tgt_pred_depth_m": tgt_depth_pred,
+                "tgt_gt_rgb_u8": tgt_u8[b : b + 1],
+                "tgt_gt_depth_m": (tgt_depth_gt_dist[b : b + 1] if has_depth_gt and tgt_depth_gt_dist is not None else zeros_tgt_depth),
+                "tgt_mask": tgt_frustum_mask,
+                "tgt_depth_mask": tgt_depth_loss_mask,
+                "tgt_apply_depth": bool(has_depth_gt),
+                "tgt_apply_grad_img": bool(has_depth_gt),
+                "tgt_grad_img_circular_h": False,
+                "tgt_apply_percep": bool(float(self.loss_fn.w.lambda_percep) > 0.0),
+                "tgt_extra_loss_terms": tgt_extra_loss_terms,
+                "aux_losses": self._aux_ray_losses(
+                    pred_rays=(
+                        out.get("unik3d_rays", None)[b : b + 1]
+                        if torch.is_tensor(out.get("unik3d_rays", None))
+                        else None
+                    ),
+                    gt_rays=(
+                        aux_ray_target_all[b : b + 1]
+                        if torch.is_tensor(aux_ray_target_all)
+                        else None
+                    ),
+                    mask=ones_mask,
+                    pred_distance=(
+                        out["unik3d_distance"][b : b + 1, 0:1]
+                        if torch.is_tensor(out.get("unik3d_distance", None))
+                        else None
+                    ),
+                    pred_distance2=src_depth2_pred,
+                    gt_distance=(
+                        src_unik3d_gt_dist[b : b + 1]
+                        if torch.is_tensor(src_unik3d_gt_dist)
+                        else None
+                    ),
+                    gt_distance2=src_depth2_gt_for_aux,
+                    depth_mask=(src_depth_gt[b : b + 1] > 0.0 if has_depth_gt else None),
+                    depth_mask2=src_depth2_mask_for_aux,
+                ),
+                "gaussian_scales": reg_inputs["gaussian_scales"],
+                "gaussian_quaternions": reg_inputs["gaussian_quaternions"],
+                "gaussian_angular_cell": reg_inputs["gaussian_angular_cell"],
+                "delta_xy": reg_inputs["delta_xy_eff"],
+                "delta_rho": reg_inputs["delta_rho_raw"],
+                "delta_grid": reg_inputs["delta_grid"],
+                "gaussian_mean_vectors": reg_inputs["gaussian_mean_vectors"],
+                "gaussian_base_mean_vectors": reg_inputs["gaussian_base_mean_vectors"],
+                "gaussian_opacities": reg_inputs["gaussian_opacities"],
+                "gauss_grid_shape": reg_inputs["gauss_grid_shape"],
+                "projected_scale_factor": reg_inputs["projected_scale_factor"],
+                "projection_model": "pinhole",
+                "projection_intrinsics": src_k_render_b,
+                "vis_payload": vis_payload,
+            }
+
+        return _ModeStrategy(
+            batch_size=int(src.shape[0]),
+            gaussians=gaussians,
+            make_world_gaussians=make_world_gaussians,
+            make_sample=make_sample,
+            collect_all_vis=bool(getattr(batch, "collect_all_vis", False)),
+        )
+
+    def _build_fisheye624_strategy(
+        self,
+        batch: Any,
+        step: int,
+        need_vis: bool = False,
+        dataset_name: str = "scannetpp_fisheye",
+    ) -> _ModeStrategy:
+        del step
+        src_u8 = self._as_bchw_rgb_u8(batch.src_rgb_u8.to(self.device, non_blocking=True))
+        tgt_u8 = self._as_bchw_rgb_u8(batch.tgt_rgb_u8.to(self.device, non_blocking=True))
+        src_depth_gt = self._clamp_distance_for_supervision(
+            self._as_b1hw_depth(batch.src_depth_m.to(self.device, non_blocking=True).to(torch.float32))
+        )
+        tgt_depth_gt = self._clamp_distance_for_supervision(
+            self._as_b1hw_depth(batch.tgt_depth_m.to(self.device, non_blocking=True).to(torch.float32))
+        )
+        src_valid_mask = self._as_b1hw_depth(batch.src_valid_mask.to(self.device, non_blocking=True).to(torch.float32))
+        tgt_valid_mask = self._as_b1hw_depth(batch.tgt_valid_mask.to(self.device, non_blocking=True).to(torch.float32))
+        src_w2c = self._as_b44_pose(batch.src_w2c.to(self.device, non_blocking=True).to(torch.float32))
+        tgt_w2c = self._as_b44_pose(batch.tgt_w2c.to(self.device, non_blocking=True).to(torch.float32))
+        src_cam_params = self._as_b16_camera_params(
+            batch.src_camera_params.to(self.device, non_blocking=True).to(torch.float32)
+        )
+        tgt_cam_params = self._as_b16_camera_params(
+            batch.tgt_camera_params.to(self.device, non_blocking=True).to(torch.float32)
+        )
+        distance_init_cap_m = self._distance_init_cap_for_dataset(dataset_name)
+
+        out = self.model(
+            image=src_u8.float().clamp(0, 255) / 255.0,
+            image_u8=src_u8,
+            camera_intrinsics=None,
+            camera_params=src_cam_params,
+            camera_model="fisheye624",
+            depth_gt=src_depth_gt,
+            distance_init_cap_m=distance_init_cap_m,
+            validity_mask=src_valid_mask,
+            return_aux=True,
+        )
+
+        gaussians = out["gaussians"]
+        src_render_cam_params = src_cam_params
+        tgt_render_cam_params = tgt_cam_params
+        src_render_valid_mask = src_valid_mask
+        tgt_render_valid_mask = tgt_valid_mask
+        aux_ray_target_all = out.get("unik3d_gt_rays", None)
+
+        def make_world_gaussians(b: int, g_b: Any) -> Any:
+            return transform_gaussians_to_world(g_b, src_w2c[b])
+
+        def make_sample(b: int, g_world: Any, enable_vis: bool) -> dict[str, Any]:
+            src_h = int(src_u8.shape[-2])
+            src_w = int(src_u8.shape[-1])
+            tgt_h = int(tgt_u8.shape[-2])
+            tgt_w = int(tgt_u8.shape[-1])
+            src_render = render_gaussians_fisheye624(
+                g_world,
+                extrinsics_w2c=src_w2c[b : b + 1],
+                camera_params=src_render_cam_params[b : b + 1],
+                image_h=src_h,
+                image_w=src_w,
+                valid_mask=src_render_valid_mask[b : b + 1],
+            )
+            tgt_render = render_gaussians_fisheye624(
+                g_world,
+                extrinsics_w2c=tgt_w2c[b : b + 1],
+                camera_params=tgt_render_cam_params[b : b + 1],
+                image_h=tgt_h,
+                image_w=tgt_w,
+                valid_mask=tgt_render_valid_mask[b : b + 1],
+            )
+            reg_inputs = self._collect_regularization_inputs(
+                out=out,
+                gaussians=gaussians,
+                b=b,
+                projected_scale_factor=None,
+            )
+            tgt_depth_for_mask = self._pick_depth_for_fisheye_frustum_mask(
+                gt_depth=tgt_depth_gt[b : b + 1],
+                pred_depth=tgt_render["depth_distance"],
+                gt_valid_mask=tgt_valid_mask[b : b + 1],
+            )
+            tgt_frustum_mask = compute_fisheye624_frustum_mask(
+                depth_distance_m=tgt_depth_for_mask,
+                tgt_w2c=tgt_w2c[b : b + 1],
+                src_w2c=src_w2c[b : b + 1],
+                tgt_camera_params=tgt_render_cam_params[b : b + 1],
+                src_camera_params=src_render_cam_params[b : b + 1],
+                src_valid_mask=src_render_valid_mask[b : b + 1] * src_render["valid_mask"],
+                source_depth_distance_m=src_depth_gt[b : b + 1],
+            )
+            src_mask = src_render_valid_mask[b : b + 1] * src_render["valid_mask"]
+            src_depth_mask = src_mask
+            tgt_mask = tgt_render_valid_mask[b : b + 1] * tgt_render["valid_mask"] * tgt_frustum_mask
+            tgt_mask_raw = tgt_mask
+            tgt_mask = self._erode_supervision_mask(
+                tgt_mask,
+                self.target_mask_erode_px,
+                circular_h=False,
+            )
+            src_depth_pred = self._clamp_distance_for_supervision(
+                out["distance_layers"][b : b + 1, 0:1],
+                clamp_max=False,
+            )
+            src_depth2_pred = (
+                self._clamp_distance_for_supervision(out["distance_layers"][b : b + 1, 1:2], clamp_max=False)
+                if out["distance_layers"] is not None and out["distance_layers"].shape[1] > 1
+                else None
+            )
+            tgt_depth_pred = self._clamp_distance_for_supervision(tgt_render["depth_distance"], clamp_max=False)
+            tgt_depth_loss_mask = self._rendered_depth_valid_for_inv_loss(tgt_depth_pred, tgt_render["alpha"])
+
+            src_loss_terms = [
+                {
+                    "pred_rgb_linear": src_render["color"],
+                    "pred_alpha": src_render["alpha"],
+                    "pred_depth_m": src_render["depth_distance"],
+                    "pred_depth2_m": None,
+                    "gt_rgb_u8": src_u8[b : b + 1],
+                    "gt_depth_m": src_depth_gt[b : b + 1],
+                    "mask": src_mask,
+                    "apply_color": True,
+                    "apply_alpha": True,
+                    "apply_depth": False,
+                    "apply_percep": False,
+                    "apply_tv": False,
+                    "apply_grad": False,
+                    "apply_grad_img": False,
+                    "grad_img_circular_h": False,
+                    "gaussian_scales": None,
+                    "gaussian_quaternions": None,
+                    "gaussian_angular_cell": None,
+                    "delta_xy": None,
+                    "gaussian_mean_vectors": None,
+                    "gaussian_opacities": None,
+                    "gauss_grid_shape": None,
+                    "projected_scale_factor": None,
+                    "apply_splat": False,
+                    "loss_scale": 1.0,
+                }
+            ]
+            src_extra_loss_terms = [
+                {
+                    "pred_rgb_linear": torch.zeros((1, 3, src_h, src_w), dtype=torch.float32, device=self.device),
+                    "pred_alpha": torch.zeros((1, 1, src_h, src_w), dtype=torch.float32, device=self.device),
+                    "pred_depth_m": src_depth_pred,
+                    "pred_depth2_m": src_depth2_pred,
+                    "gt_rgb_u8": torch.zeros((1, 3, src_h, src_w), dtype=torch.uint8, device=self.device),
+                    "gt_depth_m": src_depth_gt[b : b + 1],
+                    "mask": src_depth_mask,
+                    "apply_color": False,
+                    "apply_alpha": False,
+                    "apply_depth": False,
+                    "apply_percep": False,
+                    "apply_tv": True,
+                    "apply_grad": True,
+                    "apply_grad_img": True,
+                    "grad_img_circular_h": False,
+                    "gaussian_scales": reg_inputs["gaussian_scales"],
+                    "gaussian_quaternions": reg_inputs["gaussian_quaternions"],
+                    "gaussian_angular_cell": reg_inputs["gaussian_angular_cell"],
+                    "delta_xy": reg_inputs["delta_xy_eff"],
+                    "delta_rho": reg_inputs["delta_rho_raw"],
+                    "delta_grid": reg_inputs["delta_grid"],
+                    "gaussian_mean_vectors": reg_inputs["gaussian_mean_vectors"],
+                    "gaussian_base_mean_vectors": reg_inputs["gaussian_base_mean_vectors"],
+                    "gaussian_opacities": reg_inputs["gaussian_opacities"],
+                    "gauss_grid_shape": reg_inputs["gauss_grid_shape"],
+                    "projected_scale_factor": None,
+                    "projection_model": "fisheye624",
+                    "projection_camera_params": src_render_cam_params[b : b + 1],
+                    "apply_splat": True,
+                    "loss_scale": 1.0,
+                }
+            ]
+            tgt_extra_loss_terms = []
+
+            vis_payload = None
+            if enable_vis:
+                src_unik3d_depth = out["unik3d_distance"][b : b + 1, 0:1].detach() if torch.is_tensor(out.get("unik3d_distance", None)) else None
+                tgt_unik3d_depth = None
+                if (
+                    tgt_unik3d_depth is None
+                    and self.enable_tgt_unik3d_vis
+                ):
+                    try:
+                        with torch.no_grad():
+                            from unisharp.utils.unik3d_adapter import forward_unik3d_fisheye624
+
+                            unik_tgt = forward_unik3d_fisheye624(
+                                self._base_model().feature_extractor.unik3d,
+                                rgb_u8=tgt_u8[b : b + 1],
+                                camera_params=tgt_render_cam_params[b : b + 1],
+                                normalize=True,
+                                validity_mask=tgt_valid_mask[b : b + 1],
+                            )
+                            dist_tgt = unik_tgt.get("distance", None) if isinstance(unik_tgt, dict) else None
+                            if torch.is_tensor(dist_tgt):
+                                tgt_unik3d_depth = dist_tgt[:, 0:1].detach()
+                    except Exception:
+                        tgt_unik3d_depth = None
+                vis_payload = {
+                    "src_gt": (src_u8[b : b + 1].float() / 255.0).detach(),
+                    "src_pred": src_render["color"].clamp(0, 1).detach(),
+                    "src_alpha": src_render["alpha"].detach(),
+                    "src_gt_depth": src_depth_gt[b : b + 1].detach(),
+                    "src_pred_depth": src_render["depth_distance"].detach(),
+                    "src_unik3d_depth": src_unik3d_depth,
+                    "src_metric_mask": src_mask.detach(),
+                    "tgt_gt": (tgt_u8[b : b + 1].float() / 255.0).detach(),
+                    "tgt_pred": tgt_render["color"].clamp(0, 1).detach(),
+                    "tgt_alpha": tgt_render["alpha"].detach(),
+                    "tgt_gt_depth": tgt_depth_gt[b : b + 1].detach(),
+                    "tgt_pred_depth": tgt_depth_pred.detach(),
+                    "tgt_unik3d_depth": tgt_unik3d_depth,
+                    "dataset_name": str(dataset_name),
+                    "scene": str(self._first_item(getattr(batch, "scene", None), "unknown")),
+                    "src_idx": int(self._first_item(getattr(batch, "src_idx", None), -1)),
+                    "tgt_idx": int(self._first_item(getattr(batch, "tgt_idx", None), -1)),
+                    "src_pose_w2c": src_w2c[b : b + 1].detach(),
+                    "tgt_pose_w2c": tgt_w2c[b : b + 1].detach(),
+                    "tgt_metric_mask_raw": tgt_mask_raw.detach(),
+                    "tgt_metric_mask": tgt_mask.detach(),
+                }
+
+
+            return {
+                "src_loss_terms": src_loss_terms,
+                "src_extra_loss_terms": src_extra_loss_terms,
+                "tgt_pred_rgb_linear": tgt_render["color"],
+                "tgt_pred_alpha": tgt_render["alpha"],
+                "tgt_pred_depth_m": tgt_depth_pred,
+                "tgt_gt_rgb_u8": tgt_u8[b : b + 1],
+                "tgt_gt_depth_m": tgt_depth_gt[b : b + 1],
+                "tgt_mask": tgt_mask,
+                "tgt_depth_mask": tgt_depth_loss_mask,
+                "tgt_apply_depth": True,
+                "tgt_apply_grad_img": True,
+                "tgt_apply_splat": False,
+                "tgt_grad_img_circular_h": False,
+                "tgt_apply_percep": bool(float(self.loss_fn.w.lambda_percep) > 0.0),
+                "tgt_extra_loss_terms": tgt_extra_loss_terms,
+                "aux_losses": self._aux_ray_losses(
+                    pred_rays=(
+                        out.get("unik3d_rays", None)[b : b + 1]
+                        if torch.is_tensor(out.get("unik3d_rays", None))
+                        else None
+                    ),
+                    gt_rays=(
+                        aux_ray_target_all[b : b + 1]
+                        if torch.is_tensor(aux_ray_target_all)
+                        else None
+                    ),
+                    mask=src_render_valid_mask[b : b + 1],
+                    pred_distance=(
+                        out["unik3d_distance"][b : b + 1, 0:1]
+                        if torch.is_tensor(out.get("unik3d_distance", None))
+                        else None
+                    ),
+                    pred_distance2=None,
+                    gt_distance=src_depth_gt[b : b + 1],
+                    depth_mask=src_valid_mask[b : b + 1],
+                ),
+                "gaussian_scales": reg_inputs["gaussian_scales"],
+                "gaussian_quaternions": reg_inputs["gaussian_quaternions"],
+                "gaussian_angular_cell": reg_inputs["gaussian_angular_cell"],
+                "delta_xy": reg_inputs["delta_xy_eff"],
+                "delta_rho": reg_inputs["delta_rho_raw"],
+                "delta_grid": reg_inputs["delta_grid"],
+                "gaussian_mean_vectors": reg_inputs["gaussian_mean_vectors"],
+                "gaussian_base_mean_vectors": reg_inputs["gaussian_base_mean_vectors"],
+                "gaussian_opacities": reg_inputs["gaussian_opacities"],
+                "gauss_grid_shape": reg_inputs["gauss_grid_shape"],
+                "projected_scale_factor": reg_inputs["projected_scale_factor"],
+                "projection_model": "fisheye624",
+                "projection_camera_params": src_render_cam_params[b : b + 1],
+                "vis_payload": vis_payload,
+            }
+
+        return _ModeStrategy(
+            batch_size=int(src_u8.shape[0]),
+            gaussians=gaussians,
+            make_world_gaussians=make_world_gaussians,
+            make_sample=make_sample,
+            collect_all_vis=bool(getattr(batch, "collect_all_vis", False)),
+        )
+
+    def _build_fisheye_strategy(
+        self,
+        batch: Any,
+        step: int,
+        need_vis: bool = False,
+        dataset_name: str = "fisheye",
+    ) -> _ModeStrategy:
+        camera_model = str(getattr(batch, "camera_model", "fisheye624")).lower()
+        if camera_model != "fisheye624":
+            raise ValueError(
+                f"Unsupported fisheye camera_model={camera_model!r}; expected 'fisheye624'."
+            )
+        return self._build_fisheye624_strategy(
+            batch,
+            step,
+            need_vis=need_vis,
+            dataset_name=dataset_name,
+        )
+
+    def _build_spherical_strategy(
+        self,
+        batch: Any,
+        step: int,
+        need_vis: bool = False,
+        dataset_name: str = "hm3d",
+    ) -> _ModeStrategy:
+        src_erp_u8 = batch.src_erp_rgb_u8.to(self.device, non_blocking=True)
+        tgt_erp_u8 = batch.tgt_erp_rgb_u8.to(self.device, non_blocking=True)
+        src_erp_depth = self._clamp_distance_for_supervision(
+            batch.src_erp_depth_m.to(self.device, non_blocking=True)
+        )
+        tgt_erp_depth = self._clamp_distance_for_supervision(
+            batch.tgt_erp_depth_m.to(self.device, non_blocking=True)
+        )
+        src_cdep = self._sanitize_positive_depth(
+            batch.src_cube_depth_m.to(self.device, non_blocking=True)
+        )
+        tgt_cdep = self._sanitize_positive_depth(
+            batch.tgt_cube_depth_m.to(self.device, non_blocking=True)
+        )
+        disable_depth_gt = bool(getattr(batch, "disable_depth_gt", False))
+        
+        src_R = batch.src_R.to(self.device, non_blocking=True)
+        src_t = batch.src_t.to(self.device, non_blocking=True)
+        tgt_R = batch.tgt_R.to(self.device, non_blocking=True)
+        tgt_t = batch.tgt_t.to(self.device, non_blocking=True)
+        
+        cur_bs = int(src_erp_u8.shape[0])
+        erp_h = int(src_erp_u8.shape[-2])
+        erp_w = int(src_erp_u8.shape[-1])
+        cube_face_w = int(batch.src_cube_depth_m.shape[2]) if torch.is_tensor(batch.src_cube_depth_m) else max(1, erp_h // 2)
+        
+        use_flip_yz = str(dataset_name).lower() not in {"sim", "smx_sim_fisheye"}
+        pose_convs_per_sample = ["c2w"] * cur_bs
+        flip_yz_per_sample = [bool(use_flip_yz)] * cur_bs
+        
+        extr_src_base = torch.stack(
+            [build_extrinsics_w2c(src_R[i], src_t[i], pose_convs_per_sample[i]) for i in range(cur_bs)],
+            dim=0
+        )
+        extr_tgt_base = torch.stack(
+            [build_extrinsics_w2c(tgt_R[i], tgt_t[i], pose_convs_per_sample[i]) for i in range(cur_bs)],
+            dim=0
+        )
+        
+        with torch.autocast("cuda", enabled=False):
+            c2w_src = torch.linalg.inv(extr_src_base.to(torch.float32))
+            c2w_tgt = torch.linalg.inv(extr_tgt_base.to(torch.float32))
+            
+            flip_mask = torch.tensor(flip_yz_per_sample, device=c2w_src.device, dtype=torch.bool)
+            negate_relative_z = False
+            if bool(flip_mask.any().item()):
+                flip_mode = os.environ.get("PANO_POSE_FLIP_CONVENTION", "flip_yz_negate_rel_z").strip().lower()
+                negate_relative_z = flip_mode in {
+                    "flip_yz_negate_rel_z",
+                    "flip_yz_invert_z_translation",
+                    "flip_yz_neg_z",
+                }
+                if flip_mode in {"flip_y_only", "y", "y_only"}:
+                    diag = [1.0, -1.0, 1.0, 1.0]
+                elif flip_mode in {"none", "identity", "no_flip"}:
+                    diag = [1.0, 1.0, 1.0, 1.0]
+                else:
+                    diag = [1.0, -1.0, -1.0, 1.0]
+                D = torch.diag(torch.tensor(diag, device=c2w_src.device, dtype=torch.float32))
+                c2w_src = c2w_src.clone()
+                c2w_tgt = c2w_tgt.clone()
+                c2w_src[flip_mask] = c2w_src[flip_mask] @ D
+                c2w_tgt[flip_mask] = c2w_tgt[flip_mask] @ D
+            
+            ref_inv = torch.linalg.inv(c2w_src.to(torch.float32))
+            c2w_src = ref_inv @ c2w_src
+            c2w_tgt = ref_inv @ c2w_tgt
+            if negate_relative_z:
+                c2w_tgt = c2w_tgt.clone()
+                c2w_tgt[flip_mask, 2, 3] *= -1.0
+            
+            extr_src = torch.linalg.inv(c2w_src).to(dtype=extr_src_base.dtype)
+            extr_tgt = torch.linalg.inv(c2w_tgt).to(dtype=extr_tgt_base.dtype)
+        
+        src_erp = (src_erp_u8.float() / 255.0).clamp(0, 1)
+        distance_init_cap_m = self._distance_init_cap_for_dataset(dataset_name)
+        
+        out = self.model(
+            image=src_erp,
+            image_u8=src_erp_u8,
+            camera_intrinsics=None,
+            camera_model="spherical",
+            depth_gt=None if disable_depth_gt else src_erp_depth,
+            distance_init_cap_m=distance_init_cap_m,
+            return_aux=True,
+        )
+        gaussians = out["gaussians"]
+        aux_ray_target_all = out.get("unik3d_gt_rays", None)
+
+        def make_world_gaussians(b: int, g_b: Any) -> Any:
+            return transform_gaussians_to_world(g_b, extr_src[b])
+
+        def make_sample(b: int, g_world: Any, enable_vis: bool) -> dict[str, Any]:
+            src_rgb, src_depth, src_alpha = self._render_cubemap(g_world, extr_src[b], face_w=cube_face_w)
+            tgt_rgb, tgt_depth, tgt_alpha = self._render_cubemap(g_world, extr_tgt[b], face_w=cube_face_w)
+
+            src_erp_pred = self._cube_to_erp(src_rgb, equ_h=erp_h, equ_w=erp_w, face_w=cube_face_w)
+            tgt_erp_pred = self._cube_to_erp(tgt_rgb, equ_h=erp_h, equ_w=erp_w, face_w=cube_face_w)
+            src_erp_alpha = self._cube_to_erp(src_alpha, equ_h=erp_h, equ_w=erp_w, face_w=cube_face_w)
+            tgt_erp_alpha = self._cube_to_erp(tgt_alpha, equ_h=erp_h, equ_w=erp_w, face_w=cube_face_w)
+            src_depth_dist = self._clamp_distance_for_supervision(
+                self._cubemap_z_depth_to_distance(src_depth),
+                clamp_max=False,
+            )
+            tgt_depth_dist = self._clamp_distance_for_supervision(
+                self._cubemap_z_depth_to_distance(tgt_depth),
+                clamp_max=False,
+            )
+            src_erp_depth_render = self._cube_to_erp(
+                src_depth_dist, equ_h=erp_h, equ_w=erp_w, face_w=cube_face_w
+            ).clamp(min=1e-4)
+
+            src_erp_depth_pred = self._clamp_distance_for_supervision(
+                out["distance_layers"][b : b + 1, 0:1],
+                clamp_max=False,
+            )
+            src_erp_depth2_pred = (
+                self._clamp_distance_for_supervision(out["distance_layers"][b : b + 1, 1:2], clamp_max=False)
+                if out["distance_layers"] is not None and out["distance_layers"].shape[1] > 1
+                else None
+            )
+            tgt_erp_depth_pred = self._cube_to_erp(
+                tgt_depth_dist, equ_h=erp_h, equ_w=erp_w, face_w=cube_face_w
+            ).clamp(min=1e-4)
+            tgt_depth_loss_mask = self._rendered_depth_valid_for_inv_loss(tgt_erp_depth_pred, tgt_erp_alpha)
+            depth_novel = self._pick_depth_for_cubemap_frustum_mask(
+                gt_depth_cube=None if disable_depth_gt else (tgt_cdep[b : b + 1][0] if torch.is_tensor(tgt_cdep) else None),
+                pred_depth_cube=tgt_depth_dist,
+                face_w=cube_face_w,
+            )
+            source_depth_for_visibility = self._pick_depth_for_cubemap_frustum_mask(
+                gt_depth_cube=None if disable_depth_gt else (src_cdep[b : b + 1][0] if torch.is_tensor(src_cdep) else None),
+                pred_depth_cube=src_depth_dist,
+                face_w=cube_face_w,
+            )
+            mask_bool = view_frustum_mask_cubemap_union(
+                depth_novel=depth_novel,
+                extr_novel_w2c=extr_tgt[b],
+                extr_source_w2c=extr_src[b],
+                face_w=int(cube_face_w),
+                source_depth=source_depth_for_visibility,
+            )
+            mask_erp = self._cube_to_erp(
+                mask_bool[:, None].to(torch.float32), equ_h=erp_h, equ_w=erp_w, face_w=cube_face_w
+            )
+
+            gt_src_erp_u8 = src_erp_u8[b : b + 1]
+            gt_tgt_erp_u8 = tgt_erp_u8[b : b + 1]
+            gt_src_erp_depth = src_erp_depth[b : b + 1]
+            gt_tgt_erp_depth = tgt_erp_depth[b : b + 1]
+            gt_src_cube_u8 = batch.src_cube_rgb_u8[b].to(self.device, non_blocking=True).permute(0, 3, 1, 2).contiguous()
+            gt_tgt_cube_u8 = batch.tgt_cube_rgb_u8[b].to(self.device, non_blocking=True).permute(0, 3, 1, 2).contiguous()
+
+            src_valid = torch.ones_like(gt_src_erp_depth) if disable_depth_gt else (gt_src_erp_depth > 0.0).to(dtype=torch.float32)
+            tgt_valid = torch.ones_like(gt_tgt_erp_depth) if disable_depth_gt else (gt_tgt_erp_depth > 0.0).to(dtype=torch.float32)
+            src_mask = torch.ones_like(src_valid)
+            tgt_mask = (mask_erp.to(dtype=torch.float32) * tgt_valid).clamp(0.0, 1.0)
+            tgt_mask_raw = tgt_mask
+            tgt_mask = self._erode_supervision_mask(
+                tgt_mask,
+                self.target_mask_erode_px,
+                circular_h=True,
+            )
+            src_cube_mask = torch.ones_like(src_alpha)
+            if str(dataset_name).lower() == "hm3d" and (not disable_depth_gt) and torch.is_tensor(src_cdep):
+                src_cube_valid = (src_cdep[b : b + 1][0, ..., 0] > 0.0).to(dtype=src_alpha.dtype).unsqueeze(1)
+                if tuple(src_cube_valid.shape[-2:]) != tuple(src_alpha.shape[-2:]):
+                    src_cube_valid = F.interpolate(
+                        src_cube_valid,
+                        size=src_alpha.shape[-2:],
+                        mode="nearest",
+                    )
+                src_cube_mask = src_cube_valid.to(device=src_alpha.device, dtype=src_alpha.dtype).clamp(0.0, 1.0)
+            tgt_cube_valid = (depth_novel[..., 0] > 0.0).to(dtype=torch.float32).unsqueeze(1)
+            tgt_cube_mask = (mask_bool[:, None].to(dtype=torch.float32) * tgt_cube_valid).clamp(0.0, 1.0)
+            tgt_cube_mask = self._erode_supervision_mask(
+                tgt_cube_mask,
+                self.target_mask_erode_px,
+                circular_h=False,
+            )
+
+            src_cube_depth_zeros = torch.zeros_like(src_alpha)
+            tgt_cube_depth_zeros = torch.zeros_like(tgt_alpha)
+            src_erp_rgb_zeros = torch.zeros_like(src_erp_pred)
+            tgt_erp_rgb_zeros = torch.zeros_like(tgt_erp_pred)
+            src_erp_u8_zeros = torch.zeros_like(gt_src_erp_u8)
+            tgt_erp_u8_zeros = torch.zeros_like(gt_tgt_erp_u8)
+
+            erp_proj_scale = 0.5 * (
+                float(erp_w) / (2.0 * 3.141592653589793)
+                + float(erp_h) / 3.141592653589793
+            )
+            reg_inputs = self._collect_regularization_inputs(
+                out=out,
+                gaussians=gaussians,
+                b=b,
+                projected_scale_factor=erp_proj_scale,
+            )
+
+            vis_payload = None
+            if enable_vis:
+                src_unik3d_depth = None
+                tgt_unik3d_depth = None
+                raw_dist = out.get("unik3d_distance", None)
+                if torch.is_tensor(raw_dist):
+                    src_unik3d_depth = raw_dist[b : b + 1, 0:1].detach()
+
+                vis_payload = {
+                    "src_gt": (gt_src_erp_u8.float() / 255.0).detach(),
+                    "src_pred": src_erp_pred.clamp(0, 1).detach(),
+                    "src_alpha": src_erp_alpha.detach(),
+                    "src_gt_depth": None if disable_depth_gt else gt_src_erp_depth.detach(),
+                    "src_pred_depth": src_erp_depth_render.detach(),
+                    "src_unik3d_depth": src_unik3d_depth,
+                    "tgt_gt": (gt_tgt_erp_u8.float() / 255.0).detach(),
+                    "tgt_pred": tgt_erp_pred.clamp(0, 1).detach(),
+                    "tgt_alpha": tgt_erp_alpha.detach(),
+                    "tgt_gt_depth": None if disable_depth_gt else gt_tgt_erp_depth.detach(),
+                    "tgt_pred_depth": tgt_erp_depth_pred.detach(),
+                    "tgt_unik3d_depth": tgt_unik3d_depth,
+                    "dataset_name": str(dataset_name),
+                    "scene": str(self._item_at(getattr(batch, "scene", None), b, "unknown")),
+                    "src_idx": int(self._item_at(getattr(batch, "src_idx", None), b, -1)),
+                    "tgt_idx": int(self._item_at(getattr(batch, "tgt_idx", None), b, -1)),
+                    "src_pose_w2c": extr_src[b : b + 1].detach(),
+                    "tgt_pose_w2c": extr_tgt[b : b + 1].detach(),
+                    "src_cube_gt_u8": (
+                        batch.src_cube_rgb_u8[b].detach()
+                        if hasattr(batch, "src_cube_rgb_u8") and torch.is_tensor(batch.src_cube_rgb_u8)
+                        else None
+                    ),
+                    "tgt_cube_gt_u8": (
+                        batch.tgt_cube_rgb_u8[b].detach()
+                        if hasattr(batch, "tgt_cube_rgb_u8") and torch.is_tensor(batch.tgt_cube_rgb_u8)
+                        else None
+                    ),
+                    "src_cube_pred_linear": src_rgb.detach(),
+                    "tgt_cube_pred_linear": tgt_rgb.detach(),
+                    "src_cube_alpha": src_alpha.detach(),
+                    "tgt_cube_alpha": tgt_alpha.detach(),
+                    "tgt_metric_mask_raw": tgt_mask_raw.detach(),
+                    "tgt_metric_mask": tgt_mask.detach(),
+                }
+
+            tgt_loss_terms = [
+                {
+                    "pred_rgb_linear": tgt_rgb,
+                    "pred_alpha": tgt_alpha,
+                    "pred_depth_m": tgt_cube_depth_zeros,
+                    "pred_depth2_m": None,
+                    "gt_rgb_u8": gt_tgt_cube_u8,
+                    "gt_depth_m": tgt_cube_depth_zeros,
+                    "mask": tgt_cube_mask,
+                    "apply_color": True,
+                    "apply_alpha": True,
+                    "apply_depth": False,
+                    "apply_percep": bool(float(self.loss_fn.w.lambda_percep) > 0.0),
+                    "apply_tv": False,
+                    "apply_grad": False,
+                    "apply_grad_img": False,
+                    "grad_img_circular_h": False,
+                    "gaussian_scales": None,
+                    "gaussian_quaternions": None,
+                    "gaussian_angular_cell": None,
+                    "delta_xy": None,
+                    "gaussian_mean_vectors": None,
+                    "gaussian_opacities": None,
+                    "gauss_grid_shape": None,
+                    "projected_scale_factor": None,
+                },
+                {
+                    "pred_rgb_linear": tgt_erp_rgb_zeros,
+                    "pred_alpha": torch.zeros_like(tgt_erp_depth_pred),
+                    "pred_depth_m": tgt_erp_depth_pred,
+                    "pred_depth2_m": None,
+                    "gt_rgb_u8": tgt_erp_u8_zeros,
+                    "gt_depth_m": gt_tgt_erp_depth,
+                    "mask": tgt_mask,
+                    "depth_mask": tgt_depth_loss_mask,
+                    "apply_color": False,
+                    "apply_alpha": False,
+                    "apply_depth": not disable_depth_gt,
+                    "apply_percep": False,
+                    "apply_tv": False,
+                    "apply_grad": False,
+                    "apply_grad_img": not disable_depth_gt,
+                    "grad_img_circular_h": True,
+                    "gaussian_scales": None,
+                    "gaussian_quaternions": None,
+                    "gaussian_angular_cell": None,
+                    "delta_xy": None,
+                    "gaussian_mean_vectors": None,
+                    "gaussian_opacities": None,
+                    "gauss_grid_shape": None,
+                    "projected_scale_factor": reg_inputs["projected_scale_factor"],
+                },
+            ]
+            return {
+                "src_loss_terms": [
+                    {
+                        "pred_rgb_linear": src_rgb,
+                        "pred_alpha": src_alpha,
+                        "pred_depth_m": src_cube_depth_zeros,
+                        "pred_depth2_m": None,
+                        "gt_rgb_u8": gt_src_cube_u8,
+                        "gt_depth_m": src_cube_depth_zeros,
+                        "mask": src_cube_mask,
+                        "apply_color": True,
+                        "apply_alpha": True,
+                        "apply_depth": False,
+                        "apply_percep": False,
+                        "apply_tv": False,
+                        "apply_grad": False,
+                        "apply_grad_img": False,
+                        "grad_img_circular_h": False,
+                        "gaussian_scales": None,
+                        "gaussian_quaternions": None,
+                        "gaussian_angular_cell": None,
+                        "delta_xy": None,
+                        "gaussian_mean_vectors": None,
+                        "gaussian_opacities": None,
+                        "gauss_grid_shape": None,
+                        "projected_scale_factor": None,
+                    },
+                    {
+                        "pred_rgb_linear": src_erp_rgb_zeros,
+                        "pred_alpha": torch.zeros_like(src_erp_depth_pred),
+                        "pred_depth_m": src_erp_depth_pred,
+                        "pred_depth2_m": src_erp_depth2_pred,
+                        "gt_rgb_u8": src_erp_u8_zeros,
+                        "gt_depth_m": gt_src_erp_depth,
+                        "mask": src_mask,
+                        "apply_color": False,
+                        "apply_alpha": False,
+                        "apply_depth": False,
+                        "apply_percep": False,
+                        "apply_tv": True,
+                        "apply_grad": False,
+                        "apply_grad_img": not disable_depth_gt,
+                        "grad_img_circular_h": True,
+                        "gaussian_scales": reg_inputs["gaussian_scales"],
+                        "gaussian_quaternions": reg_inputs["gaussian_quaternions"],
+                        "gaussian_angular_cell": reg_inputs["gaussian_angular_cell"],
+                        "delta_xy": reg_inputs["delta_xy_eff"],
+                        "delta_rho": reg_inputs["delta_rho_raw"],
+                        "delta_grid": reg_inputs["delta_grid"],
+                        "gaussian_mean_vectors": reg_inputs["gaussian_mean_vectors"],
+                        "gaussian_base_mean_vectors": reg_inputs["gaussian_base_mean_vectors"],
+                        "gaussian_opacities": reg_inputs["gaussian_opacities"],
+                        "gauss_grid_shape": reg_inputs["gauss_grid_shape"],
+                        "projected_scale_factor": reg_inputs["projected_scale_factor"],
+                        "projection_model": "erp",
+                    },
+                ],
+                "tgt_loss_terms": tgt_loss_terms,
+                "gaussian_scales": reg_inputs["gaussian_scales"],
+                "gaussian_quaternions": reg_inputs["gaussian_quaternions"],
+                "gaussian_angular_cell": reg_inputs["gaussian_angular_cell"],
+                "delta_xy": reg_inputs["delta_xy_eff"],
+                "delta_rho": reg_inputs["delta_rho_raw"],
+                "delta_grid": reg_inputs["delta_grid"],
+                "gaussian_mean_vectors": reg_inputs["gaussian_mean_vectors"],
+                "gaussian_base_mean_vectors": reg_inputs["gaussian_base_mean_vectors"],
+                "gaussian_opacities": reg_inputs["gaussian_opacities"],
+                "gauss_grid_shape": reg_inputs["gauss_grid_shape"],
+                "projected_scale_factor": reg_inputs["projected_scale_factor"],
+                "projection_model": "erp",
+                "aux_losses": self._aux_ray_losses(
+                    pred_rays=(
+                        out.get("unik3d_rays", None)[b : b + 1]
+                        if torch.is_tensor(out.get("unik3d_rays", None))
+                        else None
+                    ),
+                    gt_rays=(
+                        aux_ray_target_all[b : b + 1]
+                        if torch.is_tensor(aux_ray_target_all)
+                        else None
+                    ),
+                    mask=src_valid,
+                    pred_distance=(
+                        out["unik3d_distance"][b : b + 1, 0:1]
+                        if torch.is_tensor(out.get("unik3d_distance", None))
+                        else None
+                    ),
+                    pred_distance2=src_erp_depth2_pred,
+                    gt_distance=None if disable_depth_gt else gt_src_erp_depth,
+                    depth_mask=src_valid,
+                ),
+                "vis_payload": vis_payload,
+            }
+
+        return _ModeStrategy(
+            batch_size=int(cur_bs),
+            gaussians=gaussians,
+            make_world_gaussians=make_world_gaussians,
+            make_sample=make_sample,
+            collect_all_vis=bool(getattr(batch, "collect_all_vis", False)),
+        )
+    
+    def _render_cubemap(
+        self,
+        gaussians: Any,
+        extr_w2c: torch.Tensor,
+        face_w: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        device = gaussians.mean_vectors.device
+        intr = get_pinhole_intrinsics_4x4(int(face_w)).to(device=device)[None].expand(6, -1, -1)
+        extr_faces = cubemap_face_cameras(extr_w2c, device=device)
+        out = self.renderer(
+            gaussians,
+            extrinsics=extr_faces,
+            intrinsics=intr,
+            image_width=int(face_w),
+            image_height=int(face_w),
+        )
+        return out.color.contiguous(), out.depth.contiguous(), out.alpha.contiguous()
+    
+    def _cube_to_erp(self, cube: torch.Tensor, equ_h: int, equ_w: int, face_w: int) -> torch.Tensor:
+        cube = cube.permute(1, 0, 2, 3).unsqueeze(0)
+        c2e = Cube2Equirec(face_w=int(face_w), equ_h=int(equ_h), equ_w=int(equ_w)).to(device=cube.device)
+        return c2e(cube)

unisharp/datasets/dl3dv.py

@@ -0,0 +1,305 @@
+
+from __future__ import annotations
+
+from collections import defaultdict, deque
+import json
+from pathlib import Path
+import random
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from torch.utils.data import IterableDataset
+
+from unisharp.datasets.pair_sampling import (
+    project_overlap_ratio,
+    resize_k3_align_corners_false,
+    resize_rgb_u8_chw_high_quality,
+    select_targets_for_source,
+)
+from unisharp.datasets.re10k import Re10KPairSample, re10k_collate
+from unisharp import DEFAULT_MAX_DEPTH_M
+
+
+class DL3DVDataset(IterableDataset):
+    def __init__(
+        self,
+        root: Path,
+        depth_root: Path,
+        scene_specs_file: Path | None = None,
+        min_frame_gap: int = 1,
+        max_frame_gap: int = 32,
+        pair_max_translation_m: float = 0.5,
+        pair_min_overlap: float = 0.6,
+        pair_overlap_sample_h: int = 32,
+        pair_overlap_sample_w: int = 56,
+        output_h: int | None = None,
+        output_w: int | None = None,
+        shuffle_scene: bool = True,
+        shuffle_frame: bool = False,
+        ddp_rank: int = 0,
+        ddp_world_size: int = 1,
+        batch_size_hint: int = 1,
+        depth_max_m: float = DEFAULT_MAX_DEPTH_M,
+        seed: int = 0,
+        verify_manifest_paths: bool = False,
+    ) -> None:
+        super().__init__()
+        self.root = Path(root)
+        self.depth_root = Path(depth_root)
+        self.min_frame_gap = int(min_frame_gap)
+        self.max_frame_gap = int(max_frame_gap)
+        self.pair_max_translation_m = float(pair_max_translation_m)
+        self.pair_min_overlap = float(pair_min_overlap)
+        self.pair_overlap_sample_h = int(pair_overlap_sample_h)
+        self.pair_overlap_sample_w = int(pair_overlap_sample_w)
+        self.output_h = int(output_h) if output_h is not None else None
+        self.output_w = int(output_w) if output_w is not None else None
+        self.shuffle_scene = bool(shuffle_scene)
+        self.shuffle_frame = bool(shuffle_frame)
+        self.ddp_rank = int(ddp_rank)
+        self.ddp_world_size = int(ddp_world_size)
+        self.batch_size_hint = int(max(1, batch_size_hint))
+        self.depth_max_m = float(depth_max_m)
+        self.seed = int(seed)
+        self.epoch = 0
+        self.verify_manifest_paths = bool(verify_manifest_paths)
+        self.scene_specs_file = Path(scene_specs_file) if scene_specs_file is not None else None
+        self.scene_specs = self._load_scene_specs()
+        if not self.scene_specs:
+            raise RuntimeError(f"No valid DL3DV scenes found under {self.root}")
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = int(epoch)
+
+    def _load_scene_specs(self) -> list[tuple[str, Path, Path]]:
+        if self.scene_specs_file is None:
+            return self._scan_scenes()
+        if not self.scene_specs_file.exists():
+            raise FileNotFoundError(self.scene_specs_file)
+        out: list[tuple[str, Path, Path]] = []
+        for raw in self.scene_specs_file.read_text(encoding="utf-8").splitlines():
+            line = raw.strip()
+            if not line:
+                continue
+            parts = line.split("|")
+            if len(parts) != 3:
+                continue
+            scene_name, scene_dir_raw, depth_dir_raw = parts
+            scene_dir = Path(scene_dir_raw)
+            depth_dir = Path(depth_dir_raw)
+            if (not self.verify_manifest_paths) or (scene_dir.exists() and depth_dir.exists()):
+                out.append((scene_name, scene_dir, depth_dir))
+        return out
+
+    def _scan_scenes(self) -> list[tuple[str, Path, Path]]:
+        out: list[tuple[str, Path, Path]] = []
+        for bucket_dir in sorted([p for p in self.root.iterdir() if p.is_dir()]):
+            for scene_stub in sorted([p for p in bucket_dir.iterdir() if p.is_dir()]):
+                inner_dirs = [p for p in scene_stub.iterdir() if p.is_dir()]
+                scene_dir = inner_dirs[0] if inner_dirs else scene_stub
+                transforms_path = scene_dir / "transforms.json"
+                image_dir = scene_dir / "images_4"
+                depth_dir = self.depth_root / bucket_dir.name / scene_stub.name / "exports" / "mini_npz" / "per_image"
+                if transforms_path.exists() and image_dir.exists() and depth_dir.exists():
+                    scene_name = f"{bucket_dir.name}/{scene_stub.name}"
+                    out.append((scene_name, scene_dir, depth_dir))
+        return out
+
+    @staticmethod
+    def _load_rgb_u8(path: Path) -> torch.Tensor:
+        arr = np.asarray(Image.open(path).convert("RGB"), dtype=np.uint8).copy()
+        return torch.from_numpy(arr).permute(2, 0, 1).contiguous()
+
+    def _load_depth_m(self, path: Path) -> torch.Tensor:
+        payload = np.load(path)
+        depth = payload["depth"].astype(np.float32)
+        depth[~np.isfinite(depth)] = 0.0
+        depth = np.clip(depth, a_min=0.0, a_max=self.depth_max_m)
+        return torch.from_numpy(depth).unsqueeze(0)
+
+    @staticmethod
+    def _resize_depth_to_image(depth: torch.Tensor, image_hw: tuple[int, int]) -> torch.Tensor:
+        target_h, target_w = int(image_hw[0]), int(image_hw[1])
+        if depth.shape[-2:] == (target_h, target_w):
+            return depth
+        return F.interpolate(
+            depth.unsqueeze(0),
+            size=(target_h, target_w),
+            mode="nearest",
+        ).squeeze(0)
+
+    @staticmethod
+    def _frame_id_from_name(name: str) -> int:
+        stem = Path(name).stem
+        return int(stem.split("_")[-1])
+
+    def _load_scene(
+        self,
+        scene_name: str,
+        scene_dir: Path,
+        depth_dir: Path,
+    ) -> tuple[list[int], dict[int, Path], dict[int, Path], dict[int, torch.Tensor], dict[int, torch.Tensor], torch.Tensor]:
+        meta = json.loads((scene_dir / "transforms.json").read_text())
+        orig_w = int(meta["w"])
+        orig_h = int(meta["h"])
+        k = torch.eye(3, dtype=torch.float32)
+        k[0, 0] = float(meta["fl_x"])
+        k[1, 1] = float(meta["fl_y"])
+        k[0, 2] = float(meta["cx"])
+        k[1, 2] = float(meta["cy"])
+
+        image_dir = scene_dir / "images_4"
+        image_paths = {self._frame_id_from_name(p.name): p for p in image_dir.glob("*.png")}
+        depth_paths = {self._frame_id_from_name(p.name): p for p in depth_dir.glob("*.npz")}
+        w2c_map: dict[int, torch.Tensor] = {}
+        intr_map: dict[int, torch.Tensor] = {}
+        valid_ids: list[int] = []
+
+        example_img = None
+        for frame in meta.get("frames", []):
+            rel_path = str(frame.get("file_path", ""))
+            frame_name = Path(rel_path).name
+            frame_id = self._frame_id_from_name(frame_name)
+            if frame_id not in image_paths or frame_id not in depth_paths:
+                continue
+            c2w = torch.tensor(frame["transform_matrix"], dtype=torch.float32)
+            c2w[:3, 1:3] *= -1.0
+            if example_img is None:
+                example_img = self._load_rgb_u8(image_paths[frame_id])
+            cur_h, cur_w = int(example_img.shape[1]), int(example_img.shape[2])
+            k_cur = k.clone()
+            if cur_h != orig_h or cur_w != orig_w:
+                sx = float(cur_w) / float(orig_w)
+                sy = float(cur_h) / float(orig_h)
+                k_cur = resize_k3_align_corners_false(k_cur, sx=sx, sy=sy)
+            w2c_map[frame_id] = torch.linalg.inv(c2w)
+            intr_map[frame_id] = k_cur
+            valid_ids.append(frame_id)
+        valid_ids = sorted(valid_ids)
+        return valid_ids, image_paths, depth_paths, w2c_map, intr_map, k
+
+    def __iter__(self):
+        scenes = list(self.scene_specs)
+        order_rng = random.Random(self.seed + self.epoch)
+        if self.shuffle_scene:
+            order_rng.shuffle(scenes)
+        pending_by_hw: dict[tuple[int, int], deque[Re10KPairSample]] = defaultdict(deque)
+        worker_info = torch.utils.data.get_worker_info()
+        num_workers = worker_info.num_workers if worker_info is not None else 1
+        worker_id = worker_info.id if worker_info is not None else 0
+        total_shards = max(1, self.ddp_world_size * num_workers)
+        shard_id = self.ddp_rank * num_workers + worker_id
+        src_unit_index = 0
+
+        for scene_order_idx, (scene_name, scene_dir, depth_dir) in enumerate(scenes):
+            try:
+                valid_ids, image_paths, depth_paths, w2c_map, intr_map, _ = self._load_scene(scene_name, scene_dir, depth_dir)
+            except Exception:
+                continue
+            if len(valid_ids) < 2:
+                continue
+            src_order = list(valid_ids)
+            scene_rng = random.Random(self.seed + self.epoch * 1000003 + scene_order_idx)
+            if self.shuffle_frame:
+                scene_rng.shuffle(src_order)
+            centers = torch.stack([torch.linalg.inv(w2c_map[i])[:3, 3] for i in valid_ids], dim=0)
+            frame_to_pos = {fid: pos for pos, fid in enumerate(valid_ids)}
+
+            def overlap_avg(src_pos: int, tgt_pos: int) -> float:
+                src_fid = int(valid_ids[src_pos])
+                tgt_fid = int(valid_ids[tgt_pos])
+                src_img_path = image_paths[src_fid]
+                with Image.open(src_img_path) as img:
+                    w = int(img.size[0])
+                    h = int(img.size[1])
+                return float(
+                    0.5
+                    * (
+                        project_overlap_ratio(
+                            src_w2c=w2c_map[src_fid],
+                            tgt_w2c=w2c_map[tgt_fid],
+                            src_k=intr_map[src_fid],
+                            tgt_k=intr_map[tgt_fid],
+                            h=h,
+                            w=w,
+                            sample_h=self.pair_overlap_sample_h,
+                            sample_w=self.pair_overlap_sample_w,
+                        )
+                        + project_overlap_ratio(
+                            src_w2c=w2c_map[tgt_fid],
+                            tgt_w2c=w2c_map[src_fid],
+                            src_k=intr_map[tgt_fid],
+                            tgt_k=intr_map[src_fid],
+                            h=h,
+                            w=w,
+                            sample_h=self.pair_overlap_sample_h,
+                            sample_w=self.pair_overlap_sample_w,
+                        )
+                    )
+                )
+
+            for src_idx in src_order:
+                if src_unit_index % total_shards != shard_id:
+                    src_unit_index += 1
+                    continue
+                src_unit_index += 1
+                src_pos = int(frame_to_pos[int(src_idx)])
+                tgt_pos_list = select_targets_for_source(
+                    src_idx=src_pos,
+                    candidate_indices=list(range(len(valid_ids))),
+                    centers=centers,
+                    min_index_gap=int(self.min_frame_gap),
+                    max_index_gap=int(self.max_frame_gap),
+                    pair_max_translation_m=float(self.pair_max_translation_m),
+                    pair_min_overlap=float(self.pair_min_overlap),
+                    overlap_score_fn=overlap_avg,
+                )
+                if not tgt_pos_list:
+                    continue
+                tgt_idx = int(valid_ids[scene_rng.choice(tgt_pos_list)])
+                try:
+                    src_img = self._load_rgb_u8(image_paths[int(src_idx)])
+                    tgt_img = self._load_rgb_u8(image_paths[int(tgt_idx)])
+                    src_depth = self._load_depth_m(depth_paths[int(src_idx)])
+                    tgt_depth = self._load_depth_m(depth_paths[int(tgt_idx)])
+                except Exception:
+                    continue
+                src_depth = self._resize_depth_to_image(src_depth, (int(src_img.shape[1]), int(src_img.shape[2])))
+                tgt_depth = self._resize_depth_to_image(tgt_depth, (int(tgt_img.shape[1]), int(tgt_img.shape[2])))
+                src_intr = intr_map[int(src_idx)].clone()
+                tgt_intr = intr_map[int(tgt_idx)].clone()
+                if self.output_h is not None and self.output_w is not None:
+                    oh, ow = int(src_img.shape[1]), int(src_img.shape[2])
+                    if oh != self.output_h or ow != self.output_w:
+                        sx = float(self.output_w) / float(ow)
+                        sy = float(self.output_h) / float(oh)
+                        src_img = resize_rgb_u8_chw_high_quality(src_img, size=(self.output_h, self.output_w))
+                        tgt_img = resize_rgb_u8_chw_high_quality(tgt_img, size=(self.output_h, self.output_w))
+                        src_depth = F.interpolate(src_depth[None], size=(self.output_h, self.output_w), mode="nearest")[0]
+                        tgt_depth = F.interpolate(tgt_depth[None], size=(self.output_h, self.output_w), mode="nearest")[0]
+                        src_intr = resize_k3_align_corners_false(src_intr, sx=sx, sy=sy)
+                        tgt_intr = resize_k3_align_corners_false(tgt_intr, sx=sx, sy=sy)
+                sample = Re10KPairSample(
+                    src_rgb_u8=src_img,
+                    tgt_rgb_u8=tgt_img,
+                    src_w2c=w2c_map[int(src_idx)],
+                    tgt_w2c=w2c_map[int(tgt_idx)],
+                    src_intrinsics=src_intr,
+                    tgt_intrinsics=tgt_intr,
+                    src_idx=int(src_idx),
+                    tgt_idx=int(tgt_idx),
+                    scene=scene_name,
+                    src_depth_m=src_depth,
+                    tgt_depth_m=tgt_depth,
+                )
+                hw_key = (int(sample.src_rgb_u8.shape[1]), int(sample.src_rgb_u8.shape[2]))
+                bucket = pending_by_hw[hw_key]
+                bucket.append(sample)
+                if self.batch_size_hint <= 1:
+                    yield bucket.popleft()
+                    continue
+                while len(bucket) >= self.batch_size_hint:
+                    packed = [bucket.popleft() for _ in range(self.batch_size_hint)]
+                    yield re10k_collate(packed)

unisharp/datasets/pair_sampling.py

@@ -0,0 +1,99 @@
+
+from __future__ import annotations
+
+from typing import Callable
+
+import torch
+import torch.nn.functional as F
+
+from unisharp.utils.pixel_convention import scale_intrinsics_align_corners_false
+
+
+def resize_k3_align_corners_false(k: torch.Tensor, *, sx: float, sy: float) -> torch.Tensor:
+    return scale_intrinsics_align_corners_false(k, sx=float(sx), sy=float(sy))
+
+
+def resize_rgb_u8_chw_high_quality(image: torch.Tensor, *, size: tuple[int, int]) -> torch.Tensor:
+    if not torch.is_tensor(image) or image.ndim != 3:
+        raise ValueError(f"Expected CHW tensor, got {tuple(image.shape) if torch.is_tensor(image) else type(image)}")
+    dst_h, dst_w = int(size[0]), int(size[1])
+    if tuple(image.shape[-2:]) == (dst_h, dst_w):
+        return image.contiguous()
+    resized = F.interpolate(
+        image.unsqueeze(0).to(torch.float32),
+        size=(dst_h, dst_w),
+        mode="bicubic",
+        align_corners=False,
+        antialias=True,
+    )
+    return resized[0].round().clamp(0, 255).to(torch.uint8).contiguous()
+
+
+def project_overlap_ratio(
+    src_w2c: torch.Tensor,
+    tgt_w2c: torch.Tensor,
+    src_k: torch.Tensor,
+    tgt_k: torch.Tensor,
+    h: int,
+    w: int,
+    src_hw: tuple[int, int] | None = None,
+    tgt_hw: tuple[int, int] | None = None,
+    sample_h: int = 32,
+    sample_w: int = 56,
+    proxy_depth: float = 1.0,
+) -> float:
+    device = src_w2c.device
+    src_h, src_w = tuple(int(v) for v in (src_hw or (h, w)))
+    tgt_h, tgt_w = tuple(int(v) for v in (tgt_hw or (h, w)))
+    ys = torch.linspace(0, src_h - 1, steps=sample_h, device=device)
+    xs = torch.linspace(0, src_w - 1, steps=sample_w, device=device)
+    vv, uu = torch.meshgrid(ys, xs, indexing="ij")
+    u = uu.reshape(-1)
+    v = vv.reshape(-1)
+
+    fx, fy = src_k[0, 0], src_k[1, 1]
+    cx, cy = src_k[0, 2], src_k[1, 2]
+    x = (u - cx) / fx
+    y = (v - cy) / fy
+    z = torch.ones_like(x)
+    rays = torch.stack([x, y, z], dim=-1)
+    rays = rays / torch.norm(rays, dim=-1, keepdim=True).clamp(min=1e-6)
+    pts_src = rays * float(proxy_depth)
+
+    src_c2w = torch.linalg.inv(src_w2c)
+    pts_src_h = torch.cat([pts_src, torch.ones_like(pts_src[:, :1])], dim=-1)
+    pts_w = (src_c2w @ pts_src_h.T).T
+    pts_tgt = (tgt_w2c @ pts_w.T).T
+    xt, yt, zt = pts_tgt[:, 0], pts_tgt[:, 1], pts_tgt[:, 2].clamp(min=1e-6)
+    ut = tgt_k[0, 0] * (xt / zt) + tgt_k[0, 2]
+    vt = tgt_k[1, 1] * (yt / zt) + tgt_k[1, 2]
+    inside = (zt > 0.0) & (ut >= 0.0) & (ut <= float(tgt_w - 1)) & (vt >= 0.0) & (vt <= float(tgt_h - 1))
+    return float(inside.float().mean().item())
+
+
+def select_targets_for_source(
+    *,
+    src_idx: int,
+    candidate_indices: list[int],
+    centers: torch.Tensor,
+    min_index_gap: int,
+    max_index_gap: int,
+    pair_max_translation_m: float,
+    pair_min_overlap: float,
+    overlap_score_fn: Callable[[int, int], float],
+) -> list[int]:
+    src_c = centers[int(src_idx)]
+    tgt_cands: list[int] = []
+    for j in candidate_indices:
+        j = int(j)
+        if j == int(src_idx):
+            continue
+        gap = abs(int(j) - int(src_idx))
+        if gap < int(min_index_gap) or gap > int(max_index_gap):
+            continue
+        trans = float(torch.norm(centers[j] - src_c, p=2).item())
+        if trans > float(pair_max_translation_m):
+            continue
+        if float(overlap_score_fn(int(src_idx), j)) >= float(pair_min_overlap):
+            tgt_cands.append(j)
+    return tgt_cands

unisharp/datasets/panogs.py

@@ -0,0 +1,555 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+from typing import cast
+import tarfile
+
+import numpy as np
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+
+from unisharp import DEFAULT_MAX_DEPTH_M
+
+
+MAX_DEPTH_M = DEFAULT_MAX_DEPTH_M
+
+_PAIR_RECIPE_FIXED: tuple[str, bool] = ("c2w", True)
+
+_PAIR_CONVENTIONS: tuple[str, ...] = ("c2w",)
+
+
+def _torch_load_any(path: Path) -> object:
+    try:
+        return torch.load(path, map_location="cpu", weights_only=False)
+    except TypeError:
+        return torch.load(path, map_location="cpu")
+    except (KeyError, tarfile.ReadError, EOFError, OSError, RuntimeError) as e:
+        raise RuntimeError(f"torch.load failed (possibly incomplete/corrupted): {path}") from e
+
+
+@dataclass(frozen=True)
+class PanOGSSample:
+
+    src_erp_rgb_u8: torch.Tensor
+    tgt_erp_rgb_u8: torch.Tensor
+    src_erp_depth_m: torch.Tensor
+    tgt_erp_depth_m: torch.Tensor
+
+    src_cube_rgb_u8: torch.Tensor
+    tgt_cube_rgb_u8: torch.Tensor
+    src_cube_depth_m: torch.Tensor
+    tgt_cube_depth_m: torch.Tensor
+
+    src_R: torch.Tensor
+    src_t: torch.Tensor
+    tgt_R: torch.Tensor
+    tgt_t: torch.Tensor
+
+    src_idx: int
+    tgt_idx: int
+    scene: str
+
+
+def _load_erp_rgb_u8(path: Path) -> torch.Tensor:
+    img = np.array(Image.open(path))
+    if img.ndim != 3 or img.shape[2] != 3:
+        raise ValueError(f"Expected RGB image at {path}, got shape={img.shape}")
+    return torch.from_numpy(img.astype(np.uint8)).permute(2, 0, 1).contiguous()
+
+
+def _load_depth_png(path: Path) -> torch.Tensor:
+    dep = np.array(Image.open(path))
+    return torch.from_numpy(dep)
+
+
+def _depth_to_meters(depth: torch.Tensor, max_depth_m: float = DEFAULT_MAX_DEPTH_M) -> torch.Tensor:
+    depth_f = depth.to(torch.float32)
+    maxv = float(depth_f.max().item()) if depth_f.numel() else 0.0
+    if maxv > 200.0:
+        depth_f = depth_f / 1000.0
+    depth_f[~torch.isfinite(depth_f)] = 0.0
+    return depth_f.clamp(min=0.0, max=float(max_depth_m))
+
+
+class PanOGSDataset(Dataset[PanOGSSample]):
+
+    def __init__(
+        self,
+        root: Path,
+        index_manifest_path: Path | None = None,
+        src_tgt_max_index_gap: int = 25,
+        use_cubemap_supervision: bool = True,
+        pair_sampling: bool = True,
+        pair_max_translation_m: float = 0.5,
+        pair_min_depth_overlap: float = 0.6,
+        pair_overlap_face_w: int = 64,
+        pair_overlap_margin: float = 1.05,
+        pair_max_tries: int = 48,
+        depth_max_m: float = DEFAULT_MAX_DEPTH_M,
+    ) -> None:
+        self.root = root
+        self.src_tgt_max_index_gap = int(src_tgt_max_index_gap)
+        self.use_cubemap_supervision = use_cubemap_supervision
+        self.pair_sampling = bool(pair_sampling)
+        self.pair_max_translation_m = float(pair_max_translation_m)
+        self.pair_min_depth_overlap = float(pair_min_depth_overlap)
+        self.pair_overlap_face_w = int(pair_overlap_face_w)
+        self.pair_overlap_margin = float(pair_overlap_margin)
+        self.pair_max_tries = int(pair_max_tries)
+        self.depth_max_m = float(depth_max_m)
+        self.index_manifest_path = Path(index_manifest_path) if index_manifest_path is not None else None
+
+        self._pair_valid_tgts: dict[tuple[str, int], list[int]] = {}
+        self._pair_overlap_cache: dict[tuple[str, int, int], float] = {}
+
+        if not root.exists():
+            raise FileNotFoundError(root)
+
+        self.scenes = sorted([p for p in root.iterdir() if p.is_dir()])
+        if not self.scenes:
+            raise RuntimeError(f"No scene folders found in {root}")
+
+        self._pose_cache: dict[str, tuple[np.ndarray, np.ndarray]] = {}
+        self._meta_paths: dict[str, Path] = {}
+        self._num_frames: dict[str, int] = {}
+        self._available_frames: dict[str, list[int]] = {}
+
+        if self.index_manifest_path is not None:
+            if not self.index_manifest_path.exists():
+                raise FileNotFoundError(self.index_manifest_path)
+            valid_scenes: list[Path] = []
+            for raw in self.index_manifest_path.read_text(encoding="utf-8").splitlines():
+                line = raw.strip()
+                if not line:
+                    continue
+                parts = line.split("|")
+                scene_name = parts[0].strip()
+                if not scene_name:
+                    continue
+                scene_dir = root / scene_name
+                meta_path = scene_dir / "meta.pt"
+                if not meta_path.exists():
+                    continue
+                if len(parts) >= 2:
+                    try:
+                        n_pose = int(parts[1])
+                    except ValueError:
+                        n_pose = 0
+                else:
+                    n_pose = 0
+                if n_pose <= 0:
+                    continue
+                self._meta_paths[scene_name] = meta_path
+                self._num_frames[scene_name] = n_pose
+                self._available_frames[scene_name] = list(range(n_pose))
+                valid_scenes.append(scene_dir)
+            self.scenes = valid_scenes
+
+
+        if not self._available_frames:
+
+            valid_scenes = []
+            for scene_i, scene_dir in enumerate(self.scenes):
+                meta_path = scene_dir / "meta.pt"
+                if not meta_path.exists():
+                    continue
+
+                ex = _torch_load_any(meta_path)
+                cams = ex.get("cameras", None)
+                if not isinstance(cams, torch.Tensor):
+                    raise ValueError(f"meta.pt missing 'cameras' tensor in {scene_dir}")
+                if cams.ndim != 3 or tuple(cams.shape[1:]) != (4, 4):
+                    raise ValueError(f"Bad meta.pt cameras shape {tuple(cams.shape)} in {scene_dir}")
+                n_pose = int(cams.shape[0])
+
+                frames = list(range(n_pose))
+
+                name = scene_dir.name
+                self._meta_paths[name] = meta_path
+                self._num_frames[name] = n_pose
+                self._available_frames[name] = frames
+                valid_scenes.append(scene_dir)
+
+            self.scenes = valid_scenes
+
+    def _get_pose(self, scene: str) -> tuple[np.ndarray, np.ndarray]:
+        cached = self._pose_cache.get(scene)
+        if cached is not None:
+            return cached
+
+        meta_path = self._meta_paths.get(scene)
+        if meta_path is None:
+            raise FileNotFoundError(f"meta.pt not indexed for scene={scene} under {self.root}")
+        ex = _torch_load_any(meta_path)
+        cams = ex.get("cameras", None)
+        if not isinstance(cams, torch.Tensor):
+            raise ValueError(f"meta.pt missing 'cameras' tensor for scene={scene}")
+        cams = cams.to(torch.float32)
+        if cams.ndim != 3 or tuple(cams.shape[1:]) != (4, 4):
+            raise ValueError(f"Bad meta.pt cameras shape {tuple(cams.shape)} for scene={scene}")
+        R = cams[:, :3, :3].cpu().numpy()
+        t = cams[:, :3, 3].cpu().numpy()
+        out = (R, t)
+        self._pose_cache[scene] = out
+        return out
+
+    def __len__(self) -> int:
+        return len(self._index)
+
+    def _sample_target(self, scene: str, src_idx: int) -> int:
+        frames = self._available_frames[scene]
+        if len(frames) <= 1:
+            return src_idx
+        effective_gap = self.src_tgt_max_index_gap
+        candidates = [i for i in frames if i != src_idx and abs(i - src_idx) <= effective_gap]
+        if not candidates:
+            return src_idx
+        j = int(torch.randint(low=0, high=len(candidates), size=(1,)).item())
+        return int(candidates[j])
+
+    def _candidate_targets_by_translation(self, scene: str, src_idx: int) -> list[int]:
+        frames = self._available_frames[scene]
+        if len(frames) <= 1:
+            return []
+        R_np, t_np = self._get_pose(scene)
+        if not (0 <= src_idx < len(t_np) and 0 <= src_idx < len(R_np)):
+            return []
+        th = float(self.pair_max_translation_m)
+
+        def _cam_center_from(R: np.ndarray, t: np.ndarray, conv: str) -> np.ndarray:
+            if conv in ("c2w", "w2c_t_camcenter"):
+                return t
+            if conv == "w2c":
+                return -(R.transpose(0, 2, 1) @ t[..., None])[..., 0]
+            if conv == "c2w_t_w2c":
+                return -(R @ t[..., None])[..., 0]
+            raise ValueError(conv)
+
+        def _min_dist(idxs: np.ndarray) -> np.ndarray:
+            R_sub = R_np[idxs].astype(np.float32)
+            t_sub = t_np[idxs].astype(np.float32)
+            R_src = R_np[int(src_idx) : int(src_idx) + 1].astype(np.float32)
+            t_src = t_np[int(src_idx) : int(src_idx) + 1].astype(np.float32)
+            d_min = None
+            for conv in _PAIR_CONVENTIONS:
+                C_src = _cam_center_from(R_src, t_src, conv)[0]
+                C_sub = _cam_center_from(R_sub, t_sub, conv)
+                d = np.linalg.norm(C_sub - C_src[None, :], axis=1)
+                d_min = d if (d_min is None) else np.minimum(d_min, d)
+            assert d_min is not None
+            return d_min
+
+        effective_gap = self.src_tgt_max_index_gap
+        cand0 = np.array([i for i in frames if i != src_idx and abs(i - src_idx) <= effective_gap], dtype=np.int64)
+        if cand0.size > 0:
+            d0 = _min_dist(cand0)
+            ok0 = cand0[d0 < th]
+            if ok0.size > 0:
+                return [int(x) for x in ok0.tolist()]
+        return []
+
+    def _resize_cube_depth(self, depth: torch.Tensor, face_w: int) -> torch.Tensor:
+        if depth.ndim != 4 or depth.shape[0] != 6 or depth.shape[-1] != 1:
+            raise ValueError(f"Expected cube depth shape (6,H,W,1), got {tuple(depth.shape)}")
+        H = int(depth.shape[1])
+        W = int(depth.shape[2])
+        if H == face_w and W == face_w:
+            return depth.to(dtype=torch.float32)
+        import torch.nn.functional as F
+
+        x = depth.permute(0, 3, 1, 2).to(dtype=torch.float32)
+        x = F.interpolate(x, size=(face_w, face_w), mode="bilinear", align_corners=False)
+        return x.permute(0, 2, 3, 1).contiguous()
+
+    @staticmethod
+    def _cubemap_z_depth_to_distance(depth: torch.Tensor) -> torch.Tensor:
+        if depth.ndim != 4 or depth.shape[0] != 6 or depth.shape[-1] != 1:
+            raise ValueError(f"Expected cube depth shape (6,H,W,1), got {tuple(depth.shape)}")
+        from unisharp.utils.pano import get_pinhole_intrinsics_4x4
+
+        h = int(depth.shape[1])
+        w = int(depth.shape[2])
+        if h != w:
+            raise ValueError(f"Expected square cubemap faces, got {(h, w)}")
+        depth_61hw = depth.permute(0, 3, 1, 2).to(dtype=torch.float32).contiguous()
+        intr = get_pinhole_intrinsics_4x4(w).to(device=depth_61hw.device, dtype=depth_61hw.dtype)
+        ys = torch.arange(h, device=depth_61hw.device, dtype=depth_61hw.dtype)
+        xs = torch.arange(w, device=depth_61hw.device, dtype=depth_61hw.dtype)
+        vv, uu = torch.meshgrid(ys, xs, indexing="ij")
+        x = (uu - intr[0, 2]) / intr[0, 0].clamp(min=1e-8)
+        y = (vv - intr[1, 2]) / intr[1, 1].clamp(min=1e-8)
+        ray_z = 1.0 / torch.sqrt(x * x + y * y + 1.0).clamp(min=1e-8)
+        dist = depth_61hw / ray_z.view(1, 1, h, w).clamp(min=1e-8)
+        valid = torch.isfinite(dist) & (dist > 0.0)
+        dist = torch.where(valid, dist, torch.zeros_like(dist))
+        return dist.permute(0, 2, 3, 1).contiguous()
+
+    def _pair_depth_overlap_score(
+        self,
+        *,
+        src_R: torch.Tensor,
+        src_t: torch.Tensor,
+        tgt_R: torch.Tensor,
+        tgt_t: torch.Tensor,
+        src_cube_depth_m: torch.Tensor,
+        tgt_cube_depth_m: torch.Tensor,
+    ) -> float:
+        from unisharp.utils.camera_projection import build_extrinsics_w2c, view_frustum_mask_cubemap_union  # noqa: WPS433
+
+        device = torch.device("cpu")
+        src_R = src_R.to(device=device, dtype=torch.float32)
+        src_t = src_t.to(device=device, dtype=torch.float32)
+        tgt_R = tgt_R.to(device=device, dtype=torch.float32)
+        tgt_t = tgt_t.to(device=device, dtype=torch.float32)
+
+        face_w = int(self.pair_overlap_face_w)
+        margin = float(self.pair_overlap_margin)
+        src_d = self._cubemap_z_depth_to_distance(self._resize_cube_depth(src_cube_depth_m.to(device=device), face_w=face_w))
+        tgt_d = self._cubemap_z_depth_to_distance(self._resize_cube_depth(tgt_cube_depth_m.to(device=device), face_w=face_w))
+
+        def _score_one(recipe: tuple[str, bool]) -> float:
+            pose_conv, flip_yz = recipe
+            extr_src = build_extrinsics_w2c(src_R, src_t, pose_conv)
+            extr_tgt = build_extrinsics_w2c(tgt_R, tgt_t, pose_conv)
+
+            with torch.autocast(device_type="cpu", enabled=False):
+                c2w_src = torch.linalg.inv(extr_src)
+                c2w_tgt = torch.linalg.inv(extr_tgt)
+                if bool(flip_yz):
+                    D = torch.diag(torch.tensor([1.0, -1.0, -1.0, 1.0], dtype=torch.float32, device=device))
+                    c2w_src = c2w_src @ D
+                    c2w_tgt = c2w_tgt @ D
+                ref_inv = torch.linalg.inv(c2w_src)
+                c2w_src = ref_inv @ c2w_src
+                c2w_tgt = ref_inv @ c2w_tgt
+                extr_src_n = torch.linalg.inv(c2w_src)
+                extr_tgt_n = torch.linalg.inv(c2w_tgt)
+
+            m_tgt_in_src = view_frustum_mask_cubemap_union(
+                depth_novel=tgt_d,
+                extr_novel_w2c=extr_tgt_n,
+                extr_source_w2c=extr_src_n,
+                face_w=face_w,
+                margin=margin,
+            )
+            m_src_in_tgt = view_frustum_mask_cubemap_union(
+                depth_novel=src_d,
+                extr_novel_w2c=extr_src_n,
+                extr_source_w2c=extr_tgt_n,
+                face_w=face_w,
+                margin=margin,
+            )
+            tgt_valid = torch.isfinite(tgt_d[..., 0]) & (tgt_d[..., 0] > 0.0)
+            src_valid = torch.isfinite(src_d[..., 0]) & (src_d[..., 0] > 0.0)
+            denom_t = float(tgt_valid.sum().item())
+            denom_s = float(src_valid.sum().item())
+            if denom_t < 10 or denom_s < 10:
+                return 0.0
+            a = float((m_tgt_in_src & tgt_valid).sum().item()) / denom_t
+            b = float((m_src_in_tgt & src_valid).sum().item()) / denom_s
+            return 0.5 * (a + b)
+
+        return _score_one(_PAIR_RECIPE_FIXED)
+
+    def __getitem__(self, idx: int) -> PanOGSSample:
+        src_erp: torch.Tensor | None = None
+        tgt_erp: torch.Tensor | None = None
+        src_dep: torch.Tensor | None = None
+        tgt_dep: torch.Tensor | None = None
+        src_cube: torch.Tensor | None = None
+        tgt_cube: torch.Tensor | None = None
+        src_cdep: torch.Tensor | None = None
+        tgt_cdep: torch.Tensor | None = None
+        last_err: Exception | None = None
+
+        max_outer = 16
+        for outer in range(max_outer):
+            scene, src_idx = self._index[int(idx) % len(self._index)]
+            scene_dir = self.root / scene
+            tgt_idx = self._sample_target(scene, src_idx)
+
+            max_retries = 8
+            ok = False
+            for _ in range(max_retries):
+                try:
+                    if src_erp is None:
+                        src_erp = _load_erp_rgb_u8(scene_dir / "pano" / f"{src_idx:05d}.png")
+                        src_dep = _depth_to_meters(
+                            _load_depth_png(scene_dir / "pano_depth" / f"{src_idx:05d}.png"),
+                            max_depth_m=self.depth_max_m,
+                        )
+                        if self.use_cubemap_supervision:
+                            src_cube_any = _torch_load_any(scene_dir / "cubemaps" / f"{src_idx:05d}.torch")
+                            src_cdep_any = _torch_load_any(scene_dir / "cubemaps_depth" / f"{src_idx:05d}.torch")
+                            if not all(isinstance(x, torch.Tensor) for x in [src_cube_any, src_cdep_any]):
+                                raise RuntimeError("Bad .torch payload for src (expected Tensor).")
+                            src_cube = cast(torch.Tensor, src_cube_any)
+                            src_cdep = cast(torch.Tensor, src_cdep_any).to(torch.float32).clamp(min=0.0, max=self.depth_max_m)
+                        else:
+                            src_cube = torch.zeros((6, 256, 256, 3), dtype=torch.uint8)
+                            src_cdep = torch.zeros((6, 256, 256, 1), dtype=torch.float32)
+
+                    candidates: list[int] = []
+                    if self.pair_sampling and self.use_cubemap_supervision:
+                        key = (scene, int(src_idx))
+                        cached = self._pair_valid_tgts.get(key)
+                        if cached:
+                            candidates = list(cached)
+                        else:
+                            candidates = self._candidate_targets_by_translation(scene, int(src_idx))
+                    if not candidates:
+                        candidates = [int(tgt_idx)]
+
+                    tried: set[int] = set()
+                    found = False
+                    max_try = (
+                        1
+                        if (not self.pair_sampling or not self.use_cubemap_supervision)
+                        else max(1, self.pair_max_tries)
+                    )
+                    for _try in range(max_try):
+                        pool = [
+                            c
+                            for c in candidates
+                            if int(c) not in tried and int(c) != int(src_idx)
+                        ]
+                        if not pool:
+                            break
+                        j = int(torch.randint(0, len(pool), (1,)).item())
+                        tgt_idx = int(pool[j])
+                        tried.add(int(tgt_idx))
+
+                        if self.use_cubemap_supervision:
+                            tgt_cdep_any = _torch_load_any(scene_dir / "cubemaps_depth" / f"{tgt_idx:05d}.torch")
+                            if not isinstance(tgt_cdep_any, torch.Tensor):
+                                raise RuntimeError("Bad .torch payload for tgt depth (expected Tensor).")
+                            tgt_cdep = cast(torch.Tensor, tgt_cdep_any).to(torch.float32).clamp(min=0.0, max=self.depth_max_m)
+                        else:
+                            tgt_cdep = torch.zeros((6, 256, 256, 1), dtype=torch.float32)
+
+                        if self.pair_sampling and self.use_cubemap_supervision:
+                            k = (scene, int(src_idx), int(tgt_idx))
+                            score = self._pair_overlap_cache.get(k)
+                            if score is None:
+                                R_np, t_np = self._get_pose(scene)
+                                src_R = torch.from_numpy(R_np[int(src_idx)])
+                                src_t = torch.from_numpy(t_np[int(src_idx)])
+                                tgt_R = torch.from_numpy(R_np[int(tgt_idx)])
+                                tgt_t = torch.from_numpy(t_np[int(tgt_idx)])
+                                score = self._pair_depth_overlap_score(
+                                    src_R=src_R,
+                                    src_t=src_t,
+                                    tgt_R=tgt_R,
+                                    tgt_t=tgt_t,
+                                    src_cube_depth_m=cast(torch.Tensor, src_cdep),
+                                    tgt_cube_depth_m=cast(torch.Tensor, tgt_cdep),
+                                )
+                                self._pair_overlap_cache[k] = float(score)
+                            if float(score) < float(self.pair_min_depth_overlap):
+                                continue
+                            kk = (scene, int(src_idx))
+                            self._pair_valid_tgts.setdefault(kk, []).append(int(tgt_idx))
+
+                        tgt_erp = _load_erp_rgb_u8(scene_dir / "pano" / f"{tgt_idx:05d}.png")
+                        tgt_dep = _depth_to_meters(
+                            _load_depth_png(scene_dir / "pano_depth" / f"{tgt_idx:05d}.png"),
+                            max_depth_m=self.depth_max_m,
+                        )
+                        if self.use_cubemap_supervision:
+                            tgt_cube_any = _torch_load_any(scene_dir / "cubemaps" / f"{tgt_idx:05d}.torch")
+                            if not isinstance(tgt_cube_any, torch.Tensor):
+                                raise RuntimeError("Bad .torch payload for tgt RGB cubemap (expected Tensor).")
+                            tgt_cube = cast(torch.Tensor, tgt_cube_any)
+                        else:
+                            tgt_cube = torch.zeros((6, 256, 256, 3), dtype=torch.uint8)
+
+                        found = True
+                        break
+
+                    if not found:
+                        raise RuntimeError(
+                            f"No valid tgt found for scene={scene} src={src_idx} within constraints "
+                            f"(trans<{self.pair_max_translation_m}m, overlap>{self.pair_min_depth_overlap})."
+                        )
+
+                    ok = True
+                    break
+                except (FileNotFoundError, RuntimeError, EOFError, KeyError, tarfile.ReadError, OSError) as e:
+                    last_err = e
+                    frames = self._available_frames.get(scene, [])
+                    if not frames:
+                        break
+                    src_idx = int(frames[int(torch.randint(0, len(frames), (1,)).item())])
+                    tgt_idx = self._sample_target(scene, src_idx)
+                    src_erp = None
+                    src_dep = None
+                    src_cube = None
+                    src_cdep = None
+
+            if ok:
+                break
+
+            idx = int(idx) + 9973 + outer * 13
+        else:
+            raise RuntimeError(f"PanOGS __getitem__ failed after retries. last_err={last_err}")
+
+        assert src_erp is not None and tgt_erp is not None
+        assert src_dep is not None and tgt_dep is not None
+        assert src_cube is not None and tgt_cube is not None
+        assert src_cdep is not None and tgt_cdep is not None
+
+        src_dep = src_dep.to(torch.float32).unsqueeze(0)
+        tgt_dep = tgt_dep.to(torch.float32).unsqueeze(0)
+
+        R_np, t_np = self._get_pose(scene)
+        src_R = torch.from_numpy(R_np[src_idx])
+        src_t = torch.from_numpy(t_np[src_idx])
+        tgt_R = torch.from_numpy(R_np[tgt_idx])
+        tgt_t = torch.from_numpy(t_np[tgt_idx])
+
+        return PanOGSSample(
+            src_erp_rgb_u8=src_erp,
+            tgt_erp_rgb_u8=tgt_erp,
+            src_erp_depth_m=src_dep,
+            tgt_erp_depth_m=tgt_dep,
+            src_cube_rgb_u8=src_cube,
+            tgt_cube_rgb_u8=tgt_cube,
+            src_cube_depth_m=src_cdep,
+            tgt_cube_depth_m=tgt_cdep,
+            src_R=src_R,
+            src_t=src_t,
+            tgt_R=tgt_R,
+            tgt_t=tgt_t,
+            src_idx=src_idx,
+            tgt_idx=tgt_idx,
+            scene=scene,
+        )
+
+
+def panogs_collate(batch: list[PanOGSSample]) -> PanOGSSample:
+    def stack(xs):
+        if isinstance(xs[0], torch.Tensor):
+            return torch.stack(xs, dim=0)
+        return xs
+
+    return PanOGSSample(
+        src_erp_rgb_u8=stack([b.src_erp_rgb_u8 for b in batch]),
+        tgt_erp_rgb_u8=stack([b.tgt_erp_rgb_u8 for b in batch]),
+        src_erp_depth_m=stack([b.src_erp_depth_m for b in batch]),
+        tgt_erp_depth_m=stack([b.tgt_erp_depth_m for b in batch]),
+        src_cube_rgb_u8=stack([b.src_cube_rgb_u8 for b in batch]),
+        tgt_cube_rgb_u8=stack([b.tgt_cube_rgb_u8 for b in batch]),
+        src_cube_depth_m=stack([b.src_cube_depth_m for b in batch]),
+        tgt_cube_depth_m=stack([b.tgt_cube_depth_m for b in batch]),
+        src_R=stack([b.src_R for b in batch]),
+        src_t=stack([b.src_t for b in batch]),
+        tgt_R=stack([b.tgt_R for b in batch]),
+        tgt_t=stack([b.tgt_t for b in batch]),
+        src_idx=[b.src_idx for b in batch],  # type: ignore[arg-type]
+        tgt_idx=[b.tgt_idx for b in batch],  # type: ignore[arg-type]
+        scene=[b.scene for b in batch],  # type: ignore[arg-type]
+    )
+

unisharp/datasets/re10k.py

@@ -0,0 +1,718 @@
+from __future__ import annotations
+
+from collections import defaultdict, deque
+from dataclasses import dataclass
+from io import BytesIO
+import logging
+import os
+from pathlib import Path
+import random
+import time
+
+import torch
+import torchvision.transforms as tf
+from PIL import Image
+from torch.utils.data import IterableDataset
+
+from unisharp.datasets.pair_sampling import (
+    project_overlap_ratio,
+    resize_k3_align_corners_false,
+    resize_rgb_u8_chw_high_quality,
+    select_targets_for_source,
+)
+from unisharp import DEFAULT_MAX_DEPTH_M
+from unisharp.utils.pixel_convention import normalized_intrinsics_to_integer_pixel_k
+from unisharp.utils.unik3d_adapter import infer_unik3d_pinhole, load_unik3d_model
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+def _torch_load_any(path: Path) -> object:
+    try:
+        return torch.load(path, map_location="cpu", weights_only=False)
+    except TypeError:
+        return torch.load(path, map_location="cpu")
+
+
+def _pack_re10k_batch(batch: list["Re10KPairSample"]) -> "Re10KPairSample":
+    def stack(xs):
+        if isinstance(xs[0], torch.Tensor):
+            ref_shape = tuple(xs[0].shape)
+            for idx, x in enumerate(xs[1:], start=1):
+                if tuple(x.shape) != ref_shape:
+                    raise RuntimeError(
+                        "RE10K collate got mixed tensor shapes: "
+                        f"ref={ref_shape} mismatch_idx={idx} got={tuple(x.shape)}"
+                    )
+            return torch.stack(xs, dim=0)
+        return xs
+
+    def stack_optional_depth(xs):
+        if all(torch.is_tensor(x) for x in xs):
+            ref_shape = tuple(xs[0].shape)
+            for idx, x in enumerate(xs[1:], start=1):
+                if tuple(x.shape) != ref_shape:
+                    raise RuntimeError(
+                        "RE10K collate got mixed depth shapes: "
+                        f"ref={ref_shape} mismatch_idx={idx} got={tuple(x.shape)}"
+                    )
+            return torch.stack(xs, dim=0)
+        return None
+
+    return Re10KPairSample(
+        src_rgb_u8=stack([b.src_rgb_u8 for b in batch]),
+        tgt_rgb_u8=stack([b.tgt_rgb_u8 for b in batch]),
+        src_w2c=stack([b.src_w2c for b in batch]),
+        tgt_w2c=stack([b.tgt_w2c for b in batch]),
+        src_intrinsics=stack([b.src_intrinsics for b in batch]),
+        tgt_intrinsics=stack([b.tgt_intrinsics for b in batch]),
+        src_idx=[b.src_idx for b in batch],  # type: ignore[arg-type]
+        tgt_idx=[b.tgt_idx for b in batch],  # type: ignore[arg-type]
+        scene=[b.scene for b in batch],  # type: ignore[arg-type]
+        src_depth_m=stack_optional_depth([b.src_depth_m for b in batch]),  # type: ignore[arg-type]
+        tgt_depth_m=stack_optional_depth([b.tgt_depth_m for b in batch]),  # type: ignore[arg-type]
+    )
+
+
+def re10k_passthrough(batch: "Re10KPairSample") -> "Re10KPairSample":
+    return batch
+
+
+@dataclass(frozen=True)
+class Re10KPairSample:
+    src_rgb_u8: torch.Tensor
+    tgt_rgb_u8: torch.Tensor
+    src_w2c: torch.Tensor
+    tgt_w2c: torch.Tensor
+    src_intrinsics: torch.Tensor
+    tgt_intrinsics: torch.Tensor
+    src_idx: int
+    tgt_idx: int
+    scene: str
+    src_depth_m: torch.Tensor | None = None
+    tgt_depth_m: torch.Tensor | None = None
+
+
+class Re10KDataset(IterableDataset):
+
+    def __init__(
+        self,
+        root: Path,
+        chunks_file: Path | None = None,
+        split: str = "train",
+        min_frame_gap: int = 1,
+        max_frame_gap: int = 32,
+        pair_max_translation_m: float = 0.5,
+        pair_min_overlap: float = 0.6,
+        pair_overlap_sample_h: int = 32,
+        pair_overlap_sample_w: int = 56,
+        pair_max_tries: int = 32,
+        output_h: int | None = None,
+        output_w: int | None = None,
+        shuffle_chunk: bool = True,
+        shuffle_example: bool = True,
+        ddp_rank: int = 0,
+        ddp_world_size: int = 1,
+        pseudo_depth_root: Path | None = None,
+        pseudo_depth_autogen: bool = True,
+        pseudo_depth_backbone: str = "vitl",
+        pseudo_depth_device: str = "cpu",
+        pseudo_lock_timeout_sec: float = 120.0,
+        pseudo_lock_stale_sec: float = 1800.0,
+        pseudo_wait_poll_sec: float = 0.25,
+        batch_size_hint: int = 1,
+        depth_max_m: float = DEFAULT_MAX_DEPTH_M,
+        pseudo_far_depth_invalid_m: float = 30.0,
+        seed: int = 0,
+    ) -> None:
+        super().__init__()
+        self.root = root
+        self.split = split
+        self.min_frame_gap = int(min_frame_gap)
+        self.max_frame_gap = int(max_frame_gap)
+        self.pair_max_translation_m = float(pair_max_translation_m)
+        self.pair_min_overlap = float(pair_min_overlap)
+        self.pair_overlap_sample_h = int(pair_overlap_sample_h)
+        self.pair_overlap_sample_w = int(pair_overlap_sample_w)
+        self.pair_max_tries = int(pair_max_tries)
+        self.output_h = int(output_h) if output_h is not None else None
+        self.output_w = int(output_w) if output_w is not None else None
+        self.shuffle_chunk = bool(shuffle_chunk)
+        self.shuffle_example = bool(shuffle_example)
+        self.ddp_rank = int(ddp_rank)
+        self.ddp_world_size = int(ddp_world_size)
+        self.to_tensor = tf.ToTensor()
+        self.pseudo_depth_root = Path(pseudo_depth_root) if pseudo_depth_root is not None else None
+        self.pseudo_depth_autogen = bool(pseudo_depth_autogen)
+        self.pseudo_depth_backbone = str(pseudo_depth_backbone)
+        self.pseudo_depth_device = str(pseudo_depth_device)
+        self.pseudo_lock_timeout_sec = float(max(1.0, pseudo_lock_timeout_sec))
+        self.pseudo_lock_stale_sec = float(max(30.0, pseudo_lock_stale_sec))
+        self.pseudo_wait_poll_sec = float(max(0.05, pseudo_wait_poll_sec))
+        self.batch_size_hint = int(max(1, batch_size_hint))
+        self.depth_max_m = float(depth_max_m)
+        self.pseudo_far_depth_invalid_m = float(pseudo_far_depth_invalid_m)
+        self._pseudo_model: torch.nn.Module | None = None
+        self.seed = int(seed)
+        self.epoch = 0
+
+        self.chunks_file = Path(chunks_file) if chunks_file is not None else None
+        split_dir = self.root / self.split
+        if self.chunks_file is not None:
+            if not self.chunks_file.exists():
+                raise FileNotFoundError(self.chunks_file)
+            chunks: list[Path] = []
+            for raw in self.chunks_file.read_text(encoding="utf-8").splitlines():
+                line = raw.strip()
+                if not line:
+                    continue
+                p = Path(line)
+                if not p.is_absolute():
+                    p = split_dir / p
+                if p.suffix == ".torch":
+                    chunks.append(p)
+            self.chunks = sorted(chunks)
+        else:
+            if not split_dir.exists():
+                raise FileNotFoundError(split_dir)
+            self.chunks = sorted([p for p in split_dir.iterdir() if p.suffix == ".torch"])
+        if not self.chunks:
+            source = self.chunks_file if self.chunks_file is not None else split_dir
+            raise RuntimeError(f"No .torch chunks found for {source}")
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = int(epoch)
+
+        if self.pseudo_depth_root is not None:
+            (self.pseudo_depth_root / self.split).mkdir(parents=True, exist_ok=True)
+
+    @staticmethod
+    def _decode_image_u8(image_bytes_tensor: torch.Tensor) -> torch.Tensor:
+        if image_bytes_tensor.dtype != torch.uint8:
+            raise ValueError(f"Expected uint8 bytes tensor, got {image_bytes_tensor.dtype}")
+        image = Image.open(BytesIO(image_bytes_tensor.numpy().tobytes())).convert("RGB")
+        chw_float = tf.ToTensor()(image)
+        return (chw_float * 255.0).round().to(torch.uint8)
+
+    @staticmethod
+    def _convert_pose_row_to_w2c(poses: torch.Tensor) -> torch.Tensor:
+        t = poses.shape[0]
+        w2c = torch.eye(4, dtype=torch.float32).unsqueeze(0).repeat(t, 1, 1)
+        w2c[:, :3] = poses[:, 6:].reshape(t, 3, 4).to(torch.float32)
+        return w2c
+
+    @staticmethod
+    def _convert_intrinsics_to_pixel(poses: torch.Tensor, h: int, w: int) -> torch.Tensor:
+        t = poses.shape[0]
+        fx, fy, cx, cy = poses[:, 0], poses[:, 1], poses[:, 2], poses[:, 3]
+        del t
+        return normalized_intrinsics_to_integer_pixel_k(
+            fx,
+            fy,
+            cx,
+            cy,
+            height=int(h),
+            width=int(w),
+        )
+
+    @staticmethod
+    def _sanitize_scene(scene: str) -> str:
+        s = str(scene).strip()
+        s = s.replace("\\", "__").replace("/", "__")
+        return s if len(s) > 0 else "unknown_scene"
+
+    def _pseudo_depth_path(self, scene: str, frame_idx: int) -> Path | None:
+        if self.pseudo_depth_root is None:
+            return None
+        scene_key = self._sanitize_scene(scene)
+        return self.pseudo_depth_root / self.split / scene_key / f"{int(frame_idx):05d}.pt"
+
+    @staticmethod
+    def _load_pseudo_depth(path: Path) -> tuple[torch.Tensor | None, str]:
+        if not path.exists():
+            return None, "unknown"
+        try:
+            payload = _torch_load_any(path)
+            depth_kind = "distance"
+            if isinstance(payload, dict):
+                depth = payload.get("depth_m", None)
+                depth_kind = str(payload.get("depth_kind", "distance")).strip().lower()
+                if depth_kind not in ("distance", "zdepth"):
+                    depth_kind = "distance"
+            else:
+                depth = payload
+            if not torch.is_tensor(depth):
+                return None, "unknown"
+            if depth.ndim == 3 and depth.shape[0] == 1:
+                depth = depth[0]
+            if depth.ndim != 2:
+                return None, "unknown"
+            depth = depth.to(torch.float32)
+            valid = torch.isfinite(depth) & (depth > 0.0)
+            if int(valid.sum().item()) <= 0:
+                return None, "unknown"
+            return depth.unsqueeze(0), depth_kind
+        except Exception:
+            return None, "unknown"
+
+    @staticmethod
+    def _distance_to_z_depth(depth_1hw: torch.Tensor, intrinsics_k3: torch.Tensor) -> torch.Tensor:
+        if depth_1hw.ndim != 3 or depth_1hw.shape[0] != 1:
+            raise ValueError(f"Expected depth shape (1,H,W), got {tuple(depth_1hw.shape)}")
+        d = depth_1hw.to(torch.float32)
+        h = int(d.shape[-2])
+        w = int(d.shape[-1])
+        k = intrinsics_k3.to(dtype=torch.float32, device=d.device)
+        fx = k[0, 0]
+        fy = k[1, 1]
+        cx = k[0, 2]
+        cy = k[1, 2]
+        ys = torch.arange(h, device=d.device, dtype=torch.float32)
+        xs = torch.arange(w, device=d.device, dtype=torch.float32)
+        vv, uu = torch.meshgrid(ys, xs, indexing="ij")
+        x = (uu - cx) / fx
+        y = (vv - cy) / fy
+        ray_z = 1.0 / torch.sqrt(x * x + y * y + 1.0).clamp(min=1e-8)
+        z = d[0] * ray_z
+        return z.unsqueeze(0)
+
+    @staticmethod
+    def _sanitize_pseudo_depth(
+        depth_1hw: torch.Tensor,
+        *,
+        max_depth_m: float = DEFAULT_MAX_DEPTH_M,
+        far_depth_invalid_m: float = 30.0,
+    ) -> torch.Tensor:
+        d = depth_1hw.to(torch.float32)
+        valid = torch.isfinite(d) & (d > 0.0)
+        if int(valid.sum().item()) <= 0:
+            return d
+        out = d.clone()
+        if float(far_depth_invalid_m) > 0.0:
+            valid = valid & (out <= float(far_depth_invalid_m))
+            out = torch.where(valid, out, torch.zeros_like(out))
+        out[valid] = out[valid].clamp(max=float(max_depth_m))
+        return out
+
+    def _get_or_create_pseudo_model(self) -> torch.nn.Module:
+        if self._pseudo_model is None:
+            dev = torch.device(self.pseudo_depth_device)
+            self._pseudo_model = load_unik3d_model(
+                backbone=self.pseudo_depth_backbone,
+                pretrained=True,
+                device=dev,
+            )
+            self._pseudo_model.eval()
+            LOGGER.info(
+                "Re10K pseudo-depth model loaded (split=%s, device=%s, backbone=%s)",
+                self.split,
+                str(dev),
+                self.pseudo_depth_backbone,
+            )
+        return self._pseudo_model
+
+    def _save_pseudo_depth_atomic(
+        self,
+        path: Path,
+        depth_2d: torch.Tensor,
+        scene: str,
+        frame_idx: int,
+    ) -> None:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        tmp = path.parent / f".tmp_{os.getpid()}_{int(time.time() * 1e6)}_{random.randint(0, 10_000_000)}.pt"
+        payload = {
+            "depth_m": depth_2d.to(torch.float16),
+            "depth_kind": "distance",
+            "scene": str(scene),
+            "frame_idx": int(frame_idx),
+        }
+        torch.save(payload, tmp)
+        os.replace(tmp, path)
+
+    def _acquire_lock_or_wait_for_file(self, target: Path) -> tuple[bool, bool]:
+        lock_dir = Path(str(target) + ".lock")
+        start = time.time()
+        while True:
+            if target.exists():
+                return False, True
+            try:
+                lock_dir.mkdir(parents=False, exist_ok=False)
+                meta = lock_dir / "owner.txt"
+                meta.write_text(f"pid={os.getpid()} time={time.time():.3f}\n", encoding="utf-8")
+                return True, False
+            except FileExistsError:
+                try:
+                    mtime = lock_dir.stat().st_mtime
+                    if (time.time() - float(mtime)) > self.pseudo_lock_stale_sec:
+                        for p in lock_dir.iterdir():
+                            try:
+                                p.unlink()
+                            except Exception:
+                                pass
+                        lock_dir.rmdir()
+                        continue
+                except Exception:
+                    pass
+                if (time.time() - start) >= self.pseudo_lock_timeout_sec:
+                    return False, False
+                time.sleep(self.pseudo_wait_poll_sec)
+            except Exception:
+                return False, False
+
+    def _release_lock(self, target: Path) -> None:
+        lock_dir = Path(str(target) + ".lock")
+        if not lock_dir.exists():
+            return
+        try:
+            for p in lock_dir.iterdir():
+                try:
+                    p.unlink()
+                except Exception:
+                    pass
+            lock_dir.rmdir()
+        except Exception:
+            pass
+
+    def _get_pseudo_depth_for_frame(
+        self,
+        *,
+        scene: str,
+        frame_idx: int,
+        rgb_u8: torch.Tensor,
+        intrinsics_k3: torch.Tensor,
+    ) -> torch.Tensor | None:
+        path = self._pseudo_depth_path(scene, frame_idx)
+        if path is None:
+            return None
+        depth, depth_kind = self._load_pseudo_depth(path)
+        if depth is not None:
+            if depth_kind != "zdepth":
+                try:
+                    depth = self._distance_to_z_depth(
+                        self._sanitize_pseudo_depth(
+                            depth,
+                            max_depth_m=self.depth_max_m,
+                            far_depth_invalid_m=self.pseudo_far_depth_invalid_m,
+                        ),
+                        intrinsics_k3=intrinsics_k3,
+                    )
+                except Exception:
+                    return None
+            else:
+                depth = self._sanitize_pseudo_depth(
+                    depth,
+                    max_depth_m=self.depth_max_m,
+                    far_depth_invalid_m=self.pseudo_far_depth_invalid_m,
+                )
+            return depth
+        if not self.pseudo_depth_autogen:
+            return None
+
+        acquired, ready = self._acquire_lock_or_wait_for_file(path)
+        if ready:
+            depth, depth_kind = self._load_pseudo_depth(path)
+            if depth is None:
+                return None
+            if depth_kind != "zdepth":
+                try:
+                    depth = self._distance_to_z_depth(
+                        self._sanitize_pseudo_depth(
+                            depth,
+                            max_depth_m=self.depth_max_m,
+                            far_depth_invalid_m=self.pseudo_far_depth_invalid_m,
+                        ),
+                        intrinsics_k3=intrinsics_k3,
+                    )
+                except Exception:
+                    return None
+            else:
+                depth = self._sanitize_pseudo_depth(
+                    depth,
+                    max_depth_m=self.depth_max_m,
+                    far_depth_invalid_m=self.pseudo_far_depth_invalid_m,
+                )
+            return depth
+        if not acquired:
+            depth, depth_kind = self._load_pseudo_depth(path)
+            if depth is None:
+                return None
+            if depth_kind != "zdepth":
+                try:
+                    depth = self._distance_to_z_depth(
+                        self._sanitize_pseudo_depth(
+                            depth,
+                            max_depth_m=self.depth_max_m,
+                            far_depth_invalid_m=self.pseudo_far_depth_invalid_m,
+                        ),
+                        intrinsics_k3=intrinsics_k3,
+                    )
+                except Exception:
+                    return None
+            else:
+                depth = self._sanitize_pseudo_depth(
+                    depth,
+                    max_depth_m=self.depth_max_m,
+                    far_depth_invalid_m=self.pseudo_far_depth_invalid_m,
+                )
+            return depth
+
+        try:
+            depth, depth_kind = self._load_pseudo_depth(path)
+            if depth is not None:
+                if depth_kind != "zdepth":
+                    try:
+                        depth = self._distance_to_z_depth(
+                            self._sanitize_pseudo_depth(
+                                depth,
+                                max_depth_m=self.depth_max_m,
+                                far_depth_invalid_m=self.pseudo_far_depth_invalid_m,
+                            ),
+                            intrinsics_k3=intrinsics_k3,
+                        )
+                    except Exception:
+                        return None
+                else:
+                    depth = self._sanitize_pseudo_depth(
+                        depth,
+                        max_depth_m=self.depth_max_m,
+                        far_depth_invalid_m=self.pseudo_far_depth_invalid_m,
+                    )
+                return depth
+            model = self._get_or_create_pseudo_model()
+            out = infer_unik3d_pinhole(
+                model,
+                rgb_u8=rgb_u8.unsqueeze(0),
+                intrinsics=intrinsics_k3.unsqueeze(0),
+            )
+            dist = out.get("distance", None) if isinstance(out, dict) else None
+            if not torch.is_tensor(dist) or dist.ndim != 4 or dist.shape[1] != 1:
+                return None
+            dist_1hw = self._sanitize_pseudo_depth(
+                dist[0:1, 0:1].detach().to(torch.float32).cpu()[0],
+                max_depth_m=self.depth_max_m,
+                far_depth_invalid_m=self.pseudo_far_depth_invalid_m,
+            )
+            valid = torch.isfinite(dist_1hw) & (dist_1hw > 0.0)
+            if int(valid.sum().item()) <= 0:
+                return None
+            self._save_pseudo_depth_atomic(
+                path,
+                depth_2d=dist_1hw[0],
+                scene=scene,
+                frame_idx=frame_idx,
+            )
+            return self._distance_to_z_depth(dist_1hw, intrinsics_k3=intrinsics_k3.cpu())
+        except Exception as e:
+            LOGGER.warning(
+                "Pseudo-depth generation failed scene=%s frame=%d: %s",
+                str(scene),
+                int(frame_idx),
+                str(e),
+            )
+            return None
+        finally:
+            self._release_lock(path)
+
+    def _candidate_target_indices(
+        self,
+        src_idx: int,
+        num_frames: int,
+        w2c_all: torch.Tensor,
+        intr_all: torch.Tensor,
+        h: int,
+        w: int,
+    ) -> list[int]:
+        if num_frames < 2:
+            return []
+        centers = torch.linalg.inv(w2c_all)[:, :3, 3].to(torch.float32)
+        sample_h = int(self.pair_overlap_sample_h)
+        sample_w = int(self.pair_overlap_sample_w)
+        return select_targets_for_source(
+            src_idx=int(src_idx),
+            candidate_indices=list(range(num_frames)),
+            centers=centers,
+            min_index_gap=int(self.min_frame_gap),
+            max_index_gap=int(self.max_frame_gap),
+            pair_max_translation_m=float(self.pair_max_translation_m),
+            pair_min_overlap=float(self.pair_min_overlap),
+            overlap_score_fn=lambda si, tj: float(
+                0.5
+                * (
+                    project_overlap_ratio(
+                        src_w2c=w2c_all[si],
+                        tgt_w2c=w2c_all[tj],
+                        src_k=intr_all[si],
+                        tgt_k=intr_all[tj],
+                        h=h,
+                        w=w,
+                        sample_h=sample_h,
+                        sample_w=sample_w,
+                    )
+                    + project_overlap_ratio(
+                        src_w2c=w2c_all[tj],
+                        tgt_w2c=w2c_all[si],
+                        src_k=intr_all[tj],
+                        tgt_k=intr_all[si],
+                        h=h,
+                        w=w,
+                        sample_h=sample_h,
+                        sample_w=sample_w,
+                    )
+                )
+            ),
+        )
+
+    def __iter__(self):
+        chunks = list(self.chunks)
+        order_rng = random.Random(self.seed + self.epoch)
+        if self.shuffle_chunk and self.split == "train":
+            order_rng.shuffle(chunks)
+        pending_by_hw: dict[tuple[int, int], deque[Re10KPairSample]] = defaultdict(deque)
+
+        worker_info = torch.utils.data.get_worker_info()
+        num_workers = worker_info.num_workers if worker_info is not None else 1
+        worker_id = worker_info.id if worker_info is not None else 0
+        total_shards = max(1, self.ddp_world_size * num_workers)
+        shard_id = self.ddp_rank * num_workers + worker_id
+        chunks = [chunk for i, chunk in enumerate(chunks) if i % total_shards == shard_id]
+
+        for chunk_order_idx, chunk_path in enumerate(chunks):
+            chunk = _torch_load_any(chunk_path)
+            if not isinstance(chunk, list):
+                continue
+            examples = list(chunk)
+            chunk_rng = random.Random(self.seed + self.epoch * 1000003 + chunk_order_idx)
+            if self.shuffle_example and self.split == "train":
+                chunk_rng.shuffle(examples)
+
+            for example in examples:
+                if not isinstance(example, dict):
+                    continue
+                if "cameras" not in example or "images" not in example:
+                    continue
+                poses = example["cameras"]
+                images = example["images"]
+                scene = str(example.get("key", "unknown"))
+                if not torch.is_tensor(poses) or not isinstance(images, list):
+                    continue
+                if poses.ndim != 2 or poses.shape[1] != 18:
+                    continue
+                if len(images) != int(poses.shape[0]):
+                    continue
+
+                try:
+                    src_probe = self._decode_image_u8(images[0])
+                except Exception:
+                    continue
+                h, w = int(src_probe.shape[1]), int(src_probe.shape[2])
+                w2c_all = self._convert_pose_row_to_w2c(poses)
+                intr_all = self._convert_intrinsics_to_pixel(poses, h=h, w=w)
+                src_indices = list(range(len(images)))
+                if self.shuffle_example and self.split == "train":
+                    chunk_rng.shuffle(src_indices)
+                for src_idx in src_indices:
+                    tgt_candidates = self._candidate_target_indices(
+                        int(src_idx),
+                        len(images),
+                        w2c_all=w2c_all,
+                        intr_all=intr_all,
+                        h=h,
+                        w=w,
+                    )
+                    if not tgt_candidates:
+                        continue
+                    tgt_idx = chunk_rng.choice(tgt_candidates)
+
+                    try:
+                        src_img = self._decode_image_u8(images[src_idx])
+                        tgt_img = self._decode_image_u8(images[tgt_idx])
+                    except Exception:
+                        continue
+                    if src_img.shape != tgt_img.shape:
+                        continue
+                    src_intr = intr_all[src_idx].clone()
+                    tgt_intr = intr_all[tgt_idx].clone()
+                    src_depth = self._get_pseudo_depth_for_frame(
+                        scene=scene,
+                        frame_idx=int(src_idx),
+                        rgb_u8=src_img,
+                        intrinsics_k3=intr_all[src_idx].to(torch.float32),
+                    )
+                    tgt_depth = self._get_pseudo_depth_for_frame(
+                        scene=scene,
+                        frame_idx=int(tgt_idx),
+                        rgb_u8=tgt_img,
+                        intrinsics_k3=intr_all[tgt_idx].to(torch.float32),
+                    )
+                    if self.pseudo_depth_root is not None and (
+                        (not torch.is_tensor(src_depth)) or (not torch.is_tensor(tgt_depth))
+                    ):
+                        continue
+                    if self.output_h is not None and self.output_w is not None:
+                        oh, ow = int(src_img.shape[1]), int(src_img.shape[2])
+                        if oh > 0 and ow > 0 and (oh != self.output_h or ow != self.output_w):
+                            sx = float(self.output_w) / float(ow)
+                            sy = float(self.output_h) / float(oh)
+                            src_img = resize_rgb_u8_chw_high_quality(src_img, size=(self.output_h, self.output_w))
+                            tgt_img = resize_rgb_u8_chw_high_quality(tgt_img, size=(self.output_h, self.output_w))
+                            src_intr = resize_k3_align_corners_false(src_intr, sx=sx, sy=sy)
+                            tgt_intr = resize_k3_align_corners_false(tgt_intr, sx=sx, sy=sy)
+                            if torch.is_tensor(src_depth):
+                                src_depth = (
+                                    torch.nn.functional.interpolate(
+                                        src_depth[None],
+                                        size=(self.output_h, self.output_w),
+                                        mode="bilinear",
+                                        align_corners=False,
+                                    )
+                                    .squeeze(0)
+                                    .to(torch.float32)
+                                )
+                            if torch.is_tensor(tgt_depth):
+                                tgt_depth = (
+                                    torch.nn.functional.interpolate(
+                                        tgt_depth[None],
+                                        size=(self.output_h, self.output_w),
+                                        mode="bilinear",
+                                        align_corners=False,
+                                    )
+                                    .squeeze(0)
+                                    .to(torch.float32)
+                                )
+
+                    sample = Re10KPairSample(
+                        src_rgb_u8=src_img,
+                        tgt_rgb_u8=tgt_img,
+                        src_w2c=w2c_all[src_idx],
+                        tgt_w2c=w2c_all[tgt_idx],
+                        src_intrinsics=src_intr,
+                        tgt_intrinsics=tgt_intr,
+                        src_idx=int(src_idx),
+                        tgt_idx=int(tgt_idx),
+                        scene=scene,
+                        src_depth_m=src_depth,
+                        tgt_depth_m=tgt_depth,
+                    )
+                    hw_key = (int(sample.src_rgb_u8.shape[1]), int(sample.src_rgb_u8.shape[2]))
+                    bucket = pending_by_hw[hw_key]
+                    bucket.append(sample)
+                    if self.batch_size_hint <= 1:
+                        yield bucket.popleft()
+                        continue
+                    while len(bucket) >= self.batch_size_hint:
+                        packed = [bucket.popleft() for _ in range(self.batch_size_hint)]
+                        yield _pack_re10k_batch(packed)
+
+        dropped = sum(len(bucket) for bucket in pending_by_hw.values())
+        if dropped > 0 and self.split == "train" and self.batch_size_hint > 1:
+            LOGGER.debug(
+                "Dropped %d RE10K leftover samples that could not form a same-resolution batch of size %d.",
+                int(dropped),
+                int(self.batch_size_hint),
+            )
+
+
+def re10k_collate(batch: list[Re10KPairSample]) -> Re10KPairSample:
+    return _pack_re10k_batch(batch)
+

unisharp/datasets/scannetpp_fisheye.py

@@ -0,0 +1,491 @@
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import json
+import logging
+from pathlib import Path
+import random
+
+import numpy as np
+import torch
+from PIL import Image
+from torch.utils.data import IterableDataset
+
+from unisharp import DEFAULT_MAX_DEPTH_M
+
+
+LOGGER = logging.getLogger(__name__)
+IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".JPG", ".JPEG", ".PNG"}
+DEPTH_DIR_NAMES = ("depth", "depths", "distance", "distances", "depth_maps")
+MASK_DIR_NAMES = ("masks", "mask")
+DEPTH_MAX_M = DEFAULT_MAX_DEPTH_M
+
+
+def _qvec_to_rotmat(qvec: np.ndarray) -> np.ndarray:
+    q = np.asarray(qvec, dtype=np.float64)
+    return np.array(
+        [
+            [1 - 2 * q[2] ** 2 - 2 * q[3] ** 2, 2 * q[1] * q[2] - 2 * q[0] * q[3], 2 * q[3] * q[1] + 2 * q[0] * q[2]],
+            [2 * q[1] * q[2] + 2 * q[0] * q[3], 1 - 2 * q[1] ** 2 - 2 * q[3] ** 2, 2 * q[2] * q[3] - 2 * q[0] * q[1]],
+            [2 * q[3] * q[1] - 2 * q[0] * q[2], 2 * q[2] * q[3] + 2 * q[0] * q[1], 1 - 2 * q[1] ** 2 - 2 * q[2] ** 2],
+        ],
+        dtype=np.float64,
+    )
+
+
+def _read_colmap_w2c(images_txt: Path) -> dict[str, torch.Tensor]:
+    poses: dict[str, torch.Tensor] = {}
+    if not images_txt.exists():
+        return poses
+    with images_txt.open("r", encoding="utf-8") as f:
+        for raw in f:
+            line = raw.strip()
+            if not line or line.startswith("#"):
+                continue
+            parts = line.split()
+            if len(parts) < 10:
+                continue
+            try:
+                qvec = np.asarray([float(x) for x in parts[1:5]], dtype=np.float64)
+                tvec = np.asarray([float(x) for x in parts[5:8]], dtype=np.float64)
+                image_name = parts[9]
+            except Exception:
+                continue
+            w2c = np.eye(4, dtype=np.float32)
+            w2c[:3, :3] = _qvec_to_rotmat(qvec).astype(np.float32)
+            w2c[:3, 3] = tvec.astype(np.float32)
+            poses[Path(image_name).name] = torch.from_numpy(w2c)
+    return poses
+
+
+def _opencv_fisheye_to_fisheye624_params(meta: dict[str, object]) -> torch.Tensor:
+    if str(meta.get("camera_model", "")) != "OPENCV_FISHEYE":
+        raise RuntimeError(f"Unsupported ScanNet++ camera_model={meta.get('camera_model')!r}; expected OPENCV_FISHEYE.")
+    return torch.tensor(
+        [
+            float(meta["fl_x"]),
+            float(meta["fl_y"]),
+            float(meta["cx"]),
+            float(meta["cy"]),
+            float(meta.get("k1", 0.0)),
+            float(meta.get("k2", 0.0)),
+            float(meta.get("k3", 0.0)),
+            float(meta.get("k4", 0.0)),
+            0.0,
+            0.0,
+            0.0,
+            0.0,
+            0.0,
+            0.0,
+            0.0,
+            0.0,
+        ],
+        dtype=torch.float32,
+    )
+
+
+def _camera_hw_from_meta(meta: dict[str, object]) -> tuple[int, int] | None:
+    h = meta.get("h", meta.get("height", None))
+    w = meta.get("w", meta.get("width", None))
+    if h is None or w is None:
+        return None
+    try:
+        h_i, w_i = int(h), int(w)
+    except Exception:
+        return None
+    return (h_i, w_i) if h_i > 0 and w_i > 0 else None
+
+
+def _scale_fisheye624_params(
+    params: torch.Tensor,
+    *,
+    src_hw: tuple[int, int],
+    dst_hw: tuple[int, int],
+) -> torch.Tensor:
+    if tuple(int(x) for x in src_hw) == tuple(int(x) for x in dst_hw):
+        return params.clone()
+    src_h, src_w = int(src_hw[0]), int(src_hw[1])
+    dst_h, dst_w = int(dst_hw[0]), int(dst_hw[1])
+    sx = float(dst_w) / float(max(src_w, 1))
+    sy = float(dst_h) / float(max(src_h, 1))
+    out = params.clone()
+    out[..., 0] *= sx
+    out[..., 1] *= sy
+    out[..., 2] = (out[..., 2] + 0.5) * sx - 0.5
+    out[..., 3] = (out[..., 3] + 0.5) * sy - 0.5
+    return out
+
+
+def _stack_batch(batch: list["ScannetppFisheyePairSample"]) -> "ScannetppFisheyePairSample":
+    return ScannetppFisheyePairSample(
+        src_rgb_u8=torch.stack([b.src_rgb_u8 for b in batch], dim=0),
+        tgt_rgb_u8=torch.stack([b.tgt_rgb_u8 for b in batch], dim=0),
+        src_depth_m=torch.stack([b.src_depth_m for b in batch], dim=0),
+        tgt_depth_m=torch.stack([b.tgt_depth_m for b in batch], dim=0),
+        src_valid_mask=torch.stack([b.src_valid_mask for b in batch], dim=0),
+        tgt_valid_mask=torch.stack([b.tgt_valid_mask for b in batch], dim=0),
+        src_w2c=torch.stack([b.src_w2c for b in batch], dim=0),
+        tgt_w2c=torch.stack([b.tgt_w2c for b in batch], dim=0),
+        src_camera_params=torch.stack([b.src_camera_params for b in batch], dim=0),
+        tgt_camera_params=torch.stack([b.tgt_camera_params for b in batch], dim=0),
+        src_idx=[b.src_idx for b in batch],  # type: ignore[arg-type]
+        tgt_idx=[b.tgt_idx for b in batch],  # type: ignore[arg-type]
+        scene=[b.scene for b in batch],  # type: ignore[arg-type]
+        camera_model="fisheye624",
+    )
+
+
+@dataclass(frozen=True)
+class ScannetppFisheyePairSample:
+    src_rgb_u8: torch.Tensor
+    tgt_rgb_u8: torch.Tensor
+    src_depth_m: torch.Tensor
+    tgt_depth_m: torch.Tensor
+    src_valid_mask: torch.Tensor
+    tgt_valid_mask: torch.Tensor
+    src_w2c: torch.Tensor
+    tgt_w2c: torch.Tensor
+    src_camera_params: torch.Tensor
+    tgt_camera_params: torch.Tensor
+    src_idx: int
+    tgt_idx: int
+    scene: str
+    camera_model: str = "fisheye624"
+
+
+def scannetpp_fisheye_passthrough(batch: ScannetppFisheyePairSample) -> ScannetppFisheyePairSample:
+    return batch
+
+
+class ScannetppFisheyeDataset(IterableDataset):
+
+    def __init__(
+        self,
+        root: Path,
+        scene_list_file: Path | None = None,
+        min_frame_gap: int = 1,
+        max_frame_gap: int = 10,
+        pair_max_translation_m: float = 0.5,
+        shuffle_scene: bool = True,
+        shuffle_frame: bool = True,
+        skip_bad: bool = True,
+        ddp_rank: int = 0,
+        ddp_world_size: int = 1,
+        batch_size_hint: int = 1,
+        depth_max_m: float = DEFAULT_MAX_DEPTH_M,
+        far_depth_invalid_m: float = 30.0,
+        seed: int = 0,
+    ) -> None:
+        super().__init__()
+        self.root = Path(root)
+        self.min_frame_gap = int(min_frame_gap)
+        self.max_frame_gap = int(max_frame_gap)
+        self.pair_max_translation_m = float(pair_max_translation_m)
+        self.shuffle_scene = bool(shuffle_scene)
+        self.shuffle_frame = bool(shuffle_frame)
+        self.skip_bad = bool(skip_bad)
+        self.ddp_rank = int(ddp_rank)
+        self.ddp_world_size = int(ddp_world_size)
+        self.batch_size_hint = int(max(1, batch_size_hint))
+        self.depth_max_m = float(depth_max_m)
+        self.far_depth_invalid_m = float(far_depth_invalid_m)
+        self.seed = int(seed)
+        self.epoch = 0
+        self.scene_specs = self._load_scene_specs(scene_list_file)
+        if not self.scene_specs:
+            raise RuntimeError(f"No ScanNet++ fisheye scenes found under {self.root}")
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = int(epoch)
+
+    def _load_scene_specs(self, scene_list_file: Path | None) -> list[tuple[str, Path]]:
+        specs: list[tuple[str, Path]] = []
+        if scene_list_file is not None and Path(scene_list_file).exists():
+            for raw in Path(scene_list_file).read_text(encoding="utf-8").splitlines():
+                line = raw.strip()
+                if not line:
+                    continue
+                parts = line.split("|")
+                if len(parts) == 1:
+                    scene_dir = Path(parts[0])
+                    scene_id = scene_dir.name
+                else:
+                    scene_id = parts[0]
+                    scene_dir = Path(parts[1])
+                if not scene_dir.is_absolute():
+                    scene_dir = self.root / scene_dir
+                specs.append((scene_id, scene_dir))
+            return specs
+        for transforms in sorted(self.root.glob("*/nerfstudio/transforms.json")):
+            specs.append((transforms.parent.parent.name, transforms.parent.parent))
+        for transforms in sorted(self.root.glob("*/*/nerfstudio/transforms.json")):
+            specs.append((f"{transforms.parent.parent.parent.name}/{transforms.parent.parent.name}", transforms.parent.parent))
+        return specs
+
+    @staticmethod
+    def _load_rgb(path: Path) -> torch.Tensor:
+        with Image.open(path) as image:
+            arr = np.asarray(image.convert("RGB"), dtype=np.uint8).copy()
+        return torch.from_numpy(arr).permute(2, 0, 1).contiguous()
+
+    @staticmethod
+    def _load_mask(path: Path, image_hw: tuple[int, int]) -> torch.Tensor | None:
+        if not path.exists():
+            return None
+        with Image.open(path) as image:
+            arr = np.asarray(image.convert("L"), dtype=np.uint8).copy()
+        mask = torch.from_numpy(arr).unsqueeze(0).to(torch.float32) / 255.0
+        if tuple(mask.shape[-2:]) != tuple(image_hw):
+            mask = torch.nn.functional.interpolate(mask.unsqueeze(0), size=image_hw, mode="nearest").squeeze(0)
+        return (mask > 0.5).to(torch.float32)
+
+    def _load_depth_map(self, path: Path) -> tuple[torch.Tensor, str]:
+        depth_kind = "distance"
+        if path.suffix.lower() == ".npz":
+            payload = np.load(path, allow_pickle=False)
+            for key in ("distance_m", "depth_m", "distance", "depth"):
+                if key in payload:
+                    arr = payload[key]
+                    if key in {"distance_m", "distance"}:
+                        depth_kind = "distance"
+                    elif "depth_kind" in payload:
+                        depth_kind = str(np.asarray(payload["depth_kind"]).item()).strip().lower()
+                    break
+            else:
+                raise RuntimeError(f"Unsupported ScanNet++ depth payload keys at {path}")
+        else:
+            arr = np.load(path)
+        depth = torch.from_numpy(np.asarray(arr, dtype=np.float32).copy())
+        if depth.ndim == 3 and depth.shape[0] == 1:
+            depth = depth[0]
+        if depth.ndim != 2:
+            raise RuntimeError(f"Expected 2D fisheye depth at {path}, got shape={tuple(depth.shape)}")
+        depth = depth.unsqueeze(0)
+        valid = torch.isfinite(depth) & (depth > 0.0)
+        if self.far_depth_invalid_m > 0.0:
+            valid = valid & (depth <= self.far_depth_invalid_m)
+        depth = torch.where(valid, depth, torch.zeros_like(depth))
+        if depth_kind in {"radial", "radius", "dist"}:
+            depth_kind = "distance"
+        if depth_kind not in {"distance", "z"}:
+            raise RuntimeError(f"Unsupported fisheye depth_kind={depth_kind!r} at {path}")
+        return depth.clamp(min=0.0, max=self.depth_max_m), depth_kind
+
+    @staticmethod
+    def _fisheye_z_depth_to_distance(z_depth: torch.Tensor, camera_params: torch.Tensor) -> torch.Tensor:
+        from unisharp.utils.fisheye_geer import build_fisheye624_raymap
+
+        h, w = int(z_depth.shape[-2]), int(z_depth.shape[-1])
+        rays = build_fisheye624_raymap(
+            camera_params.unsqueeze(0),
+            image_h=h,
+            image_w=w,
+            device=z_depth.device,
+            dtype=torch.float32,
+        )
+        ray_z = rays[:, 2:3].squeeze(0).to(device=z_depth.device, dtype=z_depth.dtype)
+        valid = torch.isfinite(z_depth) & (z_depth > 0.0) & torch.isfinite(ray_z) & (ray_z > 1e-4)
+        distance = z_depth / ray_z.clamp(min=1e-4)
+        return torch.where(valid, distance, torch.zeros_like(z_depth))
+
+    def _resolve_image_path(self, scene_dir: Path, image_name: str) -> Path | None:
+        rel = Path(image_name)
+        candidates = [
+            scene_dir / rel,
+            scene_dir / "images" / rel.name,
+            scene_dir / "resized_images" / rel.name,
+            scene_dir / "dslr" / rel,
+            scene_dir / "dslr" / "images" / rel.name,
+            scene_dir / "dslr" / "resized_images" / rel.name,
+        ]
+        for path in candidates:
+            if path.exists() and path.suffix in IMAGE_SUFFIXES:
+                return path
+        return None
+
+    def _resolve_depth_path(self, scene_dir: Path, image_name: str) -> Path | None:
+        stem = Path(image_name).stem
+        names = [stem, Path(image_name).name]
+        bases = [scene_dir, scene_dir / "dslr"]
+        for base in bases:
+            for depth_dir_name in DEPTH_DIR_NAMES:
+                depth_dir = base / depth_dir_name
+                for name in names:
+                    for suffix in (".npz", ".npy"):
+                        path = depth_dir / f"{name}{suffix}"
+                        if path.exists():
+                            return path
+        return None
+
+    def _resolve_mask_path(self, scene_dir: Path, image_name: str, mask_name: str | None) -> Path | None:
+        names = []
+        if mask_name:
+            names.append(Path(mask_name).name)
+        names.append(f"{Path(image_name).stem}.png")
+        bases = [scene_dir, scene_dir / "dslr"]
+        for base in bases:
+            for name in names:
+                direct = base / name
+                if direct.exists():
+                    return direct
+                for mask_dir_name in MASK_DIR_NAMES:
+                    path = base / mask_dir_name / name
+                    if path.exists():
+                        return path
+        return None
+
+    def _load_scene_frames(self, scene_id: str, scene_dir: Path) -> tuple[torch.Tensor, list[dict[str, object]]]:
+        transforms_path = scene_dir / "nerfstudio" / "transforms.json"
+        if not transforms_path.exists():
+            transforms_path = scene_dir / "dslr" / "nerfstudio" / "transforms.json"
+        meta = json.loads(transforms_path.read_text(encoding="utf-8"))
+        camera_params = _opencv_fisheye_to_fisheye624_params(meta)
+        camera_hw = _camera_hw_from_meta(meta)
+        w2c_by_name = _read_colmap_w2c(scene_dir / "colmap" / "images.txt")
+        if not w2c_by_name:
+            w2c_by_name = _read_colmap_w2c(scene_dir / "dslr" / "colmap" / "images.txt")
+
+        raw_frames = list(meta.get("frames", [])) + list(meta.get("test_frames", []))
+        frames: list[dict[str, object]] = []
+        for frame in raw_frames:
+            image_name = Path(str(frame.get("file_path", ""))).name
+            if not image_name:
+                continue
+            if self.skip_bad and bool(frame.get("is_bad", False)):
+                continue
+            image_path = self._resolve_image_path(scene_dir, image_name)
+            depth_path = self._resolve_depth_path(scene_dir, image_name)
+            if image_path is None or depth_path is None:
+                continue
+            w2c = w2c_by_name.get(image_name)
+            if w2c is None and frame.get("transform_matrix") is not None:
+                c2w = torch.tensor(frame["transform_matrix"], dtype=torch.float32)
+                w2c = torch.linalg.inv(c2w)
+            if w2c is None:
+                continue
+            center = torch.linalg.inv(w2c)[:3, 3]
+            frames.append(
+                {
+                    "image_name": image_name,
+                    "image_path": image_path,
+                    "depth_path": depth_path,
+                    "mask_path": self._resolve_mask_path(scene_dir, image_name, frame.get("mask_path")),
+                    "w2c": w2c.to(torch.float32),
+                    "center": center.to(torch.float32),
+                    "idx": len(frames),
+                    "scene": scene_id,
+                    "camera_hw": _camera_hw_from_meta(frame) or camera_hw,
+                }
+            )
+        return camera_params, sorted(frames, key=lambda x: str(x["image_name"]))
+
+    def _load_frame_tensor(self, frame: dict[str, object], camera_params: torch.Tensor) -> dict[str, torch.Tensor]:
+        rgb = self._load_rgb(frame["image_path"])  # type: ignore[arg-type]
+        rgb_hw = (int(rgb.shape[-2]), int(rgb.shape[-1]))
+        camera_hw = frame.get("camera_hw", None)
+        params = camera_params.clone()
+        if isinstance(camera_hw, tuple):
+            params = _scale_fisheye624_params(params, src_hw=camera_hw, dst_hw=rgb_hw)
+        depth, depth_kind = self._load_depth_map(frame["depth_path"])  # type: ignore[arg-type]
+        if tuple(depth.shape[-2:]) != tuple(rgb.shape[-2:]):
+            depth = torch.nn.functional.interpolate(
+                depth.unsqueeze(0),
+                size=(int(rgb.shape[-2]), int(rgb.shape[-1])),
+                mode="nearest",
+            ).squeeze(0)
+        if depth_kind == "z":
+            depth = self._fisheye_z_depth_to_distance(depth, params)
+        valid = (torch.isfinite(depth) & (depth > 0.0)).to(torch.float32)
+        mask_path = frame.get("mask_path", None)
+        if isinstance(mask_path, Path):
+            mask = self._load_mask(mask_path, (int(rgb.shape[-2]), int(rgb.shape[-1])))
+            if mask is not None:
+                valid = valid * mask
+        else:
+            valid = valid * (rgb.to(torch.float32).sum(dim=0, keepdim=True) > 1.0).to(torch.float32)
+        return {
+            "rgb_u8": rgb,
+            "depth_m": depth.clamp(min=0.0, max=self.depth_max_m),
+            "valid_mask": valid,
+            "camera_params": params,
+        }
+
+    def _iter_scene_pairs(self, scene_id: str, scene_dir: Path, rng: random.Random):
+        try:
+            camera_params, frames = self._load_scene_frames(scene_id, scene_dir)
+        except Exception as exc:
+            LOGGER.debug("Skip ScanNet++ scene %s: %s", str(scene_id), str(exc))
+            return
+        if len(frames) < 2:
+            return
+        loaded: dict[int, dict[str, torch.Tensor]] = {}
+
+        def get_loaded(pos: int) -> dict[str, torch.Tensor]:
+            if pos not in loaded:
+                loaded[pos] = self._load_frame_tensor(frames[pos], camera_params)
+            return loaded[pos]
+
+        order = list(range(len(frames)))
+        if self.shuffle_frame:
+            rng.shuffle(order)
+        for src_pos in order:
+            src_item = frames[src_pos]
+            src_center = src_item["center"]
+            assert torch.is_tensor(src_center)
+            candidates: list[int] = []
+            for tgt_pos in range(max(0, src_pos - self.max_frame_gap), min(len(frames), src_pos + self.max_frame_gap + 1)):
+                if tgt_pos == src_pos:
+                    continue
+                gap = abs(tgt_pos - src_pos)
+                if gap < self.min_frame_gap:
+                    continue
+                tgt_center = frames[tgt_pos]["center"]
+                assert torch.is_tensor(tgt_center)
+                if float(torch.norm(tgt_center - src_center, p=2).item()) > self.pair_max_translation_m:
+                    continue
+                candidates.append(tgt_pos)
+            if not candidates:
+                continue
+            tgt_pos = rng.choice(candidates)
+            try:
+                src_loaded = get_loaded(src_pos)
+                tgt_loaded = get_loaded(tgt_pos)
+            except Exception:
+                continue
+            yield ScannetppFisheyePairSample(
+                src_rgb_u8=src_loaded["rgb_u8"],
+                tgt_rgb_u8=tgt_loaded["rgb_u8"],
+                src_depth_m=src_loaded["depth_m"],
+                tgt_depth_m=tgt_loaded["depth_m"],
+                src_valid_mask=src_loaded["valid_mask"],
+                tgt_valid_mask=tgt_loaded["valid_mask"],
+                src_w2c=src_item["w2c"],  # type: ignore[arg-type]
+                tgt_w2c=frames[tgt_pos]["w2c"],  # type: ignore[arg-type]
+                src_camera_params=src_loaded["camera_params"],
+                tgt_camera_params=tgt_loaded["camera_params"],
+                src_idx=int(src_item["idx"]),
+                tgt_idx=int(frames[tgt_pos]["idx"]),
+                scene=str(scene_id),
+            )
+
+    def __iter__(self):
+        worker = torch.utils.data.get_worker_info()
+        worker_id = 0 if worker is None else int(worker.id)
+        num_workers = 1 if worker is None else int(worker.num_workers)
+        rng = random.Random(self.seed + 1009 * self.epoch + 97 * self.ddp_rank + 17 * worker_id)
+        specs = list(self.scene_specs)
+        if self.shuffle_scene:
+            rng.shuffle(specs)
+        specs = specs[self.ddp_rank :: max(self.ddp_world_size, 1)]
+        specs = specs[worker_id :: num_workers]
+        pending: dict[tuple[int, int], list[ScannetppFisheyePairSample]] = {}
+        for scene_id, scene_dir in specs:
+            for sample in self._iter_scene_pairs(scene_id, scene_dir, rng):
+                hw = (int(sample.src_rgb_u8.shape[-2]), int(sample.src_rgb_u8.shape[-1]))
+                bucket = pending.setdefault(hw, [])
+                bucket.append(sample)
+                while len(bucket) >= self.batch_size_hint:
+                    packed = bucket[: self.batch_size_hint]
+                    del bucket[: self.batch_size_hint]
+                    yield _stack_batch(packed)

unisharp/datasets/sim_panorama.py

@@ -0,0 +1,497 @@
+
+from __future__ import annotations
+
+import csv
+from dataclasses import dataclass
+import os
+from pathlib import Path
+import random
+import re
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from torch.utils.data import IterableDataset
+
+from unisharp.datasets.panogs import PanOGSSample
+from unisharp import DEFAULT_MAX_DEPTH_M
+
+try:
+    import h5py
+except ImportError:
+    h5py = None
+
+
+_NUM_RE = re.compile(r"(\d+)(?!.*\d)")
+_SIM_CACHE_VERSION = 6
+
+
+def _default_dataset_manifest_dir() -> Path:
+    repo_root = Path(__file__).resolve().parents[2]
+    parent_path = repo_root.parent / "dataset_manifests"
+    if parent_path.exists():
+        return parent_path
+    return repo_root / "dataset_manifests"
+
+
+def _frame_index_from_name(name: str) -> int | None:
+    match = _NUM_RE.search(Path(name).stem)
+    if match is None:
+        return None
+    return int(match.group(1))
+
+
+def _sim_csv_xyz_to_training_position(x: float, y: float, z: float) -> torch.Tensor:
+    return torch.tensor([float(y), -float(z), float(x)], dtype=torch.float32)
+
+
+class _EquirecToCube:
+    def __init__(self, equ_h: int, equ_w: int, face_w: int) -> None:
+        self.equ_h = int(equ_h)
+        self.equ_w = int(equ_w)
+        self.face_w = int(face_w)
+        self.grid = self._build_grid()
+        rng = torch.linspace(-0.5, 0.5, steps=self.face_w, dtype=torch.float32)
+        xx, yy = torch.meshgrid(rng, -rng, indexing="xy")
+        self.ray_z = (1.0 / torch.sqrt((2.0 * xx) ** 2 + (2.0 * yy) ** 2 + 1.0)).contiguous()
+
+    def _build_grid(self) -> torch.Tensor:
+        face_w = self.face_w
+        rng = torch.linspace(-0.5, 0.5, steps=face_w, dtype=torch.float32)
+        grid = torch.stack(torch.meshgrid(rng, -rng, indexing="xy"), dim=-1)
+        xyz = torch.zeros((6, face_w, face_w, 3), dtype=torch.float32)
+
+        xyz[0, :, :, 0] = grid[:, :, 0]
+        xyz[0, :, :, 1] = grid[:, :, 1]
+        xyz[0, :, :, 2] = 0.5
+
+        xyz[1, :, :, 2] = torch.flip(grid[:, :, 0], dims=[1])
+        xyz[1, :, :, 1] = torch.flip(grid[:, :, 1], dims=[1])
+        xyz[1, :, :, 0] = 0.5
+
+        xyz[2, :, :, 0] = torch.flip(grid[:, :, 0], dims=[1])
+        xyz[2, :, :, 1] = torch.flip(grid[:, :, 1], dims=[1])
+        xyz[2, :, :, 2] = -0.5
+
+        xyz[3, :, :, 2] = grid[:, :, 0]
+        xyz[3, :, :, 1] = grid[:, :, 1]
+        xyz[3, :, :, 0] = -0.5
+
+        xyz[4, :, :, 0] = torch.flip(grid[:, :, 0], dims=[0])
+        xyz[4, :, :, 2] = torch.flip(grid[:, :, 1], dims=[0])
+        xyz[4, :, :, 1] = 0.5
+
+        xyz[5, :, :, 0] = grid[:, :, 0]
+        xyz[5, :, :, 2] = grid[:, :, 1]
+        xyz[5, :, :, 1] = -0.5
+
+        xyz = xyz[[4, 2, 3, 0, 1, 5]]
+        x = xyz[..., 0]
+        y = xyz[..., 1]
+        z = xyz[..., 2]
+        lon = torch.atan2(x, z)
+        c = torch.sqrt(x * x + z * z).clamp(min=1e-8)
+        lat = torch.atan2(y, c)
+        grid_x = lon / np.pi
+        grid_y = (-2.0 * lat / np.pi).clamp(min=-1.0, max=1.0)
+        return torch.stack([grid_x, grid_y], dim=-1).contiguous()
+
+    def run_depth(self, depth_1hw: torch.Tensor) -> torch.Tensor:
+        depth = depth_1hw.unsqueeze(0).to(torch.float32)
+        if tuple(depth.shape[-2:]) != (self.equ_h, self.equ_w):
+            depth = F.interpolate(depth, size=(self.equ_h, self.equ_w), mode="nearest")
+        depth_faces = F.grid_sample(
+            depth.expand(6, -1, -1, -1),
+            self.grid,
+            mode="nearest",
+            padding_mode="border",
+            align_corners=True,
+        )
+        depth_faces = depth_faces[:, 0] * self.ray_z.to(depth_faces.device, depth_faces.dtype)
+        return depth_faces.unsqueeze(-1).to(torch.float32).cpu()
+
+    def run(self, rgb_chw: torch.Tensor, depth_1hw: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        rgb = rgb_chw.unsqueeze(0).to(torch.float32) / 255.0
+        if tuple(rgb.shape[-2:]) != (self.equ_h, self.equ_w):
+            rgb = F.interpolate(rgb, size=(self.equ_h, self.equ_w), mode="bilinear", align_corners=True)
+        rgb_faces = F.grid_sample(
+            rgb.expand(6, -1, -1, -1),
+            self.grid,
+            mode="bilinear",
+            padding_mode="border",
+            align_corners=True,
+        )
+        cube_rgb = (rgb_faces.permute(0, 2, 3, 1).clamp(0.0, 1.0) * 255.0).round().to(torch.uint8)
+        cube_depth = self.run_depth(depth_1hw)
+        return cube_rgb.cpu(), cube_depth
+
+    def run_rgb(self, rgb_chw: torch.Tensor) -> torch.Tensor:
+        rgb = rgb_chw.unsqueeze(0).to(torch.float32) / 255.0
+        if tuple(rgb.shape[-2:]) != (self.equ_h, self.equ_w):
+            rgb = F.interpolate(rgb, size=(self.equ_h, self.equ_w), mode="bilinear", align_corners=True)
+        rgb_faces = F.grid_sample(
+            rgb.expand(6, -1, -1, -1),
+            self.grid,
+            mode="bilinear",
+            padding_mode="border",
+            align_corners=True,
+        )
+        return (rgb_faces.permute(0, 2, 3, 1).clamp(0.0, 1.0) * 255.0).round().to(torch.uint8).cpu()
+
+
+@dataclass(frozen=True)
+class _SimFrame:
+    frame_idx: int
+    rgb_path: Path
+    depth_path: Path
+    position_xyz: torch.Tensor
+
+
+class SimPanoramaDataset(IterableDataset):
+    def __init__(
+        self,
+        root: Path,
+        pose_root: Path,
+        scene_names: list[str] | None = None,
+        scene_list_file: Path | None = None,
+        position_scale: float = 0.01,
+        max_index_gap: int = 10,
+        pair_max_translation_m: float = 0.5,
+        pair_min_depth_overlap: float = 0.6,
+        pair_overlap_margin: float = 1.05,
+        pairs_per_chunk: int = 15,
+        chunk_size: int = 30,
+        shuffle_scene: bool = True,
+        ddp_rank: int = 0,
+        ddp_world_size: int = 1,
+        depth_max_m: float = DEFAULT_MAX_DEPTH_M,
+        far_depth_invalid_m: float = 30.0,
+        far_depth_invalid_max_frac: float = 1.0,
+        max_long_edge: int = 0,
+        seed: int = 0,
+    ) -> None:
+        super().__init__()
+        self.root = Path(root)
+        self.pose_root = Path(pose_root)
+        self.scene_list_file = Path(scene_list_file) if scene_list_file is not None else None
+        requested_scene_names = [str(name).strip() for name in (scene_names or []) if str(name).strip()]
+        if self.scene_list_file is not None:
+            if not self.scene_list_file.exists():
+                raise FileNotFoundError(self.scene_list_file)
+            manifest_scene_names = [
+                line.strip()
+                for line in self.scene_list_file.read_text(encoding="utf-8").splitlines()
+                if line.strip()
+            ]
+            if requested_scene_names:
+                requested = set(requested_scene_names)
+                self.scene_names = [name for name in manifest_scene_names if name in requested]
+            else:
+                self.scene_names = manifest_scene_names
+        else:
+            self.scene_names = requested_scene_names
+        if not self.scene_names:
+            raise ValueError("SimPanoramaDataset requires scene_names or scene_list_file.")
+        self.position_scale = float(position_scale)
+        self.max_index_gap = int(max_index_gap)
+        self.pair_max_translation_m = float(pair_max_translation_m)
+        self.pair_min_depth_overlap = float(pair_min_depth_overlap)
+        self.pair_overlap_margin = float(pair_overlap_margin)
+        self.pairs_per_chunk = int(pairs_per_chunk)
+        self.chunk_size = int(chunk_size)
+        self.shuffle_scene = bool(shuffle_scene)
+        self.ddp_rank = int(ddp_rank)
+        self.ddp_world_size = int(ddp_world_size)
+        self.seed = int(seed)
+        self.depth_max_m = float(depth_max_m)
+        self.far_depth_invalid_m = float(far_depth_invalid_m)
+        self.far_depth_invalid_max_frac = float(far_depth_invalid_max_frac)
+        self.max_long_edge = max(int(max_long_edge), 0)
+        self.epoch = 0
+        self.cache_dir = _default_dataset_manifest_dir() / "sim_cache"
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self._scene_frames_cache: dict[str, list[_SimFrame]] = {}
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = int(epoch)
+
+    @staticmethod
+    def _is_depth_path(path: Path) -> bool:
+        tokens = [part.lower() for part in path.parts]
+        name = path.name.lower()
+        return ("depth" in name) or any("depth" in token for token in tokens)
+
+    @staticmethod
+    def _is_image_path(path: Path) -> bool:
+        return path.suffix.lower() in (".png", ".jpg", ".jpeg", ".webp")
+
+    @staticmethod
+    def _load_rgb(path: Path) -> torch.Tensor:
+        with Image.open(path) as img:
+            img = img.convert("RGB")
+            arr = np.asarray(img, dtype=np.uint8).copy()
+        return torch.from_numpy(arr).permute(2, 0, 1).contiguous()
+
+    @staticmethod
+    def _image_hw(path: Path) -> tuple[int, int]:
+        with Image.open(path) as img:
+            width, height = img.size
+        return int(height), int(width)
+
+    def _load_depth(self, path: Path) -> torch.Tensor:
+        suffix = path.suffix.lower()
+        if suffix == ".npy":
+            dep = np.load(path)
+        elif suffix == ".npz":
+            payload = np.load(path)
+            key = "depth" if "depth" in payload.files else payload.files[0]
+            dep = payload[key]
+        elif suffix in (".h5", ".hdf5"):
+            if h5py is None:
+                raise ImportError("h5py is required to read sim .h5 depth files but is not installed.")
+            with h5py.File(path, "r") as f:
+                keys = list(f.keys())
+                if not keys:
+                    raise RuntimeError(f"Empty sim depth file: {path}")
+                dep = f[keys[0]][()]
+        else:
+            with Image.open(path) as img:
+                dep = np.asarray(img)
+        dep = dep.astype(np.float32)
+        if dep.ndim == 3:
+            dep = dep[..., 0]
+        dep[~np.isfinite(dep)] = 0.0
+        if self.far_depth_invalid_m > 0.0:
+            far = dep > self.far_depth_invalid_m
+            if 0.0 < float(far.mean()) <= self.far_depth_invalid_max_frac:
+                dep[far] = 0.0
+        dep = np.clip(dep, a_min=0.0, a_max=self.depth_max_m)
+        return torch.from_numpy(dep).unsqueeze(0)
+
+    def _resize_erp_if_needed(self, rgb: torch.Tensor, depth: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.max_long_edge <= 0:
+            return rgb, depth
+        h = int(rgb.shape[-2])
+        w = int(rgb.shape[-1])
+        long_edge = max(h, w)
+        if long_edge <= self.max_long_edge:
+            return rgb, depth
+        scale = float(self.max_long_edge) / float(long_edge)
+        new_h = max(2, int(round(float(h) * scale)))
+        new_w = max(2, int(round(float(w) * scale)))
+        rgb_f = rgb.unsqueeze(0).to(dtype=torch.float32)
+        rgb_resized = F.interpolate(rgb_f, size=(new_h, new_w), mode="bilinear", align_corners=False)
+        rgb_out = rgb_resized[0].round().clamp(0.0, 255.0).to(dtype=torch.uint8).contiguous()
+        depth_f = depth.unsqueeze(0).to(dtype=torch.float32)
+        depth_out = F.interpolate(depth_f, size=(new_h, new_w), mode="nearest")[0].contiguous()
+        return rgb_out, depth_out
+
+    def _pose_csv_for_scene(self, scene_name: str) -> Path:
+        direct = self.pose_root / f"{scene_name}.csv"
+        if direct.exists():
+            return direct
+        matches = sorted(self.pose_root.glob(f"*{scene_name}*.csv"))
+        if matches:
+            return matches[0]
+        raise FileNotFoundError(f"No pose csv found for sim scene={scene_name} under {self.pose_root}")
+
+    def _parse_pose_csv(self, csv_path: Path) -> list[tuple[int, torch.Tensor]]:
+        with csv_path.open("r", encoding="utf-8") as f:
+            rows = list(csv.DictReader(f))
+        if not rows:
+            raise RuntimeError(f"Empty sim pose csv: {csv_path}")
+        poses: list[tuple[int, torch.Tensor]] = []
+        for row_idx, row in enumerate(rows):
+            lower = {str(k).strip().lower(): v for k, v in row.items()}
+            frame_val = None
+            for key in ("frame", "frame_idx", "idx", "index", "id", "image", "filename", "name"):
+                if key in lower and str(lower[key]).strip():
+                    frame_val = _frame_index_from_name(str(lower[key]))
+                    if frame_val is None:
+                        try:
+                            frame_val = int(float(str(lower[key]).strip()))
+                        except Exception:
+                            frame_val = None
+                    break
+            x = next((lower[k] for k in lower if k in ("x", "tx", "pos_x", "world_x")), None)
+            y = next((lower[k] for k in lower if k in ("y", "ty", "pos_y", "world_y")), None)
+            z = next((lower[k] for k in lower if k in ("z", "tz", "pos_z", "world_z")), None)
+            if x is None or y is None or z is None:
+                numeric_vals = []
+                for val in row.values():
+                    try:
+                        numeric_vals.append(float(str(val).strip()))
+                    except Exception:
+                        continue
+                if len(numeric_vals) < 3:
+                    raise ValueError(f"Failed to parse xyz from sim csv row: {row}")
+                x, y, z = numeric_vals[:3]
+            pos = _sim_csv_xyz_to_training_position(float(x), float(y), float(z)) * self.position_scale
+            poses.append((int(frame_val if frame_val is not None else row_idx), pos))
+        return poses
+
+    def _scan_scene_frames(self, scene_name: str) -> list[_SimFrame]:
+        scene_dir = self.root / scene_name
+        if not scene_dir.exists():
+            raise FileNotFoundError(scene_dir)
+        all_files = [p for p in scene_dir.rglob("*") if p.is_file()]
+        image_map: dict[int, Path] = {}
+        depth_map: dict[int, Path] = {}
+        for path in all_files:
+            idx = _frame_index_from_name(path.name)
+            if idx is None:
+                continue
+            if self._is_depth_path(path) and path.suffix.lower() in (".png", ".npy", ".npz", ".exr", ".h5", ".hdf5"):
+                depth_map.setdefault(idx, path)
+            elif self._is_image_path(path):
+                image_map.setdefault(idx, path)
+        pose_entries = self._parse_pose_csv(self._pose_csv_for_scene(scene_name))
+        frames: list[_SimFrame] = []
+        for frame_idx, pos in pose_entries:
+            rgb_path = image_map.get(int(frame_idx))
+            depth_path = depth_map.get(int(frame_idx))
+            if rgb_path is None or depth_path is None:
+                continue
+            frames.append(_SimFrame(frame_idx=int(frame_idx), rgb_path=rgb_path, depth_path=depth_path, position_xyz=pos))
+        return frames
+
+    @staticmethod
+    def _atomic_torch_save(path: Path, payload: object) -> None:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        tmp_path = path.with_suffix(path.suffix + f".tmp.{os.getpid()}")
+        torch.save(payload, tmp_path)
+        os.replace(tmp_path, path)
+
+    def _scene_index_cache_path(self, scene_name: str) -> Path:
+        scene_key = scene_name.replace("/", "__")
+        return self.cache_dir / f"{scene_key}_ps{self.position_scale:g}_frames_v{_SIM_CACHE_VERSION}.pt"
+
+    def _load_or_build_scene_frames(self, scene_name: str) -> list[_SimFrame]:
+        cached = self._scene_frames_cache.get(scene_name)
+        if cached is not None:
+            return cached
+        cache_path = self._scene_index_cache_path(scene_name)
+        frames: list[_SimFrame]
+        if cache_path.exists():
+            try:
+                payload = torch.load(cache_path, map_location="cpu")
+                frames = [
+                    _SimFrame(
+                        frame_idx=int(item["frame_idx"]),
+                        rgb_path=Path(str(item["rgb_path"])),
+                        depth_path=Path(str(item["depth_path"])),
+                        position_xyz=torch.tensor(item["position_xyz"], dtype=torch.float32),
+                    )
+                    for item in payload["frames"]
+                ]
+            except Exception:
+                frames = self._scan_scene_frames(scene_name)
+                payload = {
+                    "scene": scene_name,
+                    "frames": [
+                        {
+                            "frame_idx": int(frame.frame_idx),
+                            "rgb_path": str(frame.rgb_path),
+                            "depth_path": str(frame.depth_path),
+                            "position_xyz": frame.position_xyz.tolist(),
+                        }
+                        for frame in frames
+                    ],
+                }
+                self._atomic_torch_save(cache_path, payload)
+        else:
+            frames = self._scan_scene_frames(scene_name)
+            payload = {
+                "scene": scene_name,
+                "frames": [
+                    {
+                        "frame_idx": int(frame.frame_idx),
+                        "rgb_path": str(frame.rgb_path),
+                        "depth_path": str(frame.depth_path),
+                        "position_xyz": frame.position_xyz.tolist(),
+                    }
+                    for frame in frames
+                ],
+            }
+            self._atomic_torch_save(cache_path, payload)
+        self._scene_frames_cache[scene_name] = frames
+        return frames
+
+    def _random_chunk_pairs(self, chunk: list[_SimFrame], rng: random.Random) -> list[tuple[int, int]]:
+        if len(chunk) < self.chunk_size:
+            return []
+        indices = list(range(len(chunk)))
+        rng.shuffle(indices)
+        max_pairs = min(self.pairs_per_chunk, len(indices) // 2)
+        return [(indices[2 * i], indices[2 * i + 1]) for i in range(max_pairs)]
+
+    def __iter__(self):
+        scene_names = list(self.scene_names)
+        order_rng = random.Random(self.seed + self.epoch)
+        if self.shuffle_scene:
+            order_rng.shuffle(scene_names)
+        worker_info = torch.utils.data.get_worker_info()
+        num_workers = worker_info.num_workers if worker_info is not None else 1
+        worker_id = worker_info.id if worker_info is not None else 0
+        total_shards = max(1, self.ddp_world_size * num_workers)
+        shard_id = self.ddp_rank * num_workers + worker_id
+        pair_unit_index = 0
+
+        for scene_order_idx, scene_name in enumerate(scene_names):
+            try:
+                frames = self._load_or_build_scene_frames(scene_name)
+            except Exception:
+                continue
+            if len(frames) < self.chunk_size:
+                continue
+            for start in range(0, len(frames), self.chunk_size):
+                chunk = frames[start : start + self.chunk_size]
+                if len(chunk) < self.chunk_size:
+                    break
+                try:
+                    equ_h, equ_w = self._image_hw(chunk[0].rgb_path)
+                    if self.max_long_edge > 0 and max(equ_h, equ_w) > self.max_long_edge:
+                        scale = float(self.max_long_edge) / float(max(equ_h, equ_w))
+                        equ_h = max(2, int(round(float(equ_h) * scale)))
+                        equ_w = max(2, int(round(float(equ_w) * scale)))
+                    face_w = max(1, equ_h // 2)
+                    converter = _EquirecToCube(equ_h=equ_h, equ_w=equ_w, face_w=face_w)
+                except Exception:
+                    continue
+                def load_frame(local_idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+                    frame = chunk[local_idx]
+                    rgb = self._load_rgb(frame.rgb_path)
+                    depth = self._load_depth(frame.depth_path)
+                    rgb, depth = self._resize_erp_if_needed(rgb, depth)
+                    cube_rgb, cube_depth = converter.run(rgb, depth)
+                    return rgb, depth, cube_rgb, cube_depth
+
+                chunk_rng = random.Random(
+                    self.seed + self.epoch * 1000003 + scene_order_idx * 1009 + start
+                )
+                pairs = self._random_chunk_pairs(chunk, chunk_rng)
+                for src_local, tgt_local in pairs:
+                    if pair_unit_index % total_shards != shard_id:
+                        pair_unit_index += 1
+                        continue
+                    pair_unit_index += 1
+                    src_rgb, src_depth, src_cube_rgb, src_cube_depth = load_frame(src_local)
+                    tgt_rgb, tgt_depth, tgt_cube_rgb, tgt_cube_depth = load_frame(tgt_local)
+                    yield PanOGSSample(
+                        src_erp_rgb_u8=src_rgb,
+                        tgt_erp_rgb_u8=tgt_rgb,
+                        src_erp_depth_m=src_depth,
+                        tgt_erp_depth_m=tgt_depth,
+                        src_cube_rgb_u8=src_cube_rgb,
+                        tgt_cube_rgb_u8=tgt_cube_rgb,
+                        src_cube_depth_m=src_cube_depth,
+                        tgt_cube_depth_m=tgt_cube_depth,
+                        src_R=torch.eye(3, dtype=torch.float32),
+                        src_t=chunk[src_local].position_xyz.clone(),
+                        tgt_R=torch.eye(3, dtype=torch.float32),
+                        tgt_t=chunk[tgt_local].position_xyz.clone(),
+                        src_idx=int(chunk[src_local].frame_idx),
+                        tgt_idx=int(chunk[tgt_local].frame_idx),
+                        scene=str(scene_name),
+                    )

unisharp/datasets/wildrgbd.py

@@ -0,0 +1,352 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import json
+from pathlib import Path
+import random
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from torch.utils.data import IterableDataset
+
+from unisharp.datasets.pair_sampling import project_overlap_ratio, resize_k3_align_corners_false, resize_rgb_u8_chw_high_quality
+from unisharp import DEFAULT_MAX_DEPTH_M
+
+
+@dataclass(frozen=True)
+class WildRGBDPairSample:
+    src_rgb_u8: torch.Tensor
+    tgt_rgb_u8: torch.Tensor
+    src_depth_m: torch.Tensor
+    tgt_depth_m: torch.Tensor
+    src_w2c: torch.Tensor
+    tgt_w2c: torch.Tensor
+    src_intrinsics: torch.Tensor
+    tgt_intrinsics: torch.Tensor
+    src_idx: int
+    tgt_idx: int
+    scene: str
+
+
+class WildRGBDDataset(IterableDataset):
+
+    def __init__(
+        self,
+        root: Path | None = None,
+        scene_list_file: Path | None = None,
+        split: str = "train",
+        min_frame_gap: int = 1,
+        max_frame_gap: int = 32,
+        pair_max_translation_m: float = 0.5,
+        pair_min_overlap: float = 0.6,
+        pair_overlap_sample_h: int = 32,
+        pair_overlap_sample_w: int = 56,
+        pair_max_tries: int = 32,
+        output_h: int | None = None,
+        output_w: int | None = None,
+        shuffle_scene: bool = True,
+        shuffle_frame: bool = True,
+        ddp_rank: int = 0,
+        ddp_world_size: int = 1,
+        roots: list[Path] | None = None,
+        depth_max_m: float = DEFAULT_MAX_DEPTH_M,
+        seed: int = 0,
+        verify_manifest_paths: bool = False,
+    ) -> None:
+        super().__init__()
+        self.root = root
+        self.split = split
+        self.min_frame_gap = int(min_frame_gap)
+        self.max_frame_gap = int(max_frame_gap)
+        self.pair_max_translation_m = float(pair_max_translation_m)
+        self.pair_min_overlap = float(pair_min_overlap)
+        self.pair_overlap_sample_h = int(pair_overlap_sample_h)
+        self.pair_overlap_sample_w = int(pair_overlap_sample_w)
+        self.pair_max_tries = int(pair_max_tries)
+        self.output_h = int(output_h) if output_h is not None else None
+        self.output_w = int(output_w) if output_w is not None else None
+        self.shuffle_scene = bool(shuffle_scene)
+        self.shuffle_frame = bool(shuffle_frame)
+        self.ddp_rank = int(ddp_rank)
+        self.ddp_world_size = int(ddp_world_size)
+        self.depth_max_m = float(depth_max_m)
+        self.seed = int(seed)
+        self.epoch = 0
+        self.verify_manifest_paths = bool(verify_manifest_paths)
+        self.roots = [Path(p) for p in roots] if roots is not None else ([Path(root)] if root is not None else [])
+        if not self.roots:
+            raise ValueError("WildRGBDDataset requires at least one root path.")
+        self.scene_dirs: list[Path] = []
+        self.scene_list_file = Path(scene_list_file) if scene_list_file is not None else None
+        if self.scene_list_file is not None:
+            if not self.scene_list_file.exists():
+                raise FileNotFoundError(self.scene_list_file)
+            for raw in self.scene_list_file.read_text(encoding="utf-8").splitlines():
+                line = raw.strip()
+                if not line:
+                    continue
+                scene_dir = Path(line)
+                if (not self.verify_manifest_paths) or scene_dir.is_dir():
+                    self.scene_dirs.append(scene_dir)
+        else:
+            for ds_root in self.roots:
+                split_dir = ds_root / self.split
+                if not split_dir.exists():
+                    raise FileNotFoundError(split_dir)
+                self.scene_dirs.extend(sorted([p for p in split_dir.iterdir() if p.is_dir()]))
+        if not self.scene_dirs:
+            raise RuntimeError("No scene folders found in the configured WildRGBD roots.")
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = int(epoch)
+
+    @staticmethod
+    def _load_scene_pose_and_k(scene_dir: Path) -> tuple[np.ndarray, dict[int, np.ndarray], torch.Tensor]:
+        metadata_path = scene_dir / "metadata"
+        with metadata_path.open("r", encoding="utf-8") as f:
+            meta = json.load(f)
+        k_raw = np.asarray(meta["K"], dtype=np.float32).reshape(3, 3).T
+        k = torch.from_numpy(k_raw.copy()).to(torch.float32)
+
+        pose_path = scene_dir / "cam_poses.txt"
+        pose_rows = np.genfromtxt(str(pose_path), dtype=np.float32)
+        if pose_rows.ndim == 1:
+            pose_rows = pose_rows[None, :]
+        if pose_rows.shape[1] < 17:
+            raise ValueError(f"Bad cam_poses.txt shape={pose_rows.shape} at {pose_path}")
+        frame_ids = pose_rows[:, 0].astype(np.int64)
+        c2w = pose_rows[:, 1:17].reshape(-1, 4, 4).astype(np.float32)
+        w2c = np.linalg.inv(c2w).astype(np.float32)
+        w2c_map = {int(fid): w2c[i] for i, fid in enumerate(frame_ids.tolist())}
+        return frame_ids, w2c_map, k
+
+    @staticmethod
+    def _collect_frame_ids(folder: Path) -> set[int]:
+        ids: set[int] = set()
+        if not folder.exists():
+            return ids
+        for p in folder.iterdir():
+            if not p.is_file():
+                continue
+            if p.suffix.lower() not in (".png", ".jpg", ".jpeg"):
+                continue
+            try:
+                ids.add(int(p.stem))
+            except ValueError:
+                continue
+        return ids
+
+    @staticmethod
+    def _resolve_img_path(folder: Path, idx: int) -> Path:
+        for ext in (".png", ".jpg", ".jpeg"):
+            p = folder / f"{idx:05d}{ext}"
+            if p.exists():
+                return p
+        raise FileNotFoundError(folder / f"{idx:05d}.png")
+
+    @staticmethod
+    def _load_rgb_u8(path: Path) -> torch.Tensor:
+        img = Image.open(path).convert("RGB")
+        arr = np.asarray(img, dtype=np.uint8).copy()
+        return torch.from_numpy(arr).permute(2, 0, 1).contiguous()
+
+    def _load_depth_m(self, depth_path: Path) -> torch.Tensor:
+        dep = np.asarray(Image.open(depth_path))
+        if dep.ndim != 2:
+            raise ValueError(f"Expected single-channel depth at {depth_path}, got {dep.shape}")
+        depth = dep.astype(np.float32)
+        if float(np.nanmax(depth)) > 200.0:
+            depth = depth / 1000.0
+        depth[~np.isfinite(depth)] = 0.0
+        depth = np.clip(depth, a_min=0.0, a_max=self.depth_max_m)
+        return torch.from_numpy(depth).unsqueeze(0).to(torch.float32)
+
+    @staticmethod
+    def _scene_name(scene_dir: Path) -> str:
+        parent = scene_dir.parent.parent.name if scene_dir.parent.name == "scenes" else scene_dir.parent.name
+        return f"{parent}/{scene_dir.name}"
+
+    def _sample_target_for_src(
+        self,
+        src_idx: int,
+        valid_ids: list[int],
+        w2c_map: dict[int, np.ndarray],
+        intr: torch.Tensor,
+        h: int,
+        w: int,
+        rng: random.Random,
+    ) -> int | None:
+        src_w2c = torch.from_numpy(w2c_map[int(src_idx)]).to(torch.float32)
+        src_center = torch.linalg.inv(src_w2c)[:3, 3]
+        candidates: list[int] = []
+        for j in valid_ids:
+            if int(j) == int(src_idx):
+                continue
+            gap = abs(int(j) - int(src_idx))
+            if gap < self.min_frame_gap or gap > self.max_frame_gap:
+                continue
+            jw2c = torch.from_numpy(w2c_map[int(j)]).to(torch.float32)
+            jcenter = torch.linalg.inv(jw2c)[:3, 3]
+            trans = torch.norm(jcenter - src_center, p=2).item()
+            if trans > self.pair_max_translation_m:
+                continue
+            candidates.append(int(j))
+        if not candidates:
+            return None
+
+        rng.shuffle(candidates)
+        tries = min(self.pair_max_tries, len(candidates))
+        src_k = intr.to(torch.float32)
+        src_w2c_t = src_w2c.to(torch.float32)
+        for j in candidates[:tries]:
+            tgt_w2c_t = torch.from_numpy(w2c_map[int(j)]).to(torch.float32)
+            ov_st = project_overlap_ratio(
+                src_w2c=src_w2c_t,
+                tgt_w2c=tgt_w2c_t,
+                src_k=src_k,
+                tgt_k=src_k,
+                h=h,
+                w=w,
+                sample_h=self.pair_overlap_sample_h,
+                sample_w=self.pair_overlap_sample_w,
+            )
+            ov_ts = project_overlap_ratio(
+                src_w2c=tgt_w2c_t,
+                tgt_w2c=src_w2c_t,
+                src_k=src_k,
+                tgt_k=src_k,
+                h=h,
+                w=w,
+                sample_h=self.pair_overlap_sample_h,
+                sample_w=self.pair_overlap_sample_w,
+            )
+            if 0.5 * (ov_st + ov_ts) >= self.pair_min_overlap:
+                return int(j)
+        return None
+
+    def __iter__(self):
+        scenes = list(self.scene_dirs)
+        order_rng = random.Random(self.seed + self.epoch)
+        if self.shuffle_scene:
+            order_rng.shuffle(scenes)
+
+        worker_info = torch.utils.data.get_worker_info()
+        num_workers = worker_info.num_workers if worker_info is not None else 1
+        worker_id = worker_info.id if worker_info is not None else 0
+        total_shards = max(1, self.ddp_world_size * num_workers)
+        shard_id = self.ddp_rank * num_workers + worker_id
+        src_unit_index = 0
+
+        for scene_order_idx, scene_dir in enumerate(scenes):
+            try:
+                pose_ids_np, w2c_map, intr = self._load_scene_pose_and_k(scene_dir)
+            except Exception:
+                continue
+            pose_ids = {int(x) for x in pose_ids_np.tolist()}
+            rgb_ids = self._collect_frame_ids(scene_dir / "rgb")
+            dep_ids = self._collect_frame_ids(scene_dir / "depth")
+            valid_ids = sorted(list(pose_ids & rgb_ids & dep_ids))
+            if len(valid_ids) < 2:
+                continue
+
+            src_order = list(valid_ids)
+            scene_rng = random.Random(self.seed + self.epoch * 1000003 + scene_order_idx)
+            if self.shuffle_frame:
+                scene_rng.shuffle(src_order)
+
+            for src_idx in src_order:
+                if src_unit_index % total_shards != shard_id:
+                    src_unit_index += 1
+                    continue
+                src_unit_index += 1
+                try:
+                    rgb_src_path = self._resolve_img_path(scene_dir / "rgb", int(src_idx))
+                    dep_src_path = self._resolve_img_path(scene_dir / "depth", int(src_idx))
+                    src_img = self._load_rgb_u8(rgb_src_path)
+                    src_depth = self._load_depth_m(dep_src_path)
+                except Exception:
+                    continue
+
+                h, w = int(src_img.shape[1]), int(src_img.shape[2])
+                tgt_idx = self._sample_target_for_src(
+                    src_idx=int(src_idx),
+                    valid_ids=valid_ids,
+                    w2c_map=w2c_map,
+                    intr=intr,
+                    h=h,
+                    w=w,
+                    rng=scene_rng,
+                )
+                if tgt_idx is None:
+                    continue
+
+                try:
+                    rgb_tgt_path = self._resolve_img_path(scene_dir / "rgb", int(tgt_idx))
+                    dep_tgt_path = self._resolve_img_path(scene_dir / "depth", int(tgt_idx))
+                    tgt_img = self._load_rgb_u8(rgb_tgt_path)
+                    tgt_depth = self._load_depth_m(dep_tgt_path)
+                except Exception:
+                    continue
+
+                if src_img.shape != tgt_img.shape:
+                    continue
+
+                src_intr = intr.clone()
+                tgt_intr = intr.clone()
+                if self.output_h is not None and self.output_w is not None:
+                    oh, ow = int(src_img.shape[1]), int(src_img.shape[2])
+                    if oh > 0 and ow > 0 and (oh != self.output_h or ow != self.output_w):
+                        sx = float(self.output_w) / float(ow)
+                        sy = float(self.output_h) / float(oh)
+                        src_img = resize_rgb_u8_chw_high_quality(src_img, size=(self.output_h, self.output_w))
+                        tgt_img = resize_rgb_u8_chw_high_quality(tgt_img, size=(self.output_h, self.output_w))
+                        src_depth = F.interpolate(
+                            src_depth.unsqueeze(0),
+                            size=(self.output_h, self.output_w),
+                            mode="nearest",
+                        ).squeeze(0)
+                        tgt_depth = F.interpolate(
+                            tgt_depth.unsqueeze(0),
+                            size=(self.output_h, self.output_w),
+                            mode="nearest",
+                        ).squeeze(0)
+                        src_intr = resize_k3_align_corners_false(src_intr, sx=sx, sy=sy)
+                        tgt_intr = resize_k3_align_corners_false(tgt_intr, sx=sx, sy=sy)
+
+                yield WildRGBDPairSample(
+                    src_rgb_u8=src_img,
+                    tgt_rgb_u8=tgt_img,
+                    src_depth_m=src_depth,
+                    tgt_depth_m=tgt_depth,
+                    src_w2c=torch.from_numpy(w2c_map[int(src_idx)]).to(torch.float32),
+                    tgt_w2c=torch.from_numpy(w2c_map[int(tgt_idx)]).to(torch.float32),
+                    src_intrinsics=src_intr,
+                    tgt_intrinsics=tgt_intr,
+                    src_idx=int(src_idx),
+                    tgt_idx=int(tgt_idx),
+                    scene=self._scene_name(scene_dir),
+                )
+
+
+def wildrgbd_collate(batch: list[WildRGBDPairSample]) -> WildRGBDPairSample:
+    def stack(xs):
+        if isinstance(xs[0], torch.Tensor):
+            return torch.stack(xs, dim=0)
+        return xs
+
+    return WildRGBDPairSample(
+        src_rgb_u8=stack([b.src_rgb_u8 for b in batch]),
+        tgt_rgb_u8=stack([b.tgt_rgb_u8 for b in batch]),
+        src_depth_m=stack([b.src_depth_m for b in batch]),
+        tgt_depth_m=stack([b.tgt_depth_m for b in batch]),
+        src_w2c=stack([b.src_w2c for b in batch]),
+        tgt_w2c=stack([b.tgt_w2c for b in batch]),
+        src_intrinsics=stack([b.src_intrinsics for b in batch]),
+        tgt_intrinsics=stack([b.tgt_intrinsics for b in batch]),
+        src_idx=[b.src_idx for b in batch],  # type: ignore[arg-type]
+        tgt_idx=[b.tgt_idx for b in batch],  # type: ignore[arg-type]
+        scene=[b.scene for b in batch],  # type: ignore[arg-type]
+    )
+

unisharp/losses/__init__.py

@@ -0,0 +1,4 @@
+from .unisharp_loss import UnisharpLoss, UnisharpLossWeights
+
+__all__ = ["UnisharpLoss", "UnisharpLossWeights"]
+

unisharp/losses/unisharp_loss.py

@@ -0,0 +1,1120 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+from dataclasses import dataclass
+import math
+import torch.nn.functional as F
+
+from unisharp import DEFAULT_MAX_DEPTH_M
+from unisharp.utils import linalg
+
+
+def _masked_mean(x: torch.Tensor, m: torch.Tensor) -> torch.Tensor:
+    if m.dtype != x.dtype:
+        m = m.to(dtype=x.dtype)
+    while m.ndim < x.ndim:
+        m = m.unsqueeze(1)
+    m_expanded = m.expand_as(x)
+    return (x * m_expanded).sum() / m_expanded.sum().clamp(min=1.0)
+
+
+def _finite_masked_mean_flat(x: torch.Tensor, valid: torch.Tensor) -> torch.Tensor:
+    mask = valid.to(device=x.device, dtype=torch.bool) & torch.isfinite(x)
+    x_safe = torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)
+    safe = torch.where(mask, x_safe, torch.zeros_like(x_safe))
+    return safe.sum() / mask.to(dtype=x.dtype).sum().clamp(min=1.0)
+
+
+def _finite_abs_mean(x: torch.Tensor) -> torch.Tensor:
+    mask = torch.isfinite(x)
+    x_safe = torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)
+    safe_abs = torch.where(mask, x_safe.abs(), torch.zeros_like(x_safe))
+    return safe_abs.sum() / mask.to(dtype=x.dtype).sum().clamp(min=1.0)
+
+
+_ERP_PROJECTION_MODELS = {"erp", "spherical", "equirect", "equirectangular"}
+_FISHEYE_PROJECTION_MODELS = {"fisheye624", "opencv_fisheye"}
+
+
+def _tv_l1(img: torch.Tensor) -> torch.Tensor:
+    zero = torch.zeros((), device=img.device, dtype=img.dtype)
+    dx = (img[..., :, 1:] - img[..., :, :-1]).abs().mean() if int(img.shape[-1]) > 1 else zero
+    dy = (img[..., 1:, :] - img[..., :-1, :]).abs().mean() if int(img.shape[-2]) > 1 else zero
+    return dx + dy
+
+
+def _tv_l1_circular_h(img: torch.Tensor) -> torch.Tensor:
+    zero = torch.zeros((), device=img.device, dtype=img.dtype)
+    dx = (torch.roll(img, shifts=-1, dims=-1) - img).abs().mean() if int(img.shape[-1]) > 1 else zero
+    dy = (img[..., 1:, :] - img[..., :-1, :]).abs().mean() if int(img.shape[-2]) > 1 else zero
+    return dx + dy
+
+
+def _checkerboard_l1_5d(x: torch.Tensor, *, circular_h: bool) -> torch.Tensor:
+    if x.ndim != 5:
+        raise ValueError(f"Expected [B,C,L,H,W], got {tuple(x.shape)}")
+    if int(x.shape[-2]) < 2 or int(x.shape[-1]) < 2:
+        return torch.zeros((), device=x.device, dtype=x.dtype)
+    x = x.to(dtype=torch.float32)
+    if bool(circular_h):
+        top = x[..., :-1, :]
+        bottom = x[..., 1:, :]
+        response = top - torch.roll(top, shifts=-1, dims=-1) - bottom + torch.roll(bottom, shifts=-1, dims=-1)
+    else:
+        response = x[..., :-1, :-1] - x[..., :-1, 1:] - x[..., 1:, :-1] + x[..., 1:, 1:]
+    return _finite_abs_mean(response)
+
+
+def _delta_grid_checkerboard_loss(delta_grid: torch.Tensor, *, circular_h: bool) -> torch.Tensor:
+    if delta_grid.ndim != 5 or int(delta_grid.shape[1]) < 14:
+        raise ValueError(f"Expected delta grid [B,14,L,H,W], got {tuple(delta_grid.shape)}")
+    delta = delta_grid.to(dtype=torch.float32)
+    parts = [
+        delta[:, 3:6],
+        0.1 * delta[:, 10:13],
+        delta[:, 13:14],
+    ]
+    return torch.stack([_checkerboard_l1_5d(part, circular_h=circular_h) for part in parts]).mean()
+
+
+def _avg_pool2d_circular_h(x: torch.Tensor, kernel_size: int, stride: int) -> torch.Tensor:
+    if kernel_size <= 1 and stride <= 1:
+        return x
+    x = F.pad(x, (kernel_size - 1, 0, 0, 0), mode="circular")
+    return F.avg_pool2d(x, kernel_size=kernel_size, stride=stride)
+
+
+def _resize_max_side(img: torch.Tensor, max_side: int, *, mode: str = "bilinear") -> torch.Tensor:
+    if max_side <= 0:
+        return img
+    h, w = int(img.shape[-2]), int(img.shape[-1])
+    ms = max(h, w)
+    if ms <= max_side:
+        return img
+    scale = float(max_side) / float(ms)
+    nh = max(1, int(math.floor(h * scale)))
+    nw = max(1, int(math.floor(w * scale)))
+    if mode in ("bilinear", "bicubic"):
+        return F.interpolate(img, size=(nh, nw), mode=mode, align_corners=False)
+    return F.interpolate(img, size=(nh, nw), mode=mode)
+
+
+def _gram_matrix(fmap: torch.Tensor) -> torch.Tensor:
+    b, c, h, w = fmap.shape
+    x = fmap.reshape(b, c, h * w)
+    return x @ x.transpose(1, 2)
+
+
+class _ResNet50Perceptual(nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+        try:
+            from torchvision.models import resnet50, ResNet50_Weights
+
+            net = resnet50(weights=ResNet50_Weights.DEFAULT)
+        except Exception:
+            from torchvision.models import resnet50
+
+            net = resnet50(pretrained=True)
+
+        net.eval()
+        net.requires_grad_(False)
+
+        self.conv1 = net.conv1
+        self.bn1 = net.bn1
+        self.relu = net.relu
+        self.maxpool = net.maxpool
+        self.layer1 = net.layer1
+        self.layer2 = net.layer2
+        self.layer3 = net.layer3
+        self.layer4 = net.layer4
+
+        mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
+        std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
+        self.register_buffer("_mean", mean, persistent=False)
+        self.register_buffer("_std", std, persistent=False)
+
+    def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
+        x = x.clamp(0.0, 1.0)
+        x = (x - self._mean) / self._std
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+        f1 = self.layer1(x)
+        f2 = self.layer2(f1)
+        f3 = self.layer3(f2)
+        f4 = self.layer4(f3)
+        return [f1, f2, f3, f4]
+
+
+def _to_linear_rgb(img_srgb: torch.Tensor) -> torch.Tensor:
+    from unisharp.utils.color_space import sRGB2linearRGB
+
+    return sRGB2linearRGB(img_srgb.clamp(0.0, 1.0))
+
+
+@dataclass
+class UnisharpLossWeights:
+
+    lambda_color: float = 1.0
+    lambda_alpha: float = 1.5
+    lambda_percep: float = 3.0
+    lambda_depth: float = 0.5
+    lambda_tv: float = 1.0
+    lambda_grad: float = 1.0
+    lambda_delta: float = 0.0
+    lambda_delta_rho: float = 0.0
+    lambda_splat: float = 0.0
+    lambda_edge_splat: float = 0.0
+    lambda_grid: float = 0.0
+    lambda_grad_img: float = 0.2
+    lambda_edge_rgb: float = 0.0
+
+
+class UnisharpLoss(nn.Module):
+
+    SUPERVISION_MAX_DEPTH_M: float = DEFAULT_MAX_DEPTH_M
+
+    def __init__(
+        self,
+        weights: UnisharpLossWeights | None = None,
+        *,
+        grad_sigma: float = 1e-2,
+        grad_eps: float = 1e-2,
+        delta_clip: float = 10.0,
+        raw_delta_clip: float = 400.0,
+        raw_delta_rho_clip: float = 5.0,
+        alpha_tail_min: float = 0.99,
+        alpha_tail_weight: float = 0.0,
+        splat_sigma_min: float = 1e-1,
+        splat_sigma_max: float = 1e2,
+        edge_splat_sigma_max: float = 2.0,
+        depth_edge_log_threshold: float = 0.05,
+        depth_edge_dilate_px: int = 2,
+        percep_max_side: int = 384,
+        grad_img_scales: int = 4,
+        grad_img_circular_h: bool = True,
+    ) -> None:
+        super().__init__()
+        self.w = weights or UnisharpLossWeights()
+        self.grad_sigma = float(grad_sigma)
+        self.grad_eps = float(grad_eps)
+        self.delta_clip = float(delta_clip)
+        self.raw_delta_clip = float(raw_delta_clip)
+        self.raw_delta_rho_clip = float(raw_delta_rho_clip)
+        self.alpha_tail_min = float(alpha_tail_min)
+        self.alpha_tail_weight = float(alpha_tail_weight)
+        self.splat_sigma_min = float(splat_sigma_min)
+        self.splat_sigma_max = float(splat_sigma_max)
+        self.edge_splat_sigma_max = float(edge_splat_sigma_max)
+        self.depth_edge_log_threshold = float(depth_edge_log_threshold)
+        self.depth_edge_dilate_px = int(depth_edge_dilate_px)
+        self.percep_max_side = int(percep_max_side)
+        self.grad_img_scales = int(grad_img_scales)
+        self.grad_img_circular_h = bool(grad_img_circular_h)
+
+        sobel_kx = torch.tensor(
+            [[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]]
+        ).view(1, 1, 3, 3)
+        sobel_ky = torch.tensor(
+            [[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0]]
+        ).view(1, 1, 3, 3)
+        self.register_buffer("_sobel_kx", sobel_kx, persistent=False)
+        self.register_buffer("_sobel_ky", sobel_ky, persistent=False)
+
+        self._percep_net: nn.Module | None = None
+        if self.w.lambda_percep > 0:
+            self._percep_net = _ResNet50Perceptual()
+
+    @staticmethod
+    def _flatten_gaussian_xyz(x: torch.Tensor | None, gauss_grid_shape: tuple[int, int, int] | None = None) -> torch.Tensor | None:
+        if not torch.is_tensor(x):
+            return None
+        if x.ndim == 5:
+            return x.permute(0, 2, 3, 4, 1).flatten(1, 3)
+        if x.ndim == 3 and int(x.shape[-1]) == 3:
+            return x
+        if x.ndim == 2 and gauss_grid_shape is not None:
+            return x.unsqueeze(-1)
+        return None
+
+    @staticmethod
+    def _flatten_gaussian_quat(
+        x: torch.Tensor | None,
+        gauss_grid_shape: tuple[int, int, int] | None = None,
+    ) -> torch.Tensor | None:
+        if not torch.is_tensor(x):
+            return None
+        if x.ndim == 5 and int(x.shape[1]) == 4:
+            return x.permute(0, 2, 3, 4, 1).flatten(1, 3)
+        if x.ndim == 3 and int(x.shape[-1]) == 4:
+            return x
+        if x.ndim == 2 and gauss_grid_shape is not None:
+            return x.unsqueeze(-1)
+        return None
+
+    @staticmethod
+    def _flatten_gaussian_scalar(
+        x: torch.Tensor | None,
+        gauss_grid_shape: tuple[int, int, int] | None = None,
+    ) -> torch.Tensor | None:
+        if not torch.is_tensor(x):
+            return None
+        if x.ndim == 5:
+            return x[:, 0].flatten(1)
+        if x.ndim == 4:
+            return x.flatten(1)
+        if x.ndim == 3 and int(x.shape[-1]) == 1:
+            return x[..., 0]
+        if x.ndim == 2:
+            return x
+        return None
+
+    @staticmethod
+    def _central_disparity_gradient(inv_depth: torch.Tensor, *, circular_h: bool) -> torch.Tensor:
+        if circular_h:
+            gx = 0.5 * (torch.roll(inv_depth, shifts=-1, dims=-1) - torch.roll(inv_depth, shifts=1, dims=-1)).abs()
+        else:
+            padded_x = F.pad(inv_depth, (1, 1, 0, 0), mode="replicate")
+            gx = 0.5 * (padded_x[..., 2:] - padded_x[..., :-2]).abs()
+        padded_y = F.pad(inv_depth, (0, 0, 1, 1), mode="replicate")
+        gy = 0.5 * (padded_y[..., 2:, :] - padded_y[..., :-2, :]).abs()
+        return torch.sqrt(gx * gx + gy * gy + 1e-12)
+
+    @staticmethod
+    def _sample_map_at_uv(feat: torch.Tensor, u: torch.Tensor, v: torch.Tensor, valid: torch.Tensor) -> torch.Tensor:
+        b, _, h, w = feat.shape
+        valid_bool = valid.to(dtype=torch.bool) & torch.isfinite(u) & torch.isfinite(v)
+        u_safe = torch.where(valid_bool, u, torch.zeros_like(u)).clamp(0.0, float(max(w - 1, 0)))
+        v_safe = torch.where(valid_bool, v, torch.zeros_like(v)).clamp(0.0, float(max(h - 1, 0)))
+        grid_x = (u_safe / max(float(w - 1), 1.0)) * 2.0 - 1.0
+        grid_y = (v_safe / max(float(h - 1), 1.0)) * 2.0 - 1.0
+        grid = torch.stack([grid_x, grid_y], dim=-1).view(b, -1, 1, 2)
+        sampled = F.grid_sample(feat, grid, mode="bilinear", padding_mode="zeros", align_corners=True)
+        return sampled[:, 0, :, 0] * valid_bool.to(dtype=feat.dtype)
+
+    @staticmethod
+    def _expand_camera_params(camera_params: torch.Tensor, *, batch_size: int, device: torch.device) -> torch.Tensor:
+        params = camera_params.to(device=device, dtype=torch.float32)
+        if params.ndim == 1:
+            params = params.unsqueeze(0)
+        if int(params.shape[0]) == 1 and int(batch_size) > 1:
+            params = params.expand(int(batch_size), -1)
+        return params
+
+    @staticmethod
+    def _project_fisheye624_points_px_stable(
+        pts: torch.Tensor,
+        camera_params: torch.Tensor,
+        *,
+        image_h: int,
+        image_w: int,
+        finite: torch.Tensor,
+        require_in_bounds: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        b, n, _ = pts.shape
+        params = UnisharpLoss._expand_camera_params(camera_params, batch_size=b, device=pts.device)
+        x, y, z = pts.unbind(dim=-1)
+        radius = torch.linalg.vector_norm(pts, dim=-1).clamp(min=1e-6)
+        front = z > (radius * 1e-4).clamp(min=1e-4)
+        projectable = finite & front
+
+        safe_pts = torch.zeros_like(pts)
+        safe_pts[..., 2] = 1.0
+        pts_proj = torch.where(projectable.unsqueeze(-1), pts, safe_pts)
+        x, y, z = pts_proj.unbind(dim=-1)
+        z_safe = z.clamp(min=1e-4)
+
+        ab = torch.stack([x / z_safe, y / z_safe], dim=-1)
+        r = torch.sqrt((ab * ab).sum(dim=-1, keepdim=True) + 1e-12)
+        theta = torch.atan(r)
+        unit_ab = ab / r
+
+        coeffs = params[:, 4:10].reshape(b, 1, 6)
+        theta_powers = torch.cat([theta.pow(3 + i * 2) for i in range(6)], dim=-1)
+        theta_distorted = theta + (theta_powers * coeffs).sum(dim=-1, keepdim=True)
+        uv_dist = theta_distorted * unit_ab
+
+        p0 = params[..., -6].reshape(b, 1)
+        p1 = params[..., -5].reshape(b, 1)
+        xr = uv_dist[..., 0]
+        yr = uv_dist[..., 1]
+        xr_sq = xr.square()
+        yr_sq = yr.square()
+        rd_sq = xr_sq + yr_sq
+        uv_x = uv_dist[..., 0] + (2.0 * xr_sq + rd_sq) * p0 + 2.0 * xr * yr * p1
+        uv_y = uv_dist[..., 1] + (2.0 * yr_sq + rd_sq) * p1 + 2.0 * xr * yr * p0
+
+        s0 = params[..., -4].reshape(b, 1)
+        s1 = params[..., -3].reshape(b, 1)
+        s2 = params[..., -2].reshape(b, 1)
+        s3 = params[..., -1].reshape(b, 1)
+        rd_4 = rd_sq.square()
+        uv_x = uv_x + s0 * rd_sq + s1 * rd_4
+        uv_y = uv_y + s2 * rd_sq + s3 * rd_4
+
+        if int(params.shape[-1]) == 15:
+            fx = fy = params[..., 0:1]
+            cx = params[..., 1:2]
+            cy = params[..., 2:3]
+        else:
+            fx = params[..., 0:1]
+            fy = params[..., 1:2]
+            cx = params[..., 2:3]
+            cy = params[..., 3:4]
+        u = uv_x * fx + cx
+        v = uv_y * fy + cy
+        valid = projectable & torch.isfinite(u) & torch.isfinite(v)
+        if require_in_bounds:
+            valid = valid & (u >= 0.0) & (u <= float(image_w - 1)) & (v >= 0.0) & (v <= float(image_h - 1))
+        return u, v, valid, radius
+
+    @staticmethod
+    def _project_points_px(
+        points: torch.Tensor,
+        *,
+        projection_model: str | None,
+        image_h: int,
+        image_w: int,
+        intrinsics: torch.Tensor | None = None,
+        camera_params: torch.Tensor | None = None,
+        require_in_bounds: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        pts_raw = points.to(dtype=torch.float32)
+        finite = torch.isfinite(pts_raw).all(dim=-1)
+        pts = torch.nan_to_num(pts_raw, nan=0.0, posinf=0.0, neginf=0.0)
+        b, n, _ = pts.shape
+        x, y, z = pts.unbind(dim=-1)
+        model = (projection_model or "pinhole").lower()
+
+        if model in _ERP_PROJECTION_MODELS:
+            radius_sq_raw = (pts * pts).sum(dim=-1)
+            direction_valid = finite & (radius_sq_raw > 1e-12)
+            safe_pts = torch.zeros_like(pts)
+            safe_pts[..., 2] = 1.0
+            pts_erp = torch.where(direction_valid.unsqueeze(-1), pts, safe_pts)
+            x, y, z = pts_erp.unbind(dim=-1)
+            radius_sq = (pts_erp * pts_erp).sum(dim=-1)
+            radius = torch.sqrt(radius_sq + 1e-12)
+            horizontal_sq = x.square() + z.square()
+            horizontal = torch.sqrt(horizontal_sq + 1e-12)
+            pole_angle_eps = max(1e-4, 0.5 * math.pi / float(max(image_h, image_w, 1)))
+            lon_valid = horizontal > radius * pole_angle_eps
+            lon_x = torch.where(lon_valid, x, torch.zeros_like(x))
+            lon_z = torch.where(lon_valid, z, torch.ones_like(z))
+            lon = torch.atan2(lon_x, lon_z)
+            pitch_down = torch.atan2(y, horizontal)
+            u = (lon / (2.0 * math.pi) + 0.5) * float(max(image_w, 1)) - 0.5
+            v = (0.5 + pitch_down / math.pi) * float(max(image_h, 1)) - 0.5
+            valid = direction_valid & lon_valid
+            valid = (
+                valid
+                & torch.isfinite(u)
+                & torch.isfinite(v)
+                & (u >= 0.0)
+                & (u <= float(image_w - 1))
+                & (v >= 0.0)
+                & (v <= float(image_h - 1))
+            )
+            return u, v, valid, radius.clamp(min=1e-6)
+
+        if model in _FISHEYE_PROJECTION_MODELS and torch.is_tensor(camera_params):
+            return UnisharpLoss._project_fisheye624_points_px_stable(
+                pts,
+                camera_params,
+                image_h=image_h,
+                image_w=image_w,
+                finite=finite,
+                require_in_bounds=require_in_bounds,
+            )
+
+        valid = finite & (z > 1e-4)
+        if not torch.is_tensor(intrinsics):
+            fx = torch.full((b, 1), float(max(image_w, image_h)), device=pts.device, dtype=torch.float32)
+            fy = fx.clone()
+            cx = torch.full((b, 1), 0.5 * float(max(image_w - 1, 1)), device=pts.device, dtype=torch.float32)
+            cy = torch.full((b, 1), 0.5 * float(max(image_h - 1, 1)), device=pts.device, dtype=torch.float32)
+        else:
+            k = intrinsics.to(device=pts.device, dtype=torch.float32)
+            if k.ndim == 2:
+                k = k.unsqueeze(0)
+            if int(k.shape[0]) == 1 and b > 1:
+                k = k.expand(b, -1, -1)
+            fx = k[:, 0, 0:1]
+            fy = k[:, 1, 1:2]
+            cx = k[:, 0, 2:3]
+            cy = k[:, 1, 2:3]
+        z_safe = z.clamp(min=1e-4)
+        u = fx * (x / z_safe) + cx
+        v = fy * (y / z_safe) + cy
+        valid = valid & torch.isfinite(u) & torch.isfinite(v)
+        if require_in_bounds:
+            valid = valid & (u >= 0.0) & (u <= float(image_w - 1)) & (v >= 0.0) & (v <= float(image_h - 1))
+        return u, v, valid, z_safe
+
+    def _projected_sigma_px(
+        self,
+        *,
+        gaussian_scales: torch.Tensor,
+        gaussian_quaternions: torch.Tensor | None,
+        gaussian_mean_vectors: torch.Tensor,
+        valid: torch.Tensor,
+        projection_model: str | None,
+        image_h: int,
+        image_w: int,
+        intrinsics: torch.Tensor | None = None,
+        camera_params: torch.Tensor | None = None,
+        projected_scale_factor: float | torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        scales = self._flatten_gaussian_xyz(gaussian_scales)
+        quats = self._flatten_gaussian_quat(gaussian_quaternions)
+        means = self._flatten_gaussian_xyz(gaussian_mean_vectors)
+        if scales is None or means is None:
+            return torch.zeros_like(valid, dtype=torch.float32)
+        valid = valid.to(dtype=torch.bool) & torch.isfinite(scales).all(dim=-1) & torch.isfinite(means).all(dim=-1)
+        scales = torch.nan_to_num(scales.to(dtype=torch.float32), nan=0.0, posinf=0.0, neginf=0.0).abs()
+        means = torch.nan_to_num(means.to(dtype=torch.float32), nan=0.0, posinf=0.0, neginf=0.0)
+        model = (projection_model or "pinhole").lower()
+        if model in _ERP_PROJECTION_MODELS:
+            radius = torch.norm(means, dim=-1).clamp(min=1e-4)
+            sigma_u = scales[..., 0] / radius * (float(max(image_w, 1)) / (2.0 * math.pi))
+            sigma_v = scales[..., 1] / radius * (float(max(image_h, 1)) / math.pi)
+            sigma_px = torch.maximum(sigma_u.square(), sigma_v.square())
+            valid = valid & torch.isfinite(sigma_px)
+            sigma_px = torch.nan_to_num(sigma_px, nan=0.0, posinf=0.0, neginf=0.0)
+            return torch.where(valid, sigma_px, torch.zeros_like(sigma_px))
+
+        if quats is not None and tuple(quats.shape[:2]) == tuple(means.shape[:2]):
+            quats = torch.nan_to_num(quats.to(dtype=torch.float32), nan=0.0, posinf=0.0, neginf=0.0)
+            quat_norm = quats.norm(dim=-1, keepdim=True)
+            valid = valid & torch.isfinite(quats).all(dim=-1) & (quat_norm.squeeze(-1) > 1e-8)
+            quats = quats / quat_norm.clamp(min=1e-8)
+            rotations = linalg.rotation_matrices_from_quaternions(quats)
+            tangent_scales = scales[..., :2]
+            tangent_rotations = rotations[..., :, :2]
+            axis_offsets = (tangent_rotations * tangent_scales[..., None, :]).transpose(-1, -2)
+            axis_points = means[:, :, None, :] + axis_offsets
+            u0, v0, valid0, _ = self._project_points_px(
+                means,
+                projection_model=projection_model,
+                image_h=image_h,
+                image_w=image_w,
+                intrinsics=intrinsics,
+                camera_params=camera_params,
+                require_in_bounds=False,
+            )
+            b, n, axis_count, _ = axis_points.shape
+            u1, v1, valid1, _ = self._project_points_px(
+                axis_points.reshape(b, n * axis_count, 3),
+                projection_model=projection_model,
+                image_h=image_h,
+                image_w=image_w,
+                intrinsics=intrinsics,
+                camera_params=camera_params,
+                require_in_bounds=False,
+            )
+            u1 = u1.reshape(b, n, axis_count)
+            v1 = v1.reshape(b, n, axis_count)
+            valid1 = valid1.reshape(b, n, axis_count)
+            du = u1 - u0[..., None]
+            dv = v1 - v0[..., None]
+            if (projection_model or "pinhole").lower() in _ERP_PROJECTION_MODELS:
+                width = float(max(image_w, 1))
+                du = torch.remainder(du + 0.5 * width, width) - 0.5 * width
+            cov_xx = (du * du).sum(dim=-1)
+            cov_xy = (du * dv).sum(dim=-1)
+            cov_yy = (dv * dv).sum(dim=-1)
+            trace = cov_xx + cov_yy
+            disc = (cov_xx - cov_yy).square() + 4.0 * cov_xy.square()
+            sigma_px = 0.5 * (trace + (disc.clamp(min=0.0) + 1e-12).sqrt())
+            valid = valid & valid0 & valid1.all(dim=-1) & torch.isfinite(sigma_px)
+            sigma_px = torch.nan_to_num(sigma_px, nan=0.0, posinf=0.0, neginf=0.0)
+            return torch.where(valid, sigma_px, torch.zeros_like(sigma_px))
+
+        sigma_screen_3d = scales[..., :2].to(dtype=torch.float32).abs().amax(dim=-1).clamp(min=1e-8)
+        if model in {"fisheye624", "opencv_fisheye"} and torch.is_tensor(camera_params):
+            params = camera_params.to(device=means.device, dtype=torch.float32)
+            if params.ndim == 1:
+                params = params.unsqueeze(0)
+            if int(params.shape[0]) == 1 and int(means.shape[0]) > 1:
+                params = params.expand(int(means.shape[0]), -1)
+            if int(params.shape[-1]) == 15:
+                focal = params[:, 0:1].clamp(min=1.0)
+            else:
+                focal = 0.5 * (params[:, 0:1] + params[:, 1:2]).clamp(min=1.0)
+            radius = torch.norm(means, dim=-1).clamp(min=1e-4)
+            sigma_px = (sigma_screen_3d / radius * focal).square()
+        elif torch.is_tensor(intrinsics):
+            k = intrinsics.to(device=means.device, dtype=torch.float32)
+            if k.ndim == 2:
+                k = k.unsqueeze(0)
+            if int(k.shape[0]) == 1 and int(means.shape[0]) > 1:
+                k = k.expand(int(means.shape[0]), -1, -1)
+            focal = 0.5 * (k[:, 0, 0:1] + k[:, 1, 1:2]).clamp(min=1.0)
+            depth = means[..., 2].clamp(min=1e-4)
+            sigma_px = (sigma_screen_3d / depth * focal).square()
+        else:
+            depth = torch.norm(means, dim=-1).clamp(min=1e-4)
+            sigma_px = sigma_screen_3d / depth
+            if torch.is_tensor(projected_scale_factor):
+                sigma_px = sigma_px * projected_scale_factor.to(device=sigma_px.device, dtype=sigma_px.dtype)
+            elif projected_scale_factor is not None:
+                sigma_px = sigma_px * float(projected_scale_factor)
+            sigma_px = sigma_px.square()
+        valid = valid & torch.isfinite(sigma_px)
+        sigma_px = torch.nan_to_num(sigma_px, nan=0.0, posinf=0.0, neginf=0.0)
+        return torch.where(valid, sigma_px, torch.zeros_like(sigma_px))
+
+    def _depth_edge_band(
+        self,
+        depth_m: torch.Tensor,
+        valid_weight: torch.Tensor,
+        *,
+        circular_h: bool,
+    ) -> torch.Tensor:
+        depth = depth_m.to(dtype=torch.float32)
+        if depth.ndim == 3:
+            depth = depth.unsqueeze(1)
+        valid = torch.isfinite(depth) & (depth > 0.0) & (valid_weight[:, :1].to(dtype=torch.float32) > 0.5)
+        log_depth = torch.where(valid, depth.clamp(min=1e-4).log(), torch.zeros_like(depth))
+
+        if bool(circular_h):
+            right = torch.roll(log_depth, shifts=-1, dims=-1)
+            valid_right = valid & torch.roll(valid, shifts=-1, dims=-1)
+            edge_x = (right - log_depth).abs() > float(self.depth_edge_log_threshold)
+            edge_x = edge_x & valid_right
+        else:
+            edge_x = torch.zeros_like(valid)
+            edge_x[..., :, :-1] = (
+                (log_depth[..., :, 1:] - log_depth[..., :, :-1]).abs() > float(self.depth_edge_log_threshold)
+            ) & valid[..., :, 1:] & valid[..., :, :-1]
+
+        edge_y = torch.zeros_like(valid)
+        edge_y[..., :-1, :] = (
+            (log_depth[..., 1:, :] - log_depth[..., :-1, :]).abs() > float(self.depth_edge_log_threshold)
+        ) & valid[..., 1:, :] & valid[..., :-1, :]
+        edge = (edge_x | edge_y).to(dtype=torch.float32)
+
+        radius = max(int(self.depth_edge_dilate_px), 0)
+        if radius <= 0:
+            return edge
+        kernel = 2 * radius + 1
+        if bool(circular_h):
+            edge = F.pad(edge, (radius, radius, 0, 0), mode="circular")
+            edge = F.pad(edge, (0, 0, radius, radius), mode="constant", value=0.0)
+            return F.max_pool2d(edge, kernel_size=kernel, stride=1)
+        return F.max_pool2d(edge, kernel_size=kernel, stride=1, padding=radius)
+
+    def _ray_cell_sigma(
+        self,
+        *,
+        gaussian_scales: torch.Tensor,
+        gaussian_mean_vectors: torch.Tensor,
+        gaussian_angular_cell: torch.Tensor,
+        gauss_grid_shape: tuple[int, int, int] | None,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        scales = self._flatten_gaussian_xyz(gaussian_scales, gauss_grid_shape)
+        means = self._flatten_gaussian_xyz(gaussian_mean_vectors, gauss_grid_shape)
+        if scales is None or means is None:
+            return None, None
+        if not torch.is_tensor(gaussian_angular_cell):
+            return None, None
+        cell = gaussian_angular_cell.to(device=scales.device, dtype=torch.float32)
+        if cell.ndim != 5 or int(cell.shape[1]) != 2:
+            return None, None
+        if gauss_grid_shape is None:
+            return None, None
+        l, h, w = (int(gauss_grid_shape[0]), int(gauss_grid_shape[1]), int(gauss_grid_shape[2]))
+        if tuple(cell.shape[-2:]) != (h, w):
+            return None, None
+        if int(cell.shape[2]) == 1 and l > 1:
+            cell = cell.expand(-1, -1, l, -1, -1)
+        elif int(cell.shape[2]) != l:
+            return None, None
+        cell_flat = cell.permute(0, 2, 3, 4, 1).flatten(1, 3)
+        if int(cell_flat.shape[1]) != int(scales.shape[1]):
+            return None, None
+        radius = torch.linalg.norm(means.to(dtype=torch.float32), dim=-1, keepdim=True).clamp(min=1e-4)
+        tangent = scales[..., :2].to(dtype=torch.float32).abs()
+        sigma_cells = (tangent / radius / cell_flat.clamp(min=1e-6)).square()
+        valid = torch.isfinite(sigma_cells).all(dim=-1) & torch.isfinite(radius.squeeze(-1))
+        sigma_cells = torch.nan_to_num(sigma_cells, nan=0.0, posinf=0.0, neginf=0.0)
+        return sigma_cells, valid
+
+    def _dynamic_splat_sigma_limits(
+        self,
+        *,
+        sigma_proj: torch.Tensor,
+        projection_model: str | None,
+        image_h: int,
+        image_w: int,
+        intrinsics: torch.Tensor | None = None,
+        camera_params: torch.Tensor | None = None,
+        projected_scale_factor: float | torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        del projection_model, image_h, image_w, intrinsics, camera_params, projected_scale_factor
+        return (
+            torch.as_tensor(self.splat_sigma_min, device=sigma_proj.device, dtype=sigma_proj.dtype),
+            torch.as_tensor(self.splat_sigma_max, device=sigma_proj.device, dtype=sigma_proj.dtype),
+        )
+
+    def _sanitize_supervision_depth(self, depth_m: torch.Tensor, *, clamp_max: bool = True) -> torch.Tensor:
+        depth = depth_m.to(torch.float32)
+        valid = torch.isfinite(depth) & (depth > 0.0)
+        depth = torch.where(valid, depth, torch.zeros_like(depth))
+        if bool(valid.any().item()):
+            depth = depth.clone()
+            if bool(clamp_max):
+                depth[valid] = depth[valid].clamp(min=1e-4, max=float(self.SUPERVISION_MAX_DEPTH_M))
+            else:
+                depth[valid] = depth[valid].clamp(min=1e-4)
+        return depth
+
+    def _sobel_gradient_loss_erp(
+        self,
+        pred_depth_m: torch.Tensor,
+        gt_depth_m: torch.Tensor,
+        depth_weight: torch.Tensor,
+        circular_h: bool | None = None,
+    ) -> torch.Tensor:
+        dtype = pred_depth_m.dtype
+        device = pred_depth_m.device
+
+        kx = self._sobel_kx.to(dtype=dtype, device=device)  # type: ignore[attr-defined]
+        ky = self._sobel_ky.to(dtype=dtype, device=device)  # type: ignore[attr-defined]
+
+        log_pred = torch.log(pred_depth_m.clamp(min=1e-4))
+        log_gt = torch.log(gt_depth_m.clamp(min=1e-4))
+        log_diff = log_pred - log_gt
+
+        mask = depth_weight.to(dtype=dtype).clamp(min=0.0, max=1.0)
+        valid_mask = (mask > 0.5).to(dtype=dtype)
+        log_diff = torch.where(valid_mask > 0.5, log_diff, torch.zeros_like(log_diff))
+
+        total = torch.zeros((), device=device, dtype=dtype)
+        n_computed = 0
+
+        use_circular_h = self.grad_img_circular_h if circular_h is None else bool(circular_h)
+        ones_kernel = torch.ones((1, 1, 3, 3), device=device, dtype=dtype)
+        for _s in range(self.grad_img_scales):
+            if min(log_diff.shape[-2:]) < 4:
+                break
+
+            if use_circular_h:
+                padded = F.pad(log_diff, (1, 1, 0, 0), mode="circular")
+                padded = F.pad(padded, (0, 0, 1, 1), mode="reflect")
+                padded_mask = F.pad(valid_mask, (1, 1, 0, 0), mode="circular")
+                padded_mask = F.pad(padded_mask, (0, 0, 1, 1), mode="replicate")
+            else:
+                padded = F.pad(log_diff, (1, 1, 1, 1), mode="reflect")
+                padded_mask = F.pad(valid_mask, (1, 1, 1, 1), mode="replicate")
+
+            gx = F.conv2d(padded, kx)
+            gy = F.conv2d(padded, ky)
+            grad_mag = torch.sqrt(gx * gx + gy * gy + 1e-8)
+
+            stencil_valid = (F.conv2d(padded_mask, ones_kernel) >= 8.999).to(dtype=dtype)
+            n_valid = stencil_valid.sum().clamp(min=1.0)
+            total = total + (grad_mag * stencil_valid).sum() / n_valid
+            n_computed += 1
+
+            if _s < self.grad_img_scales - 1:
+                if use_circular_h:
+                    pooled_mask = _avg_pool2d_circular_h(valid_mask, kernel_size=2, stride=2)
+                    pooled_diff = _avg_pool2d_circular_h(log_diff * valid_mask, kernel_size=2, stride=2)
+                else:
+                    pooled_mask = F.avg_pool2d(valid_mask, kernel_size=2, stride=2)
+                    pooled_diff = F.avg_pool2d(log_diff * valid_mask, kernel_size=2, stride=2)
+                log_diff = pooled_diff / pooled_mask.clamp(min=1e-6)
+                valid_mask = (pooled_mask > 0.999).to(dtype=dtype)
+                log_diff = torch.where(valid_mask > 0.5, log_diff, torch.zeros_like(log_diff))
+
+        if n_computed == 0:
+            return torch.zeros((), device=device, dtype=dtype)
+        return total / float(n_computed)
+
+    def _sobel_xy_rgb(self, img: torch.Tensor, *, circular_h: bool) -> tuple[torch.Tensor, torch.Tensor]:
+        channels = int(img.shape[1])
+        kx = self._sobel_kx.to(dtype=img.dtype, device=img.device).expand(channels, 1, 3, 3)  # type: ignore[attr-defined]
+        ky = self._sobel_ky.to(dtype=img.dtype, device=img.device).expand(channels, 1, 3, 3)  # type: ignore[attr-defined]
+        if bool(circular_h):
+            padded = F.pad(img, (1, 1, 0, 0), mode="circular")
+            padded = F.pad(padded, (0, 0, 1, 1), mode="reflect")
+        else:
+            padded = F.pad(img, (1, 1, 1, 1), mode="reflect")
+        return (
+            F.conv2d(padded, kx, groups=channels),
+            F.conv2d(padded, ky, groups=channels),
+        )
+
+    def _edge_rgb_gradient_loss(
+        self,
+        pred_rgb_linear: torch.Tensor,
+        gt_rgb_linear: torch.Tensor,
+        valid_weight: torch.Tensor,
+        depth_edge_band: torch.Tensor | None,
+        *,
+        circular_h: bool,
+    ) -> torch.Tensor:
+        dtype = pred_rgb_linear.dtype
+        device = pred_rgb_linear.device
+        pred = pred_rgb_linear.to(dtype=torch.float32)
+        gt = gt_rgb_linear.to(device=device, dtype=torch.float32)
+        weight = valid_weight.to(device=device, dtype=torch.float32).clamp(0.0, 1.0)[:, :1]
+
+        pred_gx, pred_gy = self._sobel_xy_rgb(pred, circular_h=circular_h)
+        gt_gx, gt_gy = self._sobel_xy_rgb(gt, circular_h=circular_h)
+        gt_mag = torch.sqrt(gt_gx.square() + gt_gy.square() + 1e-8).mean(dim=1, keepdim=True)
+
+        flat = gt_mag.detach().flatten(2)
+        mean = flat.mean(dim=-1, keepdim=True)[..., None]
+        std = flat.std(dim=-1, keepdim=True, unbiased=False)[..., None]
+        rgb_edge = (gt_mag.detach() > (mean + 0.5 * std).clamp(min=0.02)).to(dtype=torch.float32)
+
+        if torch.is_tensor(depth_edge_band):
+            edge_boost = depth_edge_band.to(device=device, dtype=torch.float32).clamp(0.0, 1.0)
+            if tuple(edge_boost.shape[-2:]) != tuple(gt_mag.shape[-2:]):
+                edge_boost = F.interpolate(edge_boost, size=gt_mag.shape[-2:], mode="nearest")
+            edge_weight = rgb_edge * (1.0 + edge_boost[:, :1])
+        else:
+            edge_weight = rgb_edge
+
+        ones_kernel = torch.ones((1, 1, 3, 3), device=device, dtype=torch.float32)
+        if bool(circular_h):
+            padded_weight = F.pad(weight, (1, 1, 0, 0), mode="circular")
+            padded_weight = F.pad(padded_weight, (0, 0, 1, 1), mode="replicate")
+        else:
+            padded_weight = F.pad(weight, (1, 1, 1, 1), mode="replicate")
+        stencil_valid = (F.conv2d(padded_weight, ones_kernel) >= 8.999).to(dtype=torch.float32)
+
+        diff = (pred_gx - gt_gx).abs() + (pred_gy - gt_gy).abs()
+        diff = diff.mean(dim=1, keepdim=True)
+        final_weight = edge_weight * stencil_valid
+        return (diff * final_weight).sum().to(dtype=dtype) / final_weight.sum().clamp(min=1.0).to(dtype=dtype)
+
+    def forward(
+        self,
+        pred_rgb_linear: torch.Tensor,
+        pred_alpha: torch.Tensor,
+        pred_depth_m: torch.Tensor,
+        gt_rgb_u8: torch.Tensor,
+        gt_depth_m: torch.Tensor,
+        pred_depth2_m: torch.Tensor | None = None,
+        mask: torch.Tensor | None = None,
+        depth_mask: torch.Tensor | None = None,
+        delta_xy: torch.Tensor | None = None,
+        delta_rho: torch.Tensor | None = None,
+        delta_grid: torch.Tensor | None = None,
+        gaussian_scales: torch.Tensor | None = None,
+        gaussian_quaternions: torch.Tensor | None = None,
+        gaussian_mean_vectors: torch.Tensor | None = None,
+        gaussian_base_mean_vectors: torch.Tensor | None = None,
+        gaussian_angular_cell: torch.Tensor | None = None,
+        gaussian_opacities: torch.Tensor | None = None,
+        gauss_grid_shape: tuple[int, int, int] | None = None,
+        projected_scale_factor: float | torch.Tensor | None = None,
+        projection_model: str | None = None,
+        projection_intrinsics: torch.Tensor | None = None,
+        projection_camera_params: torch.Tensor | None = None,
+        apply_color: bool = True,
+        apply_alpha: bool = True,
+        apply_depth: bool = True,
+        apply_percep: bool = False,
+        apply_tv: bool = True,
+        apply_grad: bool = True,
+        apply_delta: bool = True,
+        apply_splat: bool = True,
+        apply_grad_img: bool = True,
+        grad_img_circular_h: bool | None = None,
+    ) -> dict[str, torch.Tensor]:
+        losses: dict[str, torch.Tensor] = {}
+        circular_h = bool(grad_img_circular_h) if grad_img_circular_h is not None else False
+
+        gt_rgb = gt_rgb_u8.to(pred_rgb_linear.device).float() / 255.0
+        gt_rgb_linear = _to_linear_rgb(gt_rgb)
+        pred_depth_m = self._sanitize_supervision_depth(pred_depth_m.to(pred_rgb_linear.device), clamp_max=False)
+        if pred_depth2_m is not None:
+            pred_depth2_m = self._sanitize_supervision_depth(pred_depth2_m.to(pred_rgb_linear.device), clamp_max=False)
+        gt_depth_raw = self._sanitize_supervision_depth(gt_depth_m.to(pred_rgb_linear.device))
+        depth_valid = torch.isfinite(gt_depth_raw) & (gt_depth_raw > 0.0)
+        gt_depth = gt_depth_raw.clamp(min=1e-4)
+
+        if mask is None:
+            m = torch.ones_like(pred_alpha)
+        else:
+            m = mask.to(pred_rgb_linear.device).to(pred_rgb_linear.dtype)
+        depth_weight = depth_valid.to(dtype=pred_depth_m.dtype) * m[:, :1].to(dtype=pred_depth_m.dtype)
+        if depth_mask is not None:
+            depth_weight = depth_weight * depth_mask.to(pred_rgb_linear.device).to(dtype=pred_depth_m.dtype)[:, :1]
+        pred_rgb_rendered = pred_rgb_linear.clamp(0.0, 1.0)
+
+        if apply_color and self.w.lambda_color > 0:
+            color_l1 = (pred_rgb_rendered - gt_rgb_linear).abs()
+            losses["color"] = _masked_mean(color_l1, m)
+        else:
+            losses["color"] = torch.zeros((), device=pred_rgb_linear.device)
+
+        if apply_alpha and self.w.lambda_alpha > 0:
+            a = pred_alpha.clamp(1e-6, 1.0 - 1e-6)
+            with torch.autocast(device_type=a.device.type, enabled=False):
+                alpha_bce = F.binary_cross_entropy(
+                    a.to(dtype=torch.float32),
+                    torch.ones_like(a, dtype=torch.float32),
+                    reduction="none",
+                )
+                alpha_loss = _masked_mean(alpha_bce, m)
+                alpha_tail_min = torch.as_tensor(
+                    self.alpha_tail_min,
+                    device=a.device,
+                    dtype=torch.float32,
+                ).clamp(min=0.0, max=1.0)
+                alpha_tail_weight = torch.as_tensor(
+                    max(0.0, self.alpha_tail_weight),
+                    device=a.device,
+                    dtype=torch.float32,
+                )
+                if self.alpha_tail_min > 0.0 and self.alpha_tail_weight > 0.0:
+                    tail = F.relu(alpha_tail_min - a.to(dtype=torch.float32))
+                    tail = tail / alpha_tail_min.clamp(min=1e-6)
+                    tail_mask = (m[:, :1].to(dtype=torch.bool)) & (tail > 0.0)
+                    alpha_loss = alpha_loss + alpha_tail_weight * _finite_masked_mean_flat(tail, tail_mask)
+            losses["alpha"] = alpha_loss.to(dtype=pred_rgb_linear.dtype)
+        else:
+            losses["alpha"] = torch.zeros((), device=pred_rgb_linear.device)
+
+        if apply_depth and self.w.lambda_depth > 0:
+            w_depth = depth_weight
+            inv_pred1 = 1.0 / pred_depth_m.clamp(min=1e-4)
+            inv_gt = torch.zeros_like(inv_pred1)
+            inv_gt[depth_valid] = 1.0 / gt_depth[depth_valid]
+            depth_abs = (inv_pred1 - inv_gt).abs()
+            losses["depth"] = _masked_mean(depth_abs, w_depth)
+        else:
+            losses["depth"] = torch.zeros((), device=pred_rgb_linear.device)
+
+        if apply_tv and self.w.lambda_tv > 0 and (pred_depth2_m is not None):
+            inv2 = 1.0 / pred_depth2_m.clamp(min=1e-4)
+            losses["tv"] = _tv_l1_circular_h(inv2) if circular_h else _tv_l1(inv2)
+        else:
+            losses["tv"] = torch.zeros((), device=pred_rgb_linear.device)
+
+        image_h, image_w = int(pred_depth_m.shape[-2]), int(pred_depth_m.shape[-1])
+        projection_points = self._flatten_gaussian_xyz(gaussian_mean_vectors, gauss_grid_shape)
+        projected_u = projected_v = None
+        projected_valid = None
+        if projection_points is not None:
+            projected_u, projected_v, projected_valid, _projected_depth = self._project_points_px(
+                projection_points,
+                projection_model=projection_model,
+                image_h=image_h,
+                image_w=image_w,
+                intrinsics=projection_intrinsics,
+                camera_params=projection_camera_params,
+            )
+
+        if apply_grad and self.w.lambda_grad > 0:
+            inv1 = 1.0 / pred_depth_m.clamp(min=1e-4)
+            op_flat = self._flatten_gaussian_scalar(gaussian_opacities, gauss_grid_shape)
+            if projected_u is not None and projected_v is not None and projected_valid is not None and op_flat is not None:
+                grad_map = self._central_disparity_gradient(inv1, circular_h=circular_h)
+                grad_at_gauss = self._sample_map_at_uv(grad_map, projected_u, projected_v, projected_valid)
+                penalty = 1.0 - torch.exp(
+                    -(1.0 / max(self.grad_sigma, 1e-8)) * F.relu(grad_at_gauss - self.grad_eps)
+                )
+                weight = projected_valid & torch.isfinite(grad_at_gauss) & torch.isfinite(op_flat)
+                mask_at_gauss = self._sample_map_at_uv(m[:, :1], projected_u, projected_v, projected_valid)
+                weight = weight & (mask_at_gauss > 0.5)
+                grad_value = op_flat.to(dtype=penalty.dtype).clamp(0, 1) * penalty
+                losses["grad"] = _finite_masked_mean_flat(grad_value, weight)
+            else:
+                raise RuntimeError(
+                    "L_grad requires gaussian_mean_vectors, gaussian_opacities, "
+                    "gauss_grid_shape, and projection metadata. The old "
+                    "pred_alpha image-space fallback is disabled for ray-local training."
+                )
+        else:
+            losses["grad"] = torch.zeros((), device=pred_rgb_linear.device)
+
+        if apply_grad_img and self.w.lambda_grad_img > 0:
+            losses["grad_img"] = self._sobel_gradient_loss_erp(
+                pred_depth_m=pred_depth_m,
+                gt_depth_m=gt_depth,
+                depth_weight=depth_weight,
+                circular_h=grad_img_circular_h,
+            )
+        else:
+            losses["grad_img"] = torch.zeros((), device=pred_rgb_linear.device)
+
+        if apply_color and self.w.lambda_edge_rgb > 0:
+            depth_edge_for_rgb = self._depth_edge_band(gt_depth, depth_weight, circular_h=circular_h)
+            losses["edge_rgb"] = self._edge_rgb_gradient_loss(
+                pred_rgb_linear=pred_rgb_rendered,
+                gt_rgb_linear=gt_rgb_linear,
+                valid_weight=m,
+                depth_edge_band=depth_edge_for_rgb,
+                circular_h=circular_h,
+            )
+        else:
+            losses["edge_rgb"] = torch.zeros((), device=pred_rgb_linear.device)
+
+        if apply_delta and self.w.lambda_delta > 0:
+            if delta_xy is not None:
+                dx = F.relu(delta_xy[:, 0:1].abs() - self.raw_delta_clip)
+                dy = F.relu(delta_xy[:, 1:2].abs() - self.raw_delta_clip)
+                losses["delta"] = (dx + dy).mean()
+            else:
+                del gaussian_base_mean_vectors
+                raise RuntimeError(
+                    "L_delta requires raw delta_xy in ray-local training. The old "
+                    "screen-space pixel displacement fallback is disabled."
+                )
+        else:
+            losses["delta"] = torch.zeros((), device=pred_rgb_linear.device)
+
+        if apply_delta and self.w.lambda_delta_rho > 0 and delta_rho is not None:
+            dz = delta_rho.to(device=pred_rgb_linear.device, dtype=pred_rgb_linear.dtype)
+            finite = torch.isfinite(dz)
+            dz_safe = torch.nan_to_num(dz, nan=0.0, posinf=0.0, neginf=0.0)
+            penalty = F.relu(dz_safe.abs() - self.raw_delta_rho_clip)
+            penalty = torch.where(finite, penalty, torch.zeros_like(penalty))
+            losses["delta_rho"] = penalty.sum() / finite.to(dtype=penalty.dtype).sum().clamp(min=1.0)
+        else:
+            losses["delta_rho"] = torch.zeros((), device=pred_rgb_linear.device)
+
+        if self.w.lambda_grid > 0 and torch.is_tensor(delta_grid):
+            losses["grid"] = _delta_grid_checkerboard_loss(
+                delta_grid.to(device=pred_rgb_linear.device),
+                circular_h=circular_h,
+            ).to(dtype=pred_rgb_linear.dtype)
+        else:
+            losses["grid"] = torch.zeros((), device=pred_rgb_linear.device)
+
+        if apply_splat and self.w.lambda_splat > 0:
+            if gaussian_scales is None:
+                raise RuntimeError("L_splat requires gaussian_scales for projected screen-space variance.")
+            if gaussian_mean_vectors is None or projected_valid is None:
+                raise RuntimeError(
+                    "L_splat requires gaussian_mean_vectors and projection metadata "
+                    "to compute projected screen-space variance."
+                )
+            sigma_proj = self._projected_sigma_px(
+                gaussian_scales=gaussian_scales,
+                gaussian_quaternions=gaussian_quaternions,
+                gaussian_mean_vectors=gaussian_mean_vectors,
+                valid=projected_valid,
+                projection_model=projection_model,
+                image_h=image_h,
+                image_w=image_w,
+                intrinsics=projection_intrinsics,
+                camera_params=projection_camera_params,
+                projected_scale_factor=projected_scale_factor,
+            )
+            valid_splat = projected_valid & torch.isfinite(sigma_proj)
+            splat_sigma_min = torch.as_tensor(
+                self.splat_sigma_min,
+                device=sigma_proj.device,
+                dtype=sigma_proj.dtype,
+            )
+            splat_sigma_max = torch.as_tensor(
+                self.splat_sigma_max,
+                device=sigma_proj.device,
+                dtype=sigma_proj.dtype,
+            )
+            lower_penalty = F.relu(splat_sigma_min - sigma_proj)
+            upper_penalty = F.relu(sigma_proj - splat_sigma_max)
+            splat_penalty = lower_penalty + upper_penalty
+            losses["splat"] = _finite_masked_mean_flat(splat_penalty, valid_splat)
+        else:
+            sigma_proj = None
+            valid_splat = None
+            losses["splat"] = torch.zeros((), device=pred_rgb_linear.device)
+
+        if apply_splat and self.w.lambda_edge_splat > 0:
+            if gaussian_scales is None:
+                raise RuntimeError("L_edge_splat requires gaussian_scales for projected screen-space variance.")
+            if gaussian_mean_vectors is None or projected_valid is None:
+                raise RuntimeError(
+                    "L_edge_splat requires gaussian_mean_vectors and projection metadata "
+                    "to sample source depth-edge bands."
+                )
+            if sigma_proj is None or valid_splat is None:
+                sigma_proj = self._projected_sigma_px(
+                    gaussian_scales=gaussian_scales,
+                    gaussian_quaternions=gaussian_quaternions,
+                    gaussian_mean_vectors=gaussian_mean_vectors,
+                    valid=projected_valid,
+                    projection_model=projection_model,
+                    image_h=image_h,
+                    image_w=image_w,
+                    intrinsics=projection_intrinsics,
+                    camera_params=projection_camera_params,
+                    projected_scale_factor=projected_scale_factor,
+                )
+                valid_splat = projected_valid & torch.isfinite(sigma_proj)
+            edge_band = self._depth_edge_band(gt_depth, depth_weight, circular_h=circular_h)
+            edge_at_gauss = self._sample_map_at_uv(edge_band, projected_u, projected_v, projected_valid)
+            edge_valid = valid_splat & torch.isfinite(edge_at_gauss) & (edge_at_gauss > 0.5)
+            edge_sigma_max = torch.as_tensor(
+                self.edge_splat_sigma_max,
+                device=sigma_proj.device,
+                dtype=sigma_proj.dtype,
+            )
+            losses["edge_splat"] = _finite_masked_mean_flat(F.relu(sigma_proj - edge_sigma_max), edge_valid)
+        else:
+            losses["edge_splat"] = torch.zeros((), device=pred_rgb_linear.device)
+
+        zero = torch.zeros((), device=pred_rgb_linear.device)
+        losses["percep_feat"] = zero
+        losses["percep_gram"] = zero
+        if apply_percep and self.w.lambda_percep > 0 and (self._percep_net is not None):
+            from unisharp.utils.color_space import linearRGB2sRGB
+
+            pred_srgb = linearRGB2sRGB(pred_rgb_rendered.to(torch.float32)).clamp(0, 1)
+            gt_srgb = gt_rgb.clamp(0, 1)
+
+            pred_srgb = _resize_max_side(pred_srgb, self.percep_max_side, mode="bilinear")
+            gt_srgb = _resize_max_side(gt_srgb, self.percep_max_side, mode="bilinear")
+
+            feats_p = self._percep_net(pred_srgb)
+            feats_g = self._percep_net(gt_srgb)
+            loss_feat_total = torch.zeros((), device=pred_rgb_linear.device)
+            loss_gram_total = torch.zeros((), device=pred_rgb_linear.device)
+            for fp, fg in zip(feats_p, feats_g):
+                d, h, w = fp.shape[1], fp.shape[2], fp.shape[3]
+                lam_gram = 10.0 / float(max(1, d * d))
+                lam_feat = 1.0 / float(max(1, d * h * w))
+                diff = (fp - fg).pow(2)
+                loss_feat = (diff.sum(dim=[1, 2, 3]) * lam_feat).mean()
+                gram_norm = float(max(1, h * w))
+                gp = _gram_matrix(fp) / gram_norm
+                gg = _gram_matrix(fg) / gram_norm
+                loss_gram = ((gp - gg).pow(2).sum(dim=[1, 2]) * lam_gram).mean()
+                loss_feat_total = loss_feat_total + loss_feat
+                loss_gram_total = loss_gram_total + loss_gram
+            layer_count = float(max(1, len(feats_p)))
+            losses["percep_feat"] = loss_feat_total / layer_count
+            losses["percep_gram"] = loss_gram_total / layer_count
+            losses["percep"] = losses["percep_feat"] + losses["percep_gram"]
+        else:
+            losses["percep"] = torch.zeros((), device=pred_rgb_linear.device)
+
+        losses["total"] = (
+            self.w.lambda_color * losses["color"]
+            + self.w.lambda_alpha * losses["alpha"]
+            + self.w.lambda_percep * losses["percep"]
+            + self.w.lambda_depth * losses["depth"]
+            + self.w.lambda_tv * losses["tv"]
+            + self.w.lambda_grad * losses["grad"]
+            + self.w.lambda_grad_img * losses["grad_img"]
+            + self.w.lambda_edge_rgb * losses["edge_rgb"]
+            + self.w.lambda_delta * losses["delta"]
+            + self.w.lambda_delta_rho * losses["delta_rho"]
+            + self.w.lambda_splat * losses["splat"]
+            + self.w.lambda_edge_splat * losses["edge_splat"]
+            + self.w.lambda_grid * losses["grid"]
+        )
+        return losses
+

unisharp/models/__init__.py

@@ -0,0 +1,23 @@
+
+from __future__ import annotations
+
+from .feature_gaussian_decoder import (
+    FeatureGaussianDecoder,
+    FeatureGaussianDecoderParams,
+    ImageFeatures,
+    create_feature_gaussian_decoder,
+)
+from .unisharp_params import PanoPredictorParams
+from .unisharp_feature import UnisharpFeatureConfig, UnisharpFeatureModel
+from .unik3d_feature_extractor import UniK3DFeatureExtractor
+
+__all__ = [
+    "PanoPredictorParams",
+    "UniK3DFeatureExtractor",
+    "FeatureGaussianDecoder",
+    "FeatureGaussianDecoderParams",
+    "ImageFeatures",
+    "create_feature_gaussian_decoder",
+    "UnisharpFeatureConfig",
+    "UnisharpFeatureModel",
+]