rpx_benchmark.vqa.outputs

  1"""Strict parsing and normalization of VQA model outputs."""
  2
  3from __future__ import annotations
  4
  5import json
  6import re
  7import unicodedata
  8from dataclasses import dataclass
  9
 10from ..exceptions import AdapterError, RPXError
 11from .contract import ATTRIBUTE_TYPES, BBOX_TYPES, BINARY_TYPES, VQASample
 12
 13_LOC_RE = re.compile(
 14    r"<loc(?P<y0>\d{4})><loc(?P<x0>\d{4})><loc(?P<y1>\d{4})><loc(?P<x1>\d{4})>"
 15)
 16_MOLMO_POINT_TAG_RE = re.compile(r"<(?:point|points)\b(?P<attrs>[^>]*)>", re.IGNORECASE)
 17_MOLMO_POINT_ATTR_RE = re.compile(
 18    r"\b(?P<axis>[xy])(?P<index>\d*)\s*=\s*[\"']"
 19    r"(?P<value>[+-]?(?:\d+(?:\.\d*)?|\.\d+))[\"']",
 20    re.IGNORECASE,
 21)
 22_NUMBER_RE = r"[+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?"
 23_INTERNVL_BBOX_RE = re.compile(
 24    rf"(?<!\[)\[\[?\s*({_NUMBER_RE})\s*,\s*({_NUMBER_RE})\s*,\s*"
 25    rf"({_NUMBER_RE})\s*,\s*({_NUMBER_RE})\s*\]\]?(?!\])"
 26)
 27
 28
 29def normalize_label(value: str) -> str:
 30    value = unicodedata.normalize("NFKC", value).strip().lower().replace("_", " ")
 31    value = re.sub(r"[^\w\s-]", "", value)
 32    return " ".join(value.split())
 33
 34
 35@dataclass(frozen=True)
 36class ParsedOutput:
 37    valid: bool
 38    label: str | None = None
 39    bbox: tuple[float, float, float, float] | None = None
 40    error: str | None = None
 41    coordinate_format: str | None = None
 42
 43
 44@dataclass(frozen=True)
 45class ParsedPoints:
 46    """A native Molmo pointing response, kept separate from bbox parsing."""
 47
 48    valid: bool
 49    points: tuple[tuple[float, float], ...] = ()
 50    error: str | None = None
 51    coordinate_format: str = "molmo_point_percent_0_100"
 52
 53
 54def parse_molmo_points(raw: str) -> ParsedPoints:
 55    """Parse Molmo's documented ``<point>``/``<points>`` markup.
 56
 57    Coordinates remain percentages on Molmo's native 0--100 grid.  This
 58    function intentionally cannot return a bbox: a point carries no object
 59    extent and must not enter the benchmark's bbox-IoU metric.
 60    """
 61    points: list[tuple[float, float]] = []
 62    errors: list[str] = []
 63    tags = list(_MOLMO_POINT_TAG_RE.finditer(raw))
 64    if not tags:
 65        return ParsedPoints(False, error="missing Molmo <point> or <points> tag")
 66    for tag in tags:
 67        by_index: dict[str, dict[str, float]] = {}
 68        for match in _MOLMO_POINT_ATTR_RE.finditer(tag.group("attrs")):
 69            index = match.group("index") or "1"
 70            axis = match.group("axis").lower()
 71            slot = by_index.setdefault(index, {})
 72            if axis in slot:
 73                errors.append(f"duplicate {axis}{index} coordinate")
 74                continue
 75            slot[axis] = float(match.group("value"))
 76        if not by_index:
 77            errors.append("point tag contains no coordinate attributes")
 78            continue
 79        for index, coordinates in sorted(
 80            by_index.items(), key=lambda item: int(item[0])
 81        ):
 82            if set(coordinates) != {"x", "y"}:
 83                errors.append(f"point {index} is missing x or y")
 84                continue
 85            x, y = coordinates["x"], coordinates["y"]
 86            if not (0 <= x <= 100 and 0 <= y <= 100):
 87                errors.append(f"point {index} is outside the 0-100 grid")
 88                continue
 89            points.append((x, y))
 90    if errors:
 91        return ParsedPoints(False, points=tuple(points), error="; ".join(errors))
 92    if not points:
 93        return ParsedPoints(False, error="Molmo point markup contains no complete point")
 94    return ParsedPoints(True, points=tuple(points))
 95
 96
 97def _json_object(raw: str) -> dict:
 98    text = raw.strip()
 99    if text.startswith("```"):
100        text = re.sub(r"^```(?:json)?\s*|\s*```$", "", text, flags=re.IGNORECASE)
101    start, end = text.find("{"), text.rfind("}")
102    if start < 0 or end < start:
103        raise AdapterError("missing JSON object", hint="return the requested bbox JSON only")
104    value = json.loads(text[start : end + 1])
105    if not isinstance(value, dict):
106        raise AdapterError("output is not a JSON object", hint="return one JSON object")
107    return value
108
109
110def _cosmos_bbox_candidates(raw: str) -> list[dict]:
111    """Extract complete native Reason2 ``bbox_2d`` objects from JSON values.
112
113    Reason2 may wrap its answer in a JSON code fence or an ``<answer>`` block.
114    Decode real JSON values instead of using a bracket regex, and retain every
115    candidate so the one-object benchmark can reject ambiguous multi-box output.
116    """
117    text = raw.strip()
118    answer = re.search(r"<answer>(.*?)</answer>", text, flags=re.IGNORECASE | re.DOTALL)
119    if answer:
120        text = answer.group(1).strip()
121    fenced = re.search(r"```(?:json)?\s*(.*?)\s*```", text, flags=re.IGNORECASE | re.DOTALL)
122    if fenced:
123        text = fenced.group(1).strip()
124
125    decoder = json.JSONDecoder()
126    candidates: list[dict] = []
127    offset = 0
128    while offset < len(text):
129        starts = [
130            index
131            for index in (text.find("[", offset), text.find("{", offset))
132            if index >= 0
133        ]
134        if not starts:
135            break
136        start = min(starts)
137        try:
138            value, consumed = decoder.raw_decode(text[start:])
139        except json.JSONDecodeError:
140            offset = start + 1
141            continue
142        values = value if isinstance(value, list) else [value]
143        candidates.extend(
144            item for item in values if isinstance(item, dict) and "bbox_2d" in item
145        )
146        offset = start + max(consumed, 1)
147    return candidates
148
149
150def parse_output(
151    sample: VQASample,
152    raw: str,
153    model_key: str,
154    *,
155    paligemma_target_only: bool = False,
156) -> ParsedOutput:
157    label: str | None
158    match: re.Match[str] | None
159    if sample.question_type in BINARY_TYPES:
160        token = normalize_label(raw)
161        if token not in {"yes", "no"}:
162            return ParsedOutput(False, error="expected exactly yes or no")
163        return ParsedOutput(True, label=token)
164    if sample.question_type in ATTRIBUTE_TYPES:
165        label = normalize_label(raw)
166        if not label or len(label.split()) > 8:
167            return ParsedOutput(False, error="invalid object-name answer")
168        return ParsedOutput(True, label=label)
169    if sample.question_type in BBOX_TYPES and model_key.startswith("cosmos-reason2-"):
170        candidates = _cosmos_bbox_candidates(raw)
171        if len(candidates) != 1:
172            return ParsedOutput(
173                False,
174                error=f"expected exactly one Cosmos bbox_2d candidate, got {len(candidates)}",
175            )
176        value = candidates[0]
177        try:
178            raw_bbox = value["bbox_2d"]
179            if not isinstance(raw_bbox, list) or len(raw_bbox) != 4:
180                raise AdapterError(
181                    "Cosmos bbox_2d must contain four numbers",
182                    hint="return one normalized XYXY box",
183                )
184            bbox = tuple(float(coordinate) for coordinate in raw_bbox)
185            if not all(float("-inf") < coordinate < float("inf") for coordinate in bbox):
186                raise AdapterError("Cosmos bbox contains a non-finite number")
187            if not all(0 <= coordinate <= 1000 for coordinate in bbox):
188                raise AdapterError("Cosmos bbox is outside the 0-1000 grid")
189            x0, y0, x1, y1 = bbox
190            if not (x0 <= x1 and y0 <= y1):
191                raise AdapterError("Cosmos bbox corners are not ordered XYXY")
192            scaled = (
193                x0 * (sample.img_w - 1) / 1000,
194                y0 * (sample.img_h - 1) / 1000,
195                x1 * (sample.img_w - 1) / 1000,
196                y1 * (sample.img_h - 1) / 1000,
197            )
198            label = normalize_label(str(value.get("label", ""))) or None
199            return ParsedOutput(
200                True,
201                label=label,
202                bbox=scaled,
203                coordinate_format="cosmos_bbox_2d_0_1000",
204            )
205        except (KeyError, TypeError, ValueError, RPXError) as exc:
206            return ParsedOutput(False, error=str(exc))
207    if sample.question_type in BBOX_TYPES and model_key.startswith("internvl3.5-"):
208        # InternVL's official RefCOCO evaluation protocol serializes one
209        # target-relative 0--1000 box.  Its own evaluator deliberately accepts
210        # either ``label[[x0,y0,x1,y1]]`` or ``[x0,y0,x1,y1]``.  Accept those
211        # two serializations before the generic JSON fallback, but never choose
212        # silently among multiple candidates in this one-object benchmark.
213        matches = list(_INTERNVL_BBOX_RE.finditer(raw))
214        if matches:
215            if len(matches) != 1:
216                return ParsedOutput(False, error="expected exactly one InternVL bbox")
217            match = matches[0]
218            bbox = tuple(float(value) for value in match.groups())
219            if not all(0 <= value <= 1000 for value in bbox):
220                return ParsedOutput(False, error="InternVL bbox is outside the 0-1000 grid")
221            x0, y0, x1, y1 = bbox
222            if not (x0 <= x1 and y0 <= y1):
223                return ParsedOutput(False, error="InternVL bbox corners are not ordered XYXY")
224            prefix = raw[: match.start()]
225            prefix = re.sub(r"</?(?:ref|box)>", "", prefix, flags=re.IGNORECASE)
226            label = normalize_label(prefix.splitlines()[-1]) or None
227            return ParsedOutput(
228                True,
229                label=label,
230                bbox=(
231                    x0 * (sample.img_w - 1) / 1000,
232                    y0 * (sample.img_h - 1) / 1000,
233                    x1 * (sample.img_w - 1) / 1000,
234                    y1 * (sample.img_h - 1) / 1000,
235                ),
236                coordinate_format="internvl_bbox_0_1000",
237            )
238        # The 1B checkpoint sometimes preserves the same 0--1000 coordinate
239        # contract but spells the single box out as labelled corners.  This is
240        # a deterministic serialization variant, not box repair: require one
241        # complete pair, validate it identically, and never infer coordinates
242        # from descriptive prose.
243        verbose_matches = re.findall(
244            r"top[\s-]*left\s+corner\s*:\s*[\[(]\s*"
245            r"(-?\d+(?:\.\d+)?)\s*,\s*(-?\d+(?:\.\d+)?)\s*[\])]"
246            r".*?bottom[\s-]*right\s+corner\s*:\s*[\[(]\s*"
247            r"(-?\d+(?:\.\d+)?)\s*,\s*(-?\d+(?:\.\d+)?)\s*[\])]",
248            raw,
249            flags=re.IGNORECASE | re.DOTALL,
250        )
251        if verbose_matches:
252            if len(verbose_matches) != 1:
253                return ParsedOutput(False, error="expected exactly one InternVL corner pair")
254            bbox = tuple(float(value) for value in verbose_matches[0])
255            if not all(0 <= value <= 1000 for value in bbox):
256                return ParsedOutput(False, error="InternVL bbox is outside the 0-1000 grid")
257            x0, y0, x1, y1 = bbox
258            if not (x0 <= x1 and y0 <= y1):
259                return ParsedOutput(False, error="InternVL bbox corners are not ordered XYXY")
260            return ParsedOutput(
261                True,
262                bbox=(
263                    x0 * (sample.img_w - 1) / 1000,
264                    y0 * (sample.img_h - 1) / 1000,
265                    x1 * (sample.img_w - 1) / 1000,
266                    y1 * (sample.img_h - 1) / 1000,
267                ),
268                coordinate_format="internvl_verbose_bbox_0_1000",
269            )
270        # Retain the strict JSON fallback for a checkpoint that follows the
271        # explicit coordinate contract but emits the generic RPX envelope.
272    if sample.question_type in BBOX_TYPES and model_key.startswith("paligemma2-"):
273        match = _LOC_RE.search(raw)
274        if match:
275            if sample.is_in_context and not paligemma_target_only:
276                return ParsedOutput(
277                    False,
278                    error=(
279                        "native PaliGemma loc tokens are relative to the labelled "
280                        "two-image composite, not Image 2; return target-relative JSON"
281                    ),
282                )
283            vals = {key: int(value) for key, value in match.groupdict().items()}
284            bbox = (
285                vals["x0"] * sample.img_w / 1024,
286                vals["y0"] * sample.img_h / 1024,
287                vals["x1"] * sample.img_w / 1024,
288                vals["y1"] * sample.img_h / 1024,
289            )
290            label = normalize_label(raw[match.end() :].split("<", 1)[0]) or None
291            return ParsedOutput(
292                True, label=label, bbox=bbox, coordinate_format="paligemma_loc_1024"
293            )
294        # The scored single-call prompt requests the same JSON contract used
295        # by the other VLMs. Fall through to that parser when PaliGemma obeys
296        # it; native <loc> remains accepted as an architecture-specific bbox
297        # serialization, but no second model call is made.
298    if sample.question_type in BBOX_TYPES:
299        try:
300            value = _json_object(raw)
301            raw_bbox = value["bbox"]
302            # Some otherwise-correct VLM responses add one redundant list
303            # dimension.  Unwrap exactly one singleton; do not guess any more
304            # complicated structure.
305            if (
306                isinstance(raw_bbox, list)
307                and len(raw_bbox) == 1
308                and isinstance(raw_bbox[0], list)
309            ):
310                raw_bbox = raw_bbox[0]
311            if not isinstance(raw_bbox, list) or len(raw_bbox) != 4:
312                raise AdapterError(
313                    "bbox must contain four numbers", hint="return normalized xyxy coordinates"
314                )
315            bbox = tuple(float(v) for v in raw_bbox)
316            if not all(float("-inf") < value < float("inf") for value in bbox):
317                raise AdapterError("bbox contains a non-finite number", hint="return finite pixels")
318            # Molmo's model-facing bbox protocol uses the same 0--100
319            # percentage grid as its documented spatial pointing interface.
320            # This is a complete four-corner box, not a point-derived box.
321            # Convert it deterministically here while retaining the untouched
322            # model text in raw_output for auditability.
323            molmo_percent = model_key in {"molmo-7b-d", "molmoe-1b"}
324            if molmo_percent and not all(0 <= coordinate <= 100 for coordinate in bbox):
325                raise AdapterError(
326                    "Molmo bbox is outside percentage 0-100 coordinates",
327                    hint="return percentage xyxy coordinates",
328                )
329
330            # JSON produced by an adapter may contain a machine-rounding
331            # residue at an exact normalized endpoint (for example
332            # 999 * 1000 / 999 -> 1000.0000000000001). Accept only that tiny
333            # numerical residue, then clamp it; genuinely out-of-range model
334            # coordinates remain invalid.
335            coordinate_epsilon = 1e-9
336            if not all(
337                -coordinate_epsilon <= coordinate <= 1000 + coordinate_epsilon
338                for coordinate in bbox
339            ):
340                raise AdapterError(
341                    "bbox is outside normalized 0-1000 coordinates",
342                    hint="return normalized xyxy coordinates",
343                )
344            bbox = tuple(min(1000.0, max(0.0, coordinate)) for coordinate in bbox)
345            # Decode the two widespread normalized conventions without using
346            # model identity. Values containing a genuine fraction and entirely
347            # within [0,1] are unit-normalized; everything else follows the
348            # requested [0,1000] protocol. This fixes the previous silent bug
349            # where [0.5,...] became a sub-pixel box at the top-left.
350            unit_normalized = not molmo_percent and all(
351                0 <= coordinate <= 1 for coordinate in bbox
352            )
353            denominator = 100 if molmo_percent else 1 if unit_normalized else 1000
354            coordinate_format = (
355                "molmo_bbox_percent_0_100"
356                if molmo_percent
357                else "normalized_0_1"
358                if unit_normalized
359                else "normalized_0_1000"
360            )
361            bbox = (
362                bbox[0] * (sample.img_w - 1) / denominator,
363                bbox[1] * (sample.img_h - 1) / denominator,
364                bbox[2] * (sample.img_w - 1) / denominator,
365                bbox[3] * (sample.img_h - 1) / denominator,
366            )
367            x0, y0, x1, y1 = bbox
368            if not (0 <= x0 <= x1 < sample.img_w and 0 <= y0 <= y1 < sample.img_h):
369                raise AdapterError(
370                    "bbox is outside the original image", hint="undo processor resizing first"
371                )
372            label = normalize_label(str(value.get("label", ""))) or None
373            return ParsedOutput(
374                True, label=label, bbox=bbox, coordinate_format=coordinate_format
375            )
376        except (KeyError, TypeError, ValueError, json.JSONDecodeError, RPXError) as exc:
377            return ParsedOutput(False, error=str(exc))
378    return ParsedOutput(False, error="unsupported task")