rpx_benchmark.tasks._video_pipeline

  1"""Per-clip pipeline for video tasks (Video Depth today, future per-clip tasks
  2later).
  3
  4The sibling of :mod:`rpx_benchmark.tasks._pipeline`, which handles
  5per-frame tasks. The two pipelines are deliberately split rather than
  6overloaded onto a single ``run_pipeline``: the iteration unit is
  7different (frame vs. clip), the dataset class is different
  8(:class:`~rpx_benchmark.loader.RPXDataset` vs.
  9:class:`~rpx_benchmark.video_loader.VideoDepthDataset`), and per-clip tasks
 10have an explicit alignment hook between ``model.predict`` and the
 11metric suite that per-frame tasks don't need. Mixing both into one
 12function obscured both code paths during the Video Depth prototype; the
 13split is the result of that prototype's lesson.
 14
 15The contract for a video model is exactly what
 16:class:`~rpx_benchmark.api.BenchmarkModel` describes: ``setup()``
 17loads weights, ``predict(batch[VideoSample]) -> list[VideoDepthPrediction]``
 18runs one inference per clip. The model declares
 19``depth_output_kind = "metric"`` (default) or ``"relative"`` on the
 20class; the pipeline applies the adapter's declared per-clip alignment
 21to the prediction *before* metrics are computed when the model is relative.
 22This matches the community convention for evaluating affine-invariant
 23video-depth models (DepthCrafter, RollingDepth, MoGe-2 video runs).
 24"""
 25
 26from __future__ import annotations
 27
 28import json
 29import os
 30import tempfile
 31import time
 32import zipfile
 33from dataclasses import dataclass
 34from pathlib import Path
 35from typing import Dict, List, Optional
 36
 37import numpy as np
 38
 39from ..api import TaskType, VideoDepthPrediction, VideoSample
 40from ..cell_log import cells_from_per_sample, write_cells
 41from ..exceptions import AdapterError, ConfigError
 42from ..hub import DEFAULT_REPO_ID, download_split
 43from ..logging_utils import get_logger
 44from ..metrics.depth_alignment import (
 45    DEPTH_MAX_M,
 46    DEPTH_MIN_M,
 47    align_pred_to_gt_pooled,
 48)
 49from ..metrics.registry import MetricSuite
 50from ..reports import format_markdown_summary, write_json
 51from ..runner import _make_cuda_sync
 52from ..video_loader import VideoDepthDataset
 53from ._pipeline import PipelineResult, TaskRunConfig, resolve_device
 54
 55log = get_logger(__name__)
 56
 57
 58# --------------------------------------------------------------------------- #
 59# Config
 60# --------------------------------------------------------------------------- #
 61
 62
 63@dataclass
 64class VideoTaskRunConfig(TaskRunConfig):
 65    """Adds Video Depth-specific knobs to :class:`TaskRunConfig`.
 66
 67    ``frame_budget`` and ``sampling`` are the temporal-resolution
 68    ablation knobs from the paper (\\S5.2). Defaults mean "use every
 69    frame on disk", which is the headline run.
 70    """
 71
 72    frame_budget: Optional[int] = None
 73    sampling: str = "all"
 74    save_predictions: bool = False
 75    resume_predictions: bool = False
 76    compute_fscore: bool = False
 77
 78
 79# --------------------------------------------------------------------------- #
 80# Alignment helper
 81# --------------------------------------------------------------------------- #
 82
 83
 84def _apply_clip_alignment(
 85    pred: VideoDepthPrediction,
 86    sample: VideoSample,
 87    *,
 88    mode: str = "ls_affine",
 89) -> VideoDepthPrediction:
 90    """Solve one ``(s, t)`` over the whole clip and apply it.
 91
 92    Relative-depth video models produce ``pred.depth_map_seq`` up to an
 93    affine transform per clip. Per the model's official convention
 94    (depth-space for most models, disparity-space for GemDepth), we fit
 95    a single ``(s, t)``
 96    over every valid pixel in the clip and apply it uniformly across
 97    all ``T`` frames. Per-frame fitting would let the model cheat
 98    temporal-inconsistency punishments; per-scene fitting (pooling all
 99    three phases) would conflate phase-stability with calibration drift.
100    """
101    gt = sample.ground_truth
102    fit_mask = (
103        gt.valid_mask_seq.astype(bool)
104        & np.isfinite(pred.depth_map_seq)
105        & np.isfinite(gt.depth_map_seq)
106        & (gt.depth_map_seq > DEPTH_MIN_M)
107        & (gt.depth_map_seq < DEPTH_MAX_M)
108    )
109    raw = pred.depth_map_seq.astype(np.float32)
110    target = gt.depth_map_seq.astype(np.float32)
111    if mode == "ls_disparity":
112        # GemDepth's official evaluator clips raw inverse depth, fits it
113        # to GT inverse depth over the entire clip, and only then takes
114        # the reciprocal.
115        raw = np.clip(raw, 1e-3, None)
116        target_inverse = np.zeros_like(target)
117        np.reciprocal(target, out=target_inverse, where=fit_mask)
118        target = target_inverse
119        aligned_inverse = align_pred_to_gt_pooled(
120            pred_seq=raw,
121            gt_seq=target,
122            mode="ls_affine",
123            valid_seq=fit_mask,
124        )
125        aligned_seq = np.reciprocal(np.clip(aligned_inverse, 1e-3, None))
126    elif mode == "ls_affine":
127        aligned_seq = align_pred_to_gt_pooled(
128            pred_seq=raw,
129            gt_seq=target,
130            mode="ls_affine",
131            valid_seq=fit_mask,
132        )
133    else:
134        raise ConfigError(
135            f"unsupported video-depth native alignment mode: {mode!r}"
136        )
137    return VideoDepthPrediction(depth_map_seq=aligned_seq)
138
139
140def _prediction_path(predictions_dir: Path, sample: VideoSample) -> Path:
141    meta = sample.metadata or {}
142    scene = meta.get("scene_id") or str(sample.id)
143    phase = meta.get("phase_idx")
144    if phase is None:
145        phase = sample.phase.value if sample.phase is not None else "unknown"
146    return predictions_dir / str(scene) / str(phase) / "depth.npz"
147
148
149def _paper_valid_mask(sample: VideoSample) -> np.ndarray:
150    gt = np.asarray(sample.ground_truth.depth_map_seq)
151    return (
152        np.asarray(sample.ground_truth.valid_mask_seq, dtype=bool)
153        & np.isfinite(gt)
154        & (gt > DEPTH_MIN_M)
155        & (gt < DEPTH_MAX_M)
156    )
157
158
159def _validate_prediction(
160    depth: np.ndarray,
161    sample: VideoSample,
162    *,
163    raise_on_error: bool,
164) -> np.ndarray | None:
165    expected_shape = np.asarray(sample.ground_truth.depth_map_seq).shape
166    error: str | None = None
167    if depth.dtype != np.dtype(np.float32):
168        error = f"dtype {depth.dtype} is not float32"
169    elif depth.ndim != 3 or depth.shape != expected_shape:
170        error = f"shape {depth.shape} does not match {expected_shape}"
171    elif not np.isfinite(depth).all():
172        error = "prediction contains non-finite values"
173    else:
174        valid = _paper_valid_mask(sample)
175        if not valid.any():
176            error = "clip has no GT pixels in 0.3 < depth < 5.0 m"
177        elif float(np.ptp(depth[valid])) <= 1e-6:
178            error = "prediction is degenerate (constant on paper-valid pixels)"
179    if error is None:
180        return depth
181    if raise_on_error:
182        raise AdapterError(
183            f"Refusing invalid Video Depth prediction for {sample.id!r}: {error}",
184            hint=(
185                "Expected one finite, non-degenerate float32 depth sequence "
186                "matching GT. Evaluation applies the documented [0.3, 5.0] "
187                "prediction clip on the strict paper-valid GT mask."
188            ),
189        )
190    return None
191
192
193def _load_cached_prediction(
194    path: Path,
195    sample: VideoSample,
196) -> tuple[VideoDepthPrediction, float] | None:
197    try:
198        with np.load(path, allow_pickle=False) as payload:
199            if set(payload.files) != {"depth", "frame_indices", "latency_ms"}:
200                return None
201            depth = np.asarray(payload["depth"])
202            frame_indices = np.asarray(payload["frame_indices"])
203            latency_raw = np.asarray(payload["latency_ms"])
204        expected_frames = np.asarray(sample.ground_truth.frame_indices, dtype=np.int32)
205        if frame_indices.dtype != np.dtype(np.int32):
206            return None
207        if not np.array_equal(frame_indices, expected_frames):
208            return None
209        if latency_raw.size != 1:
210            return None
211        latency_ms = float(latency_raw.reshape(-1)[0])
212        if not np.isfinite(latency_ms) or latency_ms < 0:
213            return None
214        depth = _validate_prediction(depth, sample, raise_on_error=False)
215        if depth is None:
216            return None
217        return VideoDepthPrediction(depth_map_seq=depth), latency_ms
218    except (OSError, ValueError, KeyError, EOFError, zipfile.BadZipFile):
219        return None
220
221
222def _atomic_save_prediction(
223    path: Path,
224    prediction: VideoDepthPrediction,
225    sample: VideoSample,
226    latency_ms: float,
227) -> None:
228    depth = np.asarray(prediction.depth_map_seq, dtype=np.float32)
229    _validate_prediction(depth, sample, raise_on_error=True)
230    path.parent.mkdir(parents=True, exist_ok=True)
231    temp_path: Path | None = None
232    try:
233        with tempfile.NamedTemporaryFile(
234            mode="w+b",
235            suffix=".npz.part",
236            dir=path.parent,
237            delete=False,
238        ) as handle:
239            temp_path = Path(handle.name)
240            np.savez_compressed(
241                handle,
242                depth=depth,
243                frame_indices=np.asarray(
244                    sample.ground_truth.frame_indices,
245                    dtype=np.int32,
246                ),
247                latency_ms=np.asarray(latency_ms, dtype=np.float64),
248            )
249            handle.flush()
250            os.fsync(handle.fileno())
251        os.replace(temp_path, path)
252    finally:
253        if temp_path is not None and temp_path.exists():
254            temp_path.unlink()
255
256
257def _atomic_write_json(path: Path, payload: dict) -> None:
258    path.parent.mkdir(parents=True, exist_ok=True)
259    temp_path: Path | None = None
260    try:
261        with tempfile.NamedTemporaryFile(
262            mode="w",
263            suffix=".json.part",
264            dir=path.parent,
265            encoding="utf-8",
266            delete=False,
267        ) as handle:
268            temp_path = Path(handle.name)
269            json.dump(payload, handle, indent=2, sort_keys=True)
270            handle.write("\n")
271            handle.flush()
272            os.fsync(handle.fileno())
273        os.replace(temp_path, path)
274    finally:
275        if temp_path is not None and temp_path.exists():
276            temp_path.unlink()
277
278
279# --------------------------------------------------------------------------- #
280# Pipeline
281# --------------------------------------------------------------------------- #
282
283
284def run_video_pipeline(
285    *,
286    task: TaskType,
287    primary_metric: str,
288    cfg: VideoTaskRunConfig,
289) -> PipelineResult:
290    """End-to-end per-clip pipeline: download → load → predict (+ align)
291    → metric suite → cell-log row → JSON / markdown summary.
292
293    The only Video Depth-specific knobs are ``frame_budget`` and ``sampling``;
294    everything else lines up with the per-frame
295    :func:`~rpx_benchmark.tasks._pipeline.run_pipeline` so the team
296    doesn't have to learn two contracts.
297    """
298    cfg.device = resolve_device(cfg.device, require_cuda=cfg.require_cuda)
299    if cfg.resume_predictions and not cfg.save_predictions:
300        raise ConfigError(
301            "resume_predictions requires save_predictions",
302            hint="Enable both flags so missing or invalid cache entries can be replaced.",
303        )
304    split_name = cfg.split.value if hasattr(cfg.split, "value") else str(cfg.split)
305    repo_id = cfg.repo_id or DEFAULT_REPO_ID
306
307    log.info(
308        "video pipeline: task=%s split=%s device=%s frame_budget=%s sampling=%s",
309        task.value,
310        split_name,
311        cfg.device,
312        cfg.frame_budget,
313        cfg.sampling,
314    )
315
316    if cfg.manifest_path:
317        manifest_path = Path(cfg.manifest_path)
318        if not manifest_path.exists():
319            raise ConfigError(
320                f"manifest_path does not exist: {manifest_path}",
321                hint="Generate per-task manifests for a local v2-webp "
322                "tree with `python -m rpx_benchmark.dataset_hub.cli "
323                "manifest ...` and pass the resulting "
324                "<staging>/manifests/video_depth/<split>.json path.",
325            )
326        log.info("using local manifest: %s", manifest_path)
327    else:
328        manifest_path = download_split(
329            task=task,
330            split=cfg.split,
331            repo_id=repo_id,
332            cache_dir=cfg.cache_dir,
333            revision=cfg.revision,
334            max_samples=cfg.max_samples,
335        )
336    dataset = VideoDepthDataset.from_manifest(
337        manifest_path,
338        batch_size=cfg.batch_size,
339        frame_budget=cfg.frame_budget,
340        sampling=cfg.sampling,
341        max_samples=cfg.max_samples,
342        compute_fscore=cfg.compute_fscore,
343    )
344    log.info("loaded %d clips from %s", len(dataset), manifest_path)
345
346    model = cfg.model
347    name = getattr(model, "name", "model")
348    is_relative = getattr(model, "depth_output_kind", "metric") == "relative"
349
350    alignment_mode = getattr(model, "native_alignment", "ls_affine")
351    if is_relative:
352        log.info(
353            "model %s is relative-depth → per-clip %s alignment will be applied",
354            name,
355            alignment_mode,
356        )
357
358    safe_name = name.replace("/", "__")
359    out_dir = Path(cfg.output_dir or f"./rpx_results/{safe_name}/{split_name}")
360    predictions_dir = out_dir / "predictions"
361    prediction_stats = {
362        "cache_hits": 0,
363        "inferred_new": 0,
364        "invalid_recomputed": 0,
365    }
366
367    # Auto-detect SystemCard (GPU / precision / batch_size) so every
368    # cell row carries the host hardware context. Required for
369    # downstream cross-host aggregation per
370    # benchmark/README.md publication gate.
371    from ..profiler import SystemCard
372
373    system_card = SystemCard.auto_detect(
374        precision=getattr(model, "native_precision", "fp32"),
375        batch_size=cfg.batch_size,
376    )
377
378    metric_suite = MetricSuite.for_task(task)
379
380    per_sample: List[dict] = []
381    n_clips = len(dataset)
382    cuda_sync = _make_cuda_sync()
383    model_is_setup = False
384    for clip_idx, batch in enumerate(dataset, start=1):
385        if not batch:
386            continue
387        predictions: List[VideoDepthPrediction | None] = [None] * len(batch)
388        latencies_ms: List[float | None] = [None] * len(batch)
389        missing_indices: List[int] = []
390        invalid_existing: set[int] = set()
391        for index, sample in enumerate(batch):
392            path = _prediction_path(predictions_dir, sample)
393            if cfg.resume_predictions and path.exists():
394                cached = _load_cached_prediction(path, sample)
395                if cached is not None:
396                    predictions[index], latencies_ms[index] = cached
397                    prediction_stats["cache_hits"] += 1
398                    continue
399                invalid_existing.add(index)
400            missing_indices.append(index)
401
402        if missing_indices:
403            if not model_is_setup:
404                model.setup()
405                model_is_setup = True
406            missing_batch = [batch[index] for index in missing_indices]
407            cuda_sync()
408            t0 = time.perf_counter()
409            inferred = model.predict(missing_batch)
410            cuda_sync()
411            elapsed_ms = (time.perf_counter() - t0) * 1000.0
412            if len(inferred) != len(missing_batch):
413                raise ConfigError(
414                    f"model returned {len(inferred)} predictions for "
415                    f"{len(missing_batch)} clips; expected one per clip",
416                )
417            per_clip_latency_ms = elapsed_ms / len(missing_batch)
418            for index, pred in zip(missing_indices, inferred, strict=True):
419                sample = batch[index]
420                if is_relative:
421                    pred = _apply_clip_alignment(
422                        pred,
423                        sample,
424                        mode=alignment_mode,
425                    )
426                canonical = np.asarray(pred.depth_map_seq, dtype=np.float32)
427                _validate_prediction(canonical, sample, raise_on_error=True)
428                pred = VideoDepthPrediction(depth_map_seq=canonical)
429                predictions[index] = pred
430                latencies_ms[index] = per_clip_latency_ms
431                if index in invalid_existing:
432                    prediction_stats["invalid_recomputed"] += 1
433                else:
434                    prediction_stats["inferred_new"] += 1
435                if cfg.save_predictions:
436                    _atomic_save_prediction(
437                        _prediction_path(predictions_dir, sample),
438                        pred,
439                        sample,
440                        per_clip_latency_ms,
441                    )
442
443        for sample, pred, latency_ms in zip(
444            batch,
445            predictions,
446            latencies_ms,
447            strict=True,
448        ):
449            if pred is None or latency_ms is None:  # pragma: no cover - invariant
450                raise ConfigError("internal error: unresolved video prediction")
451            evaluation_pred = VideoDepthPrediction(
452                depth_map_seq=np.clip(
453                    pred.depth_map_seq,
454                    DEPTH_MIN_M,
455                    DEPTH_MAX_M,
456                ).astype(np.float32, copy=False),
457            )
458            metrics: Dict[str, float] = dict(
459                metric_suite.evaluate(evaluation_pred, sample.ground_truth)
460            )
461            meta = sample.metadata or {}
462            metrics.update(
463                id=sample.id,
464                scene=meta.get("scene_id"),
465                phase=sample.phase,
466                difficulty=sample.difficulty,
467                frame_budget=int(meta.get("frame_budget", 0)),
468                latency_ms=latency_ms,
469            )
470            per_sample.append(metrics)
471
472        if clip_idx % 5 == 0 or clip_idx == n_clips:
473            log.info("video pipeline: %d/%d clips processed", clip_idx, n_clips)
474
475    result = metric_suite.build_result(per_sample)
476
477    out_dir.mkdir(parents=True, exist_ok=True)
478    json_path = out_dir / "result.json"
479    md_path = out_dir / "summary.md"
480    cells_path = out_dir / "cells.parquet"
481    metadata_path = out_dir / "run_metadata.json"
482
483    # Auto-detected metric keys exclude obvious metadata (id/scene/phase/
484    # difficulty/latency_ms via the cell_log defaults) but ``frame_budget``
485    # looks numeric, so list metrics explicitly to keep it out.
486    auto_metric_keys = sorted(
487        k for k, v in (per_sample[0].items() if per_sample else [])
488        if isinstance(v, (int, float)) and not isinstance(v, bool)
489        and k not in {"frame_budget", "latency_ms"}
490    )
491    cells = cells_from_per_sample(
492        per_sample,
493        model_name=name,
494        task=task.value,
495        metric_keys=auto_metric_keys,
496        system_card=system_card,
497    )
498    write_cells(cells, cells_path)
499    log.info("wrote %d cell rows to %s", len(cells), cells_path)
500
501    write_json(
502        json_path,
503        task=task.value,
504        model_name=name,
505        split=split_name,
506        repo_id=repo_id,
507        result=result,
508        dr_report=None,
509    )
510    md_path.write_text(
511        format_markdown_summary(
512            task=task.value,
513            model_name=name,
514            split=split_name,
515            repo_id=repo_id,
516            result=result,
517            dr_report=None,
518        ),
519        encoding="utf-8",
520    )
521    log.info("wrote %s and %s", json_path, md_path)
522    _atomic_write_json(
523        metadata_path,
524        {
525            "task": task.value,
526            "model_name": name,
527            "split": split_name,
528            "save_predictions": cfg.save_predictions,
529            "resume_predictions": cfg.resume_predictions,
530            "compute_fscore": cfg.compute_fscore,
531            "prediction_evaluation_policy": {
532                "operation": "clip",
533                "minimum_m": DEPTH_MIN_M,
534                "maximum_m": DEPTH_MAX_M,
535                "raw_cache_preserved": True,
536            },
537            "prediction_stats": prediction_stats,
538        },
539    )
540
541    paths: Dict[str, Path] = {
542        "json": json_path,
543        "markdown": md_path,
544        "cells": cells_path,
545        "out_dir": out_dir,
546        "run_metadata": metadata_path,
547    }
548    if cfg.save_predictions:
549        paths["predictions_dir"] = predictions_dir
550        paths["prediction_stats"] = prediction_stats
551
552    if cfg.upload_to_box:
553        from ..box_upload import DEFAULT_BOX_FOLDER_ID, upload_run_dir
554
555        try:
556            summary = upload_run_dir(
557                out_dir,
558                task=task.value,
559                model_name=name,
560                split=split_name,
561                root_folder_id=cfg.box_folder_id or DEFAULT_BOX_FOLDER_ID,
562            )
563            log.info(
564                "box: %d uploaded, %d skipped → box:%s",
565                summary["uploaded"],
566                summary["skipped"],
567                summary["remote_path"],
568            )
569            paths["box_remote"] = summary["remote_path"]
570        except Exception as e:  # noqa: BLE001
571            log.error(
572                "Box upload failed (%s: %s) — local artefacts at %s are intact",
573                type(e).__name__,
574                e,
575                out_dir,
576            )
577
578    return result, None, paths  # type: ignore[return-value]
579
580
581__all__ = ["VideoTaskRunConfig", "run_video_pipeline"]