#!/usr/bin/env python3 """Focused unit and CPU smoke checks for SegmentQA blue-clean v2.""" from __future__ import annotations import json import subprocess import sys import tempfile from fractions import Fraction from pathlib import Path import av from PIL import Image PROJECT_ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(PROJECT_ROOT / "scripts")) import preprocess_segmentqa_blueclean_v2 as prep # noqa: E402 PYTHON = "/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/conda-envs/medrl/bin/python" DATA_ROOT = Path("/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge") def write_test_video(path: Path, frame_count: int = 20, fps: int = 10, codec: str = "libx264") -> None: path.parent.mkdir(parents=True, exist_ok=True) with av.open(str(path), "w", format="mp4") as container: stream = container.add_stream(codec, rate=Fraction(fps, 1)) stream.width = 64 stream.height = 64 stream.pix_fmt = "yuv420p" for index in range(frame_count): frame = av.VideoFrame.from_image( Image.new("RGB", (64, 64), (index * 7 % 255, index * 11 % 255, index * 13 % 255)) ) frame.pts = index frame.time_base = Fraction(1, fps) for packet in stream.encode(frame): container.mux(packet) for packet in stream.encode(): container.mux(packet) def test_interval_helpers() -> None: assert prep.merge_intervals([(6, 8), (2, 4), (4, 5), (7, 9)]) == [(2, 5), (6, 9)] assert prep.subtract_intervals(0, 10, [(2, 4), (6, 10)]) == [(0, 2), (4, 6)] assert prep.targets_across_intervals([(0, 2), (4, 6)], 4) == [0.5, 1.5, 4.5, 5.5] assert prep.point_in_intervals(2.0, [(2.0, 4.0)]) assert not prep.point_in_intervals(4.0, [(2.0, 4.0)]) def test_full_overlap_inventory() -> None: datasets = prep.base.build_datasets(DATA_ROOT, ["train", "test"]) _, segments = prep.base.load_rows(datasets, ["train", "test"]) blue = prep.load_blue_intervals(DATA_ROOT / prep.BLUE_RELATIVE_PATH) affected = sum(bool(prep.segment_blue_intersections(segment, blue)) for segment in segments) assert len(segments) == 18_478 assert affected == 1_215 assert len(segments) - affected == 17_263 def test_low_valid_frame_train_and_test_policy() -> None: with tempfile.TemporaryDirectory(prefix="segmentqa-v2-policy-") as tmp: tmp_path = Path(tmp) source = tmp_path / "source.mp4" write_test_video(source) blocked = [(0.4, 2.0)] train = prep.base.Segment( "heico-focus-vqa", "train", source.name, 0.0, 2.0, "videos/train/heico-focus-vqa/train.mp4", str(source), ) train_item = prep.process_affected( train, blocked, {}, tmp_path / "v1", tmp_path / "out", 80, 64, 64, 16 ) assert train_item["status"] == "excluded" assert train_item["remediation_action"] == "excluded_train_insufficient_valid" assert 0 < train_item["valid_unique_source_frames"] < 16 test = prep.base.Segment( "heico-focus-vqa", "test", source.name, 0.0, 2.0, "videos/test/heico-focus-vqa/test.mp4", str(source), ) test_item = prep.process_affected( test, blocked, {}, tmp_path / "v1", tmp_path / "out", 80, 64, 64, 16 ) assert test_item["status"] == "success" assert test_item["remediation_action"] == "retained_test_insufficient_valid" assert len(test_item["selected_source_pts"]) == 80 assert len(set(test_item["selected_source_pts"])) < 16 assert all(0.0 <= pts < 0.4 for pts in test_item["selected_source_pts"]) assert prep.base.count_video_frames(tmp_path / "out" / test.relpath) == (80, 64, 64) assert prep.verified_record( test_item, test, blocked, {}, tmp_path / "out", tmp_path / "v1", 80, 64, 64, 16 ) assert not prep.verified_record( test_item, test, [(0.0, 0.5)], {}, tmp_path / "out", tmp_path / "v1", 80, 64, 64, 16 ) mpeg4 = tmp_path / "mpeg4.mp4" write_test_video(mpeg4, frame_count=80, codec="mpeg4") assert not prep.video_contract(mpeg4, 80, 64, 64) existing = tmp_path / "out" / test.relpath before = existing.read_bytes() try: prep.copy_clip(mpeg4, existing, 80, 64, 64) except RuntimeError: pass else: raise AssertionError("invalid copied media should be rejected") assert existing.read_bytes() == before def test_copy_smoke_and_resume() -> None: with tempfile.TemporaryDirectory(prefix="segmentqa-v2-copy-") as tmp: output = Path(tmp) / "clean" cmd = [ PYTHON, "scripts/preprocess_segmentqa_blueclean_v2.py", "--data-root", str(DATA_ROOT), "--output-dir", str(output), "--splits", "train", "test", "--workers", "2", "--limit", "2", ] first = subprocess.run(cmd, cwd=PROJECT_ROOT, check=True, text=True, capture_output=True) first_report = json.loads(first.stdout.strip().splitlines()[-1]) assert first_report["status"] == "pass" assert first_report["actions"] == {"copied": 2} manifest = [json.loads(line) for line in (output / "segment_manifest.jsonl").read_text().splitlines()] assert len(manifest) == 2 assert all(item["remediation_action"] == "copied" for item in manifest) assert all(str(output.resolve()) in row["videos"][0] for row in ( json.loads(line) for path in (output / "train.jsonl", output / "test_requests.jsonl") if path.exists() for line in path.read_text().splitlines() )) mtimes = {path: path.stat().st_mtime_ns for path in output.rglob("*.mp4")} resumed = subprocess.run(cmd + ["--resume"], cwd=PROJECT_ROOT, check=True, text=True, capture_output=True) resumed_report = json.loads(resumed.stdout.strip().splitlines()[-1]) assert resumed_report["status"] == "pass" assert mtimes == {path: path.stat().st_mtime_ns for path in output.rglob("*.mp4")} with (output / "segment_status.jsonl").open("a") as handle: handle.write('{"torn":') repaired = subprocess.run(cmd + ["--resume"], cwd=PROJECT_ROOT, check=True, text=True, capture_output=True) assert json.loads(repaired.stdout.strip().splitlines()[-1])["status"] == "pass" assert not (output / "segment_status.jsonl").read_text().endswith('{"torn":') status_path = output / "segment_status.jsonl" status_path.write_text(status_path.read_text().rstrip("\n")) normalized = subprocess.run(cmd + ["--resume"], cwd=PROJECT_ROOT, check=True, text=True, capture_output=True) assert json.loads(normalized.stdout.strip().splitlines()[-1])["status"] == "pass" assert status_path.read_text().endswith("\n") fresh = subprocess.run(cmd, cwd=PROJECT_ROOT, check=True, text=True, capture_output=True) assert json.loads(fresh.stdout.strip().splitlines()[-1])["status"] == "pass" assert len((output / "segment_status.jsonl").read_text().splitlines()) == 2 invalid = subprocess.run(cmd + ["--frames", "79"], cwd=PROJECT_ROOT, text=True, capture_output=True) assert invalid.returncode != 0 assert "requires --frames 80" in invalid.stderr if __name__ == "__main__": test_interval_helpers() test_full_overlap_inventory() test_low_valid_frame_train_and_test_policy() test_copy_smoke_and_resume()