demo/infer/vlm_bbox_inference.py

#!/usr/bin/env python3
"""
Shared utility module for VLM bbox-only inference.

This module provides:
- model and processor loading
- prompts for two coordinate conventions
- parsing utilities for bbox-only outputs

The model output is expected to be either:
    [[x0, y0, x1, y1], ...]
or:
    [[x_left, y_top, x_right, y_bottom], ...]

No caption is generated in this module.
"""

import ast
import re
from pathlib import Path

import torch
from transformers import AutoProcessor


# Bottom-left origin, y-axis upward.
BBOX_PROMPT_BOTTOM_LEFT = (
    "<image>This image is 1024 pixels in width and 1024 pixels in height. "
    "The coordinate origin is at the bottom-left corner: x increases to the right, y increases upward. "
    "Detect all objects (layers) in this image. "
    "Output bounding boxes as a list of [x0, y0, x1, y1] where x0=left, y0=bottom, x1=right, y1=top (pixel coordinates). "
    "Output only the list, e.g. [[x0,y0,x1,y1], ...], no other text."
)

# Top-left origin, y-axis downward.
BBOX_PROMPT_TOP_LEFT = (
    "<image>This image is 1024 pixels in width and 1024 pixels in height. "
    "The coordinate origin is at the top-left corner of the image: x increases to the right, y increases downward. "
    "Detect all objects (layers) in this image. "
    "Output bounding boxes as a list of [x_left, y_top, x_right, y_bottom] in pixel coordinates (top-left origin, y downward). "
    "Output only the list, e.g. [[x_left,y_top,x_right,y_bottom], ...], no other text."
)

# Default prompt used by the generic inference API.
BBOX_PROMPT = BBOX_PROMPT_TOP_LEFT
BBOX_SYSTEM = None


def get_model_and_processor(model_path: str, device_map: str = "auto"):
    try:
        from transformers import Qwen3VLForConditionalGeneration
        model_cls = Qwen3VLForConditionalGeneration
    except ImportError:
        try:
            from transformers import Qwen2_5_VLForConditionalGeneration
            model_cls = Qwen2_5_VLForConditionalGeneration
        except ImportError:
            from transformers import AutoModel
            model_cls = AutoModel

    base_2b = "Qwen/Qwen3-VL-2B-Instruct"
    base_8b = "Qwen/Qwen3-VL-8B-Instruct"
    base_name = base_8b if "8b" in model_path.lower() or "8B" in model_path else base_2b

    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    model_dir = Path(model_path)

    def load_config(*sources):
        from transformers import AutoConfig

        for source in sources:
            if not source:
                continue
            try:
                config = AutoConfig.from_pretrained(source, trust_remote_code=True)
                if not hasattr(config, "rope_scaling") or config.rope_scaling is None:
                    config.rope_scaling = {}
                return config
            except Exception:
                continue
        return None

    if model_dir.is_dir() and (model_dir / "adapter_config.json").exists():
        from peft import PeftConfig, PeftModel

        peft_config = PeftConfig.from_pretrained(str(model_dir))
        base_name = peft_config.base_model_name_or_path or base_name
        config = load_config(base_name, str(model_dir))

        model = model_cls.from_pretrained(
            base_name,
            config=config,
            torch_dtype=dtype,
            device_map=device_map,
            trust_remote_code=True,
        )
        model = PeftModel.from_pretrained(model, str(model_dir))

        try:
            processor = AutoProcessor.from_pretrained(str(model_dir), trust_remote_code=True)
        except Exception:
            processor = AutoProcessor.from_pretrained(base_name, trust_remote_code=True)
    else:
        config = load_config(str(model_dir) if model_dir.exists() else None, model_path, base_name)
        model = model_cls.from_pretrained(
            model_path,
            config=config,
            torch_dtype=dtype,
            device_map=device_map,
            trust_remote_code=True,
        )

        try:
            processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
        except Exception:
            processor = AutoProcessor.from_pretrained(base_name, trust_remote_code=True)

    return model, processor


def parse_bbox_output(text: str):
    """
    Parse bbox lists from model output.

    The parser extracts all standalone [a, b, c, d] patterns using regex,
    ignoring outer brackets and extra text. Duplicate boxes are removed to
    reduce the impact of repeated model outputs.
    """
    text = (text or "").strip()

    # Match standalone boxes such as [102, 611, 511, 1023].
    # Both integers and floating-point numbers are supported.
    pattern = (
        r"\[\s*-?\d+(?:\.\d+)?\s*,"
        r"\s*-?\d+(?:\.\d+)?\s*,"
        r"\s*-?\d+(?:\.\d+)?\s*,"
        r"\s*-?\d+(?:\.\d+)?\s*\]"
    )
    matches = re.findall(pattern, text)

    parsed_boxes = []
    for match_str in matches:
        try:
            box = ast.literal_eval(match_str)
            if isinstance(box, list) and len(box) == 4:
                parsed_boxes.append(box)
        except (ValueError, SyntaxError):
            continue

    # Remove duplicate boxes to avoid repeated outputs.
    unique_boxes = []
    seen = set()

    for b in parsed_boxes:
        key = tuple(float(x) for x in b)
        if key not in seen:
            seen.add(key)
            unique_boxes.append(b)

    return unique_boxes


def infer_bboxes(image_path: str, model, processor, *, prompt: str, max_new_tokens: int = 512):
    """
    Run bbox inference for a single image.

    Returns:
        A list of boxes, where each box is [a, b, c, d].
        The coordinate meaning depends on the prompt convention.
    """
    path = Path(image_path)
    if not path.exists():
        return []

    content = [
        {"type": "image", "image": str(path.absolute())},
        {"type": "text", "text": prompt},
    ]

    messages = [{"role": "user", "content": content}]

    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    )

    inputs = {k: v.to(model.device) if hasattr(v, "to") else v for k, v in inputs.items()}
    inputs.pop("token_type_ids", None)

    with torch.no_grad():
        generated = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.1,
            repetition_penalty=1.1,
            pad_token_id=processor.tokenizer.pad_token_id or processor.tokenizer.eos_token_id,
        )

    input_len = inputs["input_ids"].shape[1]
    output_ids = generated[:, input_len:]

    output_text = processor.batch_decode(
        output_ids,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=True,
    )

    raw = (output_text[0] or "").strip()
    boxes = parse_bbox_output(raw)

    result = []
    for b in boxes:
        if isinstance(b, (list, tuple)) and len(b) >= 4:
            result.append([float(b[0]), float(b[1]), float(b[2]), float(b[3])])

    return result


def detect_objects(image_path: str, model, processor, *, prompt=None, system=None, max_new_tokens=512):
    """
    Compatibility wrapper using the default bbox prompt.

    The `system` argument is kept for compatibility with older call sites.
    """
    return infer_bboxes(
        image_path,
        model,
        processor,
        prompt=prompt or BBOX_PROMPT,
        max_new_tokens=max_new_tokens,
    )