| |
| from __future__ import annotations |
|
|
| import json |
| import logging |
| import re |
| import sys |
| from pathlib import Path |
|
|
| from _gemini_runner import GeminiRunner, load_repo_module |
|
|
|
|
| BASE = load_repo_module("qwen_multi_event_base", "qwen/generate_multi_event_annotations.py") |
|
|
|
|
| def strip_detailed_caption_from_global_prompt(prompt: str) -> str: |
| cleaned = prompt |
| cleaned = cleaned.replace( |
| '- The detailed caption should be about 500-800 words.\n', |
| "", |
| ) |
| cleaned = cleaned.replace( |
| '- The detailed caption may synthesize and reuse the structured fields plus stable recurring evidence from the event summaries, but it must not turn into a segment-by-segment narration.\n', |
| "", |
| ) |
| cleaned = re.sub( |
| r',\n(\s*)"detailed_caption": "About 500-800 words\."\n', |
| "\n", |
| cleaned, |
| ) |
| return cleaned |
|
|
|
|
| _ORIGINAL_BUILD_GLOBAL_CONTEXT_PROMPT = BASE.build_global_context_prompt |
| _ORIGINAL_BUILD_GLOBAL_REFINE_PROMPT = BASE.build_global_refine_prompt |
| _ORIGINAL_COERCE_GLOBAL_PAYLOAD = BASE.coerce_global_payload |
|
|
|
|
| def _gemini_build_global_context_prompt(*args, **kwargs): |
| prompt = _ORIGINAL_BUILD_GLOBAL_CONTEXT_PROMPT(*args, **kwargs) |
| return ( |
| f"{prompt}\n\n" |
| "IMPORTANT FOR GEMINI:\n" |
| "- Return one single JSON object for the entire stitched video.\n" |
| "- Never return a JSON array.\n" |
| "- Never return one object per clip, per segment, or per event.\n" |
| "- The stitched video may contain multiple chronological event clips, but this context pass must summarize the whole stitched video into one object only.\n" |
| "- Do not return object-detection output, tracking output, grounding output, OCR output, or visual tool output.\n" |
| "- Never return keys such as box_2d, bbox, label, labels, mask, polygon, points, score, or tracking_id unless the requested schema explicitly asks for them.\n" |
| "- Never return timestamped clip-local records such as event_clip_id, start_time_seconds, or end_time_seconds in this context pass.\n" |
| "- If uncertain, still fill the requested schema directly with high-confidence whole-video annotations rather than switching to detection-style results.\n" |
| ) |
|
|
|
|
| def _gemini_build_global_refine_prompt(*args, **kwargs): |
| return strip_detailed_caption_from_global_prompt( |
| _ORIGINAL_BUILD_GLOBAL_REFINE_PROMPT(*args, **kwargs) |
| ) |
|
|
|
|
| def _gemini_coerce_global_payload(*args, **kwargs): |
| payload = _ORIGINAL_COERCE_GLOBAL_PAYLOAD(*args, **kwargs) |
| payload.pop("detailed_caption", None) |
| return payload |
|
|
|
|
| BASE.build_global_context_prompt = _gemini_build_global_context_prompt |
| BASE.build_global_refine_prompt = _gemini_build_global_refine_prompt |
| BASE.coerce_global_payload = _gemini_coerce_global_payload |
|
|
|
|
| def main() -> int: |
| args = BASE.parse_args() |
| BASE.setup_logging(args.log_level) |
|
|
| input_path = Path(args.input).expanduser().resolve() |
| output_path = Path(args.output).expanduser().resolve() |
| camera_motion_guide_path = Path(args.camera_motion_guide).expanduser().resolve() |
|
|
| suffixes = BASE.parse_video_suffixes(args.video_suffixes) |
| allow_uuids = BASE.load_uuid_allow_list(args.uuids, args.uuids_file) |
| camera_motion_guide = BASE.load_camera_motion_guide(camera_motion_guide_path) |
|
|
| groups = BASE.collect_segment_groups( |
| input_path=input_path, |
| recursive=args.recursive, |
| suffixes=suffixes, |
| allow_uuids=allow_uuids, |
| ) |
| if not groups: |
| raise RuntimeError("No valid segmented videos were found.") |
|
|
| uuids = list(groups) |
| if args.limit is not None: |
| uuids = uuids[: args.limit] |
|
|
| pending_for_print = uuids |
| if args.resume and output_path.exists(): |
| existing_videos_for_print = BASE.load_existing_output( |
| output_path=output_path, |
| overwrite=False, |
| resume=True, |
| ) |
| pending_for_print = [uuid for uuid in uuids if uuid not in existing_videos_for_print] |
|
|
| if args.print_uuids_json: |
| print(json.dumps(pending_for_print, ensure_ascii=False)) |
| return 0 |
|
|
| uuids = BASE.apply_uuid_shard( |
| uuids=pending_for_print if args.resume else uuids, |
| num_shards=args.num_shards, |
| shard_index=args.shard_index, |
| ) |
|
|
| logging.info("Discovered %d UUID groups.", len(uuids)) |
| for uuid in uuids: |
| logging.info("UUID %s has %d segments.", uuid, len(groups[uuid])) |
|
|
| if args.dry_run: |
| logging.info("Dry run complete. Gemini model was not called.") |
| return 0 |
|
|
| existing_videos = BASE.load_existing_output( |
| output_path=output_path, |
| overwrite=args.overwrite, |
| resume=args.resume, |
| ) |
| pending_uuids = [uuid for uuid in uuids if uuid not in existing_videos] |
| if not pending_uuids: |
| logging.info("Nothing to do. All requested UUIDs already exist in %s.", output_path) |
| return 0 |
|
|
| runner = GeminiRunner( |
| model_path=Path(args.model_path).expanduser().resolve(), |
| device_map=args.device_map, |
| dtype=args.dtype, |
| attn_implementation=args.attn_implementation, |
| trust_remote_code=args.trust_remote_code, |
| fps=args.fps, |
| max_pixels=args.max_pixels, |
| temperature=args.temperature, |
| top_p=args.top_p, |
| top_k=args.top_k, |
| repetition_penalty=args.repetition_penalty, |
| greedy=args.greedy, |
| ) |
|
|
| try: |
| failures: list[str] = [] |
| for index, uuid in enumerate(pending_uuids, start=1): |
| logging.info("Processing UUID %s (%d/%d)", uuid, index, len(pending_uuids)) |
| try: |
| video_payload = BASE.annotate_uuid( |
| uuid=uuid, |
| segments=groups[uuid], |
| runner=runner, |
| camera_motion_guide=camera_motion_guide, |
| args=args, |
| ) |
| existing_videos[uuid] = video_payload |
| BASE.write_output(output_path, existing_videos) |
| logging.info("Wrote %s", output_path) |
| except Exception as exc: |
| if not args.continue_on_error: |
| raise |
| logging.exception("Failed to process %s: %s", uuid, exc) |
| failures.append(uuid) |
|
|
| if failures: |
| logging.error("Completed with failures: %s", ", ".join(failures)) |
| return 1 |
| logging.info("Finished %d UUIDs.", len(pending_uuids)) |
| return 0 |
| finally: |
| runner.close() |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
