| """Metadata-only matched-video dataset with explicitly lazy teacher I/O.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| import os |
| from pathlib import Path |
| import time |
| from typing import Any, Callable, Mapping, Sequence |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import Dataset |
|
|
| from utils.matched_video_feature_cache import ( |
| MATCHED_PYAV_LAST64_DECODER_CONTRACT, |
| MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT, |
| MatchedVideoFeatureCache, |
| ) |
| from utils.matched_video_manifest import ( |
| MatchedVideoManifestEntry, |
| load_matched_video_manifest, |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class DecodedMatchedVideo: |
| """The final contiguous RGB frame window in V-JEPA2 input range.""" |
|
|
| video01: torch.Tensor |
| frame_indices: torch.Tensor |
| source_frame_count: int |
|
|
|
|
| def _frames_to_video01(frames: Sequence[Any]) -> torch.Tensor: |
| tensors: list[torch.Tensor] = [] |
| shape: tuple[int, int, int] | None = None |
| for frame in frames: |
| array = frame.to_ndarray(format="rgb24") |
| if not isinstance(array, np.ndarray): |
| raise TypeError("PyAV frame.to_ndarray did not return ndarray") |
| if array.ndim != 3 or int(array.shape[2]) != 3: |
| raise ValueError( |
| f"decoded RGB frame must be [H,W,3], got {array.shape}" |
| ) |
| if array.dtype != np.uint8: |
| raise ValueError(f"decoded RGB frame must be uint8, got {array.dtype}") |
| current_shape = tuple(int(value) for value in array.shape) |
| if shape is None: |
| shape = current_shape |
| elif current_shape != shape: |
| raise ValueError( |
| "video changes resolution inside selected frame window: " |
| f"{current_shape} != {shape}" |
| ) |
| tensors.append( |
| torch.from_numpy(np.ascontiguousarray(array)).permute(2, 0, 1) |
| ) |
| if not tensors: |
| raise ValueError("decoded frame window is empty") |
| return torch.stack(tensors).to(dtype=torch.float32).div_(255.0) |
|
|
|
|
| def decode_matched_video_last_contiguous_pyav( |
| video_path: str | os.PathLike, |
| *, |
| decoder_contract: Mapping[str, Any] | None = None, |
| ) -> DecodedMatchedVideo: |
| """Decode exactly frames 17..80 from a contractually 81-frame MP4.""" |
|
|
| contract = dict( |
| decoder_contract |
| if decoder_contract is not None |
| else MATCHED_PYAV_LAST64_DECODER_CONTRACT |
| ) |
| if contract.get("backend") != "pyav": |
| raise ValueError("decoder contract backend must be pyav") |
| if contract.get("selection") != "last_contiguous": |
| raise ValueError("decoder contract selection must be last_contiguous") |
| num_frames = int(contract.get("num_frames", 0)) |
| expected_count = int(contract.get("expected_source_frame_count", 0)) |
| if num_frames <= 0 or expected_count < num_frames: |
| raise ValueError("decoder contract frame geometry is invalid") |
| expected_bounds = [expected_count - num_frames, expected_count - 1] |
| if list(contract.get("expected_frame_indices", ())) != expected_bounds: |
| raise ValueError( |
| "decoder contract expected_frame_indices do not describe last64" |
| ) |
| try: |
| import av |
| except ImportError as exc: |
| raise ImportError( |
| "PyAV is required for a matched-video teacher cache miss; " |
| "install the pinned av==13.1.0 dependency" |
| ) from exc |
| expected_version = str(contract.get("backend_version", "")) |
| actual_version = str(getattr(av, "__version__", "")) |
| if actual_version != expected_version: |
| raise RuntimeError( |
| "PyAV version differs from the cache decoder contract: " |
| f"{actual_version!r} != {expected_version!r}" |
| ) |
|
|
| path = Path(video_path).expanduser().resolve() |
| if not path.is_file(): |
| raise FileNotFoundError(path) |
| selected: list[Any] = [] |
| decoded_count = 0 |
| start = expected_count - num_frames |
| with av.open(str(path), mode="r") as container: |
| streams = container.streams.video |
| if len(streams) != 1: |
| raise ValueError( |
| f"expected exactly one video stream in {path}, got {len(streams)}" |
| ) |
| stream = streams[0] |
| for frame_index, frame in enumerate(container.decode(stream)): |
| decoded_count = frame_index + 1 |
| if start <= frame_index < expected_count: |
| selected.append(frame) |
| if decoded_count != expected_count: |
| raise ValueError( |
| "matched video source frame count violates the last64 contract: " |
| f"sample={path.name}, actual={decoded_count}, " |
| f"expected={expected_count}" |
| ) |
| if len(selected) != num_frames: |
| raise RuntimeError( |
| f"failed to decode frames {start}..{expected_count - 1} from {path}" |
| ) |
| video01 = _frames_to_video01(selected) |
| if tuple(video01.shape[:2]) != (num_frames, 3): |
| raise RuntimeError(f"decoded video has invalid shape {tuple(video01.shape)}") |
| return DecodedMatchedVideo( |
| video01=video01, |
| frame_indices=torch.arange(start, expected_count, dtype=torch.long), |
| source_frame_count=decoded_count, |
| ) |
|
|
|
|
| DecodeFunction = Callable[[Path, Mapping[str, Any]], DecodedMatchedVideo] |
|
|
|
|
| def _default_decode( |
| path: Path, |
| contract: Mapping[str, Any], |
| ) -> DecodedMatchedVideo: |
| return decode_matched_video_last_contiguous_pyav( |
| path, decoder_contract=contract |
| ) |
|
|
|
|
| class MatchedManifestVideoDataset(Dataset): |
| """Return metadata eagerly; teacher cache/decode only on explicit request.""" |
|
|
| def __init__( |
| self, |
| manifest_path: str | os.PathLike, |
| video_root: str | os.PathLike, |
| cache_root: str | os.PathLike, |
| expected_model_sha256: str, |
| expected_feature_contract_sha256: str, |
| *, |
| expected_source_csv_sha256: str | None = None, |
| expected_count: int | None = None, |
| feature_contract: Mapping[str, Any] | None = None, |
| decoder_contract: Mapping[str, Any] | None = None, |
| decode_fn: DecodeFunction | None = None, |
| ) -> None: |
| self.manifest_path = Path(manifest_path).expanduser().resolve() |
| self.video_root = Path(video_root).expanduser().resolve() |
| if not self.video_root.is_dir(): |
| raise NotADirectoryError(self.video_root) |
| self.entries, self.manifest_metadata = load_matched_video_manifest( |
| self.manifest_path, |
| expected_source_csv_sha256=expected_source_csv_sha256, |
| expected_count=expected_count, |
| ) |
| if not self.entries: |
| raise ValueError("matched-video manifest is empty") |
| self.feature_cache = MatchedVideoFeatureCache( |
| cache_root, |
| expected_model_sha256=expected_model_sha256, |
| expected_feature_contract_sha256=expected_feature_contract_sha256, |
| feature_contract=( |
| feature_contract |
| if feature_contract is not None |
| else MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT |
| ), |
| decoder_contract=( |
| decoder_contract |
| if decoder_contract is not None |
| else MATCHED_PYAV_LAST64_DECODER_CONTRACT |
| ), |
| ) |
| self.decoder_contract = self.feature_cache.decoder_contract |
| self.decode_fn = decode_fn if decode_fn is not None else _default_decode |
| self._cache_key_to_index: dict[str, int] = {} |
| self._sample_id_to_index: dict[str, int] = {} |
| for index, entry in enumerate(self.entries): |
| cache_key = self.feature_cache.key_for(entry) |
| if cache_key in self._cache_key_to_index: |
| raise ValueError(f"duplicate matched-video cache key: {cache_key}") |
| if entry.sample_id in self._sample_id_to_index: |
| raise ValueError(f"duplicate sample_id: {entry.sample_id}") |
| self._cache_key_to_index[cache_key] = index |
| self._sample_id_to_index[entry.sample_id] = index |
|
|
| def __len__(self) -> int: |
| return len(self.entries) |
|
|
| def _entry_for_sample(self, sample_id: str) -> MatchedVideoManifestEntry: |
| try: |
| return self.entries[self._sample_id_to_index[str(sample_id)]] |
| except KeyError as exc: |
| raise KeyError(f"unknown matched-video sample_id: {sample_id!r}") from exc |
|
|
| def _video_path_and_verify_stat( |
| self, |
| entry: MatchedVideoManifestEntry, |
| ) -> Path: |
| path = self.video_root / entry.video_name |
| if not path.is_file(): |
| raise FileNotFoundError(path) |
| stat = path.stat() |
| actual = (int(stat.st_size), int(stat.st_mtime_ns)) |
| expected = (entry.video_size_bytes, entry.video_mtime_ns) |
| if actual != expected: |
| raise ValueError( |
| "matched video is stale after manifest freeze for " |
| f"sample_id={entry.sample_id}: actual={actual}, " |
| f"expected={expected}" |
| ) |
| return path |
|
|
| def cache_key_for_sample(self, sample_id: str) -> str: |
| return self.feature_cache.key_for(self._entry_for_sample(sample_id)) |
|
|
| def __getitem__(self, index: int) -> dict[str, Any]: |
| """Return only prompt/sample/path metadata; never touch teacher data.""" |
|
|
| entry = self.entries[int(index)] |
| video_path = self._video_path_and_verify_stat(entry) |
| return { |
| "idx": entry.manifest_index, |
| "prompts": entry.caption, |
| "sample_id": entry.sample_id, |
| "teacher_cache_key": self.feature_cache.key_for(entry), |
| "video_name": entry.video_name, |
| "video_path": str(video_path), |
| } |
|
|
| def load_teacher_payload(self, sample_id: str) -> dict[str, Any]: |
| """Load/decode teacher input only after suffix replay selects a sample.""" |
|
|
| entry = self._entry_for_sample(sample_id) |
| video_path = self._video_path_and_verify_stat(entry) |
| cache_key = self.feature_cache.key_for(entry) |
| cached = self.feature_cache.load(entry, cache_key=cache_key) |
| payload: dict[str, Any] = { |
| "sample_id": entry.sample_id, |
| "teacher_cache_key": cache_key, |
| "video_name": entry.video_name, |
| "video_path": str(video_path), |
| } |
| if cached is not None: |
| payload.update({ |
| "current_teacher_cache_hit": True, |
| "current_teacher_decode_seconds": 0.0, |
| "current_teacher_features": cached, |
| }) |
| return payload |
|
|
| decode_started = time.perf_counter() |
| decoded = self.decode_fn(video_path, self.decoder_contract) |
| decode_seconds = time.perf_counter() - decode_started |
| if not isinstance(decoded, DecodedMatchedVideo): |
| raise TypeError("decode_fn must return DecodedMatchedVideo") |
| expected_count = int(self.decoder_contract["expected_source_frame_count"]) |
| expected_frames = int(self.decoder_contract["num_frames"]) |
| expected_indices = torch.arange( |
| expected_count - expected_frames, |
| expected_count, |
| dtype=torch.long, |
| ) |
| if ( |
| int(decoded.source_frame_count) != expected_count |
| or not isinstance(decoded.frame_indices, torch.Tensor) |
| or decoded.frame_indices.dtype != torch.long |
| or not torch.equal( |
| decoded.frame_indices.detach().cpu(), expected_indices |
| ) |
| ): |
| raise ValueError( |
| "decode_fn violated the strict last-contiguous 17..80 contract" |
| ) |
| video = decoded.video01 |
| if ( |
| not isinstance(video, torch.Tensor) |
| or video.dtype != torch.float32 |
| or video.ndim != 4 |
| or tuple(video.shape[:2]) != (expected_frames, 3) |
| or not bool(torch.isfinite(video).all().item()) |
| or float(video.min()) < 0.0 |
| or float(video.max()) > 1.0 |
| ): |
| raise ValueError( |
| "decode_fn returned invalid video01; expected finite FP32 " |
| f"[{expected_frames},3,H,W] in [0,1]" |
| ) |
| payload.update({ |
| "current_teacher_cache_hit": False, |
| "current_teacher_decode_seconds": decode_seconds, |
| "current_teacher_frame_indices": decoded.frame_indices, |
| "current_teacher_source_frame_count": expected_count, |
| "current_teacher_video01": video.contiguous(), |
| }) |
| return payload |
|
|
| def write_feature_cache( |
| self, |
| cache_key: str, |
| feature: torch.Tensor, |
| *, |
| sample_id: str, |
| ) -> Path: |
| try: |
| index = self._cache_key_to_index[str(cache_key)] |
| except KeyError as exc: |
| raise KeyError(f"unknown teacher cache key: {cache_key!r}") from exc |
| entry = self.entries[index] |
| if entry.sample_id != str(sample_id): |
| raise ValueError( |
| "sample_id does not own teacher cache key: " |
| f"{sample_id!r} != {entry.sample_id!r}" |
| ) |
| self._video_path_and_verify_stat(entry) |
| return self.feature_cache.write( |
| entry, feature, cache_key=str(cache_key) |
| ) |
|
|
| @staticmethod |
| def collate_fn(batch: Sequence[Mapping[str, Any]]) -> dict[str, Any]: |
| return collate_matched_video_metadata(batch) |
|
|
|
|
| def collate_matched_video_metadata( |
| batch: Sequence[Mapping[str, Any]], |
| ) -> dict[str, Any]: |
| """Collate only metadata; teacher payloads intentionally have no collator.""" |
|
|
| if not batch: |
| raise ValueError("cannot collate an empty matched-video batch") |
| forbidden = { |
| "current_teacher_features", |
| "current_teacher_video01", |
| "current_teacher_cache_hit", |
| } |
| for item in batch: |
| present = forbidden.intersection(item) |
| if present: |
| raise ValueError( |
| "metadata collator received eager teacher payload fields: " |
| f"{sorted(present)}" |
| ) |
| return { |
| "idx": torch.tensor([int(item["idx"]) for item in batch]), |
| "prompts": [str(item["prompts"]) for item in batch], |
| "sample_ids": [str(item["sample_id"]) for item in batch], |
| "teacher_cache_keys": [str(item["teacher_cache_key"]) for item in batch], |
| "video_names": [str(item["video_name"]) for item in batch], |
| "video_paths": [str(item["video_path"]) for item in batch], |
| } |
|
|
|
|
| __all__ = [ |
| "DecodedMatchedVideo", |
| "MatchedManifestVideoDataset", |
| "collate_matched_video_metadata", |
| "decode_matched_video_last_contiguous_pyav", |
| ] |
|
|
