Spaces:
Paused
Paused
| """ | |
| 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"] | |
| class BoundingBox: | |
| x_min: float | |
| y_min: float | |
| x_max: float | |
| y_max: float | |
| 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) | |