| |
| """Build fixed-frame SegmentQA clips and ms-swift JSONL manifests.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import re |
| import time |
| from concurrent.futures import ThreadPoolExecutor, as_completed |
| from dataclasses import dataclass |
| from datetime import datetime, timezone |
| from fractions import Fraction |
| from pathlib import Path |
| from typing import Any |
|
|
| import av |
| import pandas as pd |
| from PIL import Image, ImageOps |
|
|
|
|
| DATASET_NAMES = ("heico-focus-vqa", "lapchole-focus-vqa") |
| STATUS_NAME = "segment_status.jsonl" |
| BASELINE_SYSTEM_PROMPT = ( |
| "You are a surgical video analysis assistant. You are shown a short clip from a laparoscopic procedure and asked a single question about it.\n\n" |
| "The clip is given as a sequence of frames in chronological order. Each frame is preceded by its timestamp in the source video, written as " |
| "<SECONDS seconds> -- a frame marked <1234.5 seconds> was taken 1234.5 seconds (00:20:34) after the video began. Consecutive frames may be several seconds " |
| "apart, so an event can happen between two of them.\n\n" |
| "A foreign object (FO) is any object fully introduced into the patient's body\n" |
| "cavity during surgery that must be retrieved or accounted for. Importantly,\n" |
| "standard surgical instruments that remain connected to the external environment\n" |
| "(e.g., graspers, scissors, trocars, staplers, cameras) are not considered foreign\n" |
| "objects. Furthermore, we exclude detachable parts of surgical instruments,\n" |
| "particularly anvil components of staplers.\n\n" |
| "The foreign object classes are exactly: Sponge, Clip, Specimen Bag, Silicone Loop, External Drain, Needle, Gallstone, Specimen, Mesh, Absorbable Hemostatic Agent.\n\n" |
| "Reply with the answer and nothing else -- no reasoning, no preamble, no\n" |
| "explanation, no restating the question. A single short line.\n\n" |
| "Rules for the answer:\n" |
| "- Write the value only. No sentence, no explanation, no units, no trailing\n" |
| " period, and never repeat the question.\n" |
| "- Asks yes or no -> write exactly: yes or no\n" |
| "- Asks how many / for a count -> write digits only, e.g. 0 or 1 or 2.\n" |
| "- Asks which foreign object class(es) -> write class names exactly as spelled\n" |
| " in the list above, comma-separated (e.g. Clip, Sponge), or exactly: none\n" |
| " Never answer with a generic description such as \"surgical instrument\".\n" |
| "- Asks when something happens -> write hh:mm:ss, read off the frame timestamps.\n" |
| " Those timestamps count from the START OF THE VIDEO, not from the start of this\n" |
| " clip, so <1234.5 seconds> is 00:20:34. Convert carefully: 3600 seconds is one\n" |
| " hour, 60 seconds is one minute. If the event falls between two frames, give\n" |
| " your best estimate between their timestamps rather than snapping to one.\n" |
| " For several time points, separate them with commas: 00:09:19, 00:12:44\n" |
| "- Asks for a percentage or a share -> write the number only, e.g. 40\n" |
| "- Lists options to choose from -> copy exactly one of those options, verbatim.\n" |
| "- Anything else -> a short phrase, at most a few words.\n\n" |
| "If you are unsure, still commit to your single best answer in the required\n" |
| "form. An empty, hedged, or explanatory answer is scored as wrong.\n" |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class Segment: |
| dataset: str |
| split: str |
| video: str |
| start: float |
| end: float |
| relpath: str |
| source_path: str |
|
|
| @property |
| def key(self) -> tuple[str, str, str, float, float]: |
| return self.dataset, self.split, self.video, self.start, self.end |
|
|
|
|
| def parse_time(value: Any) -> float: |
| if isinstance(value, (int, float)): |
| return float(value) |
| parts = str(value).split(":") |
| if len(parts) != 3: |
| raise ValueError(f"expected HH:MM:SS timestamp, got {value!r}") |
| hours, minutes, seconds = parts |
| return int(hours) * 3600 + int(minutes) * 60 + float(seconds) |
|
|
|
|
| def jsonable(value: Any) -> Any: |
| if hasattr(value, "tolist"): |
| return value.tolist() |
| if not isinstance(value, (list, tuple, dict)) and pd.isna(value): |
| return None |
| return value |
|
|
|
|
| def slug(value: str) -> str: |
| return re.sub(r"[^A-Za-z0-9._-]+", "_", value).strip("_") |
|
|
|
|
| def seconds_token(value: float) -> str: |
| return f"{value:.3f}".rstrip("0").rstrip(".").replace(".", "p") |
|
|
|
|
| def segment_relpath(dataset: str, split: str, video: str, start: float, end: float) -> str: |
| stem = slug(Path(video).stem) |
| name = f"{stem}__{seconds_token(start)}_{seconds_token(end)}.mp4" |
| return str(Path("videos") / split / dataset / name) |
|
|
|
|
| def target_times(start: float, end: float, frames: int) -> list[float]: |
| duration = end - start |
| if duration <= 0: |
| raise ValueError(f"non-positive segment duration: start={start} end={end}") |
| return [start + (i + 0.5) * duration / frames for i in range(frames)] |
|
|
|
|
| def build_datasets(data_root: Path, splits: list[str]) -> dict[str, dict[str, Path]]: |
| return { |
| name: { |
| **{split: data_root / name / f"data/segment/{split}.parquet" for split in splits}, |
| "videos": data_root / name / "videos", |
| } |
| for name in DATASET_NAMES |
| } |
|
|
|
|
| def load_rows(datasets: dict[str, dict[str, Path]], splits: list[str]) -> tuple[list[dict[str, Any]], list[Segment]]: |
| rows: list[dict[str, Any]] = [] |
| by_key: dict[tuple[str, str, str, float, float], Segment] = {} |
| for split in splits: |
| for dataset, cfg in datasets.items(): |
| df = pd.read_parquet(cfg[split]) |
| for _, row in df.iterrows(): |
| raw = {k: jsonable(v) for k, v in row.to_dict().items()} |
| start = parse_time(raw["timestamp_start"]) |
| end = parse_time(raw["timestamp_end"]) |
| video = str(raw["video"]) |
| source = cfg["videos"] / video |
| rel = segment_relpath(dataset, split, video, start, end) |
| segment = Segment(dataset, split, video, start, end, rel, str(source)) |
| by_key.setdefault(segment.key, segment) |
| raw.update( |
| { |
| "dataset": dataset, |
| "split": split, |
| "source_video_path": str(source), |
| "segment_start_seconds": start, |
| "segment_end_seconds": end, |
| "derived_video": rel, |
| } |
| ) |
| rows.append(raw) |
| return rows, list(by_key.values()) |
|
|
|
|
| def apply_limit(segments: list[Segment], limit: int | None) -> list[Segment]: |
| ordered = sorted(segments, key=lambda s: s.key) |
| if limit is None: |
| return ordered |
| buckets = {name: [s for s in ordered if s.dataset == name] for name in DATASET_NAMES} |
| picked: list[Segment] = [] |
| index = 0 |
| while len(picked) < limit: |
| before = len(picked) |
| for dataset in DATASET_NAMES: |
| if index < len(buckets[dataset]) and len(picked) < limit: |
| picked.append(buckets[dataset][index]) |
| if len(picked) == before: |
| break |
| index += 1 |
| return picked |
|
|
|
|
| def frame_seconds(frame: av.VideoFrame, stream: av.video.stream.VideoStream) -> float | None: |
| if frame.pts is None: |
| return None |
| base = frame.time_base or stream.time_base |
| return float(frame.pts * base) |
|
|
|
|
| def letterbox(frame: av.VideoFrame, width: int, height: int) -> av.VideoFrame: |
| image = frame.to_image().convert("RGB") |
| padded = ImageOps.pad(image, (width, height), method=Image.Resampling.BICUBIC, color=(0, 0, 0), centering=(0.5, 0.5)) |
| return av.VideoFrame.from_image(padded).reformat(format="yuv420p") |
|
|
|
|
| def select_frames(segment: Segment, frames: int, width: int, height: int) -> tuple[list[av.VideoFrame], list[float]]: |
| targets = target_times(segment.start, segment.end, frames) |
| selected: list[av.VideoFrame | None] = [None] * frames |
| selected_pts: list[float | None] = [None] * frames |
| last: av.VideoFrame | None = None |
| last_pts: float | None = None |
| next_idx = 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 = frame_seconds(frame, stream) |
| if pts is None: |
| continue |
| if pts < segment.start: |
| continue |
| if pts >= segment.end: |
| stop_decode = True |
| break |
| while last is not None and next_idx < frames and targets[next_idx] <= pts: |
| if abs(last_pts - targets[next_idx]) <= abs(pts - targets[next_idx]): |
| chosen, chosen_pts = last, last_pts |
| else: |
| chosen, chosen_pts = frame, pts |
| selected[next_idx] = letterbox(chosen, width, height) |
| selected_pts[next_idx] = chosen_pts |
| next_idx += 1 |
| if last is None and next_idx < frames and targets[next_idx] <= pts: |
| selected[next_idx] = letterbox(frame, width, height) |
| selected_pts[next_idx] = pts |
| next_idx += 1 |
| last, last_pts = frame, pts |
| if stop_decode or next_idx >= frames: |
| break |
|
|
| if last is None or last_pts is None: |
| raise RuntimeError("decoded zero frames inside segment") |
| for index, frame in enumerate(selected): |
| if frame is None: |
| selected[index] = letterbox(last, width, height) |
| selected_pts[index] = last_pts |
| return [f for f in selected if f is not None], [float(x) for x in selected_pts if x is not None] |
|
|
|
|
| def count_video_frames(path: Path) -> tuple[int, int, int]: |
| with av.open(str(path)) as container: |
| stream = container.streams.video[0] |
| count = 0 |
| width = height = 0 |
| for frame in container.decode(stream): |
| count += 1 |
| width, height = frame.width, frame.height |
| return count, width, height |
|
|
|
|
| def metadata_from_pts(source_pts: list[float], width: int, height: int) -> dict[str, Any]: |
| frame_indices = [int(round(value * 1000)) for value in source_pts] |
| total_num_frames = max(frame_indices) + 1 |
| return { |
| "fps": 1000, |
| "frames_indices": frame_indices, |
| "total_num_frames": total_num_frames, |
| "width": width, |
| "height": height, |
| "duration": total_num_frames / 1000, |
| "video_backend": "precomputed_absolute_pts", |
| } |
|
|
|
|
| def metadata_contract_matches(item: dict[str, Any], frames: int, width: int, height: int) -> bool: |
| metadata = item.get("video_metadata") |
| source_pts = item.get("selected_source_pts") |
| if not isinstance(metadata, dict) or not isinstance(source_pts, list) or len(source_pts) != frames: |
| return False |
| frame_indices = [int(round(float(value) * 1000)) for value in source_pts] |
| return ( |
| metadata.get("fps") == 1000 |
| and metadata.get("frames_indices") == frame_indices |
| and metadata.get("total_num_frames") == max(frame_indices) + 1 |
| and metadata.get("width") == width |
| and metadata.get("height") == height |
| and metadata.get("duration") == metadata.get("total_num_frames") / 1000 |
| and metadata.get("video_backend") == "precomputed_absolute_pts" |
| ) |
|
|
|
|
| def verified_success_item(item: dict[str, Any], output_dir: Path, frames: int, width: int, height: int) -> bool: |
| path = output_dir / item.get("relpath", "") |
| if item.get("status") != "success" or item.get("verified_frame_count") != frames: |
| return False |
| if item.get("width") != width or item.get("height") != height or not path.exists(): |
| return False |
| try: |
| actual = count_video_frames(path) |
| except Exception: |
| return False |
| return actual == (frames, width, height) and metadata_contract_matches(item, frames, width, height) |
|
|
|
|
| def encode_segment(segment: Segment, output_dir: Path, frames: int, width: int, height: int) -> dict[str, Any]: |
| out_path = output_dir / segment.relpath |
| partial = out_path.with_suffix(out_path.suffix + ".partial") |
| selected, source_pts = select_frames(segment, frames, width, height) |
| if len(selected) != frames: |
| raise RuntimeError(f"selected {len(selected)} frames, expected {frames}") |
|
|
| out_path.parent.mkdir(parents=True, exist_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) |
| verified_count, verified_width, verified_height = count_video_frames(partial) |
| if (verified_count, verified_width, verified_height) != (frames, width, height): |
| raise RuntimeError( |
| f"encoded {(verified_count, verified_width, verified_height)}, expected {(frames, width, height)}") |
| os.replace(partial, out_path) |
| except Exception: |
| try: |
| partial.unlink() |
| except FileNotFoundError: |
| pass |
| raise |
| metadata = metadata_from_pts(source_pts, width, height) |
| 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": metadata, |
| "sampling_rule": f"uniform_pts_bin_centers_[start,end)_{frames}", |
| "resize": f"letterbox_{width}x{height}", |
| "status": "success", |
| } |
|
|
|
|
| def done_segments(status_path: Path, output_dir: Path, frames: int, width: int, height: int) -> set[str]: |
| done: set[str] = set() |
| if not status_path.exists(): |
| return done |
| for line in status_path.read_text().splitlines(): |
| if not line.strip(): |
| continue |
| item = json.loads(line) |
| if verified_success_item(item, output_dir, frames, width, height): |
| done.add(item["relpath"]) |
| return done |
|
|
|
|
| def append_jsonl(path: Path, item: dict[str, Any]) -> None: |
| with path.open("a", encoding="utf-8") as handle: |
| handle.write(json.dumps(item, ensure_ascii=False) + "\n") |
| handle.flush() |
|
|
|
|
| def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: |
| with path.open("w", encoding="utf-8") as handle: |
| for row in rows: |
| handle.write(json.dumps(row, ensure_ascii=False) + "\n") |
|
|
|
|
| def ms_swift_record(row: dict[str, Any], segment_info: dict[str, Any], output_dir: Path, include_answer: bool) -> dict[str, Any]: |
| messages = [ |
| {"role": "system", "content": BASELINE_SYSTEM_PROMPT}, |
| {"role": "user", "content": f"<video>\n{row['question']}"}, |
| ] |
| if include_answer: |
| messages.append({"role": "assistant", "content": str(row["answer"])}) |
| return { |
| "messages": messages, |
| "videos": [str((output_dir / row["derived_video"]).resolve())], |
| "chat_template_kwargs": { |
| "nframes": segment_info["frame_count"], |
| "enable_thinking": False, |
| "video_metadata": segment_info["video_metadata"], |
| }, |
| "segment_metadata": { |
| k: row.get(k) |
| for k in [ |
| "id", |
| "dataset", |
| "split", |
| "video", |
| "timestamp_start", |
| "timestamp_end", |
| "answer_format", |
| "segment_start_seconds", |
| "segment_end_seconds", |
| "derived_video", |
| ] |
| }, |
| } |
|
|
|
|
| def write_outputs(output_dir: Path, rows: list[dict[str, Any]], manifest_by_rel: dict[str, dict[str, Any]], splits: list[str]) -> None: |
| selected = [row for row in rows if row["derived_video"] in manifest_by_rel] |
| if "train" in splits: |
| write_jsonl( |
| output_dir / "train.jsonl", |
| [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"] |
| write_jsonl( |
| output_dir / "test_requests.jsonl", |
| [ms_swift_record(row, manifest_by_rel[row["derived_video"]], output_dir, False) for row in test_rows], |
| ) |
| write_jsonl( |
| output_dir / "test_references.jsonl", |
| [{"id": row["id"], "dataset": row["dataset"], "answer": row["answer"], "answer_format": row["answer_format"]} for row in test_rows], |
| ) |
|
|
|
|
| def write_docs(output_dir: Path, args: argparse.Namespace, datasets: dict[str, dict[str, Path]], questions: int, segments: int, failures: int) -> None: |
| now = datetime.now(timezone.utc).isoformat() |
| provenance = { |
| "version": "fixed_frames_v1", |
| "created_at_utc": now, |
| "command": " ".join(os.sys.argv), |
| "inputs": {name: {k: str(v) for k, v in cfg.items()} for name, cfg in datasets.items()}, |
| "splits": args.splits, |
| "question_count": questions, |
| "unique_segment_count": segments, |
| "failed_count": failures, |
| "frame_count": args.frames, |
| "sampling_rule": "uniform PTS bin centers in [start,end), nearest real source PTS", |
| "resize": f"proportional letterbox {args.width}x{args.height}", |
| "encoding": {"container": "mp4", "codec": "H.264/libx264", "audio": "none", "pix_fmt": "yuv420p"}, |
| "metadata_contract": { |
| "fps": 1000, |
| "frames_indices": "round(abs_pts*1000)", |
| "total_num_frames": "max(frames_indices)+1", |
| "width": args.width, |
| "height": args.height, |
| "duration": "total_num_frames/fps", |
| "video_backend": "precomputed_absolute_pts", |
| }, |
| "system_prompt_source": { |
| "author_repo_commit": "3fc0850f9f76261c529e75c78ad2272558ff9e4e", |
| "source": "https://github.com/AzJ24/Orena/blob/3fc0850f9f76261c529e75c78ad2272558ff9e4e/orena_sft/prompts.py", |
| "function": "build_system_prompt", |
| "parameters": {"include_definitions": False, "style": "direct", "track": "segment"}, |
| "foreign_objects_source": "orena-focus 0.3.4 wheel: focus/foreign_objects.py", |
| }, |
| "system_prompt": BASELINE_SYSTEM_PROMPT, |
| "output_dir": str(output_dir), |
| "resume": bool(args.resume), |
| "workers": args.workers, |
| "limit": args.limit, |
| "seed": args.seed, |
| } |
| (output_dir / "provenance.json").write_text(json.dumps(provenance, indent=2, ensure_ascii=False) + "\n") |
| report = { |
| "status": "pass" if failures == 0 and args.success_count == segments else "fail", |
| "expected_frame_count": args.frames, |
| "expected_width": args.width, |
| "expected_height": args.height, |
| "segments": segments, |
| "success": args.success_count, |
| "missing": segments - args.success_count, |
| "failures": failures, |
| "checked_at_utc": now, |
| } |
| (output_dir / "validation_report.json").write_text(json.dumps(report, indent=2, ensure_ascii=False) + "\n") |
|
|
|
|
| def run(args: argparse.Namespace) -> int: |
| data_root = Path(args.data_root) |
| output_dir = Path(args.output_dir) |
| output_dir.mkdir(parents=True, exist_ok=True) |
| datasets = build_datasets(data_root, args.splits) |
| rows, all_segments = load_rows(datasets, args.splits) |
| segments = 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] |
| status_path = output_dir / STATUS_NAME |
| existing = done_segments(status_path, output_dir, args.frames, args.width, args.height) if args.resume else set() |
| todo = [s for s in segments if s.relpath not in existing] |
| failures = 0 |
| manifest_by_rel: dict[str, dict[str, Any]] = {} |
| print(json.dumps({"event": "start", "output": str(output_dir), "splits": args.splits, "segments": len(segments), "todo": len(todo), "python": os.sys.executable}, ensure_ascii=False), flush=True) |
|
|
| def work(segment: Segment) -> dict[str, Any]: |
| started = time.time() |
| try: |
| item = encode_segment(segment, output_dir, args.frames, args.width, args.height) |
| except Exception as exc: |
| 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": "failed", |
| "error": str(exc), |
| "elapsed_seconds": round(time.time() - started, 3), |
| } |
| item["elapsed_seconds"] = round(time.time() - started, 3) |
| return item |
|
|
| with ThreadPoolExecutor(max_workers=args.workers) as pool: |
| futures = [pool.submit(work, segment) for segment in todo] |
| for future in as_completed(futures): |
| item = future.result() |
| append_jsonl(status_path, item) |
| if item["status"] == "success": |
| manifest_by_rel[item["relpath"]] = item |
| else: |
| failures += 1 |
| print(f"relpath={item['relpath']} status={item['status']} elapsed_seconds={item.get('elapsed_seconds')}", flush=True) |
|
|
| for line in status_path.read_text().splitlines() if status_path.exists() else []: |
| if not line.strip(): |
| continue |
| item = json.loads(line) |
| if item.get("relpath") in selected_relpaths and verified_success_item( |
| item, output_dir, args.frames, args.width, args.height |
| ): |
| manifest_by_rel[item["relpath"]] = item |
| manifest = [manifest_by_rel[s.relpath] for s in segments if s.relpath in manifest_by_rel] |
| write_jsonl(output_dir / "segment_manifest.jsonl", manifest) |
| write_outputs(output_dir, selected_rows, manifest_by_rel, args.splits) |
| args.success_count = len(manifest) |
| write_docs(output_dir, args, datasets, len(selected_rows), len(segments), failures) |
| print(json.dumps({"questions": len(selected_rows), "segments": len(segments), "success": len(manifest), "failed": failures}, ensure_ascii=False)) |
| return 0 if failures == 0 and len(manifest) == len(segments) 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("--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("--limit", type=int) |
| parser.add_argument("--resume", action="store_true") |
| parser.add_argument("--seed", type=int, default=42) |
| args = parser.parse_args() |
| if args.workers < 1: |
| raise SystemExit("--workers must be >= 1") |
| if args.frames < 1: |
| raise SystemExit("--frames must be >= 1") |
| if args.width < 1 or args.height < 1: |
| raise SystemExit("--width and --height must be >= 1") |
| return run(args) |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
