"""CPU tests for B1 alignment, causal pooling, fusion, and train/infer round trips.""" from __future__ import annotations import copy import importlib.util import json import os from pathlib import Path import subprocess import sys import tempfile import unittest from types import SimpleNamespace import numpy as np import torch from PIL import Image, ImageDraw from easycom_head import TurnPredictionHead, checkpoint_head_config from extract_turn_local_features import ( FrozenQwenLocalEncoder, ego4d_clip_boxes, extract_local, validate_source_timeline, model_config_digest, box_coordinate_offset, frame_crops, parse_args, preview_boxes, ) from local_turn_features import ( LOCAL_FORMAT, causal_frame_ticks, crop_box, easycom_boxes, load_local_features, local_head_dimensions, pooled_recent_visual, reference_alignment, ) from train_turn_head import _load_examples def fixture(n=20, name="Session_4", dataset="easycom"): times = (torch.arange(n, dtype=torch.float64) + 1) / 5 reference = {"states": torch.randn(n, 4, 8), "label_indices": torch.arange(n), "times": times, "mode": "open-stream-replay", "start": 0} side = {"format": LOCAL_FORMAT, "dataset": dataset, "sample_id": name, "times": times.clone(), "label_indices": torch.arange(n), "partial": False, "visual": torch.randn(n, 6), "audio": torch.randn(n, 5), "visual_valid": torch.ones(n, dtype=torch.bool), "audio_valid": torch.ones(n, dtype=torch.bool), "config": {"visual_dim": 6, "audio_dim": 5, "encoder_id": "synthetic-test", "causal_policy": "half-open-trailing-window", "uses_speech_labels": False}} side["visual_valid"][0] = False side["visual"][0] = 0 labels = np.tile(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 1, 1]], dtype=np.float32), (n // 4 + 1, 1))[:n] return reference, side, labels class LocalTests(unittest.TestCase): def setUp(self): torch.set_num_threads(1) def test_half_open_grid_and_pooling(self): self.assertEqual(list(causal_frame_ticks(.2, 2)), [0]) self.assertEqual(list(causal_frame_ticks(1., 2)), [0, 1]) ts = torch.tensor([0., .5, 1., 1.5], dtype=torch.float64) vectors = torch.tensor([[2.], [4.], [90.], [100.]]) valid = torch.ones(4, dtype=torch.bool) value, present = pooled_recent_visual(ts, vectors, valid, 1., 1.) self.assertTrue(present) self.assertEqual(value.item(), 3.) # no frame at t=1 value, _ = pooled_recent_visual(ts, vectors, valid, 1.2, 1.) self.assertEqual(value.item(), 47.) # .5 and 1.0, not 0.0 value, present = pooled_recent_visual(ts, vectors, valid * False, 1., 1.) self.assertFalse(present) self.assertEqual(value.item(), 0.) def test_box_formats_and_crop(self): person = {"Participant_ID": -1, "x1": 2, "y1": 3, "x2": 7, "y2": 9} faces = easycom_boxes([{"Frame_Number": 1, "Participants": [person]}], "face") heads = easycom_boxes([{"Frame_Number": 0, "Participants": person}], "head") self.assertEqual(faces, heads) self.assertEqual(crop_box([202., 3., 207., 9.], 1920, 1080, offset_x=200), (2, 3, 7, 9)) self.assertIsNone(crop_box([4., 4., 3., 5.], 10, 10)) self.assertIsNone(crop_box([float("nan"), 4., 9., 5.], 10, 10)) def test_ego4d_clip_frames(self): boxes = [{"frame": 3, "clip_frame": None, "x": 1, "y": 2, "width": 3, "height": 4}] person = {"person_id": "x", "camera_wearer": False, "tracking_paths": [{"track": boxes}, {"track": boxes}]} wearer = {**person, "camera_wearer": True} actual = ego4d_clip_boxes({"persons": [wearer, person]}) self.assertEqual(actual, {3: [{"id": "x", "xyxy": [1., 2., 4., 6.]}]}) def test_dataset_specific_coordinate_defaults_and_explicit_overrides(self): required = ["--data-root", "/raw", "--feature-root", "/global", "--target-root", "/targets", "--output-root", "/local", "--samples", "4"] self.assertEqual(parse_args(required).box_coordinate_space, "encoded") self.assertEqual(parse_args(required + ["--dataset", "ego4d"]).box_coordinate_space, "camera") self.assertEqual(parse_args(required + ["--box-coordinate-space", "camera"]).box_coordinate_space, "camera") # An appended correction overrides an old LOCAL_ARGS array in the shell. self.assertEqual(parse_args(required + ["--box-coordinate-space", "camera", "--box-coordinate-space", "encoded"]).box_coordinate_space, "encoded") def test_frame_crops_subtract_each_clips_banner_before_encoding(self): expected = (100, 80, 160, 140) image = Image.new("RGB", (1920, 1080)) ImageDraw.Draw(image).rectangle(expected, fill="red") clips, annotations = [], [] for width in (2123, 2124, 2203, 1920): offset = width - 1920 clips.append(SimpleNamespace(crop=(offset, 0, 1920, 1080), dense_start=float(len(clips)))) annotations.append({0: [{"id": "3", "xyxy": [100 + offset, 80, 160 + offset, 140]}]}) media = SimpleNamespace(clips=clips, _source_frame_for_time=lambda t: (int(t), 0, 0.), frame_at_or_before=lambda t: image.copy()) args = SimpleNamespace(box_coordinate_space="encoded", crop_expansion=1.) for index, clip in enumerate(clips): with self.subTest(crop=clip.crop): _, crops, records = frame_crops(media, annotations, float(index), args) self.assertEqual(records[0]["box"], expected) self.assertEqual(records[0]["coordinate_offset"], [clip.crop[0], 0]) self.assertEqual(records[0]["annotation_box"], annotations[index][0][0]["xyxy"]) self.assertEqual(crops[0].size, (60, 60)) self.assertEqual(crops[0].getpixel((30, 30)), (255, 0, 0)) def test_cropped_camera_annotations_and_uncropped_ego4d_are_not_shifted_twice(self): clip = SimpleNamespace(crop=(204, 12, 1920, 1080), dense_start=0.) self.assertEqual(box_coordinate_offset(clip, "encoded"), (204, 12)) self.assertEqual(box_coordinate_offset(clip, "camera"), (0, 0)) self.assertEqual(box_coordinate_offset(SimpleNamespace(crop=None), "encoded"), (0, 0)) with self.assertRaisesRegex(ValueError, "coordinate space"): box_coordinate_offset(clip, "typo") media = SimpleNamespace(clips=[clip], _source_frame_for_time=lambda t: (0, 0, 0.), frame_at_or_before=lambda t: Image.new("RGB", (1920, 1080))) annotations = [{0: [{"id": "3", "xyxy": [100, 80, 160, 140]}]}] args = SimpleNamespace(box_coordinate_space="camera", crop_expansion=1.) _, _, records = frame_crops(media, annotations, .1, args) self.assertEqual(records[0]["box"], (100, 80, 160, 140)) # Encoded boxes translate on BOTH axes, then clip at the visible edge. annotations = [{0: [{"id": "3", "xyxy": [190, 4, 230, 50]}]}] args.box_coordinate_space = "encoded" _, _, records = frame_crops(media, annotations, .1, args) self.assertEqual(records[0]["box"], (0, 0, 26, 38)) def test_preview_records_the_same_transformation_used_for_extraction(self): clip = SimpleNamespace(path=Path("source.mp4"), width=2124, height=1080, crop=(204, 0, 1920, 1080), dense_start=0.) media = SimpleNamespace(clips=[clip], _source_frame_for_time=lambda t: (0, 0, 0.), frame_at_or_before=lambda t: Image.new("RGB", (1920, 1080))) annotations = [{0: [{"id": "3", "xyxy": [304, 80, 364, 140]}]}] with tempfile.TemporaryDirectory() as directory: args = SimpleNamespace(box_coordinate_space="encoded", crop_expansion=1., video_fps=2., output_root=Path(directory)) preview_boxes(media, annotations, args, "Session_4", .2) folder = Path(directory) / "Session_4_preview" audit = json.loads((folder / "boxes.json").read_text()) self.assertEqual(audit[0]["coordinate_offset"], [204, 0]) self.assertEqual(audit[0]["boxes"][0]["box"], [100, 80, 160, 140]) self.assertTrue((folder / audit[0]["image"]).is_file()) def test_sidecar_alignment_contract_and_modalities(self): reference, side, _ = fixture() with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "Session_4.pt" for mode, width in (("visual", 7), ("audio", 6), ("av", 13)): torch.save(side, path) features, contract = load_local_features(path, reference, sample_id="Session_4", dataset="easycom", modalities=mode) self.assertEqual(features.shape, (20, width)) dims = local_head_dimensions(contract) self.assertEqual(dims["local_visual_dim"], 0 if mode == "audio" else 6) for field, value in (("times", side["times"] + .2), ("label_indices", torch.arange(20) + 1), ("sample_id", "Session_5"), ("partial", True), ("visual", torch.full((20, 6), float("nan"))), ("visual_valid", torch.ones(20))): broken = {**side, field: value} torch.save(broken, path) with self.subTest(field=field), self.assertRaises(ValueError): load_local_features(path, reference, sample_id="Session_4", dataset="easycom", modalities="av") torch.save(side, path) wrong_contract = copy.deepcopy(contract) wrong_contract["config"]["encoder_id"] = "other-weights" with self.assertRaisesRegex(ValueError, "contract"): load_local_features(path, reference, sample_id="Session_4", dataset="easycom", modalities="av", expected_contract=wrong_contract) with self.assertRaisesRegex(ValueError, "times"): reference_alignment({**reference, "times": reference["times"] + .2}) def test_head_shapes_gradients_missing_and_roundtrip(self): states = torch.randn(3, 4, 8) for v, a in ((0, 0), (6, 0), (0, 5), (6, 5)): for projection in (0, 3): head = TurnPredictionHead(hidden_size=8, num_input_layers=4, layer_fusion="concat-mlp", local_visual_dim=v, local_audio_dim=a, local_projection_dim=projection, transition_auxiliary=True, dropout=0) width = sum(d + 1 for d in (v, a) if d) local = torch.randn(3, width) offset = 0 for d in (v, a): if d: local[:, offset + d] = torch.tensor([0., 1., 1.]) offset += d + 1 logits, transition = head.forward_with_transition(states, local) self.assertEqual(logits.shape, (3, 10, 3)) self.assertEqual(transition.shape, (3, 10)) (logits.square().mean() + transition.square().mean()).backward() self.assertIsNotNone(head.net[1].weight.grad) if width: # Changing invalid raw features cannot change the readout. changed = local.clone() offset = 0 for d in (v, a): if d: changed[0, offset:offset + d] = 1000 offset += d + 1 self.assertTrue(torch.equal(head(states, local)[0], head(states, changed)[0])) with self.assertRaisesRegex(ValueError, "requires aligned"): head(states) config = {"hidden_size": 8, "num_input_layers": 4, "layer_fusion": "concat-mlp", "local_visual_dim": v, "local_audio_dim": a, "local_projection_dim": projection, "transition_auxiliary": True, "dropout": 0} restored = TurnPredictionHead(**checkpoint_head_config({"head_config": config})) restored.load_state_dict(head.state_dict()) self.assertTrue(torch.equal(restored(states, local), head(states, local))) def test_legacy_head_parameter_keys_and_values(self): torch.manual_seed(42) old = torch.nn.Sequential(torch.nn.LayerNorm(8), torch.nn.Linear(8, 4), torch.nn.GELU(), torch.nn.Dropout(.1), torch.nn.Linear(4, 30)) torch.manual_seed(42) head = TurnPredictionHead(hidden_size=8) self.assertEqual(set(head.state_dict()), {f"net.{key}" for key in old.state_dict()}) for key, value in old.state_dict().items(): self.assertTrue(torch.equal(value, head.state_dict()[f"net.{key}"])) def test_sparse_selection_keeps_local_rows(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) for folder in ("features", "local", "targets"): (root / folder).mkdir() reference, side, targets = fixture() torch.save(reference, root / "features/Session_4.pt") torch.save(side, root / "local/Session_4.pt") np.save(root / "targets/Session_4.npy", targets) data = _load_examples(root / "features", root / "targets", ["Session_4"], "open-stream-replay", "easycom", {"Session_4": {1, 4}}, root / "local", "audio") self.assertTrue(torch.equal(data[-1][:, :5], side["audio"][[1, 4]])) self.assertTrue(torch.equal(data[0], reference["states"][[1, 4]])) self.assertTrue(torch.equal(data[1][:, 0], torch.from_numpy(targets[[2, 5]]))) def test_source_inventory_rejects_different_clip_order(self): source = {"path": "/old/00-00-000.mp4", "source_frames": 1200, "dense_start": 0., "duration": 60.} media = SimpleNamespace(clips=[object()], metadata=lambda: {"source_clips": [source]}) validate_source_timeline({"direct_media": {"source_clips": [source]}}, media) bad = {**source, "path": "/old/01-00-000.mp4"} with self.assertRaisesRegex(ValueError, "timeline differs"): validate_source_timeline({"direct_media": {"source_clips": [bad]}}, media) def test_config_digest_ignores_only_checkpoint_location(self): a = {"_name_or_path": "/old/model", "audio_config": {"_name_or_path": "/old/model", "output_dim": 8}} b = {"_name_or_path": "/new/model", "audio_config": {"_name_or_path": "/new/model", "output_dim": 8}} self.assertEqual(model_config_digest(a), model_config_digest(b)) b["audio_config"]["output_dim"] = 16 self.assertNotEqual(model_config_digest(a), model_config_digest(b)) def test_encoder_image_pool_splits_people_before_mean(self): encoder = FrozenQwenLocalEncoder.__new__(FrozenQwenLocalEncoder) encoder.visual_dim, encoder.device, encoder.image_batch_size, encoder.max_pixels = 2, torch.device("cpu"), 8, 50176 encoder.processor = SimpleNamespace(image_processor=lambda **kwargs: { "pixel_values": torch.ones(12, 2), "image_grid_thw": torch.tensor([[1, 2, 2], [1, 2, 4]])}) encoder.model = SimpleNamespace(visual=SimpleNamespace(dtype=torch.float32), config=SimpleNamespace(vision_config=SimpleNamespace(spatial_merge_size=2)), get_image_features=lambda *_: torch.tensor([[1., 1.], [3., 3.], [5., 5.]])) actual = encoder.images([object(), object()]) self.assertTrue(torch.equal(actual, torch.tensor([[1., 1.], [4., 4.]]))) def test_extraction_never_encodes_future_audio_and_saves_ragged_people(self): class Media: clips = [SimpleNamespace(crop=None, dense_start=0.)] def _source_frame_for_time(self, t): return 0, round(t * 20), t def frame_at_or_before(self, t): return Image.new("RGB", (16, 16)) def audio_slice(self, start, end): return np.arange(round(start * 16000), round(end * 16000), dtype=np.float32) class Encoder: visual_dim = audio_dim = 2 def __init__(self): self.ends = [] def images(self, crops): return torch.ones(len(crops), 2) def audio(self, waveform): self.ends.append(float(waveform[-1])) return torch.tensor([float(waveform[0]), float(waveform[-1])]) encoder = Encoder() args = SimpleNamespace(video_fps=2., visual_window_seconds=1., audio_window_seconds=1., box_coordinate_space="camera", crop_expansion=1., no_progress=True) annotations = [{0: [{"id": "1", "xyxy": [0., 0., 8., 8.]}]}] result = extract_local(Media(), annotations, encoder, torch.arange(1, 11, dtype=torch.float64) / 5, args) self.assertEqual(encoder.ends, [3200 * (i + 1) - 1 for i in range(10)]) self.assertEqual(result["persons"]["frame_offsets"].tolist(), [0, 1, 1, 1, 1]) self.assertTrue(result["visual_valid"][:5].all()) self.assertFalse(result["visual_valid"][5:].any()) self.assertTrue(result["audio_valid"].all()) @unittest.skipUnless(importlib.util.find_spec("transformers"), "optional: install transformers==4.52.3") class TinyQwenEncoderTests(unittest.TestCase): def test_real_encoder_apis_without_downloading_weights(self): import transformers if transformers.__version__ != "4.52.3": self.skipTest("Pinned encoder API test requires transformers==4.52.3") from transformers import ( Qwen2_5OmniThinkerConfig, Qwen2_5OmniThinkerForConditionalGeneration, Qwen2VLImageProcessor, WhisperFeatureExtractor, ) torch.set_num_threads(1) config = Qwen2_5OmniThinkerConfig( audio_config=dict(d_model=16, encoder_layers=1, encoder_attention_heads=2, encoder_ffn_dim=32, output_dim=8), vision_config=dict(depth=1, hidden_size=16, intermediate_size=32, num_heads=2, out_hidden_size=8, fullatt_block_indexes=[0]), text_config=dict(vocab_size=32, hidden_size=8, intermediate_size=16, num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=2), ) encoder = FrozenQwenLocalEncoder.__new__(FrozenQwenLocalEncoder) encoder.model = Qwen2_5OmniThinkerForConditionalGeneration(config).eval().requires_grad_(False) encoder.model.model = None encoder.model.lm_head = None encoder.visual_dim = encoder.audio_dim = 8 encoder.device, encoder.max_pixels, encoder.image_batch_size = torch.device("cpu"), 224 * 224, 8 encoder.processor = SimpleNamespace(image_processor=Qwen2VLImageProcessor(), feature_extractor=WhisperFeatureExtractor(feature_size=128)) crops = [Image.new("RGB", (56, 84), "red"), Image.new("RGB", (112, 56), "blue")] batch = encoder.images(crops) single = encoder.images(crops[:1]) self.assertEqual(batch.shape, (2, 8)) self.assertTrue(torch.allclose(batch[0], single[0], atol=1e-5)) for n in (3200, 6400, 16000): audio = encoder.audio(np.sin(np.arange(n, dtype=np.float32) / 10)) self.assertEqual(audio.shape, (8,)) self.assertTrue(torch.isfinite(audio).all()) self.assertFalse(audio.requires_grad) class RoundTripTests(unittest.TestCase): def test_train_infer_easycom_and_ego4d(self): repo = Path(__file__).resolve().parent env = dict(os.environ, OMP_NUM_THREADS="1", MKL_NUM_THREADS="1", OPENBLAS_NUM_THREADS="1") with tempfile.TemporaryDirectory() as directory: root = Path(directory) for folder in ("features", "local", "targets"): (root / folder).mkdir() train_uid = "00000000-0000-0000-0000-000000000004" test_uid = "00000000-0000-0000-0000-000000000001" manifest = root / "ego4d.tsv" manifest.write_text(f"split\tuid\ntrain\t{train_uid}\ntest\t{test_uid}\n") for name in ("Session_1", "Session_4", "Session_5", train_uid, test_uid): dataset = "easycom" if name.startswith("Session") else "ego4d" reference, side, targets = fixture(name=name, dataset=dataset) torch.save(reference, root / "features" / f"{name}.pt") torch.save(side, root / "local" / f"{name}.pt") np.save(root / "targets" / f"{name}.npy", targets) common = ["--feature-root", str(root / "features"), "--target-root", str(root / "targets"), "--device", "cpu", "--no-progress", "--expected-mode", "open-stream-replay"] for mode in ("baseline", "audio", "visual", "av", "ego4d"): checkpoint = root / f"{mode}.pth" train = [sys.executable, str(repo / "train_turn_head.py"), *common, "--epochs", "1", "--warmup-epochs", "0", "--batch-size", "4", "--layer-fusion", "concat-mlp", "--output", str(checkpoint)] infer = [sys.executable, str(repo / "infer_turn_head.py"), *common, "--checkpoint", str(checkpoint), "--output-dir", str(root / f"result_{mode}")] if mode == "ego4d": train += ["--dataset", "ego4d", "--manifest", str(manifest)] infer += ["--dataset", "ego4d", "--manifest", str(manifest)] else: train += ["--sessions", "4", "5"] infer += ["--sessions", "1"] if mode == "visual": train += ["--dev-sessions", "5", "--local-projection-dim", "0"] if mode == "av": train += ["--transition-objective", "r1-supervised"] if mode != "baseline": train += ["--local-feature-root", str(root / "local"), "--local-modalities", "av" if mode == "ego4d" else mode] infer += ["--local-feature-root", str(root / "local")] for command in (train, infer): completed = subprocess.run(command, cwd=repo, env=env, capture_output=True, text=True, timeout=60) self.assertEqual(completed.returncode, 0, f"{mode}: {completed.stdout}\n{completed.stderr}") metrics = json.loads((root / f"result_{mode}/metrics.json").read_text()) self.assertEqual(metrics["checkpoint_epoch"], 1) self.assertEqual(bool(metrics["local_contract"]), mode != "baseline") if mode == "av": self.assertTrue(metrics["transition_auxiliary"]) # The reporting command uses the same metric keys as inference. completed = subprocess.run( [sys.executable, str(repo / "summarize_local_baselines.py"), "--root", str(root)], cwd=repo, capture_output=True, text=True, timeout=20, ) self.assertEqual(completed.returncode, 0, completed.stderr) self.assertIn("Mean ± sample standard deviation", completed.stdout) if __name__ == "__main__": unittest.main()