horizonstream/data/dataloader.py

import os
import glob
from dataclasses import dataclass
from typing import List, Dict, Any, Iterator, Optional, Tuple

import cv2
import numpy as np
import torch
from PIL import Image

from horizonstream.data.undistort_utils import colmap_undistort_numpy
from horizonstream.utils.vendor.dust3r.utils.image import load_images_for_eval

dataset_metadata: Dict[str, Dict[str, Any]] = {
    "davis": {
        "img_path": "data/davis/DAVIS/JPEGImages/480p",
        "mask_path": "data/davis/DAVIS/masked_images/480p",
        "dir_path_func": lambda img_path, seq: os.path.join(img_path, seq),
        "gt_traj_func": lambda img_path, anno_path, seq: None,
        "traj_format": None,
        "seq_list": None,
        "full_seq": True,
        "mask_path_seq_func": lambda mask_path, seq: os.path.join(mask_path, seq),
        "skip_condition": None,
        "process_func": None,
    },
    "kitti": {
        "img_path": "data/kitti/sequences",
        "anno_path": "data/kitti/poses",
        "mask_path": None,
        "dir_path_func": lambda img_path, seq: os.path.join(img_path, seq, "image_2"),
        "gt_traj_func": lambda img_path, anno_path, seq: os.path.join(
            anno_path, f"{seq}.txt"
        )
        if os.path.exists(os.path.join(anno_path, f"{seq}.txt"))
        else None,
        "traj_format": "kitti",
        "seq_list": ["00", "01", "02", "03", "04", "05", "06", "07", "08", "09", "10"],
        "full_seq": True,
        "mask_path_seq_func": lambda mask_path, seq: None,
        "skip_condition": None,
        "process_func": None,
    },
    "bonn": {
        "img_path": "data/bonn/rgbd_bonn_dataset",
        "mask_path": None,
        "dir_path_func": lambda img_path, seq: os.path.join(
            img_path, f"rgbd_bonn_{seq}", "rgb_110"
        ),
        "gt_traj_func": lambda img_path, anno_path, seq: os.path.join(
            img_path, f"rgbd_bonn_{seq}", "groundtruth_110.txt"
        ),
        "traj_format": "tum",
        "seq_list": ["balloon2", "crowd2", "crowd3", "person_tracking2", "synchronous"],
        "full_seq": False,
        "mask_path_seq_func": lambda mask_path, seq: None,
        "skip_condition": None,
        "process_func": None,
    },
    "nyu": {
        "img_path": "data/nyu-v2/val/nyu_images",
        "mask_path": None,
        "process_func": None,
    },
    "scannet": {
        "img_path": "data/scannetv2",
        "mask_path": None,
        "dir_path_func": lambda img_path, seq: os.path.join(img_path, seq, "color_90"),
        "gt_traj_func": lambda img_path, anno_path, seq: os.path.join(
            img_path, seq, "pose_90.txt"
        ),
        "traj_format": "replica",
        "seq_list": None,
        "full_seq": True,
        "mask_path_seq_func": lambda mask_path, seq: None,
        "skip_condition": None,
        "process_func": None,
    },
    "tum": {
        "img_path": "data/tum",
        "mask_path": None,
        "dir_path_func": lambda img_path, seq: os.path.join(img_path, seq, "rgb_90"),
        "gt_traj_func": lambda img_path, anno_path, seq: os.path.join(
            img_path, seq, "groundtruth_90.txt"
        ),
        "traj_format": "tum",
        "seq_list": None,
        "full_seq": True,
        "mask_path_seq_func": lambda mask_path, seq: None,
        "skip_condition": None,
        "process_func": None,
    },
    "sintel": {
        "img_path": "data/sintel/training/final",
        "anno_path": "data/sintel/training/camdata_left",
        "mask_path": None,
        "dir_path_func": lambda img_path, seq: os.path.join(img_path, seq),
        "gt_traj_func": lambda img_path, anno_path, seq: os.path.join(anno_path, seq),
        "traj_format": None,
        "seq_list": [
            "alley_2",
            "ambush_4",
            "ambush_5",
            "ambush_6",
            "cave_2",
            "cave_4",
            "market_2",
            "market_5",
            "market_6",
            "shaman_3",
            "sleeping_1",
            "sleeping_2",
            "temple_2",
            "temple_3",
        ],
        "full_seq": False,
        "mask_path_seq_func": lambda mask_path, seq: None,
        "skip_condition": None,
        "process_func": None,
    },
    "waymo": {
        "img_path": "toyourdata/scatt3r_evaluation/waymo_open_dataset_v1_4_3",
        "anno_path": None,
        "mask_path": None,
        "dir_path_func": lambda img_path, seq: os.path.join(
            img_path,
            seq.split("_cam")[0] if "_cam" in seq else seq,
            "images",
            seq.split("_cam")[1] if "_cam" in seq else "00",
        ),
        "gt_traj_func": lambda img_path, anno_path, seq: os.path.join(
            img_path,
            seq.split("_cam")[0] if "_cam" in seq else seq,
            "cameras",
            seq.split("_cam")[1] if "_cam" in seq else "00",
            "extri.yml",
        ),
        "traj_format": "waymo",
        "seq_list": None,
        "full_seq": True,
        "mask_path_seq_func": lambda mask_path, seq: None,
        "skip_condition": None,
        "process_func": None,
    },
}


@dataclass
class HorizonStreamSequenceInfo:
    name: str
    scene_root: str
    image_dir: str
    image_paths: List[str]
    camera: Optional[str]


class HorizonStreamSequence:
    def __init__(
        self,
        name: str,
        images: torch.Tensor,
        image_paths: List[str],
        scene_root: Optional[str] = None,
        image_dir: Optional[str] = None,
        camera: Optional[str] = None,
    ):
        self.name = name
        self.images = images
        self.image_paths = image_paths
        self.scene_root = scene_root
        self.image_dir = image_dir
        self.camera = camera


def _read_list_file(path: str) -> List[str]:
    with open(path, "r") as f:
        lines = []
        for line in f.readlines():
            line = line.strip()
            if not line:
                continue
            if line.startswith("#"):
                continue
            lines.append(line)
    return lines


def _is_generalizable_scene_root(path: str) -> bool:
    return os.path.isdir(os.path.join(path, "images"))


def _direct_image_files(dir_path: str) -> List[str]:
    filelist = sorted(glob.glob(os.path.join(dir_path, "*.png")))
    if not filelist:
        filelist = sorted(glob.glob(os.path.join(dir_path, "*.jpg")))
    if not filelist:
        filelist = sorted(glob.glob(os.path.join(dir_path, "*.jpeg")))
    return filelist


def _frame_stems(image_paths: List[str]) -> List[str]:
    stems = [os.path.splitext(os.path.basename(p))[0] for p in image_paths]
    if len(set(stems)) == len(stems):
        return stems
    parents = [os.path.basename(os.path.dirname(p)) for p in image_paths]
    if len(set(parents)) == len(parents):
        return parents
    return stems


def _format_fps_for_name(fps: float) -> str:
    return f"{float(fps):g}".replace(".", "p")


def _safe_video_scene_name(path: str, stride: int, target_fps: Optional[float]) -> str:
    stem = os.path.splitext(os.path.basename(path))[0]
    safe = "".join(ch if ch.isalnum() or ch in "_.-" else "_" for ch in stem).strip("._")
    if not safe:
        safe = "video"
    if target_fps is not None:
        return f"{safe}_fps{_format_fps_for_name(target_fps)}"
    return f"{safe}_stride{int(stride)}"


class HorizonStreamDataLoader:
    def __init__(self, cfg: Dict[str, Any]):
        self.cfg = cfg
        self.dataset = cfg.get("dataset", None)
        meta = dataset_metadata.get(self.dataset, {})
        self.img_path = cfg.get("img_path", meta.get("img_path"))
        self.mask_path = cfg.get("mask_path", meta.get("mask_path"))
        self.dir_path_func = meta.get("dir_path_func", lambda p, s: os.path.join(p, s))
        self.mask_path_seq_func = meta.get("mask_path_seq_func", lambda p, s: None)
        self.full_seq = bool(cfg.get("full_seq", meta.get("full_seq", True)))
        self.seq_list = cfg.get("seq_list", None)
        self.stride = int(cfg.get("stride", 1))
        self.max_frames = cfg.get("max_frames", None)
        self.size = int(cfg.get("size", 518))
        self.crop = bool(cfg.get("crop", False))
        self.patch_size = int(cfg.get("patch_size", 14))
        self.format = cfg.get("format", "auto")
        self.data_roots_file = cfg.get("data_roots_file", None)
        self.split = cfg.get("split", None)
        self.camera = cfg.get("camera", None)
        self.seq_match_mode = str(cfg.get("seq_match_mode", "exact")).lower()
        self.video_path = cfg.get("video_path", None)
        self.image_paths_cfg = cfg.get("image_paths", None)
        self.image_scene_name = cfg.get("image_scene_name", None)
        self.video_stride = max(1, int(cfg.get("video_stride", 1)))
        self.video_sample_fps = cfg.get("video_sample_fps", None)
        if self.video_sample_fps in ("", "null"):
            self.video_sample_fps = None
        if self.video_sample_fps is not None:
            self.video_sample_fps = float(self.video_sample_fps)
        self.video_scene_name = cfg.get("video_scene_name", None)

        # camera VAL-aligned preprocess: undistort -> resize/crop with K sync.
        self.camera_preprocess = bool(cfg.get("camera_preprocess", True))
        self.camera_preprocess_strict = bool(
            cfg.get("camera_preprocess_strict", True)
        )
        self.undistort_backend = str(
            cfg.get("undistort_backend", "colmap")
        ).lower()
        self.undistort_dist_opt_k = bool(cfg.get("undistort_dist_opt_k", False))
        self.undistort_safe_bound = int(cfg.get("undistort_safe_bound", 4))
        self.undistort_center_crop = bool(cfg.get("undistort_center_crop", True))
        self._undistort_warned_no_pycolmap = False

    def _infer_format(self) -> str:
        if self.image_paths_cfg not in (None, "", "null"):
            return "image_list"
        if self.video_path not in (None, "", "null"):
            return "video"
        if self.format in ["relpose", "generalizable"]:
            return self.format
        if self.img_path is None:
            return "relpose"
        if _is_generalizable_scene_root(self.img_path):
            return "generalizable"
        default_list = self.data_roots_file or "data_roots.txt"
        if os.path.exists(os.path.join(self.img_path, default_list)):
            return "generalizable"
        return "relpose"

    def _resolve_seq_list_generalizable(self) -> List[str]:
        if self.seq_list is not None and self.seq_match_mode not in ["contains", "in"]:
            return list(self.seq_list)
        if self.img_path is None or not os.path.isdir(self.img_path):
            return []

        if _is_generalizable_scene_root(self.img_path):
            return self._filter_seq_candidates([self.img_path])

        candidates = []
        if isinstance(self.data_roots_file, str) and self.data_roots_file:
            candidates.append(self.data_roots_file)
        if isinstance(self.split, str) and self.split:
            split_name = self.split.lower()
            if split_name in ["val", "valid", "validate"]:
                split_name = "validate"
            candidates.append(f"{split_name}_data_roots.txt")
        candidates.append("data_roots.txt")
        candidates.append("train_data_roots.txt")
        candidates.append("validate_data_roots.txt")

        for fname in candidates:
            path = os.path.join(self.img_path, fname)
            if os.path.exists(path):
                return self._filter_seq_candidates(_read_list_file(path))

        img_dirs = sorted(
            glob.glob(os.path.join(self.img_path, "**", "images"), recursive=True)
        )
        scene_roots = [os.path.dirname(p) for p in img_dirs]

        rels = []
        for p in scene_roots:
            try:
                rels.append(os.path.relpath(p, self.img_path))
            except ValueError:
                rels.append(p)
        return self._filter_seq_candidates(sorted(set(rels)))

    def _resolve_seq_list_relpose(self) -> List[str]:
        if self.seq_list is not None and self.seq_match_mode not in ["contains", "in"]:
            return list(self.seq_list)
        meta = dataset_metadata.get(self.dataset, {})
        if self.full_seq:
            if self.img_path is None or not os.path.isdir(self.img_path):
                return []
            seqs = [
                s
                for s in os.listdir(self.img_path)
                if os.path.isdir(os.path.join(self.img_path, s))
            ]
            return self._filter_seq_candidates(sorted(seqs))
        seqs = meta.get("seq_list", []) or []
        return self._filter_seq_candidates(list(seqs))

    def _filter_seq_candidates(self, candidates: List[str]) -> List[str]:
        if self.seq_list is None or self.seq_match_mode not in ["contains", "in"]:
            return candidates
        needles = [str(x).strip() for x in self.seq_list if str(x).strip()]
        if not needles:
            return candidates
        matched = []
        for seq in candidates:
            seq_text = str(seq)
            seq_name = os.path.basename(os.path.normpath(seq_text))
            if any(needle in seq_text or needle in seq_name for needle in needles):
                matched.append(seq)
        return matched

    def _resolve_seq_list(self) -> List[str]:
        fmt = self._infer_format()
        if fmt == "generalizable":
            result = self._resolve_seq_list_generalizable()
        else:
            result = self._resolve_seq_list_relpose()
        if not result:
            print(f"[horizonstream] warning: no sequences resolved (img_path={self.img_path}, seq_list={self.seq_list}, seq_match_mode={self.seq_match_mode})", flush=True)
        return result

    def _sequence_name_from_scene_root(self, seq_entry: str, scene_root: str) -> str:
        scene_root_abs = os.path.abspath(scene_root)
        if self.img_path is not None:
            try:
                rel = os.path.relpath(scene_root_abs, os.path.abspath(self.img_path))
            except ValueError:
                rel = None
            if rel and rel not in [os.curdir, os.pardir] and not rel.startswith(os.pardir + os.path.sep):
                return rel

        if os.path.isabs(seq_entry):
            return os.path.basename(os.path.normpath(seq_entry))

        name = os.path.normpath(seq_entry)
        if name in ["", os.curdir, os.pardir] or name.startswith(os.pardir + os.path.sep):
            return os.path.basename(os.path.normpath(scene_root))
        return name

    def _resolve_scene_root(self, seq_entry: str) -> Tuple[str, str]:
        if os.path.isabs(seq_entry):
            scene_root = seq_entry
        elif os.path.sep in seq_entry:
            scene_root = os.path.join(self.img_path, seq_entry)
        else:
            scene_root = os.path.join(self.img_path, seq_entry)
        name = self._sequence_name_from_scene_root(seq_entry, scene_root)
        return name, scene_root

    def _resolve_image_dir_generalizable(self, scene_root: str) -> Optional[str]:
        images_root = os.path.join(scene_root, "images")
        if not os.path.isdir(images_root):
            return None

        if isinstance(self.camera, str) and self.camera:
            cam_dir = os.path.join(images_root, self.camera)
            if os.path.isdir(cam_dir):
                return cam_dir

        if _direct_image_files(images_root):
            return images_root

        cams = [
            d
            for d in os.listdir(images_root)
            if os.path.isdir(os.path.join(images_root, d))
        ]
        if not cams:
            return None
        cams = sorted(cams)

        frame_dirs = []
        for name in cams:
            child_dir = os.path.join(images_root, name)
            child_images = _direct_image_files(child_dir)
            if child_images:
                frame_dirs.append((name, len(child_images)))

        if (
            len(cams) > 10
            and len(frame_dirs) == len(cams)
            and max(count for _, count in frame_dirs) == 1
        ):
            return images_root

        return os.path.join(images_root, cams[0])

    def _camera_from_image_dir(self, image_dir: str) -> Optional[str]:
        parent = os.path.basename(os.path.dirname(image_dir))
        if parent != "images":
            return None
        return os.path.basename(image_dir)

    def _has_multiple_camera_dirs(self, scene_root: str) -> bool:
        images_root = os.path.join(scene_root, "images")
        if not os.path.isdir(images_root):
            return False
        camera_dirs = []
        for name in os.listdir(images_root):
            child_dir = os.path.join(images_root, name)
            if os.path.isdir(child_dir) and _direct_image_files(child_dir):
                camera_dirs.append(name)
        return len(camera_dirs) > 1

    def _collect_filelist(self, dir_path: str) -> List[str]:
        filelist = _direct_image_files(dir_path)
        if not filelist:
            nested = []
            child_dirs = sorted(
                d for d in glob.glob(os.path.join(dir_path, "*")) if os.path.isdir(d)
            )
            for child_dir in child_dirs:
                child_images = _direct_image_files(child_dir)
                if child_images:
                    nested.append(child_images[0])
            filelist = nested
        if self.stride > 1:
            filelist = filelist[:: self.stride]
        if self.max_frames is not None and self.max_frames > 0:
            filelist = filelist[: self.max_frames]
        return filelist

    def _load_intrinsics_distortion(
        self, scene_root: str, camera: Optional[str]
    ) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], str]:
        intri_candidates = []
        if isinstance(camera, str) and camera:
            intri_candidates.append(os.path.join(scene_root, "cameras", camera, "intri.yml"))
        intri_candidates.append(os.path.join(scene_root, "intri.yml"))

        intri_path = None
        for p in intri_candidates:
            if os.path.exists(p):
                intri_path = p
                break
        if intri_path is None:
            raise FileNotFoundError(
                f"intri.yml not found under scene_root={scene_root} camera={camera}"
            )

        fs = cv2.FileStorage(intri_path, cv2.FILE_STORAGE_READ)
        names_node = fs.getNode("names")
        names = []
        if not names_node.empty():
            for i in range(names_node.size()):
                names.append(names_node.at(i).string())

        K_map: Dict[str, np.ndarray] = {}
        D_map: Dict[str, np.ndarray] = {}
        for name in names:
            K = fs.getNode(f"K_{name}").mat()
            if K is None:
                continue
            D = fs.getNode(f"D_{name}").mat()
            if D is None:
                D_arr = np.zeros((5, 1), dtype=np.float64)
            else:
                D_arr = np.asarray(D, dtype=np.float64).reshape(-1)
                if D_arr.size < 5:
                    D_arr = np.pad(D_arr, (0, 5 - D_arr.size), mode="constant")
                elif D_arr.size > 5:
                    D_arr = D_arr[:5]
                D_arr = D_arr.reshape(5, 1)

            K_map[name] = np.asarray(K, dtype=np.float64)
            D_map[name] = D_arr

        fs.release()
        if not K_map:
            raise ValueError(f"No intrinsics K_* found in {intri_path}")
        return K_map, D_map, intri_path

    def _undistort_image(
        self,
        img: np.ndarray,
        K: np.ndarray,
        D: np.ndarray,
    ) -> Tuple[np.ndarray, np.ndarray]:
        K = np.asarray(K, dtype=np.float64).copy()
        D = np.asarray(D, dtype=np.float64).reshape(-1)
        if D.size < 5:
            D = np.pad(D, (0, 5 - D.size), mode="constant")
        elif D.size > 5:
            D = D[:5]

        if np.sum(np.abs(D)) == 0.0:
            return img, K

        backend = self.undistort_backend
        if backend == "colmap" and os.name != "nt":
            try:
                undist, new_K = colmap_undistort_numpy(
                    img, K, D.reshape(5, 1), blank_pixels=not self.undistort_dist_opt_k
                )
                undist = np.rint(undist).clip(0, 255).astype(np.uint8)
                return undist, np.asarray(new_K, dtype=np.float64)
            except Exception as exc:
                if not self._undistort_warned_no_pycolmap:
                    print(
                        f"[horizonstream] warning: camera preprocess colmap undistort unavailable ({exc}); fallback to cv2",
                        flush=True,
                    )
                    self._undistort_warned_no_pycolmap = True

        H, W = img.shape[:2]
        if self.undistort_dist_opt_k:
            new_K, _ = cv2.getOptimalNewCameraMatrix(K, D, (W, H), 0, (W, H))
            undist = cv2.undistort(img, K, D, newCameraMatrix=new_K)
            return undist, np.asarray(new_K, dtype=np.float64)

        undist = cv2.undistort(img, K, D)
        return undist, K

    def _resize_crop_with_k(
        self,
        img: np.ndarray,
        K: np.ndarray,
        target_h: int,
        target_w: int,
    ) -> Tuple[np.ndarray, np.ndarray]:
        H, W = img.shape[:2]
        K = np.asarray(K, dtype=np.float64).copy()

        safe_bound = max(int(self.undistort_safe_bound), 0)
        ratio = max(
            float(target_h + safe_bound) / max(H, 1),
            float(target_w + safe_bound) / max(W, 1),
        )
        if ratio <= 0:
            ratio = 1.0

        h = int(H * ratio)
        w = int(W * ratio)
        if h != H or w != W:
            K[0:1] *= float(w) / float(max(W, 1))
            K[1:2] *= float(h) / float(max(H, 1))
            img = cv2.resize(img, (w, h), interpolation=cv2.INTER_AREA)
        else:
            h, w = H, W

        # Match camera strict-center crop behavior:
        # use principal point for crop center and pad with zeros when out-of-bound.
        if self.undistort_center_crop:
            crop_h = int(round(float(K[1, 2]))) - target_h // 2
            crop_w = int(round(float(K[0, 2]))) - target_w // 2
        else:
            crop_h = 0
            crop_w = 0

        end_h = crop_h + target_h
        end_w = crop_w + target_w
        offset_h = crop_h
        offset_w = crop_w

        if self.undistort_center_crop and (
            crop_h < 0 or crop_w < 0 or end_h > h or end_w > w
        ):
            pad_top = max(0, -crop_h)
            pad_bottom = max(0, end_h - h)
            pad_left = max(0, -crop_w)
            pad_right = max(0, end_w - w)
            if pad_top or pad_bottom or pad_left or pad_right:
                img = cv2.copyMakeBorder(
                    img,
                    pad_top,
                    pad_bottom,
                    pad_left,
                    pad_right,
                    cv2.BORDER_CONSTANT,
                    value=(0, 0, 0),
                )
                crop_h += pad_top
                crop_w += pad_left
                end_h = crop_h + target_h
                end_w = crop_w + target_w

        img = img[crop_h:end_h, crop_w:end_w]

        K[0, 2] -= offset_w
        K[1, 2] -= offset_h
        return img, K

    def _load_images_preprocessed(self, info: HorizonStreamSequenceInfo) -> torch.Tensor:
        K_map, D_map, intri_path = self._load_intrinsics_distortion(
            info.scene_root, info.camera
        )

        stems = _frame_stems(info.image_paths)
        missing = [s for s in stems if s not in K_map]
        if missing:
            sample = ", ".join(missing[:5])
            msg = (
                f"[horizonstream] preprocess: stem mismatch in {intri_path}, "
                f"missing={len(missing)} sample=[{sample}]"
            )
            if self.camera_preprocess_strict:
                raise KeyError(msg)
            print(msg, flush=True)

        processed = []
        target_h = None
        target_w = None

        for i, path in enumerate(info.image_paths):
            stem = stems[i]
            if stem not in K_map:
                if self.camera_preprocess_strict:
                    raise KeyError(f"Missing K/D for frame stem={stem} path={path}")
                continue

            K = K_map[stem]
            D = D_map.get(stem, np.zeros((5, 1), dtype=np.float64))

            img_bgr = cv2.imread(path, cv2.IMREAD_COLOR)
            if img_bgr is None:
                raise IOError(f"Failed to read image: {path}")
            img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

            img, K = self._undistort_image(img, K, D)

            if target_h is None or target_w is None:
                H0, W0 = img.shape[:2]
                if self.size > 0:
                    h = int(float(H0) / float(max(W0, 1)) * float(self.size))
                    w = int(self.size)
                else:
                    h, w = H0, W0
                target_h = max((h // self.patch_size) * self.patch_size, self.patch_size)
                target_w = max((w // self.patch_size) * self.patch_size, self.patch_size)

            img, _ = self._resize_crop_with_k(img, K, target_h, target_w)
            img_t = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0
            img_t = (img_t - 0.5) / 0.5
            processed.append(img_t[None])

        if not processed:
            raise RuntimeError(f"preprocess produced 0 frames for {info.name}")

        imgs = torch.cat(processed, dim=0)
        images = imgs.unsqueeze(0)
        images = (images + 1.0) / 2.0
        return images

    def _load_images(self, info: HorizonStreamSequenceInfo) -> torch.Tensor:
        if self.camera_preprocess:
            try:
                return self._load_images_preprocessed(info)
            except Exception as exc:
                if self.camera_preprocess_strict and not isinstance(exc, FileNotFoundError):
                    raise
                print(
                    f"[horizonstream] warning: preprocess failed for {info.name}, fallback to load_images_for_eval: {exc}",
                    flush=True,
                )

        views = load_images_for_eval(
            info.image_paths,
            size=self.size,
            verbose=False,
            crop=self.crop,
            patch_size=self.patch_size,
        )
        imgs = torch.cat([view["img"] for view in views], dim=0)
        images = imgs.unsqueeze(0)
        images = (images + 1.0) / 2.0
        return images

    def _preprocess_rgb_frame_for_eval(self, rgb: np.ndarray) -> torch.Tensor:
        img = Image.fromarray(np.asarray(rgb, dtype=np.uint8), mode="RGB")
        w1, h1 = img.size
        if self.size == 224:
            long_edge = round(self.size * max(w1 / h1, h1 / w1))
        else:
            long_edge = self.size

        scale = float(long_edge) / float(max(w1, h1))
        new_size = tuple(int(round(x * scale)) for x in img.size)
        interp = Image.Resampling.LANCZOS if max(w1, h1) > long_edge else Image.Resampling.BICUBIC
        img = img.resize(new_size, interp)

        w, h = img.size
        cx, cy = w // 2, h // 2
        if self.size == 224:
            half = min(cx, cy)
            if self.crop:
                img = img.crop((cx - half, cy - half, cx + half, cy + half))
            else:
                img = img.resize((2 * half, 2 * half), Image.Resampling.LANCZOS)
        else:
            halfw = ((2 * cx) // self.patch_size) * (self.patch_size // 2)
            halfh = ((2 * cy) // self.patch_size) * (self.patch_size // 2)
            if w == h:
                halfh = int(3 * halfw / 4)
            if self.crop:
                img = img.crop((cx - halfw, cy - halfh, cx + halfw, cy + halfh))
            else:
                img = img.resize((2 * halfw, 2 * halfh), Image.Resampling.LANCZOS)

        arr = np.asarray(img, dtype=np.float32) / 255.0
        tensor = torch.from_numpy(arr).permute(2, 0, 1)
        tensor = (tensor - 0.5) / 0.5
        return tensor

    def _load_video_sequence(self) -> HorizonStreamSequence:
        video_path = os.path.abspath(os.path.expanduser(str(self.video_path)))
        if not os.path.exists(video_path):
            raise FileNotFoundError(f"Video file not found: {video_path}")
        if self.video_sample_fps is not None and self.video_sample_fps <= 0.0:
            raise ValueError(f"video_sample_fps must be positive, got {self.video_sample_fps}")

        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise RuntimeError(f"Failed to open video: {video_path}")

        source_fps = float(cap.get(cv2.CAP_PROP_FPS) or 0.0)
        if source_fps <= 0.0:
            source_fps = None
        if self.video_sample_fps is not None and source_fps is None:
            cap.release()
            raise ValueError(f"Cannot sample video by FPS because source FPS is unavailable: {video_path}")

        scene_name = self.video_scene_name or _safe_video_scene_name(
            video_path,
            self.video_stride,
            self.video_sample_fps,
        )
        effective_fps = (
            min(float(self.video_sample_fps), float(source_fps))
            if self.video_sample_fps is not None and source_fps is not None
            else (None if source_fps is None else source_fps / float(self.video_stride))
        )
        tensors = []
        source_idx = 0
        out_idx = 0

        def should_keep_frame(frame_idx: int, output_idx: int) -> bool:
            if self.video_sample_fps is None:
                return frame_idx % self.video_stride == 0
            if source_fps is None or self.video_sample_fps >= source_fps:
                return True
            target_source_idx = int(round(float(output_idx) * float(source_fps) / float(self.video_sample_fps)))
            return frame_idx >= target_source_idx

        try:
            while True:
                ok, frame_bgr = cap.read()
                if not ok:
                    break
                if should_keep_frame(source_idx, out_idx):
                    rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
                    tensors.append(self._preprocess_rgb_frame_for_eval(rgb)[None])
                    out_idx += 1
                    if self.max_frames is not None and self.max_frames > 0 and out_idx >= self.max_frames:
                        break
                source_idx += 1
        finally:
            cap.release()

        if not tensors:
            raise RuntimeError(f"No frames read from video: {video_path}")

        imgs = torch.cat(tensors, dim=0)
        images = ((imgs.unsqueeze(0)) + 1.0) / 2.0
        print(
            f"[horizonstream] loaded video {scene_name}: path={video_path} "
            f"frames={out_idx} last_source_frame={source_idx} stride={self.video_stride} "
            f"target_fps={self.video_sample_fps} source_fps={source_fps} "
            f"effective_fps={effective_fps} shape={tuple(images.shape)}",
            flush=True,
        )
        return HorizonStreamSequence(
            scene_name,
            images,
            [],
            scene_root=os.path.dirname(video_path),
            image_dir=None,
            camera="video",
        )

    def _load_image_list_sequence(self) -> HorizonStreamSequence:
        image_paths = self.image_paths_cfg
        if isinstance(image_paths, str):
            image_paths = _read_list_file(image_paths) if os.path.isfile(image_paths) else [image_paths]
        image_paths = [os.path.abspath(os.path.expanduser(str(p))) for p in (image_paths or [])]
        image_paths = [p for p in image_paths if p and os.path.isfile(p)]
        if self.max_frames is not None and self.max_frames > 0:
            image_paths = image_paths[: int(self.max_frames)]
        if not image_paths:
            raise RuntimeError("No images found for data.image_paths")
        scene_name = self.image_scene_name or "images"
        info = HorizonStreamSequenceInfo(
            name=scene_name,
            scene_root=os.path.dirname(image_paths[0]),
            image_dir=os.path.dirname(image_paths[0]),
            image_paths=image_paths,
            camera="image_list",
        )
        print(
            f"[horizonstream] loading image list {scene_name}: {len(image_paths)} frames",
            flush=True,
        )
        images = self._load_images(info)
        print(
            f"[horizonstream] loaded image list {scene_name}: {tuple(images.shape)}",
            flush=True,
        )
        return HorizonStreamSequence(
            scene_name,
            images,
            image_paths,
            scene_root=info.scene_root,
            image_dir=info.image_dir,
            camera=info.camera,
        )

    def iter_sequence_infos(self) -> Iterator[HorizonStreamSequenceInfo]:
        fmt = self._infer_format()
        seqs = self._resolve_seq_list()
        for seq_entry in seqs:
            if fmt == "generalizable":
                seq, scene_root = self._resolve_scene_root(seq_entry)
                dir_path = self._resolve_image_dir_generalizable(scene_root)
                if dir_path is None or not os.path.isdir(dir_path):
                    print(f"[horizonstream] skip seq={seq_entry}: image dir not found under {scene_root}", flush=True)
                    continue
                camera = self._camera_from_image_dir(dir_path)
                if camera is not None and self._has_multiple_camera_dirs(scene_root):
                    seq = os.path.join(seq, camera)
            else:
                seq = seq_entry
                scene_root = os.path.join(self.img_path, seq)
                dir_path = self.dir_path_func(self.img_path, seq)
                if not os.path.isdir(dir_path):
                    continue
                camera = None

            filelist = self._collect_filelist(dir_path)
            if not filelist:
                print(f"[horizonstream] skip seq={seq_entry}: no images found in {dir_path}", flush=True)
                continue
            yield HorizonStreamSequenceInfo(
                name=seq,
                scene_root=scene_root,
                image_dir=dir_path,
                image_paths=filelist,
                camera=camera,
            )

    def __iter__(self) -> Iterator[HorizonStreamSequence]:
        if self._infer_format() == "image_list":
            yield self._load_image_list_sequence()
            return
        if self._infer_format() == "video":
            yield self._load_video_sequence()
            return

        for info in self.iter_sequence_infos():
            print(
                f"[horizonstream] loading sequence {info.name}: {len(info.image_paths)} frames",
                flush=True,
            )
            images = self._load_images(info)
            print(
                f"[horizonstream] loaded sequence {info.name}: {tuple(images.shape)}",
                flush=True,
            )
            yield HorizonStreamSequence(
                info.name,
                images,
                info.image_paths,
                scene_root=info.scene_root,
                image_dir=info.image_dir,
                camera=info.camera,
            )