| |
| """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 |
|
|
|
|
| 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() |
|
|
