utils/detector.py

import cv2
import numpy as np
from ultralytics import YOLO
import torch
from typing import Dict, List, Any
import os
import gc

# Arabic letter mapping from English class names to Arabic characters
ARABIC_MAP = {
    "aleff": "ا",
    "bb": "ب",
    "ta": "ت",
    "thaa": "ث",
    "jeem": "ج",
    "haa": "ح",
    "khaa": "خ",
    "dal": "د",
    "thal": "ذ",
    "ra": "ر",
    "zay": "ز",
    "seen": "س",
    "sheen": "ش",
    "saad": "ص",
    "dhad": "ض",
    "taa": "ط",
    "dha": "ظ",
    "ain": "ع",
    "ghain": "غ",
    "fa": "ف",
    "gaaf": "ق",
    "kaaf": "ك",
    "laam": "ل",
    "la": "لا",
    "meem": "م",
    "nun": "ن",
    "ha": "ه",
    "waw": "و",
    "ya": "ي",
    "yaa": "ي",
    "toot": "ة",
    "al": "ال"
}

class ArabicSignDetector:
    def __init__(self, model_path: str = None):
        print("🔄 Initializing ArabicSignDetector...")
        
        # Check GPU status
        print(f"🎮 CUDA available: {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            print(f"🎯 GPU device: {torch.cuda.get_device_name(0)}")
            torch.cuda.empty_cache()
        else:
            print("⚡ Running on CPU")
        
        if model_path is None:
            possible_paths = ['best.pt', '/tmp/best.pt', 'utils/best.pt', './best.pt', '/app/best.pt']
            
            found_path = None
            for path in possible_paths:
                if os.path.exists(path):
                    found_path = path
                    print(f"✅ Found model at: {path}")
                    break
            
            if found_path:
                model_path = found_path
            else:
                print("❌ No best.pt found!")
                self.model = None
                return
        
        try:
            print(f"🔄 Loading YOLO model from: {model_path}")
            
            # Optimized YOLO loading for ZeroGPU
            self.model = YOLO(model_path)
            
            # Set to eval mode and optimize
            if hasattr(self.model, 'model'):
                self.model.model.eval()
            
            # Lower confidence threshold for better real-time detection
            self.confidence_threshold = 0.15
            
            # Clear memory after loading
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            print(f"✅ YOLO model loaded successfully!")
            if hasattr(self.model, 'names') and self.model.names:
                print(f"📊 Number of classes: {len(self.model.names)}")
                print(f"🎯 Confidence threshold: {self.confidence_threshold}")
            
        except Exception as e:
            print(f"❌ YOLO loading failed: {e}")
            try:
                print("🔄 Trying alternative YOLO loading...")
                checkpoint = torch.load(model_path, map_location='cpu', weights_only=False)
                self.model = YOLO(model_path)
                print("✅ YOLO model loaded with alternative method!")
            except Exception as e2:
                print(f"❌ All loading methods failed: {e2}")
                self.model = None

    def detect_letters(self, image: np.ndarray) -> Dict[str, Any]:
        """Detect Arabic letters and form text - optimized for ZeroGPU"""
        if self.model is None:
            print("❌ YOLO model is not loaded")
            return {
                'success': False,
                'error': 'YOLO model not loaded',
                'arabic_text': '',
                'letters': [],
                'total_detections': 0
            }

        try:
            # Use GPU if available, with optimizations
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
            
            # Optimized inference settings for ZeroGPU
            with torch.inference_mode():  # Use inference_mode for better performance
                results = self.model(
                    image, 
                    conf=self.confidence_threshold,
                    device=device,
                    verbose=False,  # Reduce output
                    half=torch.cuda.is_available()  # Use FP16 on GPU
                )
            
            detected_letters = []
            confidences = []
            
            for result in results:
                if hasattr(result, 'boxes') and result.boxes is not None:
                    boxes = result.boxes
                    for i in range(len(boxes.cls)):
                        class_id = int(boxes.cls[i])
                        confidence = float(boxes.conf[i])
                        letter_english = self.model.names.get(class_id, "")
                        
                        if confidence > self.confidence_threshold:
                            # Convert English class name to Arabic character
                            letter_arabic = ARABIC_MAP.get(letter_english.lower(), letter_english)
                            detected_letters.append(letter_arabic)
                            confidences.append(confidence)

            # Clear GPU memory after inference
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            if detected_letters:
                arabic_text = "".join(detected_letters)
                print(f"📝 Detected: '{arabic_text}' ({len(detected_letters)} letters)")
                return {
                    'success': True,
                    'arabic_text': arabic_text,
                    'letters': detected_letters,
                    'confidences': confidences,
                    'total_detections': len(detected_letters)
                }
            else:
                return {
                    'success': False,
                    'error': 'No Arabic sign letters detected',
                    'arabic_text': '',
                    'letters': [],
                    'total_detections': 0
                }

        except Exception as e:
            print(f"❌ Detection error: {e}")
            # Clean up on error
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            return {
                'success': False,
                'error': str(e),
                'arabic_text': '',
                'letters': [],
                'total_detections': 0
            }