inference.py

"""
inference.py

Grounding DINO + SAM 2 inference helpers.

Key points:
- Loads Grounding DINO and SAM2 lazily and keeps global singletons.
- Runs DINO first, then offloads it to CPU before running SAM2 to save VRAM.
- Uses autocast (BF16 / FP16) for SAM2.
- Splits SAM2 predictions into smaller box batches to avoid CUDA OOM.
"""

import logging
import os
import gc
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np
from PIL import Image

# Set PyTorch CUDA allocator to use expandable segments to reduce fragmentation.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import torch
import torchvision.transforms as T
import supervision as sv

from groundingdino.util.inference import load_model, predict
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor

from config import (
    DINO_CONFIG_PATH,
    DINO_CHECKPOINT_PATH,
    SAM2_CONFIG_PATH,
    SAM2_CHECKPOINT_PATH,
)
from models import Annotation, ProcessedImage

logger = logging.getLogger(__name__)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Use lower precision for SAM2 when possible
if DEVICE == "cuda":
    if torch.cuda.is_bf16_supported():
        SAM_DTYPE = torch.bfloat16
    else:
        SAM_DTYPE = torch.float16
else:
    SAM_DTYPE = torch.float32

# Safety limits to avoid CUDA OOM
MAX_SHORT_SIDE = 800           # images are resized so max(H, W) <= this for DINO
MAX_DETECTIONS = 50            # keep only top-K DINO boxes per image
MAX_SAM2_BOXES_PER_BATCH = 16  # run SAM2 on at most this many boxes at once


# Global singletons -----------------------------------------------------------
dino_model: Optional[torch.nn.Module] = None
sam2_predictor: Optional[SAM2ImagePredictor] = None


# Grounding DINO --------------------------------------------------------------
def _load_grounding_dino() -> None:
    """Lazy-load Grounding DINO."""
    global dino_model
    if dino_model is not None:
        return

    logger.info("Loading Grounding DINO from %s", DINO_CHECKPOINT_PATH)
    try:
        model = load_model(
            model_config_path=str(DINO_CONFIG_PATH),
            model_checkpoint_path=str(DINO_CHECKPOINT_PATH),
            device=DEVICE,
        )
        dino_model = model
        logger.info("Grounding DINO loaded on %s", DEVICE)
    except Exception as e:
        logger.exception("Error loading Grounding DINO: %s", e)
        dino_model = None


# SAM 2 -----------------------------------------------------------------------
def _load_sam2() -> None:
    """Lazy-load SAM2 image predictor."""
    global sam2_predictor

    if sam2_predictor is not None:
        return

    logger.info("Loading SAM2 from %s", SAM2_CHECKPOINT_PATH)
    try:
        sam2_model = build_sam2(
            config=str(SAM2_CONFIG_PATH),
            checkpoint=str(SAM2_CHECKPOINT_PATH),
            device=DEVICE,
        )

        # Move model to the chosen dtype for SAM
        sam2_model = sam2_model.to(device=DEVICE, dtype=SAM_DTYPE)

        sam2_predictor = SAM2ImagePredictor(sam2_model)
        logger.info("SAM2 loaded on %s with dtype %s", DEVICE, SAM_DTYPE)
    except Exception as e:
        logger.exception("Error loading SAM2: %s", e)
        sam2_predictor = None


def unload_models() -> None:
    """Forcefully unload DINO and SAM2 from GPU/CPU to free memory."""
    global dino_model, sam2_predictor

    try:
        if dino_model is not None:
            dino_model.cpu()
        if sam2_predictor is not None and hasattr(sam2_predictor, "model"):
            sam2_predictor.model.cpu()

        dino_model = None
        sam2_predictor = None

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()
        logger.info("Unloaded DINO and SAM2 models and cleared CUDA cache.")
    except Exception as e:
        logger.exception("Error while unloading models: %s", e)


def set_active_model_path(new_checkpoint: str) -> None:
    """
    Called from the web app when the user switches to a fine-tuned checkpoint.
    We simply update the global checkpoint path and drop the cached model;
    it will be re-loaded on next inference.
    """
    global dino_model, DINO_CHECKPOINT_PATH
    DINO_CHECKPOINT_PATH = Path(new_checkpoint)
    dino_model = None

    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    logger.info("Active Grounding DINO checkpoint switched to %s", DINO_CHECKPOINT_PATH)


# Core DINO + SAM2 pipeline ---------------------------------------------------
def run_grounded_dino_sam(
    image: Image.Image,
    prompt: str,
) -> Dict[str, Any]:
    """
    Runs Grounding DINO -> SAM2 on a single image.

    Steps:
    1. Resize image for DINO (max side = MAX_SHORT_SIDE).
    2. Run DINO to get candidate boxes.
    3. Keep only top-K boxes (MAX_DETECTIONS) to protect VRAM.
    4. Run SAM2 in smaller box batches (MAX_SAM2_BOXES_PER_BATCH).
    5. Return list of Annotation objects and original image size.
    """
    global dino_model, sam2_predictor

    try:
        width, height = image.size

        # ------------------------------------------------------------------
        # 1) Resize for DINO (we keep original image for SAM2 / output)
        # ------------------------------------------------------------------
        scale = 1.0
        if max(width, height) > MAX_SHORT_SIDE:
            scale = MAX_SHORT_SIDE / float(max(width, height))
            new_w = int(round(width * scale))
            new_h = int(round(height * scale))
            image_for_dino = image.resize(
                (new_w, new_h), resample=Image.Resampling.LANCZOS
            )
        else:
            image_for_dino = image

        # Prompt -> list of class names
        class_names = [x.strip() for x in prompt.split(",") if x.strip()]
        if not class_names:
            class_names = ["object"]
        caption = " . ".join(class_names)

        # Same transform as GroundingDINO demo
        transform = T.Compose(
            [
                T.RandomResize([800], max_size=1333),
            ]
        )
        image_tensor, _ = transform(image_for_dino.convert("RGB"), None)

        # ------------------------------------------------------------------
        # 2) Run DINO
        # ------------------------------------------------------------------
        _load_grounding_dino()
        if dino_model is None:
            logger.error("Grounding DINO model is not available.")
            return {
                "annotated_image": image,
                "annotations": [],
                "image_size": (width, height),
            }

        logger.debug("Running Grounding DINO on device=%s", DEVICE)
        with torch.inference_mode():
            boxes, logits, phrases = predict(
                model=dino_model,
                image=image_tensor,
                caption=caption,
                box_threshold=0.35,
                text_threshold=0.25,
                device=DEVICE,
            )

        # Move DINO back to CPU before running SAM2 to save VRAM
        if DEVICE == "cuda":
            dino_model = dino_model.to("cpu")
            torch.cuda.empty_cache()

        if boxes is None or boxes.numel() == 0:
            logger.info("No detections from Grounding DINO.")
            return {
                "annotated_image": image,
                "annotations": [],
                "image_size": (width, height),
            }

        # ------------------------------------------------------------------
        # 3) Convert boxes (cx, cy, w, h in normalized coords) -> xyxy
        #    and map back to original image coordinates.
        # ------------------------------------------------------------------
        W_resized, H_resized = image_for_dino.size
        # [N,4] in absolute coords relative to resized image
        boxes_cxcywh = boxes * torch.tensor(
            [W_resized, H_resized, W_resized, H_resized],
            dtype=boxes.dtype,
            device=boxes.device,
        )

        boxes_xyxy = boxes_cxcywh.clone()
        boxes_xyxy[:, 0] = boxes_cxcywh[:, 0] - 0.5 * boxes_cxcywh[:, 2]
        boxes_xyxy[:, 1] = boxes_cxcywh[:, 1] - 0.5 * boxes_cxcywh[:, 3]
        boxes_xyxy[:, 2] = boxes_cxcywh[:, 0] + 0.5 * boxes_cxcywh[:, 2]
        boxes_xyxy[:, 3] = boxes_cxcywh[:, 1] + 0.5 * boxes_cxcywh[:, 3]

        # Undo resize so boxes are in original image coordinates
        if scale != 1.0:
            boxes_xyxy /= scale

        boxes_xyxy = boxes_xyxy.clamp(min=0.0)
        boxes_xyxy = boxes_xyxy.cpu()

        # Map DINO phrases to class indices
        class_ids: List[int] = []
        for phrase in phrases:
            phrase_lower = phrase.lower()
            matched_idx = 0
            for i, name in enumerate(class_names):
                name_lower = name.lower()
                if name_lower in phrase_lower or phrase_lower in name_lower:
                    matched_idx = i
                    break
            class_ids.append(matched_idx)

        detections = sv.Detections(
            xyxy=boxes_xyxy.numpy(),
            confidence=logits.cpu().numpy(),
            class_id=np.array(class_ids, dtype=int),
        )

        if len(detections) == 0:
            logger.info("No valid detections after filtering.")
            return {
                "annotated_image": image,
                "annotations": [],
                "image_size": (width, height),
            }

        # Limit number of detections to protect SAM2 from OOM
        if len(detections) > MAX_DETECTIONS:
            order = np.argsort(-detections.confidence)
            keep = order[:MAX_DETECTIONS]
            detections = detections[keep]
            logger.info(
                "Clipped detections from %d to %d for SAM2.",
                len(order),
                MAX_DETECTIONS,
            )

        # ------------------------------------------------------------------
        # 4) Run SAM2 in smaller batches of boxes
        # ------------------------------------------------------------------
        _load_sam2()
        if sam2_predictor is None:
            logger.error("SAM2 predictor not available.")
            # Return boxes only (no masks)
            annotations: List[Annotation] = []
            for i in range(len(detections)):
                xyxy_box = detections.xyxy[i]
                score = float(detections.confidence[i])
                class_id = int(detections.class_id[i])
                label_name = (
                    class_names[class_id]
                    if 0 <= class_id < len(class_names)
                    else "object"
                )

                annotations.append(
                    Annotation(
                        label=label_name,
                        score=score,
                        bbox=[float(v) for v in xyxy_box],
                        mask=None,
                    )
                )

            return {
                "annotated_image": image,
                "annotations": annotations,
                "image_size": (width, height),
            }

        image_np = np.array(image.convert("RGB"))

        logger.debug(
            "Running SAM2 on %d boxes (batch size %d) with dtype=%s",
            len(detections),
            MAX_SAM2_BOXES_PER_BATCH,
            SAM_DTYPE,
        )

        masks_list: List[Any] = []
        scores_list: List[Any] = []

        with torch.inference_mode():
            if DEVICE == "cuda":
                ctx = torch.autocast(device_type="cuda", dtype=SAM_DTYPE)
            else:
                from contextlib import nullcontext

                ctx = nullcontext()

            with ctx:
                sam2_predictor.set_image(image_np)

                xyxy_all = detections.xyxy
                num_boxes = xyxy_all.shape[0]

                for start in range(0, num_boxes, MAX_SAM2_BOXES_PER_BATCH):
                    end = min(start + MAX_SAM2_BOXES_PER_BATCH, num_boxes)
                    box_batch = xyxy_all[start:end]

                    batch_masks, batch_scores, _ = sam2_predictor.predict(
                        point_coords=None,
                        point_labels=None,
                        box=box_batch,
                        multimask_output=False,
                    )

                    masks_list.append(batch_masks)
                    scores_list.append(batch_scores)

        # Offload SAM2 to CPU to free VRAM for next image
        if DEVICE == "cuda" and hasattr(sam2_predictor, "model"):
            sam2_predictor.model = sam2_predictor.model.to("cpu")
            torch.cuda.empty_cache()

        # ------------------------------------------------------------------
        # 5) Concatenate masks / scores and build our Annotation objects
        # ------------------------------------------------------------------
        if not masks_list:
            logger.info("SAM2 returned no masks.")
            annotations: List[Annotation] = []
            for i in range(len(detections)):
                xyxy_box = detections.xyxy[i]
                score = float(detections.confidence[i])
                class_id = int(detections.class_id[i])
                label_name = (
                    class_names[class_id]
                    if 0 <= class_id < len(class_names)
                    else "object"
                )
                annotations.append(
                    Annotation(
                        label=label_name,
                        score=score,
                        bbox=[float(v) for v in xyxy_box],
                        mask=None,
                    )
                )

            return {
                "annotated_image": image,
                "annotations": annotations,
                "image_size": (width, height),
            }

        first_mask = masks_list[0]
        if torch.is_tensor(first_mask):
            masks = torch.cat(masks_list, dim=0)
            scores = torch.cat(scores_list, dim=0)
            if masks.ndim == 4:
                masks = masks.squeeze(1)
            masks_np = masks.detach().cpu().numpy()
        else:
            masks = np.concatenate(masks_list, axis=0)
            scores = np.concatenate(scores_list, axis=0)
            if masks.ndim == 4:
                masks = masks.squeeze(1)
            masks_np = masks

        detections.mask = masks_np > 0.0

        annotations: List[Annotation] = []
        for i in range(len(detections)):
            xyxy_box = detections.xyxy[i]
            score = float(detections.confidence[i])
            class_id = int(detections.class_id[i])
            label_name = (
                class_names[class_id]
                if 0 <= class_id < len(class_names)
                else "object"
            )

            mask_i = None
            if detections.mask is not None:
                mask_raw = detections.mask[i]
                if isinstance(mask_raw, torch.Tensor):
                    mask_i = (
                        mask_raw.detach().cpu().numpy().astype(np.uint8)
                    )
                else:
                    mask_i = mask_raw.astype(np.uint8)

            annotations.append(
                Annotation(
                    label=label_name,
                    score=score,
                    bbox=[float(v) for v in xyxy_box],
                    mask=mask_i,
                )
            )

        return {
            "annotated_image": image,
            "annotations": annotations,
            "image_size": (width, height),
        }

    except torch.cuda.OutOfMemoryError as oom:
        logger.exception("CUDA OOM during Grounded SAM2 inference: %s", oom)
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()
        raise

    except Exception as e:
        logger.exception("Unexpected error during Grounded SAM2 inference: %s", e)
        return {
            "annotated_image": image,
            "annotations": [],
            "image_size": image.size,
            "error": str(e),
        }


# Batch helpers ---------------------------------------------------------------
def list_image_files(folder: Path) -> List[Path]:
    exts = {".jpg", ".jpeg", ".png", ".bmp"}
    return sorted([p for p in folder.iterdir() if p.suffix.lower() in exts])


def load_images_without_inference(folder: Path) -> List[ProcessedImage]:
    """Used when we just want to list images in the UI without running models."""
    images: List[ProcessedImage] = []
    for path in list_image_files(folder):
        try:
            img = Image.open(path).convert("RGB")
            images.append(
                ProcessedImage(
                    path=str(path),
                    original_image=img,
                    annotated_image=img,  # no overlays
                    annotations=[],
                    image_size=img.size,
                )
            )
        except Exception as e:
            logger.warning("Failed to load image %s: %s", path, e)
    return images


def process_images_with_inference(
    folder: Path,
    prompt: str,
) -> List[ProcessedImage]:
    """Run DINO+SAM2 on all images in a folder."""
    results: List[ProcessedImage] = []
    for path in list_image_files(folder):
        try:
            img = Image.open(path).convert("RGB")
            result = run_grounded_dino_sam(img, prompt)

            processed = ProcessedImage(
                path=str(path),
                original_image=img,
                annotated_image=result.get("annotated_image", img),
                annotations=result.get("annotations", []),
                image_size=result.get("image_size", img.size),
            )
            results.append(processed)
        except Exception as e:
            logger.warning("Failed to process %s: %s", path, e)

    return results