| |
| from __future__ import annotations |
|
|
| import argparse |
| import ast |
| import json |
| import re |
| import subprocess |
| from pathlib import Path |
| from typing import Any |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser( |
| description="Use Qwen3-VL to segment one video into consecutive events." |
| ) |
| parser.add_argument("video", help="Path to the input video.") |
| parser.add_argument( |
| "--model-path", |
| default="Qwen/Qwen3-VL-32B-Instruct", |
| help="Local checkpoint path or Hugging Face model id.", |
| ) |
| parser.add_argument( |
| "--output", |
| default=None, |
| help="Optional JSON output path. If omitted, prints JSON to stdout.", |
| ) |
| parser.add_argument("--fps", type=float, default=1.0, help="Video sampling FPS.") |
| parser.add_argument( |
| "--max-pixels", |
| type=int, |
| default=360000, |
| help="Optional Qwen video processor max_pixels hint. Use 0 to disable.", |
| ) |
| parser.add_argument( |
| "--max-new-tokens", |
| type=int, |
| default=1024, |
| help="Max new tokens for the model response.", |
| ) |
| parser.add_argument( |
| "--device-map", |
| default="auto", |
| help="Transformers device_map argument.", |
| ) |
| parser.add_argument( |
| "--dtype", |
| default="bfloat16", |
| choices=["auto", "float16", "bfloat16", "float32"], |
| help="Torch dtype for model loading.", |
| ) |
| parser.add_argument( |
| "--attn-implementation", |
| default="", |
| help='Optional attn_implementation, e.g. "flash_attention_2". Empty disables it.', |
| ) |
| parser.add_argument( |
| "--temperature", |
| type=float, |
| default=0.0, |
| help="Sampling temperature. 0 uses greedy decoding.", |
| ) |
| parser.add_argument( |
| "--ffprobe-bin", |
| default="ffprobe", |
| help="ffprobe executable used to read video duration.", |
| ) |
| parser.set_defaults(trust_remote_code=True) |
| parser.add_argument( |
| "--no-trust-remote-code", |
| dest="trust_remote_code", |
| action="store_false", |
| help="Disable trust_remote_code for model and processor loading.", |
| ) |
| return parser.parse_args() |
|
|
|
|
| def probe_duration(video_path: Path, ffprobe_bin: str) -> float | None: |
| try: |
| result = subprocess.run( |
| [ |
| ffprobe_bin, |
| "-v", |
| "error", |
| "-show_entries", |
| "format=duration", |
| "-of", |
| "default=noprint_wrappers=1:nokey=1", |
| str(video_path), |
| ], |
| check=True, |
| capture_output=True, |
| text=True, |
| ) |
| except (FileNotFoundError, subprocess.CalledProcessError): |
| return None |
| try: |
| duration = float(result.stdout.strip()) |
| except ValueError: |
| return None |
| return round(duration, 3) if duration > 0 else None |
|
|
|
|
| def extract_json_text(raw_text: str) -> str: |
| text = raw_text.strip() |
| fenced = re.findall(r"```(?:json)?\s*(.*?)```", text, flags=re.DOTALL | re.I) |
| candidates = [item.strip() for item in fenced if item.strip()] |
| candidates.append(text) |
| for candidate in candidates: |
| if candidate.startswith("{") and candidate.endswith("}"): |
| return candidate |
|
|
| start = text.find("{") |
| if start < 0: |
| raise ValueError(f"No JSON object found in model output: {text[:300]}") |
|
|
| depth = 0 |
| in_string = False |
| escape = False |
| for index in range(start, len(text)): |
| char = text[index] |
| if in_string: |
| if escape: |
| escape = False |
| elif char == "\\": |
| escape = True |
| elif char == '"': |
| in_string = False |
| continue |
| if char == '"': |
| in_string = True |
| elif char == "{": |
| depth += 1 |
| elif char == "}": |
| depth -= 1 |
| if depth == 0: |
| return text[start : index + 1] |
| raise ValueError(f"Unclosed JSON object in model output: {text[:300]}") |
|
|
|
|
| def parse_json_payload(raw_text: str) -> dict[str, Any]: |
| json_text = extract_json_text(raw_text) |
| try: |
| payload = json.loads(json_text) |
| except json.JSONDecodeError: |
| payload = ast.literal_eval(json_text) |
| if not isinstance(payload, dict): |
| raise ValueError(f"Expected a JSON object, got {type(payload).__name__}.") |
| return payload |
|
|
|
|
| def coerce_float(value: Any, field_name: str) -> float: |
| if isinstance(value, str): |
| value = value.strip().rstrip("s") |
| try: |
| return round(float(value), 3) |
| except (TypeError, ValueError) as exc: |
| raise ValueError(f"Invalid {field_name}: {value!r}") from exc |
|
|
|
|
| def build_prompt(duration: float | None) -> str: |
| duration_line = ( |
| f"The video duration is approximately {duration:.3f} seconds." |
| if duration is not None |
| else "The video duration is not provided; infer timestamps in seconds from the video timeline." |
| ) |
| return f""" |
| You are segmenting a single video into consecutive semantic events. |
| |
| {duration_line} |
| |
| Definition of an event: |
| - An event is a continuous temporal interval with one dominant visible action, interaction, camera-framing state, or meaningful scene/object state. |
| - Split when the dominant action changes, a new subject/object appears or disappears, the focus changes materially, or the scene state changes enough that a separate event caption would be needed. |
| - Do not over-segment tiny pose changes, ordinary motion within the same action, or brief camera shake. |
| - Do not leave gaps between events. |
| - Events must be chronological and non-overlapping. |
| - The first event should begin at 0.0 seconds. |
| - If the duration is known, the last event should end at the video duration. |
| - Use seconds as numbers, not strings. |
| - Include a short description for each event so the boundary decision is inspectable. |
| |
| Return strict JSON only. Do not use markdown fences. |
| |
| Return JSON in exactly this schema: |
| {{ |
| "event_count": 3, |
| "events": [ |
| {{ |
| "event_id": 1, |
| "begin_time": 0.0, |
| "end_time": 2.5, |
| "description": "The first visible event in one concise sentence." |
| }}, |
| {{ |
| "event_id": 2, |
| "begin_time": 2.5, |
| "end_time": 5.0, |
| "description": "The second visible event in one concise sentence." |
| }} |
| ] |
| }} |
| """.strip() |
|
|
|
|
| def normalize_events(payload: dict[str, Any], duration: float | None) -> dict[str, Any]: |
| raw_events = payload.get("events") |
| if not isinstance(raw_events, list) or not raw_events: |
| raise ValueError("Model output must contain a non-empty events list.") |
|
|
| events: list[dict[str, Any]] = [] |
| for index, item in enumerate(raw_events, start=1): |
| if not isinstance(item, dict): |
| raise ValueError(f"Event {index} is not an object.") |
| begin_time = coerce_float(item.get("begin_time"), "begin_time") |
| end_time = coerce_float(item.get("end_time"), "end_time") |
| if duration is not None: |
| begin_time = min(max(begin_time, 0.0), duration) |
| end_time = min(max(end_time, 0.0), duration) |
| if end_time <= begin_time: |
| raise ValueError( |
| f"Event {index} has non-positive duration: {begin_time} to {end_time}" |
| ) |
| description = re.sub( |
| r"\s+", " ", str(item.get("description", "")).strip() |
| ) |
| events.append( |
| { |
| "event_id": index, |
| "begin_time": begin_time, |
| "end_time": end_time, |
| "description": description, |
| } |
| ) |
|
|
| events.sort(key=lambda item: (item["begin_time"], item["end_time"])) |
| for index, event in enumerate(events, start=1): |
| event["event_id"] = index |
|
|
| if events[0]["begin_time"] != 0.0: |
| events[0]["begin_time"] = 0.0 |
| if duration is not None: |
| events[-1]["end_time"] = duration |
|
|
| return { |
| "event_count": len(events), |
| "events": events, |
| } |
|
|
|
|
| class Qwen3VLRunner: |
| def __init__( |
| self, |
| model_path: str, |
| device_map: str, |
| dtype: str, |
| attn_implementation: str, |
| trust_remote_code: bool, |
| fps: float, |
| max_pixels: int, |
| temperature: float, |
| ) -> None: |
| try: |
| import torch |
| from transformers import AutoProcessor, Qwen3VLForConditionalGeneration |
| except ModuleNotFoundError as exc: |
| raise RuntimeError( |
| "Missing dependencies. Install torch, transformers, and preferably qwen-vl-utils." |
| ) from exc |
|
|
| self.torch = torch |
| self.temperature = temperature |
| self.fps = fps if fps > 0 else None |
| self.max_pixels = max_pixels if max_pixels > 0 else None |
|
|
| try: |
| from qwen_vl_utils import process_vision_info |
| except ModuleNotFoundError: |
| process_vision_info = None |
| self.process_vision_info = process_vision_info |
|
|
| dtype_map = { |
| "auto": "auto", |
| "float16": torch.float16, |
| "bfloat16": torch.bfloat16, |
| "float32": torch.float32, |
| } |
| load_kwargs: dict[str, Any] = { |
| "device_map": device_map, |
| "torch_dtype": dtype_map[dtype], |
| "trust_remote_code": trust_remote_code, |
| } |
| if attn_implementation: |
| load_kwargs["attn_implementation"] = attn_implementation |
|
|
| self.model = Qwen3VLForConditionalGeneration.from_pretrained( |
| model_path, |
| **load_kwargs, |
| ) |
| self.processor = AutoProcessor.from_pretrained( |
| model_path, |
| trust_remote_code=trust_remote_code, |
| ) |
|
|
| def _build_messages(self, video_path: Path, prompt: str) -> list[dict[str, Any]]: |
| video_item: dict[str, Any] = { |
| "type": "video", |
| "video": str(video_path.resolve()), |
| } |
| if self.fps is not None: |
| video_item["fps"] = self.fps |
| if self.max_pixels is not None: |
| video_item["max_pixels"] = self.max_pixels |
| return [ |
| { |
| "role": "system", |
| "content": "You are a careful video event segmentation assistant. Return strict JSON only.", |
| }, |
| { |
| "role": "user", |
| "content": [video_item, {"type": "text", "text": prompt}], |
| }, |
| ] |
|
|
| def _prepare_inputs(self, messages: list[dict[str, Any]]): |
| if self.process_vision_info is not None: |
| text = self.processor.apply_chat_template( |
| messages, |
| tokenize=False, |
| add_generation_prompt=True, |
| ) |
| image_inputs, video_inputs = self.process_vision_info(messages) |
| inputs = self.processor( |
| text=[text], |
| images=image_inputs, |
| videos=video_inputs, |
| padding=True, |
| return_tensors="pt", |
| ) |
| else: |
| try: |
| inputs = self.processor.apply_chat_template( |
| messages, |
| tokenize=True, |
| add_generation_prompt=True, |
| return_dict=True, |
| return_tensors="pt", |
| ) |
| except Exception as exc: |
| raise RuntimeError( |
| "Failed to prepare video inputs. Install qwen-vl-utils or use a transformers build with Qwen3-VL video chat support." |
| ) from exc |
| return inputs.to(self.model.device) |
|
|
| def segment(self, video_path: Path, prompt: str, max_new_tokens: int) -> dict[str, Any]: |
| messages = self._build_messages(video_path, prompt) |
| inputs = self._prepare_inputs(messages) |
| generate_kwargs: dict[str, Any] = {"max_new_tokens": max_new_tokens} |
| if self.temperature and self.temperature > 0: |
| generate_kwargs.update({"do_sample": True, "temperature": self.temperature}) |
| else: |
| generate_kwargs["do_sample"] = False |
|
|
| with self.torch.inference_mode(): |
| generated_ids = self.model.generate(**inputs, **generate_kwargs) |
| generated_ids_trimmed = [ |
| out_ids[len(in_ids) :] |
| for in_ids, out_ids in zip(inputs.input_ids, generated_ids) |
| ] |
| raw_text = self.processor.batch_decode( |
| generated_ids_trimmed, |
| skip_special_tokens=True, |
| clean_up_tokenization_spaces=False, |
| )[0] |
| return parse_json_payload(raw_text) |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| video_path = Path(args.video).expanduser().resolve() |
| if not video_path.is_file(): |
| raise FileNotFoundError(f"Video not found: {video_path}") |
|
|
| duration = probe_duration(video_path, args.ffprobe_bin) |
| runner = Qwen3VLRunner( |
| model_path=args.model_path, |
| 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, |
| ) |
| payload = runner.segment( |
| video_path=video_path, |
| prompt=build_prompt(duration), |
| max_new_tokens=args.max_new_tokens, |
| ) |
| result = normalize_events(payload, duration) |
|
|
| output_text = json.dumps(result, ensure_ascii=False, indent=2) + "\n" |
| if args.output: |
| output_path = Path(args.output).expanduser().resolve() |
| output_path.parent.mkdir(parents=True, exist_ok=True) |
| output_path.write_text(output_text, encoding="utf-8") |
| else: |
| print(output_text, end="") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
