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