images.py

"""Image inputs for the vision variant: loading, resizing to the ViT token budget, and the patch-grid geometry
that the `point` head works in.

Geometry: the vision tower emits one token per (patch_size * spatial_merge_size)^2 pixel block (32x32 px for
Qwen3.5), row-major over the merged grid (hm, wm) = (h / merge, w / merge) where (t, h, w) = image_grid_thw.
"""
from __future__ import annotations

import base64
import io
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union

DEFAULT_MAX_PIXELS = 1280 * 800  # ~1000 visual tokens at 32 px/token
HIGH_DETAIL_MAX_PIXELS = 1920 * 1200  # ~2280 tokens
DEFAULT_MIN_PIXELS = 64 * 64
MAX_IMAGES_PER_REQUEST = 4
MAX_VISUAL_TOKENS = 8192


def load_image(src: Any):
    """PIL image from: PIL.Image, bytes, data-URI / base64 string, http(s) URL, or a filesystem path."""
    from PIL import Image

    if isinstance(src, Image.Image):
        return src.convert("RGB")
    if isinstance(src, (bytes, bytearray)):
        return Image.open(io.BytesIO(src)).convert("RGB")
    if hasattr(src, "shape"):  # numpy array (H, W, 3)
        return Image.fromarray(src).convert("RGB")
    if not isinstance(src, str):
        raise TypeError(f"unsupported image source {type(src)}")
    if src.startswith("data:"):
        _, _, payload = src.partition(",")
        return Image.open(io.BytesIO(base64.b64decode(payload))).convert("RGB")
    if src.startswith(("http://", "https://")):
        import urllib.request

        req = urllib.request.Request(src, headers={"User-Agent": "openthai-systemone/1.0"})
        with urllib.request.urlopen(req, timeout=20) as r:
            return Image.open(io.BytesIO(r.read())).convert("RGB")
    if len(src) > 512 and "/" not in src[:16] and not src.startswith("."):
        try:  # bare base64
            return Image.open(io.BytesIO(base64.b64decode(src, validate=True))).convert("RGB")
        except Exception:
            pass
    return Image.open(src).convert("RGB")


@dataclass
class ImageRef:
    """One image attached to a request."""

    id: str
    image: Any  # anything load_image accepts
    detail: str = "auto"  # auto | high
    role: Optional[str] = None  # free text tag, shown in the sequence (screenshot, frame, photo, document)

    @classmethod
    def parse(cls, obj: Union["ImageRef", Dict[str, Any]]) -> "ImageRef":
        if isinstance(obj, ImageRef):
            return obj
        d = dict(obj)
        src = d.get("image", d.get("data", d.get("url", d.get("path"))))
        if src is None:
            raise ValueError("image needs one of: image / data / url / path")
        return cls(id=str(d.get("id", "img")), image=src, detail=str(d.get("detail", "auto")), role=d.get("role"))


@dataclass
class ProcessedImages:
    ids: List[str]
    pixel_values: Any  # torch.Tensor (sum_patches, C*T*P*P)
    grid_thw: List[List[int]]  # per image (t, h, w) in patches
    sizes: List[Tuple[int, int]]  # original (W, H)
    n_tokens: List[int] = field(default_factory=list)
    merge: int = 2

    def tokens_for(self, i: int) -> int:
        t, h, w = self.grid_thw[i]
        return t * h * w // (self.merge * self.merge)

    def merged_grid(self, i: int) -> Tuple[int, int]:
        t, h, w = self.grid_thw[i]
        return h // self.merge, w // self.merge


def fast_patchify(im, image_processor, *, min_pixels: int, max_pixels: int, device=None):
    """Qwen2-VL preprocessing without the HF processor's overhead (~3 ms vs ~40 ms per frame on a busy CPU).

    Same maths: smart_resize to multiples of patch*merge within [min_pixels, max_pixels], bicubic resize, rescale,
    normalise with the processor's mean/std, duplicate the frame temporal_patch_size times, and flatten patches in the
    processor's (t, h/merge, w/merge, merge, merge, C, T, p, p) order. tests/test_vision.py checks equality with the HF
    processor output.
    """
    import numpy as np
    import torch
    from PIL import Image
    from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize

    patch = int(getattr(image_processor, "patch_size", 16))
    merge = int(getattr(image_processor, "merge_size", 2))
    tp = int(getattr(image_processor, "temporal_patch_size", 2))
    mean = torch.tensor(image_processor.image_mean, dtype=torch.float32).view(3, 1, 1)
    std = torch.tensor(image_processor.image_std, dtype=torch.float32).view(3, 1, 1)
    w, h = im.size
    rh, rw = smart_resize(h, w, factor=patch * merge, min_pixels=min_pixels, max_pixels=max_pixels)
    if (rw, rh) != (w, h):
        im = im.resize((rw, rh), Image.BICUBIC)
    x = torch.from_numpy(np.asarray(im, dtype=np.uint8))
    if device is not None:  # normalise + patchify on the GPU: the CPU only resizes and uploads 3 bytes/pixel
        x = x.to(device, non_blocking=True)
        mean, std = mean.to(device), std.to(device)
    x = x.permute(2, 0, 1).float() / 255.0  # (C, H, W)
    x = (x - mean) / std
    gh, gw = rh // patch, rw // patch
    x = x.unsqueeze(0).expand(tp, -1, -1, -1)  # (T, C, H, W)
    x = x.reshape(tp, 3, gh // merge, merge, patch, gw // merge, merge, patch)
    # HF order: (grid_t, temporal, C, gh//m, m, p, gw//m, m, p) -> transpose(0, 3, 6, 4, 7, 2, 1, 5, 8)
    x = x.unsqueeze(0)  # (1, T, C, gh/m, m, p, gw/m, m, p)
    x = x.permute(0, 3, 6, 4, 7, 2, 1, 5, 8).reshape(gh * gw, 3 * tp * patch * patch)
    return x, [1, gh, gw]


def process_images(image_processor, refs: Sequence[ImageRef], *, max_pixels: int = DEFAULT_MAX_PIXELS,
                   min_pixels: int = DEFAULT_MIN_PIXELS, fast: bool = True, device=None) -> ProcessedImages:
    """Preprocess each image (own detail cap each) and concatenate. fast=True uses fast_patchify (identical output)."""
    import torch

    if len(refs) > MAX_IMAGES_PER_REQUEST:
        raise ValueError(f"at most {MAX_IMAGES_PER_REQUEST} images per request")
    merge = int(getattr(image_processor, "merge_size", 2))
    pv, grids, sizes, ids, ntoks = [], [], [], [], []
    for r in refs:
        im = load_image(r.image)
        sizes.append(im.size)
        cap = HIGH_DETAIL_MAX_PIXELS if r.detail == "high" else max_pixels
        if fast:
            x, g = fast_patchify(im, image_processor, min_pixels=min_pixels, max_pixels=cap, device=device)
            pv.append(x)
        else:
            out = image_processor(images=[im], size={"shortest_edge": min_pixels, "longest_edge": cap}, return_tensors="pt")
            pv.append(out["pixel_values"])
            g = [int(v) for v in out["image_grid_thw"][0].tolist()]
        grids.append(g)
        ids.append(r.id)
        ntoks.append(g[0] * g[1] * g[2] // (merge * merge))
    if sum(ntoks) > MAX_VISUAL_TOKENS:
        raise ValueError(f"images need {sum(ntoks)} visual tokens > {MAX_VISUAL_TOKENS}; lower detail or send fewer images")
    return ProcessedImages(ids=ids, pixel_values=torch.cat(pv, 0), grid_thw=grids, sizes=sizes, n_tokens=ntoks, merge=merge)


def bbox_to_token_target(bbox: Optional[Sequence[float]], hm: int, wm: int) -> "torch.Tensor":
    """Soft target over the hm*wm tokens (+1 null slot at the end) for a normalised bbox [x0,y0,x1,y1] in 0..1.

    Cells whose area overlaps the box get weight = overlap fraction of the cell (GUI-Actor style multi-patch
    supervision); normalised to sum 1. bbox None -> all mass on the null slot ("not on screen").
    """
    import torch

    tgt = torch.zeros(hm * wm + 1)
    if bbox is None:
        tgt[-1] = 1.0
        return tgt
    x0, y0, x1, y1 = [float(v) for v in bbox]
    x0, x1 = sorted((max(0.0, min(1.0, x0)), max(0.0, min(1.0, x1))))
    y0, y1 = sorted((max(0.0, min(1.0, y0)), max(0.0, min(1.0, y1))))
    # a degenerate (point) box marks the cell that contains it
    if x1 - x0 < 1e-6:
        x0, x1 = max(0.0, x0 - 0.5 / wm), min(1.0, x0 + 0.5 / wm)
    if y1 - y0 < 1e-6:
        y0, y1 = max(0.0, y0 - 0.5 / hm), min(1.0, y0 + 0.5 / hm)
    for r in range(hm):
        cy0, cy1 = r / hm, (r + 1) / hm
        oy = max(0.0, min(cy1, y1) - max(cy0, y0))
        if oy <= 0:
            continue
        for c in range(wm):
            cx0, cx1 = c / wm, (c + 1) / wm
            ox = max(0.0, min(cx1, x1) - max(cx0, x0))
            if ox > 0:
                tgt[r * wm + c] = ox * oy * hm * wm  # fraction of the cell covered
    if tgt[:-1].sum() <= 0:
        tgt[-1] = 1.0
    else:
        tgt = tgt / tgt.sum()
    return tgt


def token_xy(index: int, hm: int, wm: int) -> Tuple[float, float]:
    """Centre of merged token `index` in normalised image coordinates."""
    r, c = divmod(index, wm)
    return (c + 0.5) / wm, (r + 0.5) / hm


def heatmap_to_candidates(probs: Sequence[float], hm: int, wm: int, *, rel_threshold: float = 0.3, max_candidates: int = 5):
    """GUI-Actor decoding: threshold at rel_threshold * max, 8-connected components, mass-weighted centroids.

    probs: hm*wm token probabilities (null slot already removed). Returns list of {x, y, p, bbox} sorted by p.
    """
    n = hm * wm
    p = [float(v) for v in probs[:n]]
    if n == 0 or max(p) <= 0:
        return []
    thr = rel_threshold * max(p)
    keep = {i for i, v in enumerate(p) if v >= thr}
    seen, comps = set(), []
    for start in sorted(keep, key=lambda i: -p[i]):
        if start in seen:
            continue
        comp, stack = [], [start]
        seen.add(start)
        while stack:
            i = stack.pop()
            comp.append(i)
            r, c = divmod(i, wm)
            for dr in (-1, 0, 1):
                for dc in (-1, 0, 1):
                    rr, cc = r + dr, c + dc
                    if 0 <= rr < hm and 0 <= cc < wm:
                        j = rr * wm + cc
                        if j in keep and j not in seen:
                            seen.add(j)
                            stack.append(j)
        mass = sum(p[i] for i in comp)
        xs = [token_xy(i, hm, wm) for i in comp]
        x = sum(px * p[i] for (px, _), i in zip(xs, comp)) / mass
        y = sum(py * p[i] for (_, py), i in zip(xs, comp)) / mass
        rows = [i // wm for i in comp]
        cols = [i % wm for i in comp]
        bbox = [min(cols) / wm, min(rows) / hm, (max(cols) + 1) / wm, (max(rows) + 1) / hm]
        comps.append({"x": round(x, 4), "y": round(y, 4), "p": round(mass, 4), "bbox": [round(v, 4) for v in bbox]})
    comps.sort(key=lambda d: -d["p"])
    return comps[:max_candidates]