| """Utilities for loading and iterating the S23DR 2026 dataset.
|
|
|
| IMPORTANT: The dataset contains at least one corrupt image (around row ~13077
|
| in the training split). All iteration helpers in this module wrap entries in
|
| try/except so a single bad scene never kills the entire run.
|
| """
|
|
|
| import io
|
| import os
|
| import tempfile
|
| import zipfile
|
| import warnings
|
| from typing import Dict, List, Optional, Tuple
|
|
|
| import numpy as np
|
| from PIL import Image, UnidentifiedImageError
|
|
|
|
|
|
|
|
|
|
|
|
|
| def validate_entry(entry: dict) -> bool:
|
| """Check whether a dataset entry has valid (non-corrupt) images.
|
|
|
| Returns False if any image column contains an unreadable PIL image.
|
| This catches the PIL.UnidentifiedImageError that occurs around row ~13077.
|
| """
|
| try:
|
| for col in ("gestalt", "ade", "depth"):
|
| imgs = entry.get(col, [])
|
| if imgs is None:
|
| return False
|
| for img in imgs:
|
| if img is None:
|
| return False
|
|
|
| np.array(img)
|
| return True
|
| except (UnidentifiedImageError, OSError, ValueError, Exception):
|
| return False
|
|
|
|
|
|
|
|
|
|
|
|
|
| def load_dataset_streaming(split: str = "train", token: Optional[str] = None):
|
| """Load the dataset in streaming mode (low memory, no download)."""
|
| from datasets import load_dataset
|
| if token is None:
|
| token = os.environ.get("HF_TOKEN")
|
| ds = load_dataset(
|
| "usm3d/hoho22k_2026_trainval",
|
| streaming=True,
|
| trust_remote_code=True,
|
| token=token,
|
| )
|
| return ds[split]
|
|
|
|
|
| def load_dataset_cached(split: str = "train", cache_dir: str = "/data/s23dr",
|
| token: Optional[str] = None):
|
| """Load dataset with caching to disk for faster repeat access."""
|
| from datasets import load_dataset
|
| if token is None:
|
| token = os.environ.get("HF_TOKEN")
|
| ds = load_dataset(
|
| "usm3d/hoho22k_2026_trainval",
|
| split=split,
|
| cache_dir=cache_dir,
|
| trust_remote_code=True,
|
| token=token,
|
| )
|
| return ds
|
|
|
|
|
| def download_dataset(cache_dir: str = "/data/s23dr", token: Optional[str] = None):
|
| """Download the full dataset to disk (both train + validation).
|
|
|
| Usage on Modal::
|
|
|
| modal run modal_download_data.py
|
|
|
| Or locally::
|
|
|
| python -c "from src.data_loader import download_dataset; download_dataset('./data')"
|
| """
|
| from datasets import load_dataset
|
| if token is None:
|
| token = os.environ.get("HF_TOKEN")
|
| print(f"Downloading train split to {cache_dir} ...")
|
| load_dataset(
|
| "usm3d/hoho22k_2026_trainval",
|
| split="train",
|
| cache_dir=cache_dir,
|
| trust_remote_code=True,
|
| token=token,
|
| )
|
| print(f"Downloading validation split to {cache_dir} ...")
|
| load_dataset(
|
| "usm3d/hoho22k_2026_trainval",
|
| split="validation",
|
| cache_dir=cache_dir,
|
| trust_remote_code=True,
|
| token=token,
|
| )
|
| print("✅ Dataset downloaded.")
|
|
|
|
|
|
|
|
|
|
|
|
|
| def iter_entries_safe(dataset, max_entries: Optional[int] = None,
|
| verbose: bool = True):
|
| """Iterate over dataset entries, **skipping corrupt rows**.
|
|
|
| Every entry is wrapped in try/except so a single broken image (e.g. the
|
| known PIL.UnidentifiedImageError near row ~13077) does not crash the run.
|
|
|
| Yields
|
| ------
|
| idx : int
|
| The original row number (before skipping).
|
| entry : dict
|
| A validated dataset row.
|
| """
|
| from tqdm import tqdm
|
|
|
| yielded = 0
|
| skipped = 0
|
| for idx, entry in enumerate(tqdm(dataset, total=max_entries,
|
| desc="Processing entries")):
|
| if max_entries and yielded >= max_entries:
|
| break
|
| try:
|
| if not validate_entry(entry):
|
| skipped += 1
|
| if verbose:
|
| oid = entry.get("order_id", f"row_{idx}")
|
| print(f" ⚠ Skipping corrupt entry {oid} (row {idx})")
|
| continue
|
| yield idx, entry
|
| yielded += 1
|
| except (UnidentifiedImageError, OSError, Exception) as exc:
|
| skipped += 1
|
| if verbose:
|
| print(f" ⚠ Skipping row {idx}: {type(exc).__name__}: {exc}")
|
|
|
| if verbose and skipped:
|
| print(f" ℹ Skipped {skipped} corrupt entries total")
|
|
|
|
|
| def iter_entries(dataset, max_entries: Optional[int] = None):
|
| """Legacy helper — yields entries only (no index)."""
|
| for _, entry in iter_entries_safe(dataset, max_entries=max_entries,
|
| verbose=False):
|
| yield entry
|
|
|
|
|
|
|
|
|
|
|
|
|
| def read_colmap_reconstruction(colmap_bytes: bytes):
|
| import pycolmap
|
| with tempfile.TemporaryDirectory() as tmpdir:
|
| with zipfile.ZipFile(io.BytesIO(colmap_bytes), "r") as zf:
|
| zf.extractall(tmpdir)
|
| rec = pycolmap.Reconstruction(tmpdir)
|
| return rec
|
|
|
|
|
| def get_colmap_points(rec) -> np.ndarray:
|
| pts = [p3d.xyz for p3d in rec.points3D.values()]
|
| return np.array(pts) if pts else np.zeros((0, 3))
|
|
|
|
|
| def get_colmap_points_with_colors(rec) -> Tuple[np.ndarray, np.ndarray]:
|
| pts, cols = [], []
|
| for p3d in rec.points3D.values():
|
| pts.append(p3d.xyz)
|
| cols.append(p3d.color)
|
| if not pts:
|
| return np.zeros((0, 3)), np.zeros((0, 3), dtype=np.uint8)
|
| return np.array(pts), np.array(cols, dtype=np.uint8)
|
|
|
|
|
|
|
|
|
|
|
|
|
| def depth_to_meters(depth_image) -> np.ndarray:
|
| return np.array(depth_image).astype(np.float32) / 1000.0
|
|
|
|
|
| def get_valid_camera_indices(entry: dict) -> List[int]:
|
| flags = entry.get("pose_only_in_colmap", [])
|
| return [i for i, f in enumerate(flags) if not f]
|
|
|
|
|
| def get_camera_params(entry: dict, cam_idx: int):
|
| K = np.array(entry["K"][cam_idx])
|
| R = np.array(entry["R"][cam_idx])
|
| t = np.array(entry["t"][cam_idx])
|
| return K, R, t
|
|
|
|
|
| def project_3d_to_2d(points_3d, K, R, t):
|
| t = t.reshape(3, 1) if t.ndim == 1 else t
|
| pts_cam = R @ points_3d.T + t
|
| depth = pts_cam[2, :]
|
| pts_2d = K @ pts_cam
|
| uv = pts_2d[:2, :] / pts_2d[2:3, :]
|
| return uv.T, depth
|
|
|
|
|
| def backproject_2d_to_3d(uv, depth, K, R, t):
|
| t = t.reshape(3, 1) if t.ndim == 1 else t
|
| fx, fy = K[0, 0], K[1, 1]
|
| cx, cy = K[0, 2], K[1, 2]
|
| x_n = (uv[:, 0] - cx) / fx
|
| y_n = (uv[:, 1] - cy) / fy
|
| pts_cam = np.stack([x_n * depth, y_n * depth, depth], axis=1)
|
| R_inv = R.T
|
| t_inv = -R.T @ t
|
| return (R_inv @ pts_cam.T + t_inv).T
|
|
|
