module.py

import os
import cv2
import numpy as np
from PIL import Image
from ultralytics import YOLO

class CVModelPipeline:
    def __init__(self, detection_model_path, classification_model_path):
        """
        Initializes the CV pipeline.
        
        Args:
            detection_model_path (str): Path to the YOLOv11 Face+Gender detection model.
            classification_model_path (str): Path to the custom 14-class Face Recognition model.
        """
        self.detection_model = None
        self.classifier_model = None
        
        # Load detection model (Face + Gender)
        if os.path.exists(detection_model_path):
            try:
                self.detection_model = YOLO(detection_model_path)
                print(f"Loaded detection model from {detection_model_path}")
            except Exception as e:
                print(f"Error loading detection model: {e}")
        else:
            print(f"Warning: Detection model not found at {detection_model_path}")

        # Load classification model (Face Recognition)
        if os.path.exists(classification_model_path):
            try:
                self.classifier_model = YOLO(classification_model_path)
                print(f"Loaded classification model from {classification_model_path}")
            except Exception as e:
                print(f"Error loading classification model: {e}")
        else:
            print(f"Warning: Classification model not found at {classification_model_path}")

        # Class mappings
        self.gender_names = {0: 'Female', 1: 'Male'}
        self.class_names = {
            0: 'Iqra Aziz',
            1: 'Hamuyun Saeed',
            2: 'Fahad Mustafa',
            3: 'Sjal Ali',
            4: 'Mahira Khan',
            5: 'Ayeza Khan',
            6: 'Bilal Abbas',
            7: 'Maya Ali',
            8: 'Wahaj Ali',
            9: 'Sarah Khan',
            10: 'Hania Aamir',
            11: 'Yumna Zaidi',
            12: 'Feroze Khan',
            13: 'Ahad Raza Mir',
            14: 'Hamza Ali Abbasi'
        }

    def process_image(self, image_input):
        """
        Processes an image, detects face+gender, crops the face, and runs recognition.
        
        Args:
            image_input (str or np.ndarray): Path to the image file or an OpenCV image array.
            
        Returns:
            list: A list of dictionaries with structured JSON-like results for each detected face.
        """
        final_results = []
        
        # Load image if it's a file path
        if isinstance(image_input, str):
            if not os.path.exists(image_input):
                print(f"Image not found at {image_input}")
                return final_results
            image = cv2.imread(image_input)
            if image is None:
                print(f"Could not read image: {image_input}")
                return final_results
        else:
            image = image_input

        if self.detection_model is None:
            print("Detection model is not loaded. Skipping inference.")
            return final_results

        # Step 1: Detect faces and predict gender
        try:
            results = self.detection_model(image, verbose=False)
        except Exception as e:
            print(f"Error during detection inference: {e}")
            return final_results
        
        face_detections = []
        if results and len(results) > 0:
            result = results[0]
            boxes = result.boxes
            if boxes is not None and len(boxes) > 0:
                for box in boxes:
                    x1, y1, x2, y2 = map(int, box.xyxy[0])
                    conf = float(box.conf[0])
                    cls = int(box.cls[0])
                    
                    # Map gender
                    gender = self.gender_names.get(cls, "Unknown")
                    
                    face_detections.append({
                        'bbox': [x1, y1, x2, y2],
                        'confidence': conf,
                        'gender': gender
                    })

        if not face_detections:
            return final_results

        # Step 2: Face recognition for each detected face
        for detection in face_detections:
            x1, y1, x2, y2 = detection['bbox']

            # Make sure bbox is strictly within image bounds
            h, w = image.shape[:2]
            x1_crop = max(0, x1)
            y1_crop = max(0, y1)
            x2_crop = min(w, x2)
            y2_crop = min(h, y2)

            face_crop = image[y1_crop:y2_crop, x1_crop:x2_crop]
            
            predicted_class = -1
            class_name = "Unknown"
            recognition_conf = 0.0

            if face_crop.size > 0 and self.classifier_model:
                try:
                    # Resize face for recognition (as standard in the notebook)
                    face_resized = cv2.resize(face_crop, (240, 240))
                    face_rgb = cv2.cvtColor(face_resized, cv2.COLOR_BGR2RGB)
                    face_pil = Image.fromarray(face_rgb)
                    
                    recognition_results = self.classifier_model(face_pil, verbose=False)
                    if recognition_results and len(recognition_results) > 0:
                        probs = recognition_results[0].probs
                        if probs is not None:
                            predicted_class = int(probs.top1)
                            recognition_conf = float(probs.top1conf)
                            class_name = self.class_names.get(predicted_class, f"Unknown_{predicted_class}")
                except Exception as e:
                    print(f"Recognition inference error: {e}")

            final_results.append({
                'bbox': detection['bbox'],
                'confidence': round(detection['confidence'], 4),
                'gender': detection['gender'],
                'predicted_class': predicted_class,
                'class_name': class_name,
                'recognition_confidence': round(recognition_conf, 4)
            })

        return final_results