src/models/realtime_morph_detector.py

"""
Real-time Morph Detection Model
Optimized for WebRTC integration with quality-aware processing
"""

import torch
import torch.nn as nn
import torchvision.transforms as transforms
import cv2
import numpy as np
from typing import Dict, List, Tuple, Optional, Any
import time
import logging
from dataclasses import dataclass
from datetime import datetime
import json

# Import existing models
from models.morph_detector_pl import MorphDetectorPL
from models.freqfacenet import FreqFaceNet

# Import face processing utilities
from src.webrtc.advanced_face_capture import FaceQualityMetrics, FaceCaptureResult

logger = logging.getLogger(__name__)

@dataclass
class MorphDetectionResult:
    """Results from real-time morph detection"""
    is_morphed: bool
    confidence: float
    morph_score: float
    processing_time_ms: float
    model_version: str
    quality_factor: float
    timestamp: datetime
    session_id: str
    frame_number: int

@dataclass
class ModelPerformanceMetrics:
    """Performance metrics for model monitoring"""
    inference_time_ms: float
    memory_usage_mb: float
    gpu_utilization: float
    accuracy_score: float
    throughput_fps: float
    error_count: int
    total_inferences: int

class QualityAwareMorphDetector:
    """
    Real-time morph detector with quality-aware processing
    Integrates with WebRTC capture system for optimal performance
    """
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        # Model configuration
        self.model_configs = {
            'primary': {
                'model_class': MorphDetectorPL,
                'model_path': config.get('primary_model_path', 'models/detector/morph_detector.ckpt'),
                'weight': 0.7
            },
            'frequency': {
                'model_class': FreqFaceNet,
                'model_path': config.get('frequency_model_path', 'models/detector/freq_detector.pt'),
                'weight': 0.3
            }
        }
        
        # Quality thresholds for processing
        self.quality_thresholds = {
            'minimum': config.get('min_quality_threshold', 0.5),
            'optimal': config.get('optimal_quality_threshold', 0.8),
            'bypass': config.get('bypass_quality_threshold', 0.3)  # Emergency processing
        }
        
        # Performance monitoring
        self.performance_metrics = ModelPerformanceMetrics(
            inference_time_ms=0,
            memory_usage_mb=0,
            gpu_utilization=0,
            accuracy_score=0,
            throughput_fps=0,
            error_count=0,
            total_inferences=0
        )
        
        # Load models
        self.models = {}
        self.transforms = {}
        self.load_models()
        
        # Processing history for temporal consistency
        self.detection_history = []
        self.max_history_length = 10
        
    def load_models(self):
        """Load and initialize all detection models"""
        logger.info("Loading real-time morph detection models...")
        
        try:
            # Load primary dual-stream model
            if self.model_configs['primary']['model_path']:
                self.models['primary'] = MorphDetectorPL.load_from_checkpoint(
                    self.model_configs['primary']['model_path']
                )
                self.models['primary'].eval()
                self.models['primary'].to(self.device)
                logger.info("Primary morph detector loaded successfully")
            
            # Load frequency domain model
            if self.model_configs['frequency']['model_path']:
                self.models['frequency'] = FreqFaceNet()
                checkpoint = torch.load(
                    self.model_configs['frequency']['model_path'],
                    map_location=self.device
                )
                self.models['frequency'].load_state_dict(checkpoint['model_state_dict'])
                self.models['frequency'].eval()
                self.models['frequency'].to(self.device)
                logger.info("Frequency domain detector loaded successfully")
            
            # Setup transforms for each model
            self.setup_transforms()
            
            # Warm up models
            self.warmup_models()
            
        except Exception as e:
            logger.error(f"Failed to load models: {e}")
            raise
    
    def setup_transforms(self):
        """Setup image transforms for each model"""
        # Standard transforms for primary model
        self.transforms['primary'] = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])
        
        # Frequency domain transforms
        self.transforms['frequency'] = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.5, 0.5, 0.5],
                std=[0.5, 0.5, 0.5]
            )
        ])
    
    def warmup_models(self):
        """Warm up models with dummy inputs"""
        logger.info("Warming up models...")
        
        dummy_input = torch.randn(1, 3, 224, 224).to(self.device)
        
        with torch.no_grad():
            for model_name, model in self.models.items():
                try:
                    _ = model(dummy_input)
                    logger.info(f"Model {model_name} warmed up successfully")
                except Exception as e:
                    logger.warning(f"Failed to warm up model {model_name}: {e}")
    
    async def detect_morph_realtime(self, 
                                  capture_result: FaceCaptureResult) -> Optional[MorphDetectionResult]:
        """
        Real-time morph detection with quality-aware processing
        """
        start_time = time.time()
        
        try:
            # Quality-based processing decision
            if not self.should_process(capture_result.quality_metrics):
                return None
            
            # Extract face region
            face_image = self.extract_face_region(capture_result)
            if face_image is None:
                return None
            
            # Run ensemble detection
            detection_scores = await self.run_ensemble_detection(face_image)
            
            # Apply quality weighting
            quality_weighted_score = self.apply_quality_weighting(
                detection_scores, 
                capture_result.quality_metrics
            )
            
            # Temporal consistency check
            final_score = self.apply_temporal_consistency(quality_weighted_score)
            
            # Determine final prediction
            is_morphed = final_score > self.config.get('detection_threshold', 0.5)
            confidence = abs(final_score - 0.5) * 2  # Convert to confidence [0, 1]
            
            processing_time = (time.time() - start_time) * 1000
            
            # Create result
            result = MorphDetectionResult(
                is_morphed=is_morphed,
                confidence=confidence,
                morph_score=final_score,
                processing_time_ms=processing_time,
                model_version=self.config.get('model_version', '2.0'),
                quality_factor=capture_result.quality_metrics.overall_score,
                timestamp=datetime.now(),
                session_id=capture_result.session_id,
                frame_number=len(self.detection_history) + 1
            )
            
            # Update history
            self.update_detection_history(result)
            
            # Update performance metrics
            self.update_performance_metrics(processing_time)
            
            return result
            
        except Exception as e:
            logger.error(f"Morph detection failed: {e}")
            self.performance_metrics.error_count += 1
            return None
    
    def should_process(self, quality_metrics: FaceQualityMetrics) -> bool:
        """Determine if frame quality is sufficient for processing"""
        overall_score = quality_metrics.overall_score
        
        # Always process if quality is good
        if overall_score >= self.quality_thresholds['minimum']:
            return True
        
        # Emergency processing for very low quality if enabled
        if (overall_score >= self.quality_thresholds['bypass'] and 
            self.config.get('enable_emergency_processing', False)):
            return True
        
        return False
    
    def extract_face_region(self, capture_result: FaceCaptureResult) -> Optional[np.ndarray]:
        """Extract and preprocess face region from capture result"""
        try:
            image = capture_result.image
            face_box = capture_result.face_box
            
            # Extract face with margin
            margin = self.config.get('face_margin', 0.2)
            x, y, w, h = face_box
            
            # Add margin
            margin_x = int(w * margin)
            margin_y = int(h * margin)
            
            x1 = max(0, x - margin_x)
            y1 = max(0, y - margin_y)
            x2 = min(image.shape[1], x + w + margin_x)
            y2 = min(image.shape[0], y + h + margin_y)
            
            # Extract face region
            face_region = image[y1:y2, x1:x2]
            
            # Quality check on extracted region
            if face_region.size == 0 or face_region.shape[0] < 64 or face_region.shape[1] < 64:
                return None
            
            return face_region
            
        except Exception as e:
            logger.error(f"Face extraction failed: {e}")
            return None
    
    async def run_ensemble_detection(self, face_image: np.ndarray) -> Dict[str, float]:
        """Run ensemble of detection models"""
        detection_scores = {}
        
        for model_name, model in self.models.items():
            try:
                # Preprocess image for this model
                transform = self.transforms[model_name]
                input_tensor = transform(face_image).unsqueeze(0).to(self.device)
                
                # Run inference
                with torch.no_grad():
                    if model_name == 'primary':
                        output = model(input_tensor)
                        score = float(output.squeeze().cpu())
                    elif model_name == 'frequency':
                        output = model(input_tensor)
                        score = float(torch.sigmoid(output).squeeze().cpu())
                    else:
                        score = 0.5  # Default fallback
                
                detection_scores[model_name] = score
                
            except Exception as e:
                logger.error(f"Model {model_name} inference failed: {e}")
                detection_scores[model_name] = 0.5  # Neutral score as fallback
        
        return detection_scores
    
    def apply_quality_weighting(self, 
                              detection_scores: Dict[str, float],
                              quality_metrics: FaceQualityMetrics) -> float:
        """Apply quality-based weighting to ensemble scores"""
        
        # Calculate quality-based weights
        quality_factor = quality_metrics.overall_score
        sharpness_factor = quality_metrics.sharpness_score
        illumination_factor = quality_metrics.illumination_score
        
        # Adjust model weights based on quality
        adjusted_weights = {}
        total_weight = 0
        
        for model_name, config in self.model_configs.items():
            base_weight = config['weight']
            
            # Frequency domain model is more robust to quality issues
            if model_name == 'frequency':
                quality_adjustment = 1.0 + (1.0 - quality_factor) * 0.3
            else:
                quality_adjustment = quality_factor
            
            adjusted_weights[model_name] = base_weight * quality_adjustment
            total_weight += adjusted_weights[model_name]
        
        # Normalize weights
        for model_name in adjusted_weights:
            adjusted_weights[model_name] /= total_weight
        
        # Calculate weighted ensemble score
        ensemble_score = 0
        for model_name, score in detection_scores.items():
            if model_name in adjusted_weights:
                ensemble_score += score * adjusted_weights[model_name]
        
        return ensemble_score
    
    def apply_temporal_consistency(self, current_score: float) -> float:
        """Apply temporal consistency using detection history"""
        if len(self.detection_history) < 3:
            return current_score
        
        # Get recent scores
        recent_scores = [result.morph_score for result in self.detection_history[-3:]]
        recent_scores.append(current_score)
        
        # Apply temporal smoothing
        smoothing_factor = self.config.get('temporal_smoothing', 0.3)
        
        # Weighted average with higher weight on recent frames
        weights = [0.1, 0.2, 0.3, 0.4]  # Current frame has highest weight
        weighted_score = sum(score * weight for score, weight in zip(recent_scores, weights))
        
        # Blend with current score
        final_score = (1 - smoothing_factor) * current_score + smoothing_factor * weighted_score
        
        return final_score
    
    def update_detection_history(self, result: MorphDetectionResult):
        """Update detection history for temporal consistency"""
        self.detection_history.append(result)
        
        # Maintain history length
        if len(self.detection_history) > self.max_history_length:
            self.detection_history.pop(0)
    
    def update_performance_metrics(self, processing_time_ms: float):
        """Update model performance metrics"""
        self.performance_metrics.total_inferences += 1
        self.performance_metrics.inference_time_ms = processing_time_ms
        
        # Update throughput (simplified calculation)
        if processing_time_ms > 0:
            self.performance_metrics.throughput_fps = 1000.0 / processing_time_ms
        
        # Update GPU utilization if available
        if torch.cuda.is_available():
            try:
                self.performance_metrics.gpu_utilization = torch.cuda.utilization()
                self.performance_metrics.memory_usage_mb = torch.cuda.memory_allocated() / 1024 / 1024
            except:
                pass
    
    def get_performance_stats(self) -> Dict[str, Any]:
        """Get current performance statistics"""
        return {
            'inference_time_ms': self.performance_metrics.inference_time_ms,
            'memory_usage_mb': self.performance_metrics.memory_usage_mb,
            'gpu_utilization': self.performance_metrics.gpu_utilization,
            'throughput_fps': self.performance_metrics.throughput_fps,
            'total_inferences': self.performance_metrics.total_inferences,
            'error_count': self.performance_metrics.error_count,
            'error_rate': self.performance_metrics.error_count / max(self.performance_metrics.total_inferences, 1)
        }
    
    def get_model_health(self) -> Dict[str, Any]:
        """Get model health status"""
        health_score = 1.0
        issues = []
        
        # Check inference time
        if self.performance_metrics.inference_time_ms > 500:
            health_score -= 0.3
            issues.append("High inference latency")
        
        # Check error rate
        error_rate = self.performance_metrics.error_count / max(self.performance_metrics.total_inferences, 1)
        if error_rate > 0.1:
            health_score -= 0.4
            issues.append("High error rate")
        
        # Check GPU memory
        if self.performance_metrics.memory_usage_mb > 2048:  # 2GB threshold
            health_score -= 0.2
            issues.append("High memory usage")
        
        # Check model availability
        if not self.models:
            health_score = 0
            issues.append("No models loaded")
        
        return {
            'health_score': max(0, health_score),
            'status': 'healthy' if health_score > 0.8 else 'degraded' if health_score > 0.5 else 'unhealthy',
            'issues': issues,
            'models_loaded': list(self.models.keys()),
            'timestamp': datetime.now().isoformat()
        }
    
    def optimize_for_quality(self, average_quality: float):
        """Dynamically adjust model parameters based on input quality"""
        if average_quality < 0.6:
            # Lower quality: favor frequency domain model
            self.model_configs['frequency']['weight'] = 0.6
            self.model_configs['primary']['weight'] = 0.4
        elif average_quality > 0.8:
            # Higher quality: favor primary model
            self.model_configs['primary']['weight'] = 0.8
            self.model_configs['frequency']['weight'] = 0.2
        else:
            # Balanced quality: default weights
            self.model_configs['primary']['weight'] = 0.7
            self.model_configs['frequency']['weight'] = 0.3
    
    def cleanup(self):
        """Cleanup model resources"""
        for model in self.models.values():
            del model
        
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        logger.info("Model resources cleaned up")