""" Inference pipeline for emotion recognition. """ import numpy as np from pathlib import Path from typing import Dict, List, Optional, Tuple, Union import cv2 from PIL import Image import tensorflow as tf from tensorflow.keras.models import Model import sys sys.path.append(str(Path(__file__).parent.parent.parent)) from src.config import ( IMAGE_SIZE, IMAGE_SIZE_TRANSFER, EMOTION_CLASSES, IDX_TO_EMOTION, INTENSITY_HIGH_THRESHOLD, INTENSITY_MEDIUM_THRESHOLD, CUSTOM_CNN_PATH, MOBILENET_PATH, VGG_PATH ) from src.preprocessing.face_detector import FaceDetector from src.models.model_utils import load_model class EmotionPredictor: """ Unified prediction interface for emotion recognition. """ def __init__( self, model_name: str = "custom_cnn", model_path: Optional[Path] = None, use_face_detection: bool = True ): """ Initialize the predictor. Args: model_name: Name of the model ('custom_cnn', 'mobilenet', 'vgg19') model_path: Optional custom model path use_face_detection: Whether to detect faces before prediction """ self.model_name = model_name self.model = None self.face_detector = FaceDetector() if use_face_detection else None # Determine model path if model_path: self.model_path = Path(model_path) else: paths = { "custom_cnn": CUSTOM_CNN_PATH, "mobilenet": MOBILENET_PATH, "vgg19": VGG_PATH } self.model_path = paths.get(model_name) # Set preprocessing based on model type self.is_transfer_model = model_name in ["mobilenet", "vgg19"] self.target_size = IMAGE_SIZE_TRANSFER if self.is_transfer_model else IMAGE_SIZE self.use_rgb = self.is_transfer_model def load(self) -> bool: """ Load the model. Returns: True if model loaded successfully """ try: if self.model_path and self.model_path.exists(): self.model = load_model(self.model_path) return True else: print(f"Model file not found: {self.model_path}") return False except Exception as e: print(f"Error loading model: {e}") return False def preprocess_image( self, image: np.ndarray, detect_face: bool = True ) -> Tuple[Optional[np.ndarray], List[dict]]: """ Preprocess an image for prediction. Args: image: Input image (BGR or RGB format) detect_face: Whether to detect and extract face Returns: Tuple of (preprocessed image, face info) """ faces_info = [] if detect_face and self.face_detector: # Detect and extract face face, faces_info = self.face_detector.detect_and_extract( image, target_size=self.target_size, to_grayscale=not self.use_rgb ) if face is None: return None, faces_info processed = face else: # Resize directly processed = cv2.resize(image, self.target_size) # Convert color if needed if self.use_rgb: if len(processed.shape) == 2: processed = cv2.cvtColor(processed, cv2.COLOR_GRAY2RGB) elif processed.shape[2] == 1: processed = np.repeat(processed, 3, axis=2) else: if len(processed.shape) == 3 and processed.shape[2] == 3: processed = cv2.cvtColor(processed, cv2.COLOR_BGR2GRAY) # Normalize processed = processed.astype(np.float32) / 255.0 # Add channel dimension if grayscale if len(processed.shape) == 2: processed = np.expand_dims(processed, axis=-1) # Add batch dimension processed = np.expand_dims(processed, axis=0) return processed, faces_info def predict( self, image: Union[np.ndarray, str, Path], detect_face: bool = True, return_all_scores: bool = True ) -> Dict: """ Predict emotion from an image. Args: image: Input image (array, file path, or PIL Image) detect_face: Whether to detect face first return_all_scores: Whether to return all class scores Returns: Prediction result dictionary """ if self.model is None: success = self.load() if not success: return {"error": "Model not loaded"} # Load image if path provided if isinstance(image, (str, Path)): image = cv2.imread(str(image)) if image is None: return {"error": f"Could not load image: {image}"} elif isinstance(image, Image.Image): image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) # Preprocess processed, faces_info = self.preprocess_image(image, detect_face) if processed is None: return { "error": "No face detected", "face_detected": False, "faces_info": faces_info } # Predict predictions = self.model.predict(processed, verbose=0) # Get top prediction pred_idx = int(np.argmax(predictions[0])) confidence = float(predictions[0][pred_idx]) emotion = IDX_TO_EMOTION[pred_idx] # Calculate intensity intensity = self._calculate_intensity(confidence) result = { "emotion": emotion, "confidence": confidence, "intensity": intensity, "face_detected": len(faces_info) > 0, "faces_info": faces_info, "model_used": self.model_name } if return_all_scores: result["all_probabilities"] = { EMOTION_CLASSES[i]: float(predictions[0][i]) for i in range(len(EMOTION_CLASSES)) } return result def predict_batch( self, images: List[Union[np.ndarray, str, Path]], detect_face: bool = True ) -> Dict: """ Predict emotions for multiple images. Args: images: List of images detect_face: Whether to detect faces Returns: Batch prediction results """ results = [] emotion_counts = {e: 0 for e in EMOTION_CLASSES} successful_predictions = 0 for i, image in enumerate(images): result = self.predict(image, detect_face) result["image_index"] = i results.append(result) if "error" not in result: emotion_counts[result["emotion"]] += 1 successful_predictions += 1 # Calculate distribution if successful_predictions > 0: emotion_distribution = { e: count / successful_predictions for e, count in emotion_counts.items() } else: emotion_distribution = {e: 0.0 for e in EMOTION_CLASSES} # Find dominant emotion dominant_emotion = max(emotion_counts.items(), key=lambda x: x[1]) return { "results": results, "summary": { "total_images": len(images), "successful_predictions": successful_predictions, "failed_predictions": len(images) - successful_predictions, "emotion_counts": emotion_counts, "emotion_distribution": emotion_distribution, "dominant_emotion": dominant_emotion[0], "dominant_emotion_count": dominant_emotion[1] }, "model_used": self.model_name } def _calculate_intensity(self, confidence: float) -> str: """ Calculate emotion intensity based on confidence. Args: confidence: Prediction confidence Returns: Intensity level ('high', 'medium', 'low') """ if confidence >= INTENSITY_HIGH_THRESHOLD: return "high" elif confidence >= INTENSITY_MEDIUM_THRESHOLD: return "medium" else: return "low" def visualize_prediction( self, image: np.ndarray, prediction: Dict ) -> np.ndarray: """ Visualize prediction on image. Args: image: Original image prediction: Prediction result Returns: Image with visualizations """ result = image.copy() if self.face_detector and prediction.get("faces_info"): # Draw face detection and emotion label result = self.face_detector.draw_detections( result, prediction["faces_info"], emotions=[prediction.get("emotion", "Unknown")], confidences=[prediction.get("confidence", 0)] ) return result @staticmethod def get_available_models() -> Dict[str, bool]: """ Get available trained models. Returns: Dictionary of model name -> availability """ return { "custom_cnn": CUSTOM_CNN_PATH.exists(), "mobilenet": MOBILENET_PATH.exists(), "vgg19": VGG_PATH.exists() } def create_predictor( model_name: str = "custom_cnn", auto_load: bool = True ) -> Optional[EmotionPredictor]: """ Factory function to create a predictor. Args: model_name: Name of the model auto_load: Whether to automatically load the model Returns: EmotionPredictor instance or None if loading fails """ predictor = EmotionPredictor(model_name) if auto_load: if not predictor.load(): return None return predictor if __name__ == "__main__": # Show available models print("Available models:") for name, available in EmotionPredictor.get_available_models().items(): status = "✓" if available else "✗" print(f" {status} {name}")