""" Model Manager - Handles loading and inference for Grounding DINO + SAM 2 """ import os import torch import numpy as np from PIL import Image from typing import List, Dict, Tuple, Optional from dataclasses import dataclass @dataclass class DetectionResult: """Single detection result""" label: str confidence: float bbox: np.ndarray # [x1, y1, x2, y2] mask: Optional[np.ndarray] = None # H x W binary mask class ModelManager: """Manages Grounding DINO + SAM 2 pipeline""" def __init__(self, config): self.config = config self.device = config.model.device self.gdino_model = None self.gdino_processor = None self.sam2_predictor = None self.sam2_video_predictor = None self._loaded = False def load_models(self, progress_callback=None): """Load all models into memory""" if self._loaded: return if progress_callback: progress_callback(0.1, "Loading Grounding DINO...") self._load_grounding_dino() if progress_callback: progress_callback(0.5, "Loading SAM 2...") self._load_sam2() self._loaded = True if progress_callback: progress_callback(1.0, "Models loaded ✅") def _load_grounding_dino(self): """Load Grounding DINO model""" try: # Try HuggingFace Transformers first (easier setup) from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection model_id = self.config.model.gdino_model_id print(f"📥 Loading {model_id} (cached in ~/.cache/huggingface after first download)") self.gdino_processor = AutoProcessor.from_pretrained(model_id) self.gdino_model = AutoModelForZeroShotObjectDetection.from_pretrained( model_id ).to(self.device) if self.config.model.use_fp16 and self.device == "cuda": self.gdino_model = self.gdino_model.half() self.gdino_model.eval() print(f"✅ Grounding DINO loaded from {model_id}") self._gdino_backend = "transformers" except Exception as e: print(f"⚠️ Transformers loading failed ({e}), trying GroundingDINO package...") self._load_grounding_dino_native() def _load_grounding_dino_native(self): """Fallback: Load Grounding DINO from official package""" try: from groundingdino.util.inference import load_model, predict from huggingface_hub import hf_hub_download # Download checkpoint ckpt_path = hf_hub_download( repo_id="ShilongLiu/GroundingDINO", filename="groundingdino_swinb_cogcoor.pth" ) config_path = hf_hub_download( repo_id="ShilongLiu/GroundingDINO", filename="GroundingDINO_SwinB.cfg.py" ) self.gdino_model = load_model(config_path, ckpt_path, device=self.device) self._gdino_backend = "native" print("✅ Grounding DINO loaded (native)") except ImportError: raise RuntimeError( "❌ Grounding DINO failed to load!\n" "The HuggingFace Transformers backend failed, and the native package is not installed.\n" "Fix: pip install git+https://github.com/IDEA-Research/GroundingDINO.git\n" "Or check that 'transformers' is up to date: pip install -U transformers" ) def _load_sam2(self): """Load SAM 2 model""" try: # Try importing from sam2 (PyPI: pip install sam-2) try: from sam2.build_sam import build_sam2, build_sam2_video_predictor from sam2.sam2_image_predictor import SAM2ImagePredictor except ImportError: # Older versions may have different import paths from sam2.build_sam import build_sam2 from sam2.automatic_mask_generator import SAM2ImagePredictor build_sam2_video_predictor = None from huggingface_hub import hf_hub_download checkpoint = self.config.model.sam2_checkpoint model_cfg = self.config.model.sam2_model_cfg # Download checkpoint from HuggingFace ckpt_map = { "facebook/sam2.1-hiera-base-plus": "sam2.1_hiera_base_plus.pt", "facebook/sam2.1-hiera-small": "sam2.1_hiera_small.pt", "facebook/sam2.1-hiera-large": "sam2.1_hiera_large.pt", "facebook/sam2.1-hiera-tiny": "sam2.1_hiera_tiny.pt", } ckpt_file = ckpt_map.get(checkpoint, "sam2.1_hiera_base_plus.pt") try: ckpt_path = hf_hub_download( repo_id=checkpoint, filename=ckpt_file ) except Exception: # Try without version suffix alt_file = ckpt_file.replace("sam2.1_", "sam2_") ckpt_path = hf_hub_download( repo_id=checkpoint, filename=alt_file ) # Build image predictor sam2_model = build_sam2(model_cfg, ckpt_path, device=self.device) self.sam2_predictor = SAM2ImagePredictor(sam2_model) # Build video predictor (may not be available in all versions) if build_sam2_video_predictor is not None: try: self.sam2_video_predictor = build_sam2_video_predictor( model_cfg, ckpt_path, device=self.device ) except Exception as e: print(f"⚠️ Video predictor not available: {e}") print(" Will use frame-by-frame mode only.") self.sam2_video_predictor = None else: self.sam2_video_predictor = None print(f"✅ SAM 2 loaded from {checkpoint}") except Exception as e: print(f"❌ SAM 2 loading failed: {e}") print(" Make sure sam-2 is installed: pip install sam-2>=1.1.0") raise def detect_objects(self, image: np.ndarray, text_prompt: str) -> List[DetectionResult]: """ Detect objects in image using text prompt via Grounding DINO Args: image: BGR numpy array (H, W, 3) text_prompt: Text description of objects to detect (e.g., "face. hand. text.") Returns: List of DetectionResult with bounding boxes """ # Normalize prompt - ensure it ends with period for GDINO prompt = text_prompt.strip() if not prompt.endswith("."): prompt += "." pil_image = Image.fromarray(image[..., ::-1]) # BGR -> RGB -> PIL if self._gdino_backend == "transformers": return self._detect_transformers(pil_image, prompt) else: return self._detect_native(image, prompt) def _detect_transformers(self, pil_image: Image.Image, prompt: str) -> List[DetectionResult]: """Detection using HuggingFace Transformers (auto-detects API version)""" import inspect inputs = self.gdino_processor( images=pil_image, text=prompt, return_tensors="pt" ).to(self.device) with torch.no_grad(): if self.config.model.use_fp16 and self.device == "cuda": with torch.autocast("cuda"): outputs = self.gdino_model(**inputs) else: outputs = self.gdino_model(**inputs) target_sizes = [pil_image.size[::-1]] # (H, W) threshold = self.config.model.gdino_box_threshold # Inspect the actual function signature to know which params it accepts post_fn = self.gdino_processor.post_process_grounded_object_detection sig = inspect.signature(post_fn) param_names = list(sig.parameters.keys()) kwargs = {"target_sizes": target_sizes} args = [outputs] # Add threshold with correct name if "threshold" in param_names: kwargs["threshold"] = threshold elif "box_threshold" in param_names: kwargs["box_threshold"] = threshold kwargs["text_threshold"] = self.config.model.gdino_text_threshold # Add input_ids if accepted if "input_ids" in param_names: args.append(inputs.get("input_ids", None)) results = post_fn(*args, **kwargs)[0] detections = [] # Handle both 'text_labels' (new) and 'labels' (old) keys labels = results.get("text_labels", results.get("labels", [])) for bbox, score, label in zip( results["boxes"].cpu().numpy(), results["scores"].cpu().numpy(), labels ): label_str = str(label) if not isinstance(label, str) else label detections.append(DetectionResult( label=label_str, confidence=float(score), bbox=bbox )) return detections def _detect_native(self, image: np.ndarray, prompt: str) -> List[DetectionResult]: """Detection using native GroundingDINO""" from groundingdino.util.inference import predict from groundingdino.util.utils import get_phrases_from_posmap import groundingdino.datasets.transforms as T transform = T.Compose([ T.RandomResize([800], max_size=1333), T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) pil_image = Image.fromarray(image[..., ::-1]) transformed, _ = transform(pil_image, None) boxes, logits, phrases = predict( model=self.gdino_model, image=transformed, caption=prompt, box_threshold=self.config.model.gdino_box_threshold, text_threshold=self.config.model.gdino_text_threshold, device=self.device ) h, w = image.shape[:2] detections = [] for box, score, label in zip(boxes, logits, phrases): # Convert from [cx, cy, w, h] normalized to [x1, y1, x2, y2] pixels cx, cy, bw, bh = box.cpu().numpy() x1 = (cx - bw/2) * w y1 = (cy - bh/2) * h x2 = (cx + bw/2) * w y2 = (cy + bh/2) * h detections.append(DetectionResult( label=label, confidence=float(score), bbox=np.array([x1, y1, x2, y2]) )) return detections def segment_with_boxes(self, image: np.ndarray, boxes: np.ndarray) -> np.ndarray: """ Generate segmentation masks from bounding boxes using SAM 2 Args: image: BGR numpy array (H, W, 3) boxes: Array of boxes [N, 4] in [x1, y1, x2, y2] format Returns: Combined binary mask (H, W) uint8 """ rgb_image = image[..., ::-1] # BGR -> RGB self.sam2_predictor.set_image(rgb_image) if len(boxes) == 0: return np.zeros(image.shape[:2], dtype=np.uint8) input_boxes = torch.tensor(boxes, dtype=torch.float32, device=self.device) with torch.no_grad(): if self.config.model.use_fp16 and self.device == "cuda": with torch.autocast("cuda"): masks, scores, _ = self.sam2_predictor.predict( box=input_boxes, multimask_output=False, ) else: masks, scores, _ = self.sam2_predictor.predict( box=input_boxes, multimask_output=False, ) # Combine all masks into single mask if isinstance(masks, torch.Tensor): masks = masks.cpu().numpy() combined_mask = np.zeros(image.shape[:2], dtype=np.uint8) for mask in masks: if mask.ndim == 3: mask = mask[0] # Take first mask if multimask combined_mask = np.maximum(combined_mask, (mask > 0.5).astype(np.uint8) * 255) return combined_mask def init_video_tracking(self, frames_dir: str, detections: List[DetectionResult]) -> dict: """ Initialize SAM 2 video tracking from first-frame detections Args: frames_dir: Directory containing numbered JPEG frames detections: Detection results from first frame Returns: SAM 2 inference state """ state = self.sam2_video_predictor.init_state(video_path=frames_dir) # Add each detection as a tracking target for idx, det in enumerate(detections): box = det.bbox _, _, mask_logits = self.sam2_video_predictor.add_new_points_or_box( inference_state=state, frame_idx=0, obj_id=idx + 1, box=box, ) return state def propagate_video(self, state, num_frames: int, progress_callback=None): """ Propagate masks through all video frames Args: state: SAM 2 inference state num_frames: Total number of frames progress_callback: Optional callback(frame_idx, total_frames) Returns: Dict mapping frame_idx -> combined binary mask (H, W) """ frame_masks = {} for frame_idx, obj_ids, mask_logits in self.sam2_video_predictor.propagate_in_video(state): # Combine all object masks masks = (mask_logits > 0.0).cpu().numpy() # [N, 1, H, W] combined = np.zeros(masks.shape[2:], dtype=np.uint8) for mask in masks: combined = np.maximum(combined, (mask[0] > 0).astype(np.uint8) * 255) frame_masks[frame_idx] = combined if progress_callback: progress_callback(frame_idx, num_frames) return frame_masks def detect_and_segment_frame(self, frame: np.ndarray, text_prompt: str) -> np.ndarray: """ Full pipeline: detect + segment on a single frame Args: frame: BGR numpy array text_prompt: What to detect Returns: Binary mask (H, W) uint8, 0 or 255 """ detections = self.detect_objects(frame, text_prompt) if not detections: return np.zeros(frame.shape[:2], dtype=np.uint8) boxes = np.array([d.bbox for d in detections]) mask = self.segment_with_boxes(frame, boxes) return mask def unload_models(self): """Free GPU memory""" self.gdino_model = None self.gdino_processor = None self.sam2_predictor = None self.sam2_video_predictor = None self._loaded = False if torch.cuda.is_available(): torch.cuda.empty_cache() print("🗑️ Models unloaded")