Blame for qwen/classify_camera_motion_qwen3vl.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import ast
import json
import re
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parent
DEFAULT_CAMERA_GUIDE = ROOT / "camera_motion.md"

CAMERA_MOTION_LABELS = [
    "static",
    "dolly-in",
    "dolly-out",
    "pedestal-up",
    "pedestal-down",
    "truck-left",
    "truck-right",
    "pan-left",
    "pan-right",
    "tilt-up",
    "tilt-down",
    "roll-CW",
    "roll-CCW",
]
LABEL_MAP = {label.lower(): label for label in CAMERA_MOTION_LABELS}
LABEL_ALIASES = {
    "none": "static",
    "no motion": "static",
    "no-motion": "static",
    "near-static": "static",
    "stationary": "static",
    "clockwise roll": "roll-CW",
    "counterclockwise roll": "roll-CCW",
    "counter-clockwise roll": "roll-CCW",
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Classify camera motion for one video with Qwen3-VL."
    )
    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(
        "--camera-motion-guide",
        default=str(DEFAULT_CAMERA_GUIDE),
        help="Path to camera_motion.md.",
    )
    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=512,
        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.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 normalize_label(value: Any) -> str:
    text = re.sub(r"\s+", " ", str(value or "")).strip()
    if not text:
        return "static"
    lowered = text.lower()
    if lowered in LABEL_ALIASES:
        return LABEL_ALIASES[lowered]
    return LABEL_MAP.get(lowered, "")


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 load_camera_motion_guide(path: Path) -> str:
    if path.exists():
        return path.read_text(encoding="utf-8").strip()
    labels = "\n".join(f"- {label}" for label in CAMERA_MOTION_LABELS)
    return f"Allowed labels:\n{labels}"


def build_prompt(camera_motion_guide: str) -> str:
    labels = ", ".join(CAMERA_MOTION_LABELS)
    return f"""
You are classifying the dominant camera motion in a single video.

Camera motion guide:
{camera_motion_guide}

Allowed output labels:
{labels}

Rules:
- Return strict JSON only. Do not use markdown fences.
- Choose exactly one label.
- Use "static" if there is no clear camera motion or the apparent motion is mostly subject/object motion.
- Prefer rotation labels pan/tilt/roll when the camera pivots in place.
- Prefer translation labels dolly/truck/pedestal when the camera physically changes position.
- If multiple motions occur, choose the most visually dominant camera motion.
- The description must be one concise English sentence explaining the visual evidence.

Return JSON in exactly this schema:
{{
  "label": "pan-left",
  "description": "The camera pans left as the scene shifts horizontally across the frame."
}}
""".strip()


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  # type: ignore
        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 annotation 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 classify(self, video_path: Path, prompt: str, max_new_tokens: int) -> dict[str, str]:
        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]
        payload = parse_json_payload(raw_text)
        label = normalize_label(payload.get("label"))
        if not label:
            raise ValueError(
                f"Model returned unsupported label {payload.get('label')!r}; allowed labels are {CAMERA_MOTION_LABELS}"
            )
        description = re.sub(r"\s+", " ", str(payload.get("description", ""))).strip()
        if not description:
            description = "The video does not provide enough clear evidence for a more specific camera motion."
        return {"label": label, "description": description}


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

    camera_guide = load_camera_motion_guide(
        Path(args.camera_motion_guide).expanduser().resolve()
    )
    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,
    )
    result = runner.classify(
        video_path=video_path,
        prompt=build_prompt(camera_guide),
        max_new_tokens=args.max_new_tokens,
    )

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