#!/usr/bin/env python3 """Build SegmentQA fixed-frame v2 while excluding annotated blue-screen time.""" from __future__ import annotations import argparse import csv import filecmp import json import os import subprocess import threading import time from collections import Counter from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime, timezone from fractions import Fraction from pathlib import Path from typing import Any import av import preprocess_segmentqa_fixed_frames as base V1_RELATIVE_DIR = Path("derived/orena_segmentqa_fixed80_420x224_absmeta_v1/clean") BLUE_RELATIVE_PATH = Path("heico-focus-vqa/blue_screen_temporal_annotations/data/intervals.csv") EXPECTED_FULL = { "segments": 18_478, "affected": 1_215, "unaffected": 17_263, "test_questions": 6_254, "test_segments": 5_661, } def merge_intervals(intervals: list[tuple[float, float]]) -> list[tuple[float, float]]: merged: list[list[float]] = [] for start, end in sorted(intervals): if end <= start: continue if merged and start <= merged[-1][1]: merged[-1][1] = max(merged[-1][1], end) else: merged.append([start, end]) return [(start, end) for start, end in merged] def load_blue_intervals(path: Path) -> dict[str, list[tuple[float, float]]]: by_video: dict[str, list[tuple[float, float]]] = {} with path.open(newline="", encoding="utf-8") as handle: for row in csv.DictReader(handle): by_video.setdefault(row["video_filename"], []).append( (float(row["start_seconds"]), float(row["end_seconds"])) ) return {video: merge_intervals(intervals) for video, intervals in by_video.items()} def segment_blue_intersections( segment: base.Segment, blue_by_video: dict[str, list[tuple[float, float]]] ) -> list[tuple[float, float]]: if segment.dataset != "heico-focus-vqa": return [] return [ (max(segment.start, start), min(segment.end, end)) for start, end in blue_by_video.get(segment.video, []) if segment.start < end and segment.end > start ] def subtract_intervals( start: float, end: float, blocked: list[tuple[float, float]] ) -> list[tuple[float, float]]: valid: list[tuple[float, float]] = [] cursor = start for blocked_start, blocked_end in merge_intervals(blocked): if cursor < blocked_start: valid.append((cursor, blocked_start)) cursor = max(cursor, blocked_end) if cursor < end: valid.append((cursor, end)) return valid def targets_across_intervals(intervals: list[tuple[float, float]], frames: int) -> list[float]: total = sum(end - start for start, end in intervals) if total <= 0: return [] targets: list[float] = [] interval_index = 0 consumed = 0.0 for index in range(frames): offset = (index + 0.5) * total / frames while interval_index + 1 < len(intervals) and offset >= consumed + ( intervals[interval_index][1] - intervals[interval_index][0] ): consumed += intervals[interval_index][1] - intervals[interval_index][0] interval_index += 1 start, _ = intervals[interval_index] targets.append(start + offset - consumed) return targets def point_in_intervals(value: float, intervals: list[tuple[float, float]]) -> bool: return any(start <= value < end for start, end in intervals) def select_non_blue_frames( segment: base.Segment, blocked: list[tuple[float, float]], frames: int, width: int, height: int, ) -> tuple[list[av.VideoFrame], list[float], int]: valid_intervals = subtract_intervals(segment.start, segment.end, blocked) targets = targets_across_intervals(valid_intervals, frames) if not targets: return [], [], 0 selected: list[av.VideoFrame | None] = [None] * frames selected_pts: list[float | None] = [None] * frames unique_pts: set[int] = set() last: av.VideoFrame | None = None last_pts: float | None = None next_index = 0 blocked_index = 0 with av.open(segment.source_path) as container: stream = container.streams.video[0] seek_offset = int(max(segment.start - 2.0, 0.0) / float(stream.time_base)) container.seek(seek_offset, stream=stream, backward=True, any_frame=False) stop_decode = False for packet in container.demux(stream): for frame in packet.decode(): pts = base.frame_seconds(frame, stream) if pts is None or pts < segment.start: continue if pts >= segment.end: stop_decode = True break while blocked_index < len(blocked) and blocked[blocked_index][1] <= pts: blocked_index += 1 if blocked_index < len(blocked) and blocked[blocked_index][0] <= pts < blocked[blocked_index][1]: continue unique_pts.add(int(frame.pts)) while next_index < frames and targets[next_index] <= pts: if last is not None and abs(last_pts - targets[next_index]) <= abs(pts - targets[next_index]): chosen, chosen_pts = last, last_pts else: chosen, chosen_pts = frame, pts selected[next_index] = base.letterbox(chosen, width, height) selected_pts[next_index] = chosen_pts next_index += 1 last, last_pts = frame, pts if stop_decode: break if last is None or last_pts is None: return [], [], 0 for index, frame in enumerate(selected): if frame is None: selected[index] = base.letterbox(last, width, height) selected_pts[index] = last_pts return ( [frame for frame in selected if frame is not None], [float(pts) for pts in selected_pts if pts is not None], len(unique_pts), ) def encode_frames( path: Path, selected: list[av.VideoFrame], frames: int, width: int, height: int ) -> tuple[int, int, int]: partial = path.with_suffix(path.suffix + ".partial") path.parent.mkdir(parents=True, exist_ok=True) partial.unlink(missing_ok=True) try: with av.open(str(partial), "w", format="mp4") as container: stream = container.add_stream("libx264", rate=Fraction(20, 1)) stream.width = width stream.height = height stream.pix_fmt = "yuv420p" stream.options = {"crf": "18", "preset": "veryfast", "an": "1"} for index, frame in enumerate(selected): frame.pts = index frame.time_base = Fraction(1, 20) for packet in stream.encode(frame): container.mux(packet) for packet in stream.encode(): container.mux(packet) stats = base.count_video_frames(partial) if stats != (frames, width, height): raise RuntimeError(f"encoded {stats}, expected {(frames, width, height)}") os.replace(partial, path) return stats except Exception: partial.unlink(missing_ok=True) raise def copy_clip(source: Path, destination: Path, frames: int, width: int, height: int) -> None: partial = destination.with_suffix(destination.suffix + ".partial") destination.parent.mkdir(parents=True, exist_ok=True) partial.unlink(missing_ok=True) try: subprocess.run( ["cp", "--reflink=auto", "--preserve=mode,timestamps", "--", str(source), str(partial)], check=True, ) if not filecmp.cmp(source, partial, shallow=False): raise RuntimeError(f"copied file differs from v1: {source}") if not video_contract(partial, frames, width, height): raise RuntimeError(f"copied file violates the fixed media contract: {source}") os.replace(partial, destination) except Exception: partial.unlink(missing_ok=True) raise def blue_fields( segment: base.Segment, blocked: list[tuple[float, float]], valid_unique_source_frames: int | None, action: str, ) -> dict[str, Any]: overlap = sum(end - start for start, end in blocked) duration = segment.end - segment.start return { "blue_screen_intervals": [ {"start_seconds": start, "end_seconds": end} for start, end in blocked ], "blue_overlap_seconds": round(overlap, 6), "blue_overlap_ratio": round(overlap / duration, 9), "valid_duration_seconds": round(duration - overlap, 6), "valid_unique_source_frames": valid_unique_source_frames, "remediation_action": action, } def copy_v1_item( segment: base.Segment, v1_item: dict[str, Any], v1_dir: Path, output_dir: Path, blocked: list[tuple[float, float]], action: str, valid_unique_source_frames: int | None, frames: int, width: int, height: int, ) -> dict[str, Any]: source = v1_dir / segment.relpath destination = output_dir / segment.relpath if not base.metadata_contract_matches(v1_item, frames, width, height): raise RuntimeError(f"v1 clip violates the fixed media contract: {source}") copy_clip(source, destination, frames, width, height) item = dict(v1_item) item.update( blue_fields(segment, blocked, valid_unique_source_frames, action) | {"source_v1_path": str(source), "status": "success"} ) return item def process_affected( segment: base.Segment, blocked: list[tuple[float, float]], v1_item: dict[str, Any], v1_dir: Path, output_dir: Path, frames: int, width: int, height: int, min_train_valid_frames: int, ) -> dict[str, Any]: selected, source_pts, unique_count = select_non_blue_frames( segment, blocked, frames, width, height ) if segment.split == "train" and unique_count < min_train_valid_frames: (output_dir / segment.relpath).unlink(missing_ok=True) return { "dataset": segment.dataset, "split": segment.split, "video": segment.video, "source_video_path": segment.source_path, "start_seconds": segment.start, "end_seconds": segment.end, "relpath": segment.relpath, "status": "excluded", **blue_fields(segment, blocked, unique_count, "excluded_train_insufficient_valid"), } if unique_count == 0: return copy_v1_item( segment, v1_item, v1_dir, output_dir, blocked, "irrecoverable_all_blue", 0, frames, width, height, ) if len(selected) != frames or len(source_pts) != frames: raise RuntimeError(f"selected {len(selected)} frames and {len(source_pts)} pts, expected {frames}") if any(point_in_intervals(pts, blocked) for pts in source_pts): raise RuntimeError("selected source PTS intersects a blue-screen interval") verified_count, verified_width, verified_height = encode_frames( output_dir / segment.relpath, selected, frames, width, height ) action = ( "retained_test_insufficient_valid" if segment.split == "test" and unique_count < min_train_valid_frames else "resampled_excluding_blue" ) return { "dataset": segment.dataset, "split": segment.split, "video": segment.video, "source_video_path": segment.source_path, "start_seconds": segment.start, "end_seconds": segment.end, "relpath": segment.relpath, "frame_count": frames, "verified_frame_count": verified_count, "width": verified_width, "height": verified_height, "selected_source_pts": source_pts, "video_metadata": base.metadata_from_pts(source_pts, width, height), "sampling_rule": f"uniform_bin_centers_over_non_blue_time_in_[start,end)_{frames}", "resize": f"letterbox_{width}x{height}", "status": "success", **blue_fields(segment, blocked, unique_count, action), } def video_contract(path: Path, frames: int, width: int, height: int) -> bool: try: with av.open(str(path)) as container: videos = container.streams.video if len(videos) != 1 or container.streams.audio: return False stream = videos[0] if stream.codec_context.name != "h264" or (stream.width, stream.height) != (width, height): return False if stream.frames: return stream.frames == frames return base.count_video_frames(path) == (frames, width, height) except Exception: return False def verified_record( item: dict[str, Any], segment: base.Segment, blocked: list[tuple[float, float]], v1_item: dict[str, Any], output_dir: Path, v1_dir: Path, frames: int, width: int, height: int, min_train_valid_frames: int, ) -> bool: action = item.get("remediation_action") unique_count = item.get("valid_unique_source_frames") identity_matches = all( item.get(key) == value for key, value in { "dataset": segment.dataset, "split": segment.split, "video": segment.video, "source_video_path": segment.source_path, "start_seconds": segment.start, "end_seconds": segment.end, "relpath": segment.relpath, }.items() ) if not identity_matches or any( item.get(key) != value for key, value in blue_fields(segment, blocked, unique_count, action).items() ): return False if item.get("status") == "excluded": return ( action == "excluded_train_insufficient_valid" and segment.split == "train" and blocked and isinstance(unique_count, int) and unique_count < min_train_valid_frames and not (output_dir / segment.relpath).exists() ) if action in {"copied", "irrecoverable_all_blue"}: source, destination = v1_dir / segment.relpath, output_dir / segment.relpath action_matches = (action == "copied" and not blocked and unique_count is None) or ( action == "irrecoverable_all_blue" and bool(blocked) and segment.split == "test" and unique_count == 0 ) v1_matches = all( key == "elapsed_seconds" or item.get(key) == value for key, value in v1_item.items() ) and item.get("source_v1_path") == str(source) return ( item.get("status") == "success" and action_matches and v1_matches and source.exists() and destination.exists() and filecmp.cmp(source, destination, shallow=False) and base.metadata_contract_matches(item, frames, width, height) and video_contract(destination, frames, width, height) ) if action not in {"resampled_excluding_blue", "retained_test_insufficient_valid"}: return False if not blocked or not isinstance(unique_count, int) or unique_count < 1: return False if action == "retained_test_insufficient_valid": if segment.split != "test" or unique_count >= min_train_valid_frames: return False elif unique_count < min_train_valid_frames: return False source_pts = item.get("selected_source_pts") return ( isinstance(source_pts, list) and len(source_pts) == frames and source_pts == sorted(source_pts) and all(segment.start <= pts < segment.end for pts in source_pts) and not any(point_in_intervals(pts, blocked) for pts in source_pts) and base.verified_success_item(item, output_dir, frames, width, height) and video_contract(output_dir / segment.relpath, frames, width, height) ) def resume_records( status_path: Path, segments: list[base.Segment], blocked_by_rel: dict[str, list[tuple[float, float]]], v1_manifest: dict[str, dict[str, Any]], output_dir: Path, v1_dir: Path, frames: int, width: int, height: int, min_train_valid_frames: int, ) -> dict[str, dict[str, Any]]: latest: dict[str, dict[str, Any]] = {} if status_path.exists(): raw_status = status_path.read_text() lines = [line for line in raw_status.splitlines() if line.strip()] valid_lines: list[str] = [] repair_tail = bool(raw_status) and not raw_status.endswith("\n") for index, line in enumerate(lines): try: item = json.loads(line) except json.JSONDecodeError: if index == len(lines) - 1: repair_tail = True continue raise valid_lines.append(line) latest[item["relpath"]] = item if repair_tail: partial = status_path.with_name(status_path.name + ".repair") partial.write_text("\n".join(valid_lines) + ("\n" if valid_lines else "")) os.replace(partial, status_path) segment_by_rel = {segment.relpath: segment for segment in segments} return { relpath: item for relpath, item in latest.items() if relpath in segment_by_rel and verified_record( item, segment_by_rel[relpath], blocked_by_rel[relpath], v1_manifest[relpath], output_dir, v1_dir, frames, width, height, min_train_valid_frames, ) } def load_v1_manifest(path: Path) -> dict[str, dict[str, Any]]: return { item["relpath"]: item for item in (json.loads(line) for line in path.read_text().splitlines() if line.strip()) } def publish_outputs( output_dir: Path, rows: list[dict[str, Any]], manifest: list[dict[str, Any]], splits: list[str], ) -> None: manifest_by_rel = {item["relpath"]: item for item in manifest} selected = [row for row in rows if row["derived_video"] in manifest_by_rel] payloads: dict[Path, list[dict[str, Any]]] = { output_dir / "segment_manifest.jsonl": manifest, } if "train" in splits: payloads[output_dir / "train.jsonl"] = [ base.ms_swift_record(row, manifest_by_rel[row["derived_video"]], output_dir, True) for row in selected if row["split"] == "train" ] if "test" in splits: test_rows = [row for row in selected if row["split"] == "test"] payloads[output_dir / "test_requests.jsonl"] = [ base.ms_swift_record(row, manifest_by_rel[row["derived_video"]], output_dir, False) for row in test_rows ] payloads[output_dir / "test_references.jsonl"] = [ { "id": row["id"], "dataset": row["dataset"], "answer": row["answer"], "answer_format": row["answer_format"], } for row in test_rows ] staged: list[tuple[Path, Path]] = [] try: for path, content in payloads.items(): partial = path.with_name(path.name + ".partial") base.write_jsonl(partial, content) staged.append((partial, path)) for partial, path in staged: os.replace(partial, path) except Exception: for partial, _ in staged: partial.unlink(missing_ok=True) raise def write_json_atomic(path: Path, value: dict[str, Any]) -> None: partial = path.with_name(path.name + ".partial") partial.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n") os.replace(partial, path) def write_reports( output_dir: Path, args: argparse.Namespace, records: dict[str, dict[str, Any]], segments: list[base.Segment], selected_rows: list[dict[str, Any]], affected_count: int, blue_path: Path, v1_dir: Path, ) -> bool: now = datetime.now(timezone.utc).isoformat() actions = Counter(item.get("remediation_action", "failed") for item in records.values()) manifest_relpaths = { relpath for relpath, item in records.items() if item.get("status") == "success" } excluded_relpaths = { relpath for relpath, item in records.items() if item.get("status") == "excluded" } failed_relpaths = { relpath for relpath, item in records.items() if item.get("status") == "failed" } expected_relpaths = {segment.relpath for segment in segments} missing_relpaths = expected_relpaths - manifest_relpaths - excluded_relpaths - failed_relpaths test_rows = [row for row in selected_rows if row["split"] == "test"] output_test_rows = [row for row in test_rows if row["derived_video"] in manifest_relpaths] output_train_rows = [ row for row in selected_rows if row["split"] == "train" and row["derived_video"] in manifest_relpaths ] excluded_train_questions = sum( row["split"] == "train" and row["derived_video"] in excluded_relpaths for row in selected_rows ) full_contract = True if args.limit is None and set(args.splits) == {"train", "test"}: full_contract = ( len(segments) == EXPECTED_FULL["segments"] and affected_count == EXPECTED_FULL["affected"] and len(segments) - affected_count == EXPECTED_FULL["unaffected"] and len(output_test_rows) == EXPECTED_FULL["test_questions"] and len({row["derived_video"] for row in output_test_rows}) == EXPECTED_FULL["test_segments"] ) passed = ( not failed_relpaths and not missing_relpaths and len(output_test_rows) == len(test_rows) and full_contract ) provenance = { "version": "fixed_frames_blueclean_v2", "created_at_utc": now, "command": " ".join(os.sys.argv), "data_root": str(Path(args.data_root).resolve()), "v1_dir": str(v1_dir.resolve()), "blue_intervals": str(blue_path.resolve()), "output_dir": str(output_dir.resolve()), "splits": args.splits, "frame_count": args.frames, "resize": f"proportional letterbox {args.width}x{args.height}", "encoding": {"container": "mp4", "codec": "H.264/libx264", "audio": "none"}, "sampling_rule": "80 uniform bin centers over the non-blue complement inside the original [start,end)", "min_train_valid_frames": args.min_train_valid_frames, "unaffected_copy": "cp --reflink=auto; byte-identical verification; no hardlinks", "test_answers_used_for_decisions": False, "resume": bool(args.resume), "workers": args.workers, "limit": args.limit, } report = { "status": "ready_to_publish" if passed else "fail", "expected_frame_count": args.frames, "expected_width": args.width, "expected_height": args.height, "segments_total": len(segments), "affected_temporal_overlap": affected_count, "unaffected": len(segments) - affected_count, "included_segments": len(manifest_relpaths), "excluded_train_segments": len(excluded_relpaths), "excluded_train_questions": excluded_train_questions, "train_questions": len(output_train_rows), "test_questions": len(output_test_rows), "test_unique_segments": len({row["derived_video"] for row in output_test_rows}), "actions": dict(sorted(actions.items())), "missing": len(missing_relpaths), "failures": len(failed_relpaths), "full_dataset_contract": full_contract, "checked_at_utc": now, } write_json_atomic(output_dir / "provenance.json", provenance) write_json_atomic(output_dir / "validation_report.json", report) return passed def run(args: argparse.Namespace) -> int: data_root = Path(args.data_root) output_dir = Path(args.output_dir) v1_dir = Path(args.v1_dir) if args.v1_dir else data_root / V1_RELATIVE_DIR blue_path = Path(args.blue_intervals) if args.blue_intervals else data_root / BLUE_RELATIVE_PATH if output_dir.resolve() == v1_dir.resolve(): raise ValueError("v2 output directory must differ from v1") output_dir.mkdir(parents=True, exist_ok=True) datasets = base.build_datasets(data_root, args.splits) rows, all_segments = base.load_rows(datasets, args.splits) segments = base.apply_limit(all_segments, args.limit) selected_relpaths = {segment.relpath for segment in segments} selected_rows = [row for row in rows if row["derived_video"] in selected_relpaths] v1_manifest = load_v1_manifest(v1_dir / "segment_manifest.jsonl") missing_v1 = selected_relpaths - v1_manifest.keys() if missing_v1: raise RuntimeError(f"v1 manifest is missing {len(missing_v1)} selected segments") blue_by_video = load_blue_intervals(blue_path) blocked_by_rel = { segment.relpath: segment_blue_intersections(segment, blue_by_video) for segment in segments } affected_count = sum(bool(intervals) for intervals in blocked_by_rel.values()) status_path = output_dir / base.STATUS_NAME records = ( resume_records( status_path, segments, blocked_by_rel, v1_manifest, output_dir, v1_dir, args.frames, args.width, args.height, args.min_train_valid_frames, ) if args.resume else {} ) todo = [segment for segment in segments if segment.relpath not in records] print( json.dumps( { "event": "start", "output": str(output_dir), "segments": len(segments), "affected": affected_count, "unaffected": len(segments) - affected_count, "todo": len(todo), "python": os.sys.executable, }, ensure_ascii=False, ), flush=True, ) def work(segment: base.Segment) -> dict[str, Any]: started = time.time() blocked = blocked_by_rel[segment.relpath] try: if blocked: item = process_affected( segment, blocked, v1_manifest[segment.relpath], v1_dir, output_dir, args.frames, args.width, args.height, args.min_train_valid_frames, ) else: item = copy_v1_item( segment, v1_manifest[segment.relpath], v1_dir, output_dir, [], "copied", None, args.frames, args.width, args.height, ) except Exception as exc: item = { "dataset": segment.dataset, "split": segment.split, "video": segment.video, "source_video_path": segment.source_path, "start_seconds": segment.start, "end_seconds": segment.end, "relpath": segment.relpath, "status": "failed", "error": str(exc), } item["elapsed_seconds"] = round(time.time() - started, 3) return item unaffected_todo = [segment for segment in todo if not blocked_by_rel[segment.relpath]] affected_by_video: dict[str, list[base.Segment]] = {} for segment in todo: if blocked_by_rel[segment.relpath]: affected_by_video.setdefault(segment.source_path, []).append(segment) affected_groups = [ group for lane in range(2) for video in sorted(affected_by_video) if (group := affected_by_video[video][lane::2]) ] with status_path.open("a" if args.resume else "w", encoding="utf-8") as status_handle: status_lock = threading.Lock() def record(item: dict[str, Any]) -> None: with status_lock: status_handle.write(json.dumps(item, ensure_ascii=False) + "\n") status_handle.flush() records[item["relpath"]] = item print( f"relpath={item['relpath']} status={item['status']} " f"action={item.get('remediation_action')} elapsed_seconds={item['elapsed_seconds']}", flush=True, ) if unaffected_todo: print(json.dumps({"event": "phase", "name": "copy", "todo": len(unaffected_todo)}), flush=True) with ThreadPoolExecutor(max_workers=args.workers) as pool: for future in as_completed([pool.submit(work, segment) for segment in unaffected_todo]): record(future.result()) if affected_groups: print( json.dumps( { "event": "phase", "name": "resample", "segments": sum(len(group) for group in affected_groups), "source_videos": len(affected_by_video), "max_parallel_per_video": 2, } ), flush=True, ) def work_group(group: list[base.Segment]) -> None: for segment in group: record(work(segment)) with ThreadPoolExecutor(max_workers=min(args.workers, len(affected_groups))) as pool: futures = [pool.submit(work_group, group) for group in affected_groups] for future in as_completed(futures): future.result() manifest = [ records[segment.relpath] for segment in segments if records.get(segment.relpath, {}).get("status") == "success" ] passed = write_reports( output_dir, args, records, segments, selected_rows, affected_count, blue_path, v1_dir, ) if passed: publish_outputs(output_dir, selected_rows, manifest, args.splits) report_path = output_dir / "validation_report.json" report = json.loads(report_path.read_text()) if passed: report["status"] = "pass" write_json_atomic(report_path, report) print(json.dumps(report, ensure_ascii=False), flush=True) return 0 if passed else 1 def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("--data-root", required=True) parser.add_argument("--output-dir", required=True) parser.add_argument("--v1-dir") parser.add_argument("--blue-intervals") parser.add_argument("--splits", nargs="+", choices=["train", "test"], required=True) parser.add_argument("--frames", type=int, default=80) parser.add_argument("--width", type=int, default=420) parser.add_argument("--height", type=int, default=224) parser.add_argument("--workers", type=int, default=1) parser.add_argument("--min-train-valid-frames", type=int, default=16) parser.add_argument("--limit", type=int) parser.add_argument("--resume", action="store_true") args = parser.parse_args() if min(args.frames, args.width, args.height, args.workers, args.min_train_valid_frames) < 1: raise SystemExit("frames, dimensions, workers, and min valid frames must be >= 1") if (args.frames, args.width, args.height) != (80, 420, 224): raise SystemExit("v2 contract requires --frames 80 --width 420 --height 224") return run(args) if __name__ == "__main__": raise SystemExit(main())