tools/localize.py

# cxr-agent/tools/localize.py
"""
tools/localize.py
=================
Spatial grounding for CXR findings via **NV-Locate-Anything-3B**.

The orchestrator's ``localize_findings`` tool calls :func:`run_localization`
with the image path and the reasoning report; this module turns the report's
finding labels into open-vocabulary *categories*, asks the LocateAnything-3B
service to locate them, and returns a typed :class:`LocalizationResult` whose
boxes are in **original-image pixel space**.

Transport
---------
LocateAnything-3B is deployed as a FastMCP server (``locate_objects``,
``ground_phrase``, ``detect_text``, ``ground_gui``, ``point_at``) in front of a
Modal A10G worker. The server resizes the image and scales boxes back to the
original resolution, so the caller gets pixel coordinates directly. We reach it
through the Agents SDK's own MCP client (``MCPServerStreamableHttp``) — the same
streamable-HTTP transport the reasoning server uses — rather than hand-rolling
JSON-RPC/SSE.

Design notes
------------
* **Flat module.** Kept as a single ``tools/localize.py`` (the project keeps
  ``tools/`` flat), matching ``tools/voice.py`` and ``tools/triage.py``.
* **Never fatal.** Any infra error (URL down, cold-start timeout, unparseable
  payload) degrades to an *empty* ``LocalizationResult`` and a log line — the
  ``@function_tool`` still returns a valid object so the agent can proceed (the
  overlay simply has nothing to draw). It does not raise.

Environment:
    CXR_LOCALIZE_MCP_URL     localization MCP URL (else built-in default)
    CXR_LOCALIZE_TIMEOUT     per-call read timeout, seconds (default 60)
    CXR_LOCALIZE_MAX_CATS    max categories, one inference each (default 8)
    CXR_LOCALIZE_GEN_MODE    generation mode: hybrid|fast|slow (default hybrid)
"""

from __future__ import annotations

import base64
import json
import logging
import os
from pathlib import Path
from typing import Any

from models.pipeline import (
    BoundingBox,
    LocalizationResult,
    LocalizedRegion,
    ReasoningReport,
)

logger = logging.getLogger(__name__)

DEFAULT_LOCALIZE_MCP_URL = "https://nv-locate-anything-3b.fastmcp.app/mcp"
DEFAULT_TIMEOUT = 60.0
DEFAULT_MAX_CATEGORIES = 8
DEFAULT_GEN_MODE = "hybrid"
LOCATE_TOOL = "locate_objects"


# ── Inputs: report → categories, image → base64 ──────────────────────────────

def _derive_categories(reasoning_report_json: str | dict | None,
                       limit: int) -> list[str]:
    """Pull unique finding labels from a reasoning report to localize.

    Tolerant of: a ``ReasoningReport`` JSON object, a bare ``{"findings": [...]}``
    dict, a bare list of findings, or a comma/newline-separated string. Returns a
    de-duplicated, order-preserving list of category phrases (capped at ``limit``).
    """
    findings: list = []
    data: Any = reasoning_report_json

    if isinstance(data, str):
        s = data.strip()
        if s.startswith("{") or s.startswith("["):
            try:
                data = json.loads(s)
            except Exception:  # noqa: BLE001
                data = s
        if isinstance(data, str):
            # Plain text → split on commas / newlines as a last resort.
            parts = [p.strip() for chunk in data.splitlines() for p in chunk.split(",")]
            return _dedup([p for p in parts if p])[:limit]

    if isinstance(data, dict):
        findings = data.get("findings") or []
    elif isinstance(data, list):
        findings = data

    labels: list[str] = []
    for f in findings:
        if isinstance(f, dict):
            label = (f.get("label") or "").strip()
        else:
            label = str(getattr(f, "label", "") or "").strip()
        if label:
            labels.append(label)
    return _dedup(labels)[:limit]


def _dedup(items: list[str]) -> list[str]:
    seen: set[str] = set()
    out: list[str] = []
    for it in items:
        key = it.lower()
        if key not in seen:
            seen.add(key)
            out.append(it)
    return out


def _encode_image_b64(image_path: str) -> tuple[str, str]:
    """Return (base64 of the original image bytes, image_id stem)."""
    p = Path(image_path)
    data = base64.b64encode(p.read_bytes()).decode("utf-8")
    return data, p.stem


# ── Output: tool JSON → LocalizationResult ───────────────────────────────────

def _as_xyxy(box: Any) -> tuple[float, float, float, float] | None:
    """Coerce a box in various shapes to (x1, y1, x2, y2)."""
    if isinstance(box, dict):
        if all(k in box for k in ("x1", "y1", "x2", "y2")):
            return float(box["x1"]), float(box["y1"]), float(box["x2"]), float(box["y2"])
        if all(k in box for k in ("x", "y", "w", "h")):
            x, y, w, h = (float(box["x"]), float(box["y"]),
                          float(box["w"]), float(box["h"]))
            return x, y, x + w, y + h
        # {"left":..,"top":..,"right":..,"bottom":..}
        if all(k in box for k in ("left", "top", "right", "bottom")):
            return (float(box["left"]), float(box["top"]),
                    float(box["right"]), float(box["bottom"]))
    if isinstance(box, (list, tuple)) and len(box) == 4:
        return float(box[0]), float(box[1]), float(box[2]), float(box[3])
    return None


def _iter_detections(payload: Any):
    """Yield (label, box, score) tuples from a variety of response shapes."""
    # Unwrap common envelopes.
    if isinstance(payload, dict):
        for key in ("detections", "boxes", "objects", "regions", "results"):
            if key in payload:
                inner = payload[key]
                # {"results": {category: [box, ...]}}
                if isinstance(inner, dict):
                    for label, boxes in inner.items():
                        if isinstance(boxes, list):
                            for b in boxes:
                                yield label, _box_of(b), _score_of(b)
                        else:
                            yield label, _box_of(boxes), _score_of(boxes)
                    return
                if isinstance(inner, list):
                    for d in inner:
                        yield _label_of(d), _box_of(d), _score_of(d)
                    return
    if isinstance(payload, list):
        for d in payload:
            yield _label_of(d), _box_of(d), _score_of(d)


def _label_of(d: Any) -> str:
    if isinstance(d, dict):
        return str(d.get("label") or d.get("category") or d.get("phrase")
                   or d.get("name") or "finding").strip()
    return "finding"


def _box_of(d: Any) -> Any:
    if isinstance(d, dict):
        return d.get("box") or d.get("bbox") or d.get("xyxy") or d
    return d


def _score_of(d: Any) -> float:
    if isinstance(d, dict):
        for k in ("score", "confidence", "conf"):
            if k in d:
                try:
                    return max(0.0, min(1.0, float(d[k])))
                except (TypeError, ValueError):
                    pass
    return 0.5  # neutral default; LocalizedRegion.score is ge=0,le=1


def _to_regions(payload: Any) -> list[LocalizedRegion]:
    regions: list[LocalizedRegion] = []
    for label, box, score in _iter_detections(payload):
        xyxy = _as_xyxy(box)
        if xyxy is None:
            continue
        x1, y1, x2, y2 = xyxy
        # Normalize ordering and clamp to non-negative width/height.
        left, right = sorted((x1, x2))
        top, bottom = sorted((y1, y2))
        bbox = BoundingBox(
            x=int(round(left)),
            y=int(round(top)),
            w=max(0, int(round(right - left))),
            h=max(0, int(round(bottom - top))),
        )
        regions.append(LocalizedRegion(finding=label or "finding",
                                       bbox=bbox, score=score))
    return regions


def _extract_tool_text(result: Any) -> str:
    """Concatenate text from an MCP CallToolResult's content blocks."""
    content = getattr(result, "content", None) or []
    out: list[str] = []
    for block in content:
        text = getattr(block, "text", None)
        if text:
            out.append(text)
    return "\n".join(out).strip()


# ── Public entrypoint ────────────────────────────────────────────────────────

async def run_localization(
    image_path: str,
    reasoning_report_json: str,
    *,
    categories: list[str] | None = None,
) -> LocalizationResult:
    """
    Localize the findings from a reasoning report onto the CXR image.

    Each finding category is grounded in its **own** ``locate_objects`` call so
    every returned box carries the correct finding label. (A single multi-category
    call returns a flat, *unlabeled* box list — the category↔box association lives
    only in the raw model text — which would force every box to a generic
    "finding" label.) Calls share one MCP connection and run sequentially.

    Args:
        image_path: Path to the CXR image (JPEG/PNG).
        reasoning_report_json: JSON-serialised ``ReasoningReport`` (or any shape
            ``_derive_categories`` understands).
        categories: Optional explicit category override. When omitted, categories
            are derived from the report's finding labels.

    Returns:
        ``LocalizationResult`` with one region per located box, labeled by finding.
        ``score`` is None — LocateAnything-3B does not emit a detection confidence.
        Empty on infrastructure error (logged, never raised).
    """
    url = os.environ.get("CXR_LOCALIZE_MCP_URL", DEFAULT_LOCALIZE_MCP_URL)
    try:
        timeout = float(os.environ.get("CXR_LOCALIZE_TIMEOUT", DEFAULT_TIMEOUT))
    except ValueError:
        timeout = DEFAULT_TIMEOUT
    try:
        max_cats = int(os.environ.get("CXR_LOCALIZE_MAX_CATS", DEFAULT_MAX_CATEGORIES))
    except ValueError:
        max_cats = DEFAULT_MAX_CATEGORIES
    gen_mode = os.environ.get("CXR_LOCALIZE_GEN_MODE", DEFAULT_GEN_MODE)

    cats = categories if categories else _derive_categories(reasoning_report_json, max_cats)
    cats = _dedup([c for c in (cats or []) if c.strip()])[:max_cats]
    if not cats:
        logger.info("localization: no categories derived from report; nothing to locate")
        return LocalizationResult(regions=[])

    try:
        image_b64, image_id = _encode_image_b64(image_path)
    except Exception as e:  # noqa: BLE001
        logger.warning("localization: could not read image %s — %s", image_path, e)
        return LocalizationResult(regions=[])

    # Reach the service through the SDK's MCP client (lazy import keeps this
    # module importable without the SDK present, e.g. for unit tests).
    try:
        from agents.mcp import MCPServerStreamableHttp
    except Exception as e:  # noqa: BLE001
        logger.warning("localization: Agents SDK MCP client unavailable — %s", e)
        return LocalizationResult(regions=[])

    server = MCPServerStreamableHttp(
        name="cxr-localization",
        params={"url": url},
        cache_tools_list=True,
        client_session_timeout_seconds=timeout,
        max_retry_attempts=2,
    )

    regions: list[LocalizedRegion] = []
    try:
        await server.connect()
        try:
            for cat in cats:
                try:
                    result = await server.call_tool(
                        LOCATE_TOOL,
                        {
                            "image_b64": image_b64,
                            "image_id": image_id,
                            "categories": [cat],   # one category → labeled boxes
                            "generation_mode": gen_mode,
                        },
                    )
                except Exception as e:  # noqa: BLE001 — skip this finding, keep going
                    logger.warning("localization: '%s' call failed — %s", cat, e)
                    continue
                regions.extend(_regions_for_category(result, cat))
        finally:
            await server.cleanup()
    except Exception as e:  # noqa: BLE001
        logger.warning("localization: MCP connection failed (%s) — %s", url, e)
        return LocalizationResult(regions=[])

    logger.info("localization: %d region(s) across %d categor(y/ies)", len(regions), len(cats))
    return LocalizationResult(regions=regions)


def _regions_for_category(result: Any, category: str) -> list[LocalizedRegion]:
    """Parse one ``locate_objects`` result and label every box with ``category``."""
    if getattr(result, "isError", False) or getattr(result, "is_error", False):
        return []
    text = _extract_tool_text(result)
    if not text:
        return []
    try:
        payload = json.loads(text)
    except Exception:  # noqa: BLE001
        return []
    if isinstance(payload, dict) and payload.get("success") is False:
        return []

    # Known shape: {"boxes": [{x1,y1,x2,y2}, ...], ...}. Fall back to the tolerant
    # multi-shape walker for anything else.
    raw_boxes = payload.get("boxes") if isinstance(payload, dict) else None
    out: list[LocalizedRegion] = []
    if isinstance(raw_boxes, list):
        for b in raw_boxes:
            bbox = _bbox_from(b)
            if bbox is not None:
                out.append(LocalizedRegion(finding=category, bbox=bbox, score=None))
        return out
    # Fallback: tolerate other envelopes; force the known category as the label.
    for _label, box, _score in _iter_detections(payload):
        bbox = _bbox_from(box)
        if bbox is not None:
            out.append(LocalizedRegion(finding=category, bbox=bbox, score=None))
    return out


def _bbox_from(box: Any) -> "BoundingBox | None":
    """Coerce a box of any supported shape to a pixel-space ``BoundingBox``."""
    xyxy = _as_xyxy(box)
    if xyxy is None:
        return None
    x1, y1, x2, y2 = xyxy
    left, right = sorted((x1, x2))
    top, bottom = sorted((y1, y2))
    return BoundingBox(
        x=int(round(left)),
        y=int(round(top)),
        w=max(0, int(round(right - left))),
        h=max(0, int(round(bottom - top))),
    )