Blame for backend/models/grounding.py

"""
Text-guided region grounding specialist model.

Phase 6 (P0). Satisfies PS requirement:
- "Each solution must additionally implement either captioning/scene
   description OR text-guided region grounding." — we chose grounding
   (see README §4 rationale: stronger map/demo fit, produces the
   bounding-box evidence the PS explicitly calls for: "Highlight the
   water body referred to in the query.")

MODEL CHOICE — google/owlvit-base-patch32
- Verified supports: open-vocabulary object detection. Give it an image
  + arbitrary text phrase(s), it returns bounding boxes + a match score
  per phrase. This is exactly "text-guided region grounding" — no
  detection-head retraining needed to add new object classes.
- Modality: input = 1 RGB image + 1-N text phrases. Output = boxes
  (pixel coords) + confidence score per phrase.
- Integration complexity: LOW. Pure `transformers`
  (OwlViTProcessor/OwlViTForObjectDetection), same dependency footprint
  as Phase 5's VQA model — no new packages required.
- License: Apache 2.0 (OWL-ViT / Google). Free for hackathon/research use.
- NOT remote-sensing-native. It was trained on natural-image/web data,
  so expect it to do well on generic categories ("water", "building",
  "road", "vehicle", "vegetation") and poorly on RS-specific phrasing
  ("built-up area", "paddy field") until adapted. This model is a
  strong Phase-6 candidate for the SAME LoRA-style domain adaptation
  approach as Phase 5's VQA model — not yet built; would be a Phase 6.5
  follow-up if time allows, not required for the PS baseline.
- Confidence returned is the detector's own softmax match score for the
  best box, NOT a validated grounding-accuracy benchmark figure.
"""
from __future__ import annotations

import re
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from preprocessing.raster_to_rgb import load_as_rgb

MODEL_ID = "google/owlvit-base-patch32"

# Same confidence-bucket policy as Phase 5 (master brief §10) — kept
# consistent across specialist models so the UI doesn't need per-task
# thresholds.
_HIGH_THRESHOLD = 0.30   # OWL-ViT match scores run lower than VQA token
_MEDIUM_THRESHOLD = 0.12  # probs by nature of open-vocab detection; these
                           # are a prototype policy, not a validated cutoff.

# Rule-based phrase extraction — strips grounding-intent scaffolding off
# the user's natural-language query so OWL-ViT gets a clean noun phrase.
# Deterministic on purpose, consistent with agent/router.py's approach.
_STRIP_PREFIXES = [
    "where is the", "where is", "highlight the", "highlight", "locate the",
    "locate", "point to the", "point to", "show me the", "show me",
    "draw a box around the", "draw a box around", "mark the", "mark",
    "find the", "find",
]
_STRIP_SUFFIXES = ["referred to in the query", "in this image", "in the image"]


@dataclass
class BoundingBox:
    x_min: float
    y_min: float
    x_max: float
    y_max: float


@dataclass
class GroundingResult:
    phrase: str
    found: bool
    box: Optional[BoundingBox]
    score: float
    confidence_level: str  # "high" | "medium" | "insufficient"
    confidence_note: str
    image_width: int
    image_height: int
    model_id: str
    adapted: bool


def extract_referring_phrase(query: str) -> str:
    """Turn a grounding question into a short noun phrase for the detector.

    e.g. "Where is the water body?" -> "water body"
         "Highlight the flooded area." -> "flooded area"

    Purely rule-based (strip known prefixes/suffixes + punctuation) —
    no LLM call, consistent with the router's deterministic classification.
    If nothing matches, falls back to the raw query so grounding still
    attempts something rather than failing silently.
    """
    q = query.strip().rstrip("?.!").lower()

    for prefix in sorted(_STRIP_PREFIXES, key=len, reverse=True):
        if q.startswith(prefix):
            q = q[len(prefix):].strip()
            break

    for suffix in _STRIP_SUFFIXES:
        if q.endswith(suffix):
            q = q[: -len(suffix)].strip()

    q = re.sub(r"^(the|a|an)\s+", "", q).strip()
    return q or query.strip().rstrip("?.!")


class _GroundingModel:
    """Lazy-loaded singleton, mirrors models/vqa.py's pattern."""

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._processor = None
        self._model = None
        self._device = "cpu"
        self._adapted = False  # no grounding adapter exists yet — see docstring

    def _ensure_loaded(self) -> None:
        if self._model is not None:
            return
        with self._lock:
            if self._model is not None:
                return

            import torch
            from transformers import OwlViTForObjectDetection, OwlViTProcessor

            self._device = "cuda" if torch.cuda.is_available() else "cpu"
            self._processor = OwlViTProcessor.from_pretrained(MODEL_ID)
            self._model = OwlViTForObjectDetection.from_pretrained(MODEL_ID).to(self._device).eval()

    def ground(self, image_path: Path, phrase: str) -> GroundingResult:
        self._ensure_loaded()
        import torch

        image = load_as_rgb(image_path)
        width, height = image.size

        inputs = self._processor(text=[[phrase]], images=image, return_tensors="pt").to(self._device)
        with torch.no_grad():
            outputs = self._model(**inputs)

        target_sizes = torch.tensor([[height, width]])
        results = self._processor.post_process_grounded_object_detection(
            outputs=outputs, target_sizes=target_sizes, threshold=0.0
        )[0]

        boxes, scores = results["boxes"], results["scores"]
        if len(scores) == 0:
            level, note = "insufficient", "Detector returned no candidate region at all for this phrase."
            return GroundingResult(
                phrase=phrase, found=False, box=None, score=0.0,
                confidence_level=level, confidence_note=note,
                image_width=width, image_height=height,
                model_id=self._model_id_label(), adapted=self._adapted,
            )

        best_idx = int(torch.argmax(scores))
        best_score = float(scores[best_idx])
        x_min, y_min, x_max, y_max = (float(v) for v in boxes[best_idx])
        level, note = _bucket_confidence(best_score)

        return GroundingResult(
            phrase=phrase,
            found=True,
            box=BoundingBox(x_min=x_min, y_min=y_min, x_max=x_max, y_max=y_max),
            score=best_score,
            confidence_level=level,
            confidence_note=note,
            image_width=width,
            image_height=height,
            model_id=self._model_id_label(),
            adapted=self._adapted,
        )

    def _model_id_label(self) -> str:
        return MODEL_ID + (" + RS LoRA adapter" if self._adapted else " (base, not RS-adapted)")


def _bucket_confidence(score: float) -> tuple[str, str]:
    if score >= _HIGH_THRESHOLD:
        return "high", "Detector was strongly confident this region matches the phrase."
    if score >= _MEDIUM_THRESHOLD:
        return "medium", "Detector found a plausible region but with moderate uncertainty — verify visually."
    return "insufficient", "Best candidate region had low match confidence — treat as unreliable."


_singleton: Optional[_GroundingModel] = None
_singleton_lock = threading.Lock()


def get_grounding_model() -> _GroundingModel:
    global _singleton
    if _singleton is None:
        with _singleton_lock:
            if _singleton is None:
                _singleton = _GroundingModel()
    return _singleton


def run_grounding(image_path: Path, query: str) -> GroundingResult:
    phrase = extract_referring_phrase(query)
    return get_grounding_model().ground(image_path, phrase)