#!/usr/bin/env python3 from __future__ import annotations import argparse import csv from dataclasses import dataclass import json import logging import math import os import random import sys from pathlib import Path from types import SimpleNamespace from typing import Any, Callable, Iterable, Iterator import numpy as np import torch import torch.nn.functional as F from PIL import Image from tqdm import tqdm REPO_ROOT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(REPO_ROOT)) from unisharp.datasets.panogs import panogs_collate # noqa: E402 from unisharp.datasets.scannetpp_fisheye import ScannetppFisheyeDataset # noqa: E402 from unisharp.datasets.sim_panorama import _EquirecToCube, SimPanoramaDataset # noqa: E402 from unisharp.datasets.wildrgbd import WildRGBDDataset # noqa: E402 from unisharp.losses import UnisharpLoss, UnisharpLossWeights # noqa: E402 from unisharp.models.unisharp_feature import UnisharpFeatureConfig, UnisharpFeatureModel # noqa: E402 from unisharp.utils.color_space import linearRGB2sRGB # noqa: E402 from unisharp import DEFAULT_MAX_DEPTH_M # noqa: E402 from unisharp.utils.io import save_image # noqa: E402 from unisharp.utils.metrics import ( # noqa: E402 MetricsCalculator, compute_masked_rgb_metrics, default_metric_mask_cache_dir, metric_mask_from_pinhole_batch, ) from unisharp.utils.vis import colorize_alpha, colorize_scalar_map # noqa: E402 from unisharp.validation.io_common import ( # noqa: E402 decode_rgb_u8 as _decode_rgb_u8, distance_to_z_depth_pinhole as _distance_to_z_depth_pinhole, colmap_image_dir as _colmap_image_dir, colmap_scene_roots as _colmap_scene_roots, load_colmap_entries as _load_colmap_entries, load_hm3d_pose as _load_hm3d_pose, load_scaled_colmap_entries as _load_scaled_colmap_entries, load_png_depth_m as _load_png_depth_m, load_png_rgb_u8 as _load_png_rgb_u8, load_validation_pseudo_depth as _load_validation_pseudo_depth, load_validation_pseudo_distance as _load_validation_pseudo_distance, nerf_c2w_to_opencv_c2w as _nerf_c2w_to_opencv_c2w, normalize_depth_kind as _normalize_depth_kind, read_manifest_lines as _read_manifest_lines, resolve_replica_test_root as _resolve_replica_test_root, resize_k3_align_corners_false as _resize_k3_align_corners_false, torch_load_any as _torch_load_any, wild_validation_roots as _wild_validation_roots, ) logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") LOGGER = logging.getLogger(__name__) ValidationTag = str | list[str] ValidationItem = tuple[str, Any, ValidationTag, str] METRIC_FIELDS = ["psnr", "ssim", "lpips"] 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 _training_config_for_checkpoint(checkpoint_path: Path) -> dict[str, Any]: payload: dict[str, Any] = {} try: ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False) except TypeError: ckpt = torch.load(checkpoint_path, map_location="cpu") except Exception: ckpt = None if isinstance(ckpt, dict): cfg = ckpt.get("config", None) if isinstance(cfg, dict): payload.update(cfg) config_path = Path(checkpoint_path).parent / "config.json" if not config_path.exists(): return payload try: json_payload = json.loads(config_path.read_text(encoding="utf-8")) except Exception: return payload if isinstance(json_payload, dict): payload.update(json_payload) return payload def _fill_arg_from_config(args: argparse.Namespace, attr: str, config: dict[str, Any], key: str, default: float) -> None: if getattr(args, attr, None) is not None: return value = config.get(key, default) try: setattr(args, attr, float(value)) except Exception: setattr(args, attr, float(default)) def _fill_int_arg_from_config(args: argparse.Namespace, attr: str, config: dict[str, Any], key: str, default: int) -> None: if getattr(args, attr, None) is not None: return value = config.get(key, default) try: setattr(args, attr, int(value)) except Exception: setattr(args, attr, int(default)) def _apply_training_depth_config_defaults(args: argparse.Namespace) -> None: config = _training_config_for_checkpoint(Path(args.checkpoint)) _fill_arg_from_config(args, "max_depth_m", config, "max_depth_m", DEFAULT_MAX_DEPTH_M) _fill_arg_from_config(args, "sim_far_depth_invalid_m", config, "sim_far_depth_invalid_m", 30.0) _fill_arg_from_config(args, "sim_far_depth_invalid_max_frac", config, "sim_far_depth_invalid_max_frac", 1.0) _fill_arg_from_config(args, "re10k_pseudo_far_depth_invalid_m", config, "re10k_pseudo_far_depth_invalid_m", 30.0) _fill_arg_from_config(args, "scanetpp_fisheye_far_depth_invalid_m", config, "scanetpp_fisheye_far_depth_invalid_m", 30.0) _fill_arg_from_config(args, "low_pass_filter_eps", config, "render_low_pass_filter_eps", 1e-2) def _append_metrics_row(csv_path: Path, row: dict[str, float]) -> None: fieldnames = list(METRIC_FIELDS) row_out = {k: row.get(k, float("nan")) for k in fieldnames} if csv_path.exists(): try: with csv_path.open("r", newline="") as f: reader = csv.reader(f) existing_header = next(reader, []) if existing_header: if fieldnames != existing_header: with csv_path.open("r", newline="") as f: old_rows = list(csv.DictReader(f)) with csv_path.open("w", newline="") as f: writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writeheader() for r in old_rows: writer.writerow({k: r.get(k, float("nan")) for k in fieldnames}) except Exception: pass write_header = not csv_path.exists() or csv_path.stat().st_size == 0 with csv_path.open("a", newline="") as f: writer = csv.DictWriter(f, fieldnames=fieldnames) if write_header: writer.writeheader() writer.writerow(row_out) def _append_sample_metrics_row(csv_path: Path, group_key: str, tag: str, row: dict[str, float]) -> None: fieldnames = ["group", "tag", *METRIC_FIELDS] row_out = {"group": group_key, "tag": tag, **{k: row.get(k, float("nan")) for k in METRIC_FIELDS}} write_header = not csv_path.exists() or csv_path.stat().st_size == 0 with csv_path.open("a", newline="") as f: writer = csv.DictWriter(f, fieldnames=fieldnames) if write_header: writer.writeheader() writer.writerow(row_out) def _feature_config_from_checkpoint(checkpoint_path: Path, ckpt: dict[str, Any]) -> UnisharpFeatureConfig: cfg = UnisharpFeatureConfig() merged: dict[str, Any] = {} cfg_payload = ckpt.get("config", {}) if isinstance(cfg_payload, dict): merged.update(cfg_payload) for key in cfg.__dict__.keys(): if key in ckpt: merged[key] = ckpt[key] config_path = Path(checkpoint_path).parent / "config.json" if config_path.exists(): try: sidecar = json.loads(config_path.read_text(encoding="utf-8")) except Exception: sidecar = None if isinstance(sidecar, dict): merged.update({k: v for k, v in sidecar.items() if k in cfg.__dict__}) for k in cfg.__dict__.keys(): if k in merged: setattr(cfg, k, merged[k]) return cfg def _load_model(checkpoint_path: Path, device: torch.device) -> tuple[UnisharpFeatureModel, int]: try: ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False) except TypeError: ckpt = torch.load(checkpoint_path, map_location="cpu") if not isinstance(ckpt, dict): raise ValueError(f"Expected feature-only checkpoint dict, got {type(ckpt)} from {checkpoint_path}") cfg = _feature_config_from_checkpoint(checkpoint_path, ckpt) model = UnisharpFeatureModel(cfg).to(device) model.load_from_checkpoint(str(checkpoint_path), strict=True) model.eval() return model, int(ckpt.get("step", 0)) def _build_trainer(model: UnisharpFeatureModel, device: torch.device, args: argparse.Namespace) -> Any: from unisharp.cli.unified_trainer import UnifiedTrainer zero_w = UnisharpLossWeights( lambda_color=0.0, lambda_alpha=0.0, lambda_percep=0.0, lambda_depth=0.0, lambda_tv=0.0, lambda_grad=0.0, lambda_grad_img=0.0, lambda_delta=0.0, lambda_splat=0.0, ) loss_fn = UnisharpLoss(zero_w).to(device) max_depth_m = float(getattr(args, "max_depth_m", getattr(model.config, "max_distance_m", DEFAULT_MAX_DEPTH_M))) loss_fn.SUPERVISION_MAX_DEPTH_M = max_depth_m from unisharp.utils.gsplat import GSplatRenderer renderer = GSplatRenderer( color_space="sRGB", background_color="black", low_pass_filter_eps=float(getattr(args, "low_pass_filter_eps", 1e-2)), ).to(device) return UnifiedTrainer( model=model, renderer=renderer, loss_fn=loss_fn, device=device, enable_tgt_unik3d_vis=False, max_depth_m=max_depth_m, sim_far_depth_invalid_m=float(getattr(args, "sim_far_depth_invalid_m", 30.0)), re10k_pseudo_far_depth_invalid_m=float(getattr(args, "re10k_pseudo_far_depth_invalid_m", 30.0)), scanetpp_fisheye_far_depth_invalid_m=float(getattr(args, "scanetpp_fisheye_far_depth_invalid_m", 30.0)), ) def _make_pinhole_batch( *, src_img: torch.Tensor, tgt_img: torch.Tensor, src_w2c: torch.Tensor, tgt_w2c: torch.Tensor, src_k: torch.Tensor, tgt_k: torch.Tensor, scene: str, src_idx: int | list[int], tgt_idx: int | list[int], src_depth: torch.Tensor | None = None, tgt_depth: torch.Tensor | None = None, src_img_orig: torch.Tensor | None = None, tgt_img_orig: torch.Tensor | None = None, src_k_orig: torch.Tensor | None = None, tgt_k_orig: torch.Tensor | None = None, src_depth_orig: torch.Tensor | None = None, tgt_depth_orig: torch.Tensor | None = None, ) -> SimpleNamespace: batch_size = int(tgt_img.shape[0]) if torch.is_tensor(tgt_img) and tgt_img.ndim == 4 else 1 scene_values = [scene] * batch_size if isinstance(scene, str) else list(scene) src_idx_values = [int(src_idx)] * batch_size if isinstance(src_idx, int) else [int(x) for x in src_idx] tgt_idx_values = [int(tgt_idx)] if isinstance(tgt_idx, int) else [int(x) for x in tgt_idx] return SimpleNamespace( src_rgb_u8=src_img, tgt_rgb_u8=tgt_img, src_depth_m=src_depth, tgt_depth_m=tgt_depth, src_rgb_u8_orig=(src_img if src_img_orig is None else src_img_orig), tgt_rgb_u8_orig=(tgt_img if tgt_img_orig is None else tgt_img_orig), src_depth_m_orig=(src_depth if src_depth_orig is None else src_depth_orig), tgt_depth_m_orig=(tgt_depth if tgt_depth_orig is None else tgt_depth_orig), src_w2c=src_w2c, tgt_w2c=tgt_w2c, src_intrinsics=src_k, tgt_intrinsics=tgt_k, src_intrinsics_orig=(src_k if src_k_orig is None else src_k_orig), tgt_intrinsics_orig=(tgt_k if tgt_k_orig is None else tgt_k_orig), scene=scene_values, src_idx=torch.tensor(src_idx_values, dtype=torch.long), tgt_idx=torch.tensor(tgt_idx_values, dtype=torch.long), share_src_forward=True, collect_all_vis=True, ) def _make_scanetpp_fisheye_batch( *, scene: str, src_pos: int, tgt_positions: list[int], src_frame: dict[str, Any], tgt_frames: list[dict[str, Any]], src_loaded: dict[str, torch.Tensor], tgt_loaded: list[dict[str, torch.Tensor]], ) -> SimpleNamespace: n = int(len(tgt_positions)) src_rgb = src_loaded["rgb_u8"].unsqueeze(0).repeat(n, 1, 1, 1) src_depth = src_loaded["depth_m"].unsqueeze(0).repeat(n, 1, 1, 1) src_mask = src_loaded["valid_mask"].unsqueeze(0).repeat(n, 1, 1, 1) src_w2c = src_frame["w2c"].to(torch.float32).unsqueeze(0).repeat(n, 1, 1) src_cam = src_loaded["camera_params"].to(torch.float32).unsqueeze(0).repeat(n, 1) return SimpleNamespace( src_rgb_u8=src_rgb, tgt_rgb_u8=torch.stack([item["rgb_u8"] for item in tgt_loaded], dim=0), src_depth_m=src_depth, tgt_depth_m=torch.stack([item["depth_m"] for item in tgt_loaded], dim=0), src_valid_mask=src_mask, tgt_valid_mask=torch.stack([item["valid_mask"] for item in tgt_loaded], dim=0), src_w2c=src_w2c, tgt_w2c=torch.stack([frame["w2c"].to(torch.float32) for frame in tgt_frames], dim=0), src_camera_params=src_cam, tgt_camera_params=torch.stack([item["camera_params"].to(torch.float32) for item in tgt_loaded], dim=0), src_idx=torch.full((n,), int(src_pos), dtype=torch.long), tgt_idx=torch.tensor([int(x) for x in tgt_positions], dtype=torch.long), scene=[str(scene)] * n, camera_model="fisheye624", share_src_forward=True, collect_all_vis=True, ) @dataclass class _PinholeTargetAdapter: idx: int img: torch.Tensor w2c: torch.Tensor k: torch.Tensor depth: torch.Tensor | None = None @dataclass class _PinholeGroupAdapter: scene: str group_key: str src_idx: int src_img: torch.Tensor src_w2c: torch.Tensor src_k: torch.Tensor tgt_indices: list[int] load_target: Callable[[int], _PinholeTargetAdapter | None] src_depth: torch.Tensor | None = None def _iter_manifest_parts( args: argparse.Namespace, *, expected_parts: int, ) -> Iterator[tuple[int, list[str]]]: manifest_in = _read_manifest_lines(getattr(args, "manifest_file", None), max_lines=_manifest_max_groups(args)) for group_idx, raw in enumerate(manifest_in): parts = raw.split("|") if len(parts) == int(expected_parts): yield group_idx, parts def _yield_pinhole_group_batches( dataset: str, adapter: _PinholeGroupAdapter, args: argparse.Namespace, ) -> Iterator[ValidationItem]: batch_size = max(1, int(getattr(args, "validation_batch_size", 1))) pending: list[_PinholeTargetAdapter] = [] def _flush(targets: list[_PinholeTargetAdapter]) -> Iterator[ValidationItem]: if not targets: return n = len(targets) src_img_orig = adapter.src_img.clone() if n > 1: src_img_orig = src_img_orig.repeat(n, 1, 1, 1) tgt_img_orig = torch.cat([t.img.clone() for t in targets], dim=0) src_k_orig = adapter.src_k.clone() if n > 1: src_k_orig = src_k_orig.repeat(n, 1, 1) tgt_k_orig = torch.cat([t.k.clone() for t in targets], dim=0) src_w2c = adapter.src_w2c if n > 1: src_w2c = src_w2c.repeat(n, 1, 1) tgt_w2c = torch.cat([t.w2c for t in targets], dim=0) src_depth_orig = None if adapter.src_depth is None else adapter.src_depth.clone() if src_depth_orig is not None and n > 1: src_depth_orig = src_depth_orig.repeat(n, 1, 1, 1) tgt_depth_values = [t.depth for t in targets] tgt_depth_orig = None if all(torch.is_tensor(d) for d in tgt_depth_values): tgt_depth_orig = torch.cat([d for d in tgt_depth_values if torch.is_tensor(d)], dim=0) batch = _make_pinhole_batch( src_img=src_img_orig, tgt_img=tgt_img_orig, src_w2c=src_w2c, tgt_w2c=tgt_w2c, src_k=src_k_orig, tgt_k=tgt_k_orig, scene=adapter.scene, src_idx=[adapter.src_idx] * n, tgt_idx=[int(t.idx) for t in targets], src_depth=src_depth_orig, tgt_depth=tgt_depth_orig, src_img_orig=src_img_orig, tgt_img_orig=tgt_img_orig, src_k_orig=src_k_orig, tgt_k_orig=tgt_k_orig, src_depth_orig=src_depth_orig, tgt_depth_orig=tgt_depth_orig, ) tags = [f"{adapter.group_key}_t{int(t.idx):05d}" for t in targets] yield (dataset, batch, tags[0] if len(tags) == 1 else tags, adapter.group_key) for tgt_idx in adapter.tgt_indices: tgt = adapter.load_target(int(tgt_idx)) if tgt is None: continue pending.append(tgt) if len(pending) >= batch_size: yield from _flush(pending) pending = [] if pending: yield from _flush(pending) def _yield_panogs_group_batches( dataset: str, group_key: str, samples: list[Any], tags: list[str], args: argparse.Namespace, ) -> Iterator[ValidationItem]: batch_size = max(1, int(getattr(args, "validation_batch_size", 1))) if len(samples) != len(tags): raise ValueError(f"Expected samples/tags length match, got {len(samples)} vs {len(tags)}") for start in range(0, len(samples), batch_size): end = min(len(samples), start + batch_size) batch_tags = tags[start:end] batch = panogs_collate(samples[start:end]) object.__setattr__(batch, "collect_all_vis", True) yield ( dataset, batch, batch_tags[0] if len(batch_tags) == 1 else batch_tags, group_key, ) def _mask_connected_to_border(mask_2d: torch.Tensor) -> torch.Tensor: h, w = int(mask_2d.shape[0]), int(mask_2d.shape[1]) if h <= 0 or w <= 0: return torch.zeros_like(mask_2d, dtype=torch.bool) border = torch.zeros_like(mask_2d, dtype=torch.bool) border[0, :] = True border[-1, :] = True border[:, 0] = True border[:, -1] = True frontier = mask_2d & border visited = frontier.clone() kernel = torch.tensor( [[[[0.0, 1.0, 0.0], [1.0, 1.0, 1.0], [0.0, 1.0, 0.0]]]], device=mask_2d.device, dtype=torch.float32, ) for _ in range(h + w): if not bool(frontier.any()): break neigh = F.conv2d(frontier[None, None].to(torch.float32), kernel, padding=1)[0, 0] > 0.0 new_frontier = neigh & mask_2d & (~visited) visited = visited | new_frontier frontier = new_frontier return visited def _compute_metrics_from_vis( vis: dict[str, Any], metrics_calc: MetricsCalculator, ) -> dict[str, float]: tgt_gt = vis["tgt_gt"].detach().to(torch.float32).clamp(0, 1) tgt_alpha = vis["tgt_alpha"].detach().to(torch.float32).clamp(0.0, 1.0) tgt_pred = linearRGB2sRGB( (vis["tgt_pred"].detach().to(torch.float32) / tgt_alpha.clamp(min=1e-4)).clamp(0.0, 1.0) ).clamp(0, 1) geom_mask = vis.get("tgt_metric_mask", None) if torch.is_tensor(geom_mask): geom_mask = geom_mask.detach().to(device=tgt_pred.device, dtype=torch.float32) if geom_mask.ndim == 3: geom_mask = geom_mask.unsqueeze(1) if tuple(geom_mask.shape[-2:]) != tuple(tgt_pred.shape[-2:]): geom_mask = F.interpolate(geom_mask, size=tgt_pred.shape[-2:], mode="nearest") tgt_geom = compute_masked_rgb_metrics( pred=tgt_pred, gt=tgt_gt, mask=geom_mask, metrics_calc=metrics_calc, ) else: tgt_geom = metrics_calc.compute_rgb_metrics(tgt_pred, tgt_gt) return { "psnr": float(tgt_geom["psnr"]), "ssim": float(tgt_geom["ssim"]), "lpips": float(tgt_geom["lpips"]), } def _save_vis_from_payload(vis: dict[str, Any], vis_dir: Path, tag: str, step: int) -> None: from unisharp.utils.unified_vis import save_pair_visualization save_pair_visualization( vis_dir / f"step_{int(step):07d}_{tag}.png", src_gt=vis["src_gt"], src_pred=vis["src_pred"], src_alpha=vis["src_alpha"], tgt_gt=vis["tgt_gt"], tgt_pred=vis["tgt_pred"], tgt_alpha=vis["tgt_alpha"], src_gt_depth=vis.get("src_gt_depth", None), tgt_gt_depth=vis.get("tgt_gt_depth", None), src_pred_depth=vis.get("src_pred_depth", None), tgt_pred_depth=vis.get("tgt_pred_depth", None), src_unik3d_depth=vis.get("src_unik3d_depth", None), tgt_unik3d_depth=vis.get("tgt_unik3d_depth", None), dataset_name=str(vis.get("dataset_name", "unknown")), scene=str(vis.get("scene", "unknown")), step=int(step), src_idx=int(vis.get("src_idx", -1)), tgt_idx=int(vis.get("tgt_idx", -1)), src_pose_w2c=vis.get("src_pose_w2c", None), tgt_pose_w2c=vis.get("tgt_pose_w2c", None), src_cube_gt_u8=vis.get("src_cube_gt_u8", None), src_cube_pred_linear=vis.get("src_cube_pred_linear", None), src_cube_alpha=vis.get("src_cube_alpha", None), tgt_cube_gt_u8=vis.get("tgt_cube_gt_u8", None), tgt_cube_pred_linear=vis.get("tgt_cube_pred_linear", None), tgt_cube_alpha=vis.get("tgt_cube_alpha", None), ) def _save_group_pair_pngs(group_dir: Path, group_items: list[dict[str, Any]]) -> None: visual_items = [item for item in group_items if isinstance(item.get("vis", None), dict)] if not visual_items: return group_dir.mkdir(parents=True, exist_ok=True) def _save_mask(mask: torch.Tensor | None, path: Path) -> None: if not torch.is_tensor(mask): return mask_rgb = mask.detach().to(torch.float32).clamp(0.0, 1.0) if mask_rgb.ndim == 4: mask_rgb = mask_rgb[0] if mask_rgb.ndim == 3 and int(mask_rgb.shape[0]) == 1: mask_rgb = mask_rgb.repeat(3, 1, 1) if mask_rgb.ndim == 2: mask_rgb = mask_rgb[None].repeat(3, 1, 1) if mask_rgb.ndim == 3: save_image(_to_u8_hwc(mask_rgb), path) def _save_alpha(alpha: torch.Tensor | None, path: Path) -> None: if not torch.is_tensor(alpha): return a = alpha.detach().to(torch.float32).clamp(0.0, 1.0) if a.ndim == 4: a = a[0] if a.ndim == 3 and int(a.shape[0]) == 1: a = a.repeat(3, 1, 1) if a.ndim == 2: a = a[None].repeat(3, 1, 1) if a.ndim == 3: save_image(_to_u8_hwc(a), path) def _noextrap_mask_from_vis(vis: dict[str, Any], which: str) -> torch.Tensor | None: pred = vis.get(f"{which}_pred", None) alpha = vis.get(f"{which}_alpha", None) if not (torch.is_tensor(pred) and torch.is_tensor(alpha)): return None pred_pm = linearRGB2sRGB(pred.detach().to(torch.float32).clamp(min=0.0)).clamp(0.0, 1.0) alpha = alpha.detach().to(torch.float32).clamp(0.0, 1.0) masks: list[torch.Tensor] = [] for bi in range(int(pred_pm.shape[0])): black = pred_pm[bi : bi + 1].max(dim=1, keepdim=True).values <= float(2.0 / 255.0) low_alpha = alpha[bi : bi + 1] <= float(0.02) extrap_border = _mask_connected_to_border((black & low_alpha)[0, 0]) masks.append((~extrap_border)[None, None].to(torch.float32)) return torch.cat(masks, dim=0) src_row = _build_perspective_row(visual_items[0]["vis"], "src") save_image(src_row[0], group_dir / "src_gt.png") save_image(src_row[1], group_dir / "src_pred.png") _save_alpha(visual_items[0]["vis"].get("src_alpha", None), group_dir / "src_alpha.png") _save_mask(_noextrap_mask_from_vis(visual_items[0]["vis"], "src"), group_dir / "src_noextrap_mask.png") _save_mask(visual_items[0]["vis"].get("src_metric_mask", None), group_dir / "src_mask.png") for idx, item in enumerate(visual_items): tgt_row = _build_perspective_row(item["vis"], "tgt") save_image(tgt_row[0], group_dir / f"tgt_{idx:03d}_gt.png") save_image(tgt_row[1], group_dir / f"tgt_{idx:03d}_pred.png") _save_alpha(item["vis"].get("tgt_alpha", None), group_dir / f"tgt_{idx:03d}_alpha.png") _save_mask(_noextrap_mask_from_vis(item["vis"], "tgt"), group_dir / f"tgt_{idx:03d}_noextrap_mask.png") _save_mask(item["vis"].get("tgt_training_mask", None), group_dir / f"tgt_{idx:03d}_training_mask.png") _save_mask(item["vis"].get("tgt_metric_mask", None), group_dir / f"tgt_{idx:03d}_mask.png") def _aggregate_rows(rows: list[dict[str, float]]) -> dict[str, float]: agg: dict[str, float] = {} if not rows: return agg keys = sorted(set().union(*[set(r.keys()) for r in rows])) for k in keys: arr = np.array([r.get(k, np.nan) for r in rows], dtype=np.float64) agg[k] = _safe_nanmean(arr) agg["num_samples"] = float(len(rows)) return agg def _safe_nanmean(values: Any) -> float: arr = np.asarray(values, dtype=np.float64) if arr.size == 0: return float("nan") if not np.isfinite(arr).any(): return float("nan") return float(np.nanmean(arr)) def _to_u8_hwc(img_chw: torch.Tensor) -> np.ndarray: if img_chw.dtype == torch.uint8: return img_chw.permute(1, 2, 0).detach().cpu().numpy() x = img_chw.detach().to(torch.float32).clamp(0.0, 1.0) return (x * 255.0).round().to(torch.uint8).permute(1, 2, 0).cpu().numpy() def _concat_grid(rows: list[list[np.ndarray]], pad: int = 6, pad_value: int = 0) -> np.ndarray: row_imgs: list[np.ndarray] = [] for r in rows: padded: list[np.ndarray] = [] for i, im in enumerate(r): padded.append(im) if i != len(r) - 1 and pad > 0: padded.append(np.full((im.shape[0], pad, 3), pad_value, dtype=np.uint8)) row_imgs.append(np.concatenate(padded, axis=1)) merged: list[np.ndarray] = [] for i, im in enumerate(row_imgs): merged.append(im) if i != len(row_imgs) - 1 and pad > 0: merged.append(np.full((pad, im.shape[1], 3), pad_value, dtype=np.uint8)) return np.concatenate(merged, axis=0) def _resize_panel_np(panel: np.ndarray, out_h: int, out_w: int) -> np.ndarray: if panel.shape[0] == out_h and panel.shape[1] == out_w: return panel return np.asarray(Image.fromarray(panel).resize((int(out_w), int(out_h)), resample=Image.BILINEAR)) def _normalize_rows_for_grid(rows: list[list[np.ndarray]]) -> list[list[np.ndarray]]: if not rows or not rows[0]: return rows ref_h, ref_w = rows[0][0].shape[:2] return [[_resize_panel_np(panel, ref_h, ref_w) for panel in row] for row in rows] def _save_gif(frames: list[np.ndarray], out_file: Path, duration_ms: int = 250) -> None: if not frames: return out_file.parent.mkdir(parents=True, exist_ok=True) pil_frames = [Image.fromarray(frame) for frame in frames] pil_frames[0].save( out_file, save_all=True, append_images=pil_frames[1:], duration=int(duration_ms), loop=0, disposal=2, ) def _depth_range(depth: torch.Tensor | None, fallback: tuple[float, float] = (0.0, 10.0)) -> tuple[float, float]: if not torch.is_tensor(depth): return fallback valid = depth[torch.isfinite(depth) & (depth > 0.0)] if int(valid.numel()) < 8: return fallback valid = valid.to(torch.float32).flatten() if int(valid.numel()) > 262144: step = max(1, int(valid.numel()) // 262144) valid = valid[::step] vmin = float(torch.quantile(valid, 0.01).item()) vmax = float(torch.quantile(valid, 0.99).item()) vmin = max(0.0, vmin) vmax = max(vmin + 1e-3, vmax) return (vmin, vmax) def _depth_panel(depth: torch.Tensor | None, val_min: float, val_max: float, blank: np.ndarray) -> np.ndarray: if not torch.is_tensor(depth): return blank d = depth.detach().to(torch.float32) valid = torch.isfinite(d) & (d > 0.0) if int(valid.sum().item()) < 8: return blank valid_vals = d[valid].flatten() if int(valid_vals.numel()) > 262144: step = max(1, int(valid_vals.numel()) // 262144) valid_vals = valid_vals[::step] fill = float(torch.quantile(valid_vals, 0.5).item()) d_safe = torch.where(valid, d, torch.full_like(d, fill)).clamp(min=val_min, max=val_max) panel = colorize_scalar_map(d_safe[0, 0], val_min=val_min, val_max=val_max, color_map="turbo") out = _to_u8_hwc(panel) out[~valid[0, 0].detach().cpu().numpy()] = 0 return out def _build_perspective_row(vis: dict[str, Any], which: str) -> list[np.ndarray]: if which not in ("src", "tgt"): raise ValueError(f"which must be src/tgt, got {which}") gt = vis[f"{which}_gt"].detach().to(torch.float32).clamp(0.0, 1.0) pred = vis[f"{which}_pred"].detach().to(torch.float32) alpha = vis[f"{which}_alpha"].detach().to(torch.float32).clamp(0.0, 1.0) pred_vis = linearRGB2sRGB((pred / alpha.clamp(min=1e-4)).clamp(0.0, 1.0)).clamp(0.0, 1.0) err = (pred_vis - gt).abs().mean(dim=1, keepdim=True) err_vals = err.flatten() if int(err_vals.numel()) > 262144: step = max(1, int(err_vals.numel()) // 262144) err_vals = err_vals[::step] vmax = float(max(1e-3, min(float(torch.quantile(err_vals, 0.99).item()), 0.5))) err_u8 = _to_u8_hwc(colorize_scalar_map(err[0, 0], val_min=0.0, val_max=vmax, color_map="turbo")) alpha_u8 = _to_u8_hwc(colorize_alpha(alpha)[0]) blank = np.zeros_like(_to_u8_hwc(gt[0])) gt_depth = vis.get(f"{which}_gt_depth", None) pred_depth = vis.get(f"{which}_pred_depth", None) dmin, dmax = _depth_range(gt_depth) if not torch.is_tensor(gt_depth): dmin, dmax = _depth_range(pred_depth) ref_d = _depth_panel(gt_depth, dmin, dmax, blank) pred_d = _depth_panel(pred_depth, dmin, dmax, blank) return [_to_u8_hwc(gt[0]), _to_u8_hwc(pred_vis[0]), err_u8, alpha_u8, ref_d, pred_d] def _build_perspective_gif_frame(vis: dict[str, Any], which: str) -> np.ndarray: pred = vis[f"{which}_pred"].detach().to(torch.float32) alpha = vis[f"{which}_alpha"].detach().to(torch.float32).clamp(0.0, 1.0) pred_vis = linearRGB2sRGB((pred / alpha.clamp(min=1e-4)).clamp(0.0, 1.0)).clamp(0.0, 1.0) return _to_u8_hwc(pred_vis[0]) def _cube_faces_u8(cube_img: torch.Tensor, face_count: int = 6) -> list[np.ndarray]: x = cube_img if x.ndim == 5 and x.shape[0] == 1: x = x[0] if x.ndim != 4: return [] faces = [] if x.shape[0] == face_count and x.shape[1] == 3: for i in range(face_count): faces.append(_to_u8_hwc(x[i])) elif x.shape[0] == face_count and x.shape[-1] == 3: for i in range(face_count): faces.append(_to_u8_hwc(x[i].permute(2, 0, 1).contiguous())) return faces def _build_hm3d_front_gif_frame(vis: dict[str, Any], which: str) -> np.ndarray: cube_pred = vis.get(f"{which}_cube_pred_linear", None) cube_alpha = vis.get(f"{which}_cube_alpha", None) if torch.is_tensor(cube_pred) and torch.is_tensor(cube_alpha): pred = linearRGB2sRGB( (cube_pred.detach().to(torch.float32) / cube_alpha.detach().to(torch.float32).clamp(min=1e-4)).clamp(0.0, 1.0) ).clamp(0.0, 1.0) faces = _cube_faces_u8(pred) if len(faces) == 6: return faces[3] return _build_perspective_gif_frame(vis, which) def _save_group_gif( *, dataset: str, group_dir: Path, group_key: str, step: int, group_items: list[dict[str, Any]], ) -> None: visual_items = [item for item in group_items if isinstance(item.get("vis", None), dict)] if not visual_items: return if dataset in {"hm3d", "replica"}: frames = [_build_hm3d_front_gif_frame(visual_items[0]["vis"], "src")] frames.extend(_build_hm3d_front_gif_frame(item["vis"], "tgt") for item in visual_items[:10]) else: frames = [_build_perspective_gif_frame(visual_items[0]["vis"], "src")] frames.extend(_build_perspective_gif_frame(item["vis"], "tgt") for item in visual_items[:10]) _save_gif(frames, group_dir / f"step_{int(step):07d}_{group_key}.gif") def _save_perspective_group_grid( *, group_dir: Path, group_key: str, step: int, group_items: list[dict[str, Any]], ) -> None: visual_items = [item for item in group_items if isinstance(item.get("vis", None), dict)] if not visual_items: return group_dir.mkdir(parents=True, exist_ok=True) first_vis = visual_items[0]["vis"] rows: list[list[np.ndarray]] = [_build_perspective_row(first_vis, "src")] for item in visual_items[:10]: rows.append(_build_perspective_row(item["vis"], "tgt")) while len(rows) < 11: rows.append(list(rows[-1])) rows = _normalize_rows_for_grid(rows) grid = _concat_grid(rows=rows, pad=6, pad_value=0) out_file = group_dir / f"step_{int(step):07d}_{group_key}_erp_11x6.png" save_image(grid, out_file) def _manifest_max_groups(args: argparse.Namespace) -> int: return max(0, int(getattr(args, "manifest_max_groups", 0))) def _validation_pseudo_root(args: argparse.Namespace) -> Path | None: root = getattr(args, "validation_pseudo_depth_root", None) return Path(root) if root is not None else Path("/media/team_data/ML4_team/datasets/sharp/validation_unik3d_pseudo_depth") def _re10k_pseudo_scene_key(scene: Any) -> str: key = str(scene).strip().replace("\\", "__").replace("/", "__") return key if key else "unknown_scene" def _re10k_training_pseudo_depth_path(args: argparse.Namespace, scene: Any, frame_idx: Any) -> Path | None: root = getattr(args, "re10k_pseudo_depth_root", None) if root is None: return None root = Path(root) split = str(getattr(args, "split", "test")) base = root if root.name == split else root / split try: frame_key = f"{int(frame_idx):05d}" except Exception: frame_key = str(frame_idx) return base / _re10k_pseudo_scene_key(scene) / f"{frame_key}.pt" def _load_re10k_training_pseudo_depth( args: argparse.Namespace, *, scene: Any, frame_idx: Any, intrinsics_k3: torch.Tensor, ) -> torch.Tensor | None: path = _re10k_training_pseudo_depth_path(args, scene, frame_idx) if path is None or not path.exists(): return None try: payload = _torch_load_any(path) depth_kind = "distance" if isinstance(payload, dict): depth = payload.get("z_depth_m", None) if torch.is_tensor(depth): depth_kind = "zdepth" else: depth = payload.get("distance_m", None) if torch.is_tensor(depth): depth_kind = "distance" else: depth = payload.get("depth_m", None) depth_kind = _normalize_depth_kind(payload.get("depth_kind", "distance"), default="distance") else: depth = payload depth_kind = "distance" if isinstance(depth, np.ndarray): depth = torch.from_numpy(depth) if not torch.is_tensor(depth): return None if depth.ndim == 2: depth = depth.unsqueeze(0) if depth.ndim != 3 or int(depth.shape[0]) != 1: return None depth = depth.to(torch.float32) max_depth_m = float(getattr(args, "max_depth_m", DEFAULT_MAX_DEPTH_M)) far_invalid_m = float(getattr(args, "re10k_pseudo_far_depth_invalid_m", 30.0)) valid = torch.isfinite(depth) & (depth > 0.0) if far_invalid_m > 0.0: valid = valid & (depth <= far_invalid_m) depth = torch.where(valid, depth, torch.zeros_like(depth)) if int(valid.sum().item()) <= 0: return None depth[valid] = depth[valid].clamp(max=max_depth_m) if _normalize_depth_kind(depth_kind, default="distance") != "zdepth": depth = _distance_to_z_depth_pinhole(depth, intrinsics_k3=intrinsics_k3) return depth except Exception: return None def _load_val_pseudo_depth( args: argparse.Namespace, *, dataset: str, scene: Any, frame_idx: Any, intrinsics_k3: torch.Tensor, ) -> torch.Tensor | None: if str(dataset) == "re10k": depth = _load_re10k_training_pseudo_depth( args, scene=scene, frame_idx=frame_idx, intrinsics_k3=intrinsics_k3, ) if torch.is_tensor(depth): return depth return _load_validation_pseudo_depth( _validation_pseudo_root(args), dataset=dataset, scene=scene, frame_idx=frame_idx, intrinsics_k3=intrinsics_k3, ) def _load_val_pseudo_depth_b1hw( args: argparse.Namespace, *, dataset: str, scene: Any, frame_idx: Any, intrinsics_k3: torch.Tensor, ) -> torch.Tensor | None: depth = _load_val_pseudo_depth( args, dataset=dataset, scene=scene, frame_idx=frame_idx, intrinsics_k3=intrinsics_k3, ) return depth.unsqueeze(0) if torch.is_tensor(depth) else None def _load_val_pseudo_distance_b1hw( args: argparse.Namespace, *, dataset: str, scene: Any, frame_idx: Any, ) -> torch.Tensor | None: depth = _load_validation_pseudo_distance( _validation_pseudo_root(args), dataset=dataset, scene=scene, frame_idx=frame_idx, ) return depth.unsqueeze(0) if torch.is_tensor(depth) else None def _iter_re10k_manifest_items(args: argparse.Namespace) -> Iterator[ValidationItem]: for group_idx, parts in _iter_manifest_parts(args, expected_parts=4): chunk_path = Path(parts[0]) scene = str(parts[1]) src_idx = int(parts[2]) tgt_indices = [int(x) for x in parts[3].split(",") if x.strip()] payload = _torch_load_any(chunk_path) if not isinstance(payload, list): continue example = next((ex for ex in payload if isinstance(ex, dict) and str(ex.get("key", chunk_path.stem)) == scene), None) if not isinstance(example, dict): continue poses = example.get("cameras", None) images = example.get("images", None) if not torch.is_tensor(poses) or not isinstance(images, list) or poses.ndim != 2 or poses.shape[1] != 18: continue if not (0 <= src_idx < len(images)): continue src_probe = _decode_rgb_u8(images[0]) h0, w0 = int(src_probe.shape[1]), int(src_probe.shape[2]) intr_all = torch.eye(3, dtype=torch.float32).unsqueeze(0).repeat(int(poses.shape[0]), 1, 1) intr_all[:, 0, 0] = poses[:, 0] * float(w0) intr_all[:, 1, 1] = poses[:, 1] * float(h0) intr_all[:, 0, 2] = poses[:, 2] * float(w0) - 0.5 intr_all[:, 1, 2] = poses[:, 3] * float(h0) - 0.5 w2c_all = torch.eye(4, dtype=torch.float32).unsqueeze(0).repeat(int(poses.shape[0]), 1, 1) w2c_all[:, :3] = poses[:, 6:].reshape(-1, 3, 4).to(torch.float32) group_key = f"{scene}_g{group_idx:05d}" adapter = _PinholeGroupAdapter( scene=scene, group_key=group_key, src_idx=int(src_idx), src_img=_decode_rgb_u8(images[src_idx]).unsqueeze(0), src_depth=_load_val_pseudo_depth_b1hw( args, dataset="re10k", scene=scene, frame_idx=src_idx, intrinsics_k3=intr_all[src_idx] ), src_w2c=w2c_all[src_idx].unsqueeze(0), src_k=intr_all[src_idx].unsqueeze(0).clone(), tgt_indices=tgt_indices, load_target=lambda tgt_idx, images=images, intr_all=intr_all, w2c_all=w2c_all: None if not (0 <= int(tgt_idx) < len(images)) else _PinholeTargetAdapter( idx=int(tgt_idx), img=_decode_rgb_u8(images[int(tgt_idx)]).unsqueeze(0), w2c=w2c_all[int(tgt_idx)].unsqueeze(0), k=intr_all[int(tgt_idx)].unsqueeze(0).clone(), depth=_load_val_pseudo_depth_b1hw( args, dataset="re10k", scene=scene, frame_idx=int(tgt_idx), intrinsics_k3=intr_all[int(tgt_idx)] ), ), ) yield from _yield_pinhole_group_batches("re10k", adapter, args) def _iter_wildrgbd_manifest_items(args: argparse.Namespace) -> Iterator[ValidationItem]: root_map: dict[str, Path] = {} for root in _wild_validation_roots(Path(args.data_root)): scene_parent = root / "scenes" for scene_dir in sorted([p for p in scene_parent.iterdir() if p.is_dir()]) if scene_parent.exists() else []: root_map[f"{root.name}/{scene_dir.name}"] = scene_dir for group_idx, parts in _iter_manifest_parts(args, expected_parts=3): scene_name = str(parts[0]) src_idx = int(parts[1]) tgt_indices = [int(x) for x in parts[2].split(",") if x.strip()] scene_dir = root_map.get(scene_name) if scene_dir is None: continue pose_ids_np, w2c_map, intr = WildRGBDDataset._load_scene_pose_and_k(scene_dir) pose_ids = {int(x) for x in pose_ids_np.tolist()} rgb_ids = WildRGBDDataset._collect_frame_ids(scene_dir / "rgb") dep_ids = WildRGBDDataset._collect_frame_ids(scene_dir / "depth") valid_ids = pose_ids & rgb_ids & dep_ids if int(src_idx) not in valid_ids: continue ds_loader = WildRGBDDataset(root=scene_dir.parent.parent, split="scenes", scene_list_file=None) group_key = f"{scene_name}_g{group_idx:05d}" adapter = _PinholeGroupAdapter( scene=scene_name, group_key=group_key, src_idx=src_idx, src_img=WildRGBDDataset._load_rgb_u8(WildRGBDDataset._resolve_img_path(scene_dir / "rgb", src_idx)).unsqueeze(0), src_depth=ds_loader._load_depth_m(WildRGBDDataset._resolve_img_path(scene_dir / "depth", src_idx)).unsqueeze(0), src_w2c=torch.from_numpy(w2c_map[src_idx]).to(torch.float32).unsqueeze(0), src_k=intr.to(torch.float32).unsqueeze(0).clone(), tgt_indices=tgt_indices, load_target=lambda tgt_idx, scene_dir=scene_dir, valid_ids=valid_ids, ds_loader=ds_loader, intr=intr, w2c_map=w2c_map: None if int(tgt_idx) not in valid_ids else _PinholeTargetAdapter( idx=int(tgt_idx), img=WildRGBDDataset._load_rgb_u8( WildRGBDDataset._resolve_img_path(scene_dir / "rgb", int(tgt_idx)) ).unsqueeze(0), depth=ds_loader._load_depth_m( WildRGBDDataset._resolve_img_path(scene_dir / "depth", int(tgt_idx)) ).unsqueeze(0), w2c=torch.from_numpy(w2c_map[int(tgt_idx)]).to(torch.float32).unsqueeze(0), k=intr.to(torch.float32).unsqueeze(0).clone(), ), ) yield from _yield_pinhole_group_batches("wildrgbd", adapter, args) def _iter_hm3d_manifest_items(args: argparse.Namespace) -> Iterator[ValidationItem]: root = Path(args.data_root) manifest_in = _read_manifest_lines(getattr(args, "manifest_file", None), max_lines=_manifest_max_groups(args)) for group_idx, raw in enumerate(manifest_in): parts = raw.split("|") if len(parts) != 3: continue scene_name = str(parts[0]) src_idx = int(parts[1]) tgt_indices = [int(x) for x in parts[2].split(",") if x.strip()] scene_dir = root / scene_name pano_dir = scene_dir / "pano" depth_dir = scene_dir / "pano_depth" cube_dir = scene_dir / "cubemaps" cube_depth_dir = scene_dir / "cubemaps_depth" if not (pano_dir.exists() and depth_dir.exists() and cube_dir.exists() and cube_depth_dir.exists()): continue R_np, t_np = _load_hm3d_pose(scene_dir) group_key = f"{scene_dir.name}_g{group_idx:05d}" src_rgb = _load_png_rgb_u8(pano_dir / f"{src_idx:05d}.png") src_dep = _load_png_depth_m(depth_dir / f"{src_idx:05d}.png") src_cube = _torch_load_any(cube_dir / f"{src_idx:05d}.torch") src_cdep = _torch_load_any(cube_depth_dir / f"{src_idx:05d}.torch") src_R = torch.from_numpy(R_np[src_idx]) src_t = torch.from_numpy(t_np[src_idx]) batch_size = max(1, int(getattr(args, "validation_batch_size", 1))) samples: list[Any] = [] tags: list[str] = [] for tgt_idx in tgt_indices: tgt_rgb = _load_png_rgb_u8(pano_dir / f"{tgt_idx:05d}.png") tgt_dep = _load_png_depth_m(depth_dir / f"{tgt_idx:05d}.png") tgt_cube = _torch_load_any(cube_dir / f"{tgt_idx:05d}.torch") tgt_cdep = _torch_load_any(cube_depth_dir / f"{tgt_idx:05d}.torch") sample = SimpleNamespace( src_erp_rgb_u8=src_rgb, tgt_erp_rgb_u8=tgt_rgb, 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=torch.from_numpy(R_np[tgt_idx]), tgt_t=torch.from_numpy(t_np[tgt_idx]), src_idx=src_idx, tgt_idx=tgt_idx, scene=scene_dir.name, ) samples.append(sample) tags.append(f"{group_key}_t{tgt_idx:05d}") if len(samples) >= batch_size: yield from _yield_panogs_group_batches("hm3d", group_key, samples, tags, args) samples = [] tags = [] yield from _yield_panogs_group_batches("hm3d", group_key, samples, tags, args) def _iter_replica_manifest_items(args: argparse.Namespace) -> Iterator[ValidationItem]: root = _resolve_replica_test_root(Path(args.data_root)) manifest_in = _read_manifest_lines(getattr(args, "manifest_file", None), max_lines=_manifest_max_groups(args)) for group_idx, raw in enumerate(manifest_in): parts = raw.split("|") if len(parts) != 3: continue scene_name = str(parts[0]) src_idx = int(parts[1]) tgt_indices = [int(x) for x in parts[2].split(",") if x.strip()] scene_dir = root / scene_name pano_dir = scene_dir / "pano" depth_dir = scene_dir / "pano_depth" cube_dir = scene_dir / "cubemaps" cube_depth_dir = scene_dir / "cubemaps_depth" if not (pano_dir.exists() and depth_dir.exists() and cube_dir.exists() and cube_depth_dir.exists()): continue R_np, t_np = _load_hm3d_pose(scene_dir) group_key = f"replica_{scene_dir.name}_g{group_idx:05d}" src_rgb = _load_png_rgb_u8(pano_dir / f"{src_idx:05d}.png") src_dep = _load_png_depth_m(depth_dir / f"{src_idx:05d}.png") src_cube = _torch_load_any(cube_dir / f"{src_idx:05d}.torch") src_cdep = _torch_load_any(cube_depth_dir / f"{src_idx:05d}.torch") src_R = torch.from_numpy(R_np[src_idx]) src_t = torch.from_numpy(t_np[src_idx]) batch_size = max(1, int(getattr(args, "validation_batch_size", 1))) samples: list[Any] = [] tags: list[str] = [] for tgt_idx in tgt_indices: tgt_rgb = _load_png_rgb_u8(pano_dir / f"{tgt_idx:05d}.png") tgt_dep = _load_png_depth_m(depth_dir / f"{tgt_idx:05d}.png") tgt_cube = _torch_load_any(cube_dir / f"{tgt_idx:05d}.torch") tgt_cdep = _torch_load_any(cube_depth_dir / f"{tgt_idx:05d}.torch") sample = SimpleNamespace( src_erp_rgb_u8=src_rgb, tgt_erp_rgb_u8=tgt_rgb, 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=torch.from_numpy(R_np[tgt_idx]), tgt_t=torch.from_numpy(t_np[tgt_idx]), src_idx=src_idx, tgt_idx=tgt_idx, scene=scene_dir.name, ) samples.append(sample) tags.append(f"{group_key}_t{tgt_idx:05d}") if len(samples) >= batch_size: yield from _yield_panogs_group_batches("replica", group_key, samples, tags, args) samples = [] tags = [] yield from _yield_panogs_group_batches("replica", group_key, samples, tags, args) def _iter_sim_manifest_items(args: argparse.Namespace) -> Iterator[ValidationItem]: root = Path(args.data_root) pose_root = Path(getattr(args, "sim_pose_root", root / "30cm")) manifest_in = _read_manifest_lines(getattr(args, "manifest_file", None), max_lines=_manifest_max_groups(args)) dataset = SimPanoramaDataset( root=root, pose_root=pose_root, scene_names=["AI_vol3_03"], scene_list_file=None, max_index_gap=10, pair_max_translation_m=0.5, pair_min_depth_overlap=0.0, chunk_size=30, shuffle_scene=False, depth_max_m=float(getattr(args, "max_depth_m", DEFAULT_MAX_DEPTH_M)), far_depth_invalid_m=float(getattr(args, "sim_far_depth_invalid_m", 30.0)), far_depth_invalid_max_frac=float(getattr(args, "sim_far_depth_invalid_max_frac", 1.0)), seed=int(args.seed), ) def _load_scene(scene_name: str) -> tuple[dict[int, Any], _EquirecToCube] | None: try: frames = dataset._load_or_build_scene_frames(scene_name) if not frames: return None first_rgb = dataset._load_rgb(frames[0].rgb_path) equ_h, equ_w = int(first_rgb.shape[1]), int(first_rgb.shape[2]) converter = _EquirecToCube(equ_h=equ_h, equ_w=equ_w, face_w=max(1, equ_h // 2)) frame_map = {int(frame.frame_idx): frame for frame in frames} except Exception as exc: LOGGER.warning("Skip SIM scene=%s: %s", scene_name, str(exc)) return None return frame_map, converter def _load_frame(frame: Any, converter: _EquirecToCube) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: rgb = dataset._load_rgb(frame.rgb_path) depth = dataset._load_depth(frame.depth_path) cube_rgb, cube_depth = converter.run(rgb, depth) return rgb, depth, cube_rgb, cube_depth for group_idx, raw in enumerate(manifest_in): parts = raw.split("|") if len(parts) != 3: continue scene_name = str(parts[0]) src_idx = int(parts[1]) tgt_indices = [int(x) for x in parts[2].split(",") if x.strip()] loaded_scene = _load_scene(scene_name) if loaded_scene is None: continue frame_map, converter = loaded_scene src_frame = frame_map.get(src_idx) if src_frame is None: continue try: src_rgb, src_dep, src_cube, src_cdep = _load_frame(src_frame, converter) except Exception as exc: LOGGER.warning("Skip SIM src scene=%s src=%d: %s", scene_name, int(src_idx), str(exc)) continue group_key = f"sim_{scene_name}_g{group_idx:05d}" batch_size = max(1, int(getattr(args, "validation_batch_size", 1))) samples: list[Any] = [] tags: list[str] = [] for tgt_idx in tgt_indices: tgt_frame = frame_map.get(int(tgt_idx)) if tgt_frame is None: continue try: tgt_rgb, tgt_dep, tgt_cube, tgt_cdep = _load_frame(tgt_frame, converter) except Exception: continue sample = SimpleNamespace( src_erp_rgb_u8=src_rgb, tgt_erp_rgb_u8=tgt_rgb, 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=torch.eye(3, dtype=torch.float32), src_t=src_frame.position_xyz.clone(), tgt_R=torch.eye(3, dtype=torch.float32), tgt_t=tgt_frame.position_xyz.clone(), src_idx=src_idx, tgt_idx=int(tgt_idx), scene=scene_name, ) samples.append(sample) tags.append(f"{group_key}_t{int(tgt_idx):05d}") if len(samples) >= batch_size: yield from _yield_panogs_group_batches("sim", group_key, samples, tags, args) samples = [] tags = [] yield from _yield_panogs_group_batches("sim", group_key, samples, tags, args) getattr(dataset, "_scene_frames_cache", {}).pop(scene_name, None) getattr(dataset, "_scene_pair_cache", {}).pop(scene_name, None) def _iter_scannetpp_manifest_items(args: argparse.Namespace) -> Iterator[ValidationItem]: for group_idx, parts in _iter_manifest_parts(args, expected_parts=4): tf = Path(parts[0]) sample_key = str(parts[1]) src_idx = int(parts[2]) tgt_indices = [int(x) for x in parts[3].split(",") if x.strip()] payload = _torch_load_any(tf) sample_raw = payload[0] if isinstance(payload, list) else payload if str(sample_raw.get("key", tf.stem)) != sample_key: continue cameras = sample_raw["cameras"].to(torch.float32) images = sample_raw["images"] if not isinstance(images, list) or int(cameras.shape[0]) < 11: continue w2c_all = [] intr_all = [] for i in range(int(cameras.shape[0])): cam = cameras[i] fx_n, fy_n, cx_n, cy_n, w0, h0 = cam[:6] w2c = torch.eye(4, dtype=torch.float32) w2c[:3, :] = cam[6:].reshape(3, 4) k = torch.eye(3, dtype=torch.float32) k[0, 0] = fx_n * w0 k[1, 1] = fy_n * h0 k[0, 2] = cx_n * w0 k[1, 2] = cy_n * h0 w2c_all.append(w2c) intr_all.append(k) w2c_t = torch.stack(w2c_all, dim=0) intr_t = torch.stack(intr_all, dim=0) group_key = f"{sample_key}_g{group_idx:05d}" adapter = _PinholeGroupAdapter( scene=sample_key, group_key=group_key, src_idx=int(src_idx), src_img=_decode_rgb_u8(images[src_idx]).unsqueeze(0), src_depth=_load_val_pseudo_depth_b1hw( args, dataset="scannetpp", scene=sample_key, frame_idx=src_idx, intrinsics_k3=intr_t[src_idx] ), src_w2c=w2c_t[src_idx].unsqueeze(0), src_k=intr_t[src_idx].unsqueeze(0), tgt_indices=tgt_indices, load_target=lambda tgt_idx, images=images, intr_t=intr_t, w2c_t=w2c_t: _PinholeTargetAdapter( idx=int(tgt_idx), img=_decode_rgb_u8(images[int(tgt_idx)]).unsqueeze(0), w2c=w2c_t[int(tgt_idx)].unsqueeze(0), k=intr_t[int(tgt_idx)].unsqueeze(0), depth=_load_val_pseudo_depth_b1hw( args, dataset="scannetpp", scene=sample_key, frame_idx=int(tgt_idx), intrinsics_k3=intr_t[int(tgt_idx)], ), ), ) yield from _yield_pinhole_group_batches("scannetpp", adapter, args) def _iter_scanetpp_fisheye_manifest_items(args: argparse.Namespace) -> Iterator[ValidationItem]: root = Path(args.data_root) loader = ScannetppFisheyeDataset( root=root, scene_list_file=None, min_frame_gap=1, max_frame_gap=10, pair_max_translation_m=float(args.pair_max_translation_m), shuffle_scene=False, shuffle_frame=False, skip_bad=True, batch_size_hint=1, depth_max_m=float(getattr(args, "max_depth_m", DEFAULT_MAX_DEPTH_M)), far_depth_invalid_m=float(getattr(args, "scanetpp_fisheye_far_depth_invalid_m", 30.0)), seed=int(args.seed), ) batch_size = max(1, int(getattr(args, "validation_batch_size", 1))) for group_idx, parts in _iter_manifest_parts(args, expected_parts=4): scene_id = str(parts[0]) scene_dir = Path(parts[1]) if not scene_dir.is_absolute(): scene_dir = root / scene_dir src_pos = int(parts[2]) tgt_positions = [int(x) for x in parts[3].split(",") if x.strip()] try: camera_params, frames = loader._load_scene_frames(scene_id, scene_dir) except Exception as exc: LOGGER.warning("Skip ScanNet++ fisheye scene=%s: %s", scene_id, str(exc)) continue if not (0 <= src_pos < len(frames)): continue try: src_loaded = loader._load_frame_tensor(frames[src_pos], camera_params) except Exception as exc: LOGGER.warning("Skip ScanNet++ fisheye src scene=%s src=%d: %s", scene_id, int(src_pos), str(exc)) continue pending_pos: list[int] = [] pending_frames: list[dict[str, Any]] = [] pending_loaded: list[dict[str, torch.Tensor]] = [] def _flush() -> Iterator[tuple[str, Any, str | list[str], str]]: if not pending_pos: return group_key = f"scanetpp_fisheye_{scene_id}_g{group_idx:05d}" batch = _make_scanetpp_fisheye_batch( scene=scene_id, src_pos=src_pos, tgt_positions=list(pending_pos), src_frame=frames[src_pos], tgt_frames=list(pending_frames), src_loaded=src_loaded, tgt_loaded=list(pending_loaded), ) tags = [f"{group_key}_t{int(t):05d}" for t in pending_pos] yield ("scanetpp_fisheye", batch, tags[0] if len(tags) == 1 else tags, group_key) for tgt_pos in tgt_positions: if not (0 <= int(tgt_pos) < len(frames)): continue try: tgt_loaded = loader._load_frame_tensor(frames[int(tgt_pos)], camera_params) except Exception: continue pending_pos.append(int(tgt_pos)) pending_frames.append(frames[int(tgt_pos)]) pending_loaded.append(tgt_loaded) if len(pending_pos) >= batch_size: yield from _flush() pending_pos = [] pending_frames = [] pending_loaded = [] if pending_pos: yield from _flush() def _load_smx_sim_fisheye_scene(scene_dir: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]: meta = json.loads((scene_dir / "transforms.json").read_text(encoding="utf-8")) raw_frames = list(meta.get("frames", [])) frames: list[dict[str, Any]] = [] for local_idx, frame in enumerate(raw_frames): rel = Path(str(frame.get("file_path", ""))) image_path = scene_dir / rel if not image_path.exists(): image_path = Path(str(frame.get("source_image", ""))) source_image = Path(str(frame.get("source_image", ""))) if not image_path.exists() or frame.get("transform_matrix") is None: continue c2w = torch.tensor(frame["transform_matrix"], dtype=torch.float32) frames.append( { "image_name": image_path.name, "image_path": image_path, "source_image": source_image, "w2c": torch.linalg.inv(c2w), "idx": int(frame.get("source_image_index", local_idx)), "pos": int(local_idx), "yaw_pitch_roll_deg": list(frame.get("yaw_pitch_roll_deg", [0.0, 0.0, 0.0])), } ) return meta, frames def _smx_sim_fisheye_valid_mask(rgb_u8: torch.Tensor, meta: dict[str, Any]) -> torch.Tensor: h, w = int(rgb_u8.shape[-2]), int(rgb_u8.shape[-1]) yy, xx = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij") cx = float(meta.get("cx", w * 0.5)) cy = float(meta.get("cy", h * 0.5)) radius = float(meta.get("valid_radius_px", min(h, w) * 0.5)) circle = ((xx.to(torch.float32) - cx) ** 2 + (yy.to(torch.float32) - cy) ** 2) <= radius * radius nonblack = rgb_u8.to(torch.float32).sum(dim=0) > 1.0 return (circle & nonblack).to(torch.float32).unsqueeze(0) def _smx_sim_rotation_yaw_pitch_roll(yaw_deg: float, pitch_deg: float, roll_deg: float) -> torch.Tensor: yaw = math.radians(float(yaw_deg)) pitch = math.radians(float(pitch_deg)) roll = math.radians(float(roll_deg)) cy, sy = math.cos(yaw), math.sin(yaw) cp, sp = math.cos(pitch), math.sin(pitch) cr, sr = math.cos(roll), math.sin(roll) r_yaw = torch.tensor([[cy, 0.0, sy], [0.0, 1.0, 0.0], [-sy, 0.0, cy]], dtype=torch.float32) r_pitch = torch.tensor([[1.0, 0.0, 0.0], [0.0, cp, sp], [0.0, -sp, cp]], dtype=torch.float32) r_roll = torch.tensor([[cr, -sr, 0.0], [sr, cr, 0.0], [0.0, 0.0, 1.0]], dtype=torch.float32) return (r_yaw @ r_pitch @ r_roll).to(torch.float32) def _smx_sim_fisheye_grid( *, meta: dict[str, Any], frame: dict[str, Any], erp_h: int, erp_w: int, fish_h: int, fish_w: int, ) -> tuple[torch.Tensor, torch.Tensor]: yy, xx = torch.meshgrid( torch.arange(fish_h, dtype=torch.float32) + 0.5, torch.arange(fish_w, dtype=torch.float32) + 0.5, indexing="ij", ) fx = float(meta.get("fl_x", fish_w / max(math.radians(float(meta.get("fov_deg", 130.0))), 1e-6))) fy = float(meta.get("fl_y", fish_h / max(math.radians(float(meta.get("fov_deg", 130.0))), 1e-6))) cx = float(meta.get("cx", fish_w * 0.5)) cy = float(meta.get("cy", fish_h * 0.5)) fov_rad = math.radians(float(meta.get("fov_deg", 130.0))) half_fov = 0.5 * fov_rad dx = (xx - cx) / max(fx, 1e-6) dy = (yy - cy) / max(fy, 1e-6) theta = torch.sqrt(dx * dx + dy * dy) valid = theta <= float(half_fov) scale = torch.zeros_like(theta) nonzero = theta > 1e-8 scale[nonzero] = torch.sin(theta[nonzero]) / theta[nonzero] rays = torch.stack([dx * scale, -dy * scale, torch.cos(theta)], dim=-1) rays[~nonzero] = torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32) ypr = list(frame.get("yaw_pitch_roll_deg", [0.0, 0.0, 0.0])) rot = _smx_sim_rotation_yaw_pitch_roll(float(ypr[0]), float(ypr[1]), float(ypr[2])) rays = rays @ rot.T rays = rays / rays.norm(dim=-1, keepdim=True).clamp_min(1e-8) lon = torch.atan2(rays[..., 0], rays[..., 2]) lat = torch.atan2(rays[..., 1], torch.sqrt(rays[..., 0] ** 2 + rays[..., 2] ** 2)) map_x = (lon / (2.0 * math.pi) + 0.5) * float(erp_w) - 0.5 map_y = (0.5 - lat / math.pi) * float(erp_h) - 0.5 map_x = torch.remainder(map_x + 0.5, float(erp_w)) - 0.5 x_norm = 2.0 * (map_x + 0.5) / float(erp_w) - 1.0 y_norm = 2.0 * (map_y + 0.5) / float(erp_h) - 1.0 grid = torch.stack([x_norm, y_norm], dim=-1).to(torch.float32) return grid, valid.to(torch.float32).unsqueeze(0) def _smx_sim_project_erp_tensor_to_fisheye( tensor: torch.Tensor, grid: torch.Tensor, valid_mask: torch.Tensor, *, mode: str = "bilinear", ) -> torch.Tensor: x = tensor.detach().to(torch.float32) if x.ndim == 3: x = x.unsqueeze(0) if x.ndim != 4: raise ValueError(f"Expected 3D/4D tensor for SMX projection, got shape={tuple(x.shape)}") device = x.device grid_b = grid.to(device=device, dtype=torch.float32).unsqueeze(0).expand(int(x.shape[0]), -1, -1, -1) mask_b = valid_mask.to(device=device, dtype=x.dtype) if mask_b.ndim == 3: mask_b = mask_b.unsqueeze(0) if mask_b.ndim != 4: raise ValueError(f"Expected 3D/4D SMX valid mask, got shape={tuple(mask_b.shape)}") mask_b = mask_b.expand(int(x.shape[0]), -1, -1, -1) out = F.grid_sample(x, grid_b, mode=mode, padding_mode="zeros", align_corners=False) return out * mask_b def _project_smx_sim_fisheye_vis(vis: dict[str, Any], batch: Any, sample_idx: int) -> dict[str, Any]: if not hasattr(batch, "smx_fisheye_src_grid"): return vis i = int(sample_idx) src_grid = batch.smx_fisheye_src_grid[i] tgt_grid = batch.smx_fisheye_tgt_grid[i] src_mask = batch.smx_fisheye_src_valid_mask[i : i + 1] tgt_mask = batch.smx_fisheye_tgt_valid_mask[i : i + 1] def _batch_value(name: str, idx: int, device: torch.device) -> torch.Tensor | None: value = getattr(batch, name, None) if not torch.is_tensor(value) or idx >= int(value.shape[0]): return None return value[idx : idx + 1].to(device=device) device = vis["tgt_pred"].device if torch.is_tensor(vis.get("tgt_pred", None)) else torch.device("cpu") src_gt = _batch_value("smx_fisheye_src_rgb_u8", i, device) tgt_gt = _batch_value("smx_fisheye_tgt_rgb_u8", i, device) if src_gt is not None: vis["src_gt"] = (src_gt.to(torch.float32) / 255.0).clamp(0.0, 1.0) if tgt_gt is not None: vis["tgt_gt"] = (tgt_gt.to(torch.float32) / 255.0).clamp(0.0, 1.0) for prefix, grid, mask in (("src", src_grid, src_mask), ("tgt", tgt_grid, tgt_mask)): pred_key = f"{prefix}_pred" alpha_key = f"{prefix}_alpha" if torch.is_tensor(vis.get(pred_key, None)): vis[pred_key] = _smx_sim_project_erp_tensor_to_fisheye(vis[pred_key], grid, mask, mode="bilinear") if torch.is_tensor(vis.get(alpha_key, None)): vis[alpha_key] = _smx_sim_project_erp_tensor_to_fisheye(vis[alpha_key], grid, mask, mode="bilinear") for depth_key in (f"{prefix}_gt_depth", f"{prefix}_pred_depth", f"{prefix}_unik3d_depth"): if torch.is_tensor(vis.get(depth_key, None)): vis[depth_key] = _smx_sim_project_erp_tensor_to_fisheye(vis[depth_key], grid, mask, mode="bilinear") vis["src_metric_mask"] = src_mask.to(device=device, dtype=torch.float32) vis["tgt_metric_mask"] = tgt_mask.to(device=device, dtype=torch.float32) vis["dataset_name"] = "smx_sim_fisheye" vis["projection_pipeline"] = "source_fisheye_to_source_pano_infer_target_pano_to_target_fisheye" return vis def _yield_smx_sim_fisheye_pano_batches( *, group_key: str, samples: list[Any], tags: list[str], args: argparse.Namespace, ) -> Iterator[ValidationItem]: batch_size = max(1, int(getattr(args, "validation_batch_size", 1))) if len(samples) != len(tags): raise ValueError(f"Expected samples/tags length match, got {len(samples)} vs {len(tags)}") for start in range(0, len(samples), batch_size): end = min(len(samples), start + batch_size) chunk = samples[start:end] batch_tags = tags[start:end] batch = panogs_collate(chunk) object.__setattr__(batch, "collect_all_vis", True) object.__setattr__(batch, "disable_depth_gt", True) for attr in ( "smx_fisheye_src_rgb_u8", "smx_fisheye_tgt_rgb_u8", "smx_fisheye_src_valid_mask", "smx_fisheye_tgt_valid_mask", "smx_fisheye_src_grid", "smx_fisheye_tgt_grid", ): object.__setattr__(batch, attr, torch.stack([getattr(s, attr) for s in chunk], dim=0)) yield ("smx_sim_fisheye", batch, batch_tags[0] if len(batch_tags) == 1 else batch_tags, group_key) def _smx_frame_position_from_w2c(frame: dict[str, Any], meta: dict[str, Any]) -> torch.Tensor: w2c = frame["w2c"] if torch.is_tensor(w2c): w2c_t = w2c.detach().clone().to(torch.float32) else: w2c_t = torch.as_tensor(w2c, dtype=torch.float32) raw_xyz = torch.linalg.inv(w2c_t)[:3, 3].clone() raw_scale = float(meta.get("position_scale", 1.0)) if abs(raw_scale) > 1e-8: raw_xyz = raw_xyz / raw_scale return torch.stack([raw_xyz[1], -raw_xyz[2], raw_xyz[0]], dim=0).to(torch.float32) * 0.01 def _iter_smx_sim_fisheye_manifest_items(args: argparse.Namespace) -> Iterator[ValidationItem]: root = Path(args.data_root) batch_size = max(1, int(getattr(args, "validation_batch_size", 1))) for group_idx, parts in _iter_manifest_parts(args, expected_parts=4): scene_id = str(parts[0]) scene_dir = Path(parts[1]) if not scene_dir.is_absolute(): scene_dir = root / scene_dir src_pos = int(parts[2]) tgt_positions = [int(x) for x in parts[3].split(",") if x.strip()] try: meta, frames = _load_smx_sim_fisheye_scene(scene_dir) except Exception as exc: LOGGER.warning("Skip SMX SIM fisheye scene=%s: %s", scene_id, str(exc)) continue if not (0 <= src_pos < len(frames)): continue def _load_frame(frame: dict[str, Any]) -> dict[str, torch.Tensor]: source_image = Path(str(frame.get("source_image", ""))) if not source_image.exists(): raise FileNotFoundError(source_image) erp_rgb = _load_png_rgb_u8(source_image) converter = _EquirecToCube( equ_h=int(erp_rgb.shape[-2]), equ_w=int(erp_rgb.shape[-1]), face_w=max(1, int(erp_rgb.shape[-2]) // 2), ) cube_rgb = converter.run_rgb(erp_rgb) erp_depth = torch.zeros((1, int(erp_rgb.shape[-2]), int(erp_rgb.shape[-1])), dtype=torch.float32) cube_depth = torch.zeros((6, int(converter.face_w), int(converter.face_w), 1), dtype=torch.float32) fish_rgb = _load_png_rgb_u8(Path(frame["image_path"])) grid, valid = _smx_sim_fisheye_grid( meta=meta, frame=frame, erp_h=int(erp_rgb.shape[-2]), erp_w=int(erp_rgb.shape[-1]), fish_h=int(fish_rgb.shape[-2]), fish_w=int(fish_rgb.shape[-1]), ) valid = (valid * _smx_sim_fisheye_valid_mask(fish_rgb, meta)).clamp(0.0, 1.0) return { "erp_rgb_u8": erp_rgb, "erp_depth_m": erp_depth, "cube_rgb_u8": cube_rgb, "cube_depth_m": cube_depth, "fish_rgb_u8": fish_rgb, "fish_valid_mask": valid, "fish_grid": grid, } try: src_loaded = _load_frame(frames[src_pos]) except Exception as exc: LOGGER.warning("Skip SMX SIM fisheye src scene=%s src=%d: %s", scene_id, int(src_pos), str(exc)) continue pending_pos: list[int] = [] pending_samples: list[Any] = [] pending_tags: list[str] = [] def _flush() -> Iterator[tuple[str, Any, str | list[str], str]]: if not pending_pos: return group_key = f"smx_sim_fisheye_{scene_id}_g{group_idx:05d}" yield from _yield_smx_sim_fisheye_pano_batches( group_key=group_key, samples=list(pending_samples), tags=list(pending_tags), args=args, ) for tgt_pos in tgt_positions: if not (0 <= int(tgt_pos) < len(frames)): continue try: tgt_loaded = _load_frame(frames[int(tgt_pos)]) except Exception: continue group_key = f"smx_sim_fisheye_{scene_id}_g{group_idx:05d}" sample = SimpleNamespace( src_erp_rgb_u8=src_loaded["erp_rgb_u8"], tgt_erp_rgb_u8=tgt_loaded["erp_rgb_u8"], src_erp_depth_m=src_loaded["erp_depth_m"], tgt_erp_depth_m=tgt_loaded["erp_depth_m"], src_cube_rgb_u8=src_loaded["cube_rgb_u8"], tgt_cube_rgb_u8=tgt_loaded["cube_rgb_u8"], src_cube_depth_m=src_loaded["cube_depth_m"], tgt_cube_depth_m=tgt_loaded["cube_depth_m"], src_R=torch.eye(3, dtype=torch.float32), src_t=_smx_frame_position_from_w2c(frames[src_pos], meta), tgt_R=torch.eye(3, dtype=torch.float32), tgt_t=_smx_frame_position_from_w2c(frames[int(tgt_pos)], meta), src_idx=int(frames[src_pos].get("idx", src_pos)), tgt_idx=int(frames[int(tgt_pos)].get("idx", tgt_pos)), scene=scene_id, smx_fisheye_src_rgb_u8=src_loaded["fish_rgb_u8"], smx_fisheye_tgt_rgb_u8=tgt_loaded["fish_rgb_u8"], smx_fisheye_src_valid_mask=src_loaded["fish_valid_mask"], smx_fisheye_tgt_valid_mask=tgt_loaded["fish_valid_mask"], smx_fisheye_src_grid=src_loaded["fish_grid"], smx_fisheye_tgt_grid=tgt_loaded["fish_grid"], ) pending_pos.append(int(tgt_pos)) pending_samples.append(sample) pending_tags.append(f"{group_key}_t{int(tgt_pos):05d}") if len(pending_pos) >= batch_size: yield from _flush() pending_pos = [] pending_samples = [] pending_tags = [] if pending_pos: yield from _flush() def _iter_tat_manifest_items(args: argparse.Namespace) -> Iterator[ValidationItem]: root = Path(args.data_root) scene_roots = _colmap_scene_roots(root) scene_root_map = {scene_root.name: scene_root for scene_root in scene_roots} for group_idx, parts in _iter_manifest_parts(args, expected_parts=3): scene_name = str(parts[0]) scene_root = scene_root_map.get(scene_name) if scene_root is None: continue image_dir = _colmap_image_dir(scene_root) image_paths = sorted([p for p in image_dir.iterdir() if p.suffix.lower() in (".png", ".jpg", ".jpeg")]) image_map = {p.name: p for p in image_paths} colmap_entries = _load_scaled_colmap_entries(scene_root) if not colmap_entries: continue image_paths = [p for p in image_paths if p.name in colmap_entries] image_map = {p.name: p for p in image_paths} src_name = str(parts[1]) tgt_names = [x for x in parts[2].split(",") if x.strip()] if src_name not in image_map: continue group_key = f"tat_{scene_name}_g{group_idx:05d}" src_img = _load_png_rgb_u8(image_map[src_name]).unsqueeze(0) src_meta = colmap_entries[src_name] src_k = src_meta["k"].unsqueeze(0).clone() src_w2c = src_meta["w2c"].unsqueeze(0).clone() ref_h = int(src_meta["height"]) ref_w = int(src_meta["width"]) if (int(src_img.shape[-2]) != ref_h) or (int(src_img.shape[-1]) != ref_w): sx0 = float(int(src_img.shape[-1])) / float(ref_w) sy0 = float(int(src_img.shape[-2])) / float(ref_h) src_k = _resize_k3_align_corners_false(src_k, sx=sx0, sy=sy0) name_to_idx = {p.name: i for i, p in enumerate(image_paths)} src_idx = int(name_to_idx[src_name]) def _load_tat_target(tgt_idx: int) -> _PinholeTargetAdapter | None: tgt_name = image_paths[int(tgt_idx)].name if tgt_name not in image_map or tgt_name not in colmap_entries: return None tgt_img = _load_png_rgb_u8(image_map[tgt_name]).unsqueeze(0) tgt_meta = colmap_entries[tgt_name] tgt_k = tgt_meta["k"].unsqueeze(0).clone() if (int(tgt_img.shape[-2]) != int(tgt_meta["height"])) or (int(tgt_img.shape[-1]) != int(tgt_meta["width"])): sx0 = float(int(tgt_img.shape[-1])) / float(int(tgt_meta["width"])) sy0 = float(int(tgt_img.shape[-2])) / float(int(tgt_meta["height"])) tgt_k = _resize_k3_align_corners_false(tgt_k, sx=sx0, sy=sy0) return _PinholeTargetAdapter( idx=int(tgt_idx), img=tgt_img, w2c=tgt_meta["w2c"].unsqueeze(0).clone(), k=tgt_k, depth=_load_val_pseudo_depth_b1hw( args, dataset="tat", scene=scene_name, frame_idx=int(tgt_idx), intrinsics_k3=tgt_k[0] ), ) adapter = _PinholeGroupAdapter( scene=scene_name, group_key=group_key, src_idx=src_idx, src_img=src_img, src_depth=_load_val_pseudo_depth_b1hw( args, dataset="tat", scene=scene_name, frame_idx=src_idx, intrinsics_k3=src_k[0] ), src_w2c=src_w2c, src_k=src_k, tgt_indices=[int(name_to_idx[tgt_name]) for tgt_name in tgt_names if tgt_name in name_to_idx], load_target=_load_tat_target, ) yield from _yield_pinhole_group_batches("tat", adapter, args) def _dl3dv_frame_id_from_name(name: str) -> int: return int(Path(name).stem.split("_")[-1]) def _load_dl3dv_scene(scene_dir: Path) -> tuple[dict[int, Path], dict[int, torch.Tensor], dict[int, torch.Tensor]] | None: transforms_path = scene_dir / "transforms.json" image_dir = scene_dir / "images_4" if not (transforms_path.exists() and image_dir.exists()): return None meta = json.loads(transforms_path.read_text(encoding="utf-8")) image_paths = {int(_dl3dv_frame_id_from_name(p.name)): p for p in image_dir.glob("*.png")} if not image_paths: return None 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"]) example_path = next(iter(image_paths.values())) with Image.open(example_path) as img: cur_w, cur_h = int(img.size[0]), int(img.size[1]) k_cur = k.clone() if cur_h != orig_h or cur_w != orig_w: k_cur = _resize_k3_align_corners_false( k_cur.unsqueeze(0), sx=float(cur_w) / float(orig_w), sy=float(cur_h) / float(orig_h), )[0] w2c_map: dict[int, torch.Tensor] = {} intr_map: dict[int, torch.Tensor] = {} for frame in meta.get("frames", []): rel_path = str(frame.get("file_path", "")) try: frame_id = int(_dl3dv_frame_id_from_name(Path(rel_path).name)) except Exception: continue if frame_id not in image_paths: continue c2w = _nerf_c2w_to_opencv_c2w(frame["transform_matrix"]) w2c_map[frame_id] = torch.linalg.inv(c2w) intr_map[frame_id] = k_cur.clone() return image_paths, w2c_map, intr_map def _resolve_dl3dv_scene_dir(args: argparse.Namespace, scene_name: str, scene_dir_raw: str) -> Path | None: scene_dir = Path(scene_dir_raw) if scene_dir.exists(): return scene_dir root = Path(args.data_root) parts = scene_name.split("/", 1) if len(parts) == 2: candidate = root / parts[0] / parts[1] / parts[1] if candidate.exists(): return candidate candidate = root / parts[0] / parts[1] if candidate.exists(): return candidate return None def _iter_dl3dv_manifest_items(args: argparse.Namespace) -> Iterator[ValidationItem]: for group_idx, parts in _iter_manifest_parts(args, expected_parts=4): scene_name = str(parts[0]) scene_dir = _resolve_dl3dv_scene_dir(args, scene_name=scene_name, scene_dir_raw=parts[1]) if scene_dir is None: continue src_idx = int(parts[2]) tgt_indices = [int(x) for x in parts[3].split(",") if x.strip()] loaded = _load_dl3dv_scene(scene_dir) if loaded is None: continue image_paths, w2c_map, intr_map = loaded if src_idx not in image_paths or src_idx not in w2c_map: continue group_key = f"dl3dv_{scene_name.replace('/', '_')}_g{group_idx:05d}" adapter = _PinholeGroupAdapter( scene=scene_name, group_key=group_key, src_idx=src_idx, src_img=_load_png_rgb_u8(image_paths[src_idx]).unsqueeze(0), src_depth=_load_val_pseudo_depth_b1hw( args, dataset="dl3dv", scene=scene_name, frame_idx=src_idx, intrinsics_k3=intr_map[src_idx] ), src_w2c=w2c_map[src_idx].unsqueeze(0), src_k=intr_map[src_idx].unsqueeze(0).clone(), tgt_indices=tgt_indices, load_target=lambda tgt_idx, image_paths=image_paths, w2c_map=w2c_map, intr_map=intr_map: None if int(tgt_idx) not in image_paths or int(tgt_idx) not in w2c_map else _PinholeTargetAdapter( idx=int(tgt_idx), img=_load_png_rgb_u8(image_paths[int(tgt_idx)]).unsqueeze(0), w2c=w2c_map[int(tgt_idx)].unsqueeze(0), k=intr_map[int(tgt_idx)].unsqueeze(0).clone(), depth=_load_val_pseudo_depth_b1hw( args, dataset="dl3dv", scene=scene_name, frame_idx=int(tgt_idx), intrinsics_k3=intr_map[int(tgt_idx)], ), ), ) yield from _yield_pinhole_group_batches("dl3dv", adapter, args) def _iter_dataset_items(args: argparse.Namespace) -> Iterable[ValidationItem]: if getattr(args, "manifest_file", None) is None: raise ValueError("Validation requires --manifest-file. Build manifests first with scripts/build_validation_manifests.py.") dataset = str(args.dataset) if dataset == "re10k": return _iter_re10k_manifest_items(args) if dataset == "dl3dv": return _iter_dl3dv_manifest_items(args) if dataset == "replica": return _iter_replica_manifest_items(args) if dataset == "sim": return _iter_sim_manifest_items(args) if dataset == "wildrgbd": return _iter_wildrgbd_manifest_items(args) if dataset == "scannetpp": return _iter_scannetpp_manifest_items(args) if dataset == "scanetpp_fisheye": return _iter_scanetpp_fisheye_manifest_items(args) if dataset == "smx_sim_fisheye": return _iter_smx_sim_fisheye_manifest_items(args) if dataset == "tat": return _iter_tat_manifest_items(args) return _iter_hm3d_manifest_items(args) def _finalize_validation_group( *, dataset: str, step: int, vis_dir: Path, group_key: str, group_items: list[dict[str, Any]], ) -> dict[str, float] | None: if not group_items: return None group_row = { "psnr": _safe_nanmean([float(e["row"]["psnr"]) for e in group_items]), "ssim": _safe_nanmean([float(e["row"]["ssim"]) for e in group_items]), "lpips": _safe_nanmean([float(e["row"]["lpips"]) for e in group_items]), } group_dir = vis_dir / group_key _save_group_pair_pngs(group_dir, group_items) _save_perspective_group_grid( group_dir=group_dir, group_key=group_key, step=int(step), group_items=group_items, ) _save_group_gif( dataset=dataset, group_dir=group_dir, group_key=group_key, step=int(step), group_items=group_items, ) visual_items = [item for item in group_items if isinstance(item.get("vis", None), dict)] if dataset in {"hm3d", "replica"}: for j, item in enumerate(visual_items[:10], start=1): _save_vis_from_payload( item["vis"], vis_dir=group_dir, tag=f"{group_key}_t{j:02d}", step=int(step), ) return group_row def run_validation(args: argparse.Namespace) -> None: random.seed(int(args.seed)) np.random.seed(int(args.seed)) torch.manual_seed(int(args.seed)) dev = torch.device(args.device) model, step = _load_model(Path(args.checkpoint), dev) trainer = _build_trainer(model, dev, args) metrics_calc = MetricsCalculator(device=dev, compute_lpips=not bool(getattr(args, "fast_metrics", False))) dataset = str(args.dataset) items = _iter_dataset_items(args) if getattr(args, "out_dir", None) is not None: out_dir = Path(args.out_dir) else: out_dir = Path(args.checkpoint).parent / f"validation_{dataset}" out_dir.mkdir(parents=True, exist_ok=True) vis_dir = out_dir / "vis" vis_dir.mkdir(parents=True, exist_ok=True) sample_csv = out_dir / f"validation_sample_metrics_{dataset}.csv" group_rows: list[dict[str, float]] = [] failure_rows: list[dict[str, Any]] = [] current_group_key: str | None = None current_group_items: list[dict[str, Any]] = [] num_rows = 0 LOGGER.info("Validation start: dataset=%s checkpoint=%s", dataset, str(args.checkpoint)) pbar = tqdm(items, desc=f"validate_{dataset}", leave=False, disable=True) for i, (dataset_name, batch, tag, group_key) in enumerate(pbar): if current_group_key is not None and group_key != current_group_key: group_row = _finalize_validation_group( dataset=dataset, step=int(step), vis_dir=vis_dir, group_key=current_group_key, group_items=current_group_items, ) if group_row is not None: group_rows.append(group_row) num_rows += len(current_group_items) pbar.set_postfix(groups=len(group_rows), targets=num_rows, refresh=False) current_group_items = [] current_group_key = group_key try: with torch.no_grad(): result = trainer.process_batch( batch, dataset_name=dataset_name, step=int(step), need_vis=True, ) vis_payloads = result.get("vis_payloads", None) if isinstance(vis_payloads, list) and vis_payloads: vis_list = [v for v in vis_payloads if isinstance(v, dict)] else: vis = result.get("vis_payload", None) vis_list = [vis] if isinstance(vis, dict) else [] if not vis_list: continue if str(dataset_name) == "smx_sim_fisheye": vis_list = [_project_smx_sim_fisheye_vis(vis, batch, j) for j, vis in enumerate(vis_list)] tags = tag if isinstance(tag, list) else [tag] metric_mask = metric_mask_from_pinhole_batch( batch, dataset=str(dataset_name), cache_dir=Path(args.metric_mask_cache_dir) if args.metric_mask_cache_dir is not None else None, device=dev, ) for j, vis in enumerate(vis_list): if torch.is_tensor(vis.get("tgt_metric_mask", None)): vis["tgt_training_mask"] = vis["tgt_metric_mask"].detach() if torch.is_tensor(metric_mask): vis_b = int(vis["tgt_gt"].shape[0]) if torch.is_tensor(vis.get("tgt_gt", None)) else 1 if len(vis_list) == 1 and vis_b == int(metric_mask.shape[0]): vis["tgt_metric_mask"] = metric_mask.detach() elif j < int(metric_mask.shape[0]): vis["tgt_metric_mask"] = metric_mask[j : j + 1].detach() row = _compute_metrics_from_vis( vis, metrics_calc=metrics_calc, ) item_tag = str(tags[j]) if j < len(tags) else str(tag) _append_sample_metrics_row(sample_csv, str(group_key), item_tag, row) item = { "dataset_name": dataset_name, "tag": item_tag, "row": row, } if len(current_group_items) < 10: item["vis"] = vis current_group_items.append(item) pbar.set_postfix( groups=len(group_rows), targets=num_rows + len(current_group_items), refresh=False, ) except Exception as e: LOGGER.warning("Skip %s sample idx=%d: %s", dataset, int(i), str(e)) failure_rows.append( { "step": int(step), "sample_idx": int(i), "dataset": str(dataset), "tag": str(tag), "group_key": str(group_key), "error": str(e), } ) if "cuda" in str(e).lower(): raise RuntimeError( f"CUDA error during {dataset} validation at sample idx={int(i)}; " "the CUDA context may be corrupted, so this validation round must fail." ) from e if failure_rows: fail_csv = out_dir / f"validation_failures_{dataset}_step_{int(step):07d}.csv" with fail_csv.open("w", newline="") as f: writer = csv.DictWriter( f, fieldnames=["step", "sample_idx", "dataset", "tag", "group_key", "error"], ) writer.writeheader() writer.writerows(failure_rows) if current_group_key is not None: group_row = _finalize_validation_group( dataset=dataset, step=int(step), vis_dir=vis_dir, group_key=current_group_key, group_items=current_group_items, ) if group_row is not None: group_rows.append(group_row) num_rows += len(current_group_items) if not group_rows: if dataset in ("scannetpp", "scanetpp_fisheye", "tat"): LOGGER.warning("No validation samples processed for dataset=%s; skip this round.", dataset) return raise RuntimeError(f"No validation samples processed for dataset={dataset}") agg = _aggregate_rows(group_rows) agg["step"] = float(step) csv_main = out_dir / f"validation_metrics_{dataset}.csv" _append_metrics_row(csv_main, agg) LOGGER.info( "Validation done: dataset=%s groups=%d samples=%d psnr=%.3f ssim=%.4f lpips=%.4f", dataset, int(len(group_rows)), int(num_rows), float(agg.get("psnr", float("nan"))), float(agg.get("ssim", float("nan"))), float(agg.get("lpips", float("nan"))), ) def build_parser() -> argparse.ArgumentParser: p = argparse.ArgumentParser(description="Unified UniSharp validation") p.add_argument("--checkpoint", type=Path, required=True) p.add_argument( "--dataset", type=str, required=True, choices=[ "re10k", "dl3dv", "hm3d", "replica", "sim", "wildrgbd", "scannetpp", "scanetpp_fisheye", "smx_sim_fisheye", "tat", ], ) p.add_argument("--data-root", type=Path, required=True) p.add_argument("--device", type=str, default="cuda:0") p.add_argument("--seed", type=int, default=42) p.add_argument("--max-index-gap", type=int, default=10) p.add_argument("--pair-max-translation-m", type=float, default=0.5) p.add_argument("--pair-min-overlap", type=float, default=0.6) p.add_argument("--split", type=str, default="test") p.add_argument("--manifest-file", type=Path, default=None) p.add_argument("--manifest-max-groups", type=int, default=0) p.add_argument("--validation-batch-size", type=int, default=10) p.add_argument("--out-dir", type=Path, default=None) p.add_argument("--fast-metrics", action="store_true", help="Skip LPIPS during validation; keep PSNR/SSIM/depth metrics.") p.add_argument("--metric-mask-cache-dir", type=Path, default=default_metric_mask_cache_dir()) p.add_argument("--max-depth-m", type=float, default=None) p.add_argument("--sim-far-depth-invalid-m", type=float, default=None) p.add_argument("--sim-far-depth-invalid-max-frac", type=float, default=None) p.add_argument("--re10k-pseudo-far-depth-invalid-m", type=float, default=None) p.add_argument("--scanetpp-fisheye-far-depth-invalid-m", type=float, default=None) p.add_argument("--low-pass-filter-eps", type=float, default=None) p.add_argument( "--validation-pseudo-depth-root", type=Path, default=Path("/media/team_data/ML4_team/datasets/sharp/validation_unik3d_pseudo_depth"), ) p.add_argument("--sim-pose-root", type=Path, default=Path("/media/team_data/ML4_team/datasets/smx_sim/30cm")) p.add_argument( "--re10k-pseudo-depth-root", type=Path, default=Path("/media/team_data/ML4_team/datasets/nopose/re10k_unik3d_pseudo_depth/test"), ) return p def main() -> None: _configure_torchhub_cache() args = build_parser().parse_args() _apply_training_depth_config_defaults(args) run_validation(args) if __name__ == "__main__": main()