grounding_dino2.py

# grounding_dino2.py
# Lightweight Grounding DINO wrapper for box detection + cropping + visualization.
from __future__ import annotations

import os
import threading
from pathlib import Path
from typing import List, Dict, Any, Tuple, Optional

import torch
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection

# ---- Writable caches (HF Spaces / containers) ----
CACHE_DIR = os.getenv("HF_CACHE_DIR", "/tmp/hf-cache")
Path(CACHE_DIR).mkdir(parents=True, exist_ok=True)
os.environ.setdefault("HOME", "/tmp")
os.environ.setdefault("XDG_CACHE_HOME", CACHE_DIR)
os.environ.setdefault("HF_HOME", CACHE_DIR)
os.environ.setdefault("HUGGINGFACE_HUB_CACHE", CACHE_DIR)
os.environ.setdefault("TRANSFORMERS_CACHE", CACHE_DIR)
os.environ.setdefault("HF_DATASETS_CACHE", f"{CACHE_DIR}/datasets")
os.environ.setdefault("TORCH_HOME", CACHE_DIR)
os.environ.setdefault("PYTHONPYCACHEPREFIX", "/tmp/pycache")


def _clamp_xyxy(box: List[float], w: int, h: int) -> Tuple[int, int, int, int]:
    x0, y0, x1, y1 = box
    x0 = max(0, min(int(round(x0)), w - 1))
    y0 = max(0, min(int(round(y0)), h - 1))
    x1 = max(0, min(int(round(x1)), w - 1))
    y1 = max(0, min(int(round(y1)), h - 1))
    if x1 < x0:
        x0, x1 = x1, x0
    if y1 < y0:
        y0, y1 = y1, y0
    return x0, y0, x1, y1


def _pad_box(box: Tuple[int, int, int, int], w: int, h: int, frac: float = 0.06) -> Tuple[int, int, int, int]:
    x0, y0, x1, y1 = box
    bw, bh = x1 - x0, y1 - y0
    dx, dy = int(bw * frac), int(bh * frac)
    return max(0, x0 - dx), max(0, y0 - dy), min(w - 1, x1 + dx), min(h - 1, y1 + dy)


def crop_from_box(img: Image.Image, box_xyxy: Tuple[int, int, int, int]) -> Image.Image:
    return img.crop(box_xyxy)


def _parse_to_flat_labels(labels: List[str] | str) -> List[str]:
    """
    Accepts a comma-separated string or a list of strings and returns a flat list of non-empty labels.
    """
    if isinstance(labels, str):
        items = [x.strip() for x in labels.split(",") if x.strip()]
    else:
        items = [str(x).strip() for x in labels if str(x).strip()]
    if not items:
        raise ValueError("No labels provided.")
    return items


def _build_dot_separated_prompt(items: List[str]) -> str:
    """
    Builds the recommended GroundingDINO text prompt: "a man . a dog ."
    """
    return " . ".join(items) + " ."


class GroundingDINORunner:
    """
    Minimal singleton-style wrapper for Grounding DINO zero-shot detector.
    """

    def __init__(self, model_id: Optional[str] = None, device: Optional[str] = None):
        self.model_id = model_id or os.getenv("GDINO_MODEL_ID", "IDEA-Research/grounding-dino-tiny")
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self._lock = threading.Lock()

        self.processor = AutoProcessor.from_pretrained(self.model_id, cache_dir=CACHE_DIR)
        self.model = AutoModelForZeroShotObjectDetection.from_pretrained(
            self.model_id, cache_dir=CACHE_DIR
        ).to(self.device)
        self.model.eval()

    def detect(
        self,
        image: Image.Image,
        labels: List[str] | str,
        box_threshold: float = 0.4,
        text_threshold: float = 0.3,
        pad_frac: float = 0.06,
    ) -> List[Dict[str, Any]]:
        """
        Runs zero-shot detection and returns:
          [{ 'label': str, 'score': float, 'box_xyxy': (x0,y0,x1,y1), 'crop': PIL.Image }, ...]
        """
        w, h = image.size

        # ---- FIX: use dot-separated string or flat list; avoid nested lists ----
        items = _parse_to_flat_labels(labels)
        text_prompt = _build_dot_separated_prompt(items)  # "a man . a dog ."

        # Prepare inputs
        inputs = self.processor(images=image, text=text_prompt, return_tensors="pt").to(self.device)

        # Inference
        with self._lock, torch.no_grad():
            outputs = self.model(**inputs)

            # transformers>=4.51 uses "threshold", older expects "box_threshold"
            try:
                post = self.processor.post_process_grounded_object_detection(
                    outputs=outputs,
                    input_ids=inputs.input_ids,
                    threshold=float(box_threshold),
                    text_threshold=float(text_threshold),
                    target_sizes=[(h, w)],
                )
            except TypeError:
                post = self.processor.post_process_grounded_object_detection(
                    outputs=outputs,
                    input_ids=inputs.input_ids,
                    box_threshold=float(box_threshold),
                    text_threshold=float(text_threshold),
                    target_sizes=[(h, w)],
                )

        det = post[0]
        boxes = det.get("boxes", [])
        scores = det.get("scores", [])
        # Newer transformers populate "text_labels"; else "labels"
        labels_out = det.get("text_labels", det.get("labels", []))

        results: List[Dict[str, Any]] = []
        for b, s, lab in zip(boxes, scores, labels_out):
            b = b.tolist() if hasattr(b, "tolist") else list(b)
            bx = _clamp_xyxy(b, w, h)
            bx = _pad_box(bx, w, h, pad_frac)
            crop = crop_from_box(image, bx)
            score = float(s.item()) if torch.is_tensor(s) else float(s)
            results.append({"label": lab, "score": score, "box_xyxy": bx, "crop": crop})

        return results


# --- Visualization helper ------------------------------------------------------
def visualize_detections(
    image: Image.Image,
    detections: list[dict],
    *,
    box_color: tuple[int, int, int] = (0, 255, 0),
    text_color: tuple[int, int, int] = (0, 0, 0),
    box_width: int = 3,
) -> Image.Image:
    """
    Draw boxes + labels on a copy of `image`.
    Each detection item expects: {'label': str, 'score': float, 'box_xyxy': (x0,y0,x1,y1)}
    """
    vis = image.copy()
    draw = ImageDraw.Draw(vis)
    try:
        font = ImageFont.truetype("DejaVuSans.ttf", 16)
    except Exception:
        font = None

    for det in detections:
        x0, y0, x1, y1 = det["box_xyxy"]
        lab = det.get("label", "")
        sc = det.get("score", 0.0)
        draw.rectangle((x0, y0, x1, y1), outline=box_color, width=box_width)
        text = f"{lab} {sc:.2f}"
        # textlength fallback
        try:
            text_w = draw.textlength(text, font=font)  # type: ignore[attr-defined]
        except Exception:
            text_w = len(text) * 8
        pad = 4
        draw.rectangle((x0, max(0, y0 - 20), x0 + int(text_w) + pad * 2, y0), fill=box_color)
        draw.text((x0 + pad, max(0, y0 - 18)), text, fill=text_color, font=font)
    return vis


# convenience singleton
_runner_singleton: GroundingDINORunner | None = None

def get_runner() -> GroundingDINORunner:
    global _runner_singleton
    if _runner_singleton is None:
        _runner_singleton = GroundingDINORunner()
    return _runner_singleton