icassp2027-experiment/test_ego4d_box_frames.py

"""CPU regressions for bounded, explicitly authorized Ego4D bbox tails."""

from __future__ import annotations

import contextlib
import copy
import io
import json
from pathlib import Path
import sys
import tempfile
from types import SimpleNamespace
import unittest
from unittest.mock import patch

import numpy as np
from PIL import Image
import torch

import extract_turn_local_features as extract
from local_turn_features import load_local_features
from test_local_turn_features import fixture


UID = "31488ff6-05ef-4362-9786-35981b06f307"
TWO_FRAME_UID = "ac7d3d16-d74d-4fe3-a1e0-cab0b1adf492"
THIRTEEN_FRAME_UID = "d2cb4045-8ef1-4288-8398-ab6d4d6a0709"


def terminal_clip(*, uid=UID, count=8999, tail_frames=1):
    num_steps = count // 6  # 30 FPS source, 5 Hz targets.
    num_video_frames = (num_steps * 2 + 4) // 5  # ceil(duration * 2 FPS).
    last_grid_frame = (num_video_frames - 1) * 15
    return {
        "clip_uid": uid, "clip_start_frame": 0, "clip_end_frame": count,
        "clip_start_sec": 0, "clip_end_sec": count / 30.,
        "persons": [{"person_id": "p1", "tracking_paths": [{"track": [
            {"frame": frame, "clip_frame": None, "x": 0, "y": 0, "width": 2, "height": 2}
            for frame in (0, last_grid_frame, count - 1, *range(count, count + tail_frames))
        ]}]}],
    }


def terminal_annotation(**kwargs):
    clip = terminal_clip(**kwargs)
    return {"frames": extract.ego4d_clip_boxes(clip),
            "bounds": {key: clip[key] for key in (
                "clip_start_frame", "clip_end_frame", "clip_start_sec", "clip_end_sec",
            )}}


class Ego4DBoxFrameTests(unittest.TestCase):
    def setUp(self):
        torch.set_num_threads(1)
        self.source = SimpleNamespace(num_frames=8999, source_fps=30.)

    def test_confirmed_terminal_only_is_removed_without_mutating_input(self):
        annotation = terminal_annotation()
        original = copy.deepcopy(annotation)
        with contextlib.redirect_stdout(io.StringIO()) as log:
            frames, audit = extract.validate_ego4d_box_frames(UID, annotation, self.source)
        self.assertEqual(frames, {frame: original["frames"][frame] for frame in (0, 8985, 8998)})
        self.assertEqual(annotation, original)
        self.assertEqual(audit["ignored_frame_indices"], [8999])
        self.assertEqual(audit["ignored_box_count"], 1)
        self.assertEqual(audit["annotation_frame_range"], [0, 8999])
        self.assertIn("No frame shift or feature/target row removal", log.getvalue())

    def test_valid_and_empty_mappings_are_unchanged_even_without_bounds(self):
        for frames in ({}, {0: [], 8998: [{"id": "valid"}]}):
            with self.subTest(frames=frames):
                result, audit = extract.validate_ego4d_box_frames(
                    UID, {"frames": frames, "bounds": {}}, self.source)
                self.assertIs(result, frames)
                self.assertEqual(audit["ignored_frame_indices"], [])

    def test_confirmed_two_frame_tail_preserves_last_valid_box_and_counts_people(self):
        annotation = terminal_annotation(uid=TWO_FRAME_UID, count=8998, tail_frames=2)
        annotation["frames"][8999].append({"id": "p2", "xyxy": [1, 1, 2, 2]})
        original = copy.deepcopy(annotation)
        with contextlib.redirect_stdout(io.StringIO()):
            frames, audit = extract.validate_ego4d_box_frames(
                TWO_FRAME_UID, annotation, SimpleNamespace(num_frames=8998, source_fps=30.))
        self.assertEqual(frames, {frame: original["frames"][frame] for frame in (0, 8985, 8997)})
        self.assertEqual(annotation, original)
        self.assertEqual(audit["ignored_frame_indices"], [8998, 8999])
        self.assertEqual(audit["ignored_box_count"], 3)
        self.assertEqual(audit["max_terminal_frame_overrun"], 2)

    def test_sparse_tail_is_bounded_by_position_not_outlier_count(self):
        annotation = terminal_annotation(count=8998, tail_frames=2)
        del annotation["frames"][8998]
        source = SimpleNamespace(num_frames=8998, source_fps=30.)
        with contextlib.redirect_stdout(io.StringIO()):
            frames, audit = extract.validate_ego4d_box_frames(TWO_FRAME_UID, annotation, source)
        self.assertEqual(audit["ignored_frame_indices"], [8999])
        self.assertEqual(sorted(frames), [0, 8985, 8997])
        annotation["frames"][9999] = annotation["frames"].pop(8999)
        with self.assertRaisesRegex(ValueError, "out_of_range_count=1"):
            extract.validate_ego4d_box_frames(TWO_FRAME_UID, annotation, source)

    def test_larger_overrun_negative_and_truncated_source_still_fail(self):
        for extra in (-1, 9001, 12000):
            annotation = terminal_annotation()
            annotation["frames"][extra] = []
            with self.subTest(frame=extra), self.assertRaisesRegex(ValueError, "source_frames=8999"):
                extract.validate_ego4d_box_frames(UID, annotation, self.source)
        with self.assertRaisesRegex(ValueError, "source_frames=8998"):
            extract.validate_ego4d_box_frames(
                UID, terminal_annotation(), SimpleNamespace(num_frames=8998, source_fps=30.))

    def test_terminal_requires_zero_origin_matching_declared_end_and_duration(self):
        for key, value in (("clip_start_frame", 1), ("clip_end_frame", 9000),
                           ("clip_start_sec", .1), ("clip_end_sec", 300.),
                           ("clip_end_sec", float("nan"))):
            annotation = terminal_annotation()
            annotation["bounds"][key] = value
            with self.subTest(key=key, value=value), self.assertRaisesRegex(ValueError, "declared_bounds"):
                extract.validate_ego4d_box_frames(UID, annotation, self.source)
        for key in terminal_annotation()["bounds"]:
            annotation = terminal_annotation()
            del annotation["bounds"][key]
            with self.subTest(missing=key), self.assertRaises(ValueError):
                extract.validate_ego4d_box_frames(UID, annotation, self.source)
        for fps in (29.97, 15., 0., float("nan")):
            with self.subTest(fps=fps), self.assertRaises(ValueError):
                extract.validate_ego4d_box_frames(
                    UID, terminal_annotation(), SimpleNamespace(num_frames=8999, source_fps=fps))

    def test_thirteen_frame_tail_requires_explicit_limit_and_keeps_valid_boxes(self):
        annotation = terminal_annotation(uid=THIRTEEN_FRAME_UID, count=8987, tail_frames=13)
        original = copy.deepcopy(annotation)
        source = SimpleNamespace(num_frames=8987, source_fps=30.)
        with self.assertRaisesRegex(ValueError, "terminal_overrun_frames=13"):
            extract.validate_ego4d_box_frames(THIRTEEN_FRAME_UID, annotation, source)
        with contextlib.redirect_stdout(io.StringIO()):
            frames, audit = extract.validate_ego4d_box_frames(
                THIRTEEN_FRAME_UID, annotation, source, max_terminal_frame_overrun=13)
        self.assertEqual(annotation, original)
        self.assertEqual(frames, {frame: boxes for frame, boxes in original["frames"].items() if frame < 8987})
        self.assertIn(8986, frames)  # Never discard/replace the last valid frame.
        self.assertEqual(audit["ignored_frame_indices"], list(range(8987, 9000)))
        self.assertEqual(audit["ignored_box_count"], 13)
        self.assertEqual(audit["max_terminal_frame_overrun"], 13)
        self.assertEqual(audit["annotation_terminal_overrun_frames"], 13)

    def test_explicit_limit_does_not_bypass_endpoint_duration_or_far_outliers(self):
        source = SimpleNamespace(num_frames=8987, source_fps=30.)
        for key, value in (("clip_start_frame", 1), ("clip_end_frame", 9000),
                           ("clip_start_sec", .1), ("clip_end_sec", 300.),
                           ("clip_end_sec", float("nan"))):
            annotation = terminal_annotation(count=8987, tail_frames=13)
            annotation["bounds"][key] = value
            with self.subTest(key=key, value=value), self.assertRaises(ValueError):
                extract.validate_ego4d_box_frames(
                    THIRTEEN_FRAME_UID, annotation, source, max_terminal_frame_overrun=13)
        for frame in (-1, 9000, 12000):
            annotation = terminal_annotation(count=8987, tail_frames=0)
            annotation["frames"][frame] = []
            with self.subTest(frame=frame), self.assertRaises(ValueError):
                extract.validate_ego4d_box_frames(
                    THIRTEEN_FRAME_UID, annotation, source, max_terminal_frame_overrun=13)
        with self.assertRaises(ValueError):
            extract.validate_ego4d_box_frames(
                THIRTEEN_FRAME_UID, terminal_annotation(count=8987, tail_frames=13),
                SimpleNamespace(num_frames=8986, source_fps=30.), max_terminal_frame_overrun=13)

    def test_zero_limit_disables_tail_compatibility_and_invalid_limits_fail(self):
        with self.assertRaisesRegex(ValueError, "Current tail limit is 0"):
            extract.validate_ego4d_box_frames(
                UID, terminal_annotation(), self.source, max_terminal_frame_overrun=0)
        annotation = terminal_annotation(tail_frames=0)
        frames, audit = extract.validate_ego4d_box_frames(
            UID, annotation, self.source, max_terminal_frame_overrun=0)
        self.assertIs(frames, annotation["frames"])
        self.assertEqual(audit["ignored_frame_indices"], [])
        for limit in (-1, 1.5, True):
            with self.subTest(limit=limit), self.assertRaises(ValueError):
                extract.validate_ego4d_box_frames(
                    UID, annotation, self.source, max_terminal_frame_overrun=limit)

    def test_full_extraction_keeps_1499_rows_saves_audit_and_reuses(self):
        self._assert_full_extraction_and_reuse(UID, 8999, 1)

    def test_two_frame_tail_extraction_matches_manual_removal_of_unsampled_boxes(self):
        self._assert_full_extraction_and_reuse(TWO_FRAME_UID, 8998, 2)

    def test_thirteen_frame_tail_extraction_keeps_target_grid_and_reuses(self):
        self._assert_full_extraction_and_reuse(THIRTEEN_FRAME_UID, 8987, 13, explicit_limit=13)

    def _assert_full_extraction_and_reuse(self, uid, count, tail_frames, explicit_limit=None):
        num_steps = count // 6
        last_video_time = ((num_steps * 2 + 4) // 5 - 1) / 2
        image = Image.new("RGB", (2, 2))
        source = SimpleNamespace(num_frames=count, source_fps=30., crop=None, dense_start=0.,
                                 path=Path(f"{uid}.mp4"), width=2, height=2)
        media = SimpleNamespace(
            clips=[source], metadata=lambda: {},
            _source_frame_for_time=lambda t: (0, int(t * 30 + 1e-8), int(t * 30 + 1e-8) / 30),
            frame_at_or_before=lambda t: image,
            audio_slice=lambda start, end: np.array([start, end]),
        )
        fake_module = SimpleNamespace(CausalSourceAV=lambda *a, **kw: media,
                                      SourceClipSpec=lambda path: path, _easycom_directory_specs=None)
        encoder = SimpleNamespace(
            visual_dim=6, audio_dim=5, transformers_version="4.57.6", model_config_sha256="test",
            images=lambda crops: torch.ones(len(crops), 6),
            audio=lambda waveform: torch.full((5,), float(waveform[-1])),
        )
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            for folder in ("global", "targets", "v1/clips", "v1/annotations"):
                (root / folder).mkdir(parents=True)
            reference, _, targets = fixture(n=num_steps, name=uid, dataset="ego4d")
            torch.save(reference, root / "global" / f"{uid}.pt")
            np.save(root / "targets" / f"{uid}.npy", targets)
            (root / "v1/clips" / f"{uid}.mp4").write_bytes(b"fake-video")
            clip = terminal_clip(uid=uid, count=count, tail_frames=tail_frames)
            for split, clips in (("train", [clip]), ("val", [])):
                (root / f"v1/annotations/av_{split}.json").write_text(json.dumps({"videos": [{"clips": clips}]}))
            loaded = extract.load_ego4d_boxes([root / "v1/annotations/av_train.json"], [uid])
            self.assertEqual(loaded[uid], terminal_annotation(uid=uid, count=count, tail_frames=tail_frames))
            argv = ["--dataset", "ego4d", "--data-root", str(root), "--feature-root", str(root / "global"),
                    "--target-root", str(root / "targets"), "--output-root", str(root / "local"),
                    "--samples", uid, "--no-progress"]
            if explicit_limit is not None:
                before = extract.requested_local_config(extract.parse_args(argv))
                argv += ["--ego4d-max-terminal-overrun-frames", str(explicit_limit)]
                self.assertEqual(extract.requested_local_config(extract.parse_args(argv)), before)
            with patch.dict(sys.modules, {"transformers": SimpleNamespace(__version__="4.57.6"),
                                          "streaming_turn_model_directly": fake_module}), \
                 patch.object(extract, "FrozenQwenLocalEncoder", return_value=encoder) as model, \
                 contextlib.redirect_stdout(io.StringIO()):
                extract.main(argv)
                model.assert_called_once()
                destination = root / "local" / f"{uid}.pt"
                result = torch.load(destination, weights_only=False)
                self.assertEqual(result["visual"].shape, (num_steps, 6))
                self.assertEqual(result["audio"].shape, (num_steps, 5))
                self.assertTrue(torch.equal(result["times"], reference["times"]))
                self.assertTrue(torch.equal(result["label_indices"], reference["label_indices"]))
                self.assertAlmostEqual(result["times"][-1].item(), num_steps / 5)
                self.assertLessEqual(num_steps / 5, count / 30)
                self.assertFalse(result["partial"])
                ignored = list(range(count, count + tail_frames))
                self.assertEqual(result["box_frame_audit"]["ignored_frame_indices"], ignored)
                self.assertEqual(result["persons"]["frame_times"][-1].item(), last_video_time)
                self.assertEqual(result["persons"]["records"][-1]["source_frame_index"], int(last_video_time * 30))
                load_local_features(destination, reference, sample_id=uid, dataset="ego4d", modalities="av")
                # Previously completed valid sidecars have no audit. Their
                # vectors/config remain reusable: only unsampled boxes changed.
                del result["box_frame_audit"]
                torch.save(result, destination)
                model.reset_mock()
                extract.main(argv + ["--reuse-existing"])
                model.assert_not_called()
                extract.main(argv + ["--preview-only", "--preview-count", "1"])
                audit = json.loads((root / "local" / f"{uid}_preview/box_frame_audit.json").read_text())
                self.assertEqual(audit["ignored_frame_indices"], ignored)
                # The automatic tail handling must produce the same vectors,
                # masks, timestamps and source-frame records as a clean fixture.
                track = clip["persons"][0]["tracking_paths"][0]
                track["track"] = [box for box in track["track"] if box["frame"] < count]
                (root / "v1/annotations/av_train.json").write_text(json.dumps({"videos": [{"clips": [clip]}]}))
                extract.main(argv + ["--output-root", str(root / "clean")])
                clean = torch.load(root / "clean" / f"{uid}.pt", weights_only=False)
                for key in ("visual", "audio", "visual_valid", "audio_valid", "times", "label_indices"):
                    self.assertTrue(torch.equal(result[key], clean[key]), key)
                for key in ("features", "frame_times", "frame_offsets"):
                    self.assertTrue(torch.equal(result["persons"][key], clean["persons"][key]), key)
                self.assertEqual(result["persons"]["records"], clean["persons"]["records"])
                self.assertEqual(clean["box_frame_audit"]["ignored_frame_indices"], [])


if __name__ == "__main__":
    unittest.main()