models/model_manager.py

"""
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")