utils/detector.py

import torch
import torch.nn as nn
import cv2
import numpy as np
from PIL import Image
import mediapipe as mp
from facenet_pytorch import MTCNN
import time
import warnings
warnings.filterwarnings('ignore')

class DeepfakeDetector:
    def __init__(self, device='cuda' if torch.cuda.is_available() else 'cpu'):
        self.device = device
        self.face_detector = MTCNN(keep_all=True, device=device)
        self.mp_face_mesh = mp.solutions.face_mesh
        self.face_mesh = self.mp_face_mesh.FaceMesh(
            static_image_mode=True,
            max_num_faces=1,
            refine_landmarks=True,
            min_detection_confidence=0.5
        )
        
        # Initialize models
        self.models = self.load_models()
        self.threshold = 0.7
        
    def load_models(self):
        """Load pretrained models"""
        models = {}
        
        # Load EfficientNet-B4
        from efficientnet_pytorch import EfficientNet
        models['efficientnet'] = EfficientNet.from_pretrained('efficientnet-b4')
        models['efficientnet']._fc = nn.Linear(1792, 2)
        
        # Load Xception
        from torchvision.models import xception
        models['xception'] = xception(pretrained=False)
        models['xception'].fc = nn.Linear(2048, 2)
        
        # Move to device and set to eval mode
        for name, model in models.items():
            model_path = f"models/{name}.pth"
            try:
                model.load_state_dict(torch.load(model_path, map_location=self.device))
                print(f"Loaded {name}")
            except:
                print(f"Using pretrained {name} without fine-tuning")
            model.to(self.device)
            model.eval()
        
        return models
    
    def detect_image(self, image):
        """Detect deepfake in image"""
        start_time = time.time()
        
        # Convert to numpy if PIL
        if isinstance(image, Image.Image):
            image = np.array(image)
        
        # Run all detection methods
        results = {}
        
        # Frequency analysis
        results['frequency_score'] = self.analyze_frequency(image)
        
        # Face artifact detection
        face_results = self.analyze_faces(image)
        results['face_score'] = face_results['confidence']
        results['num_faces'] = face_results['num_faces']
        
        # Model predictions
        model_predictions = []
        for name, model in self.models.items():
            pred = self.predict_with_model(image, model)
            model_predictions.append(pred)
        
        # Ensemble voting
        final_score = np.mean([
            results['frequency_score'],
            results['face_score'],
            *model_predictions
        ])
        
        results['is_fake'] = final_score > self.threshold
        results['confidence'] = final_score
        results['quality_score'] = self.assess_quality(image)
        results['processing_time'] = time.time() - start_time
        
        return results
    
    def detect_video(self, video_path, sample_frames=30):
        """Detect deepfake in video"""
        start_time = time.time()
        
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise ValueError(f"Cannot open video: {video_path}")
        
        # Get video info
        fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        # Sample frames
        frame_indices = np.linspace(0, total_frames-1, min(sample_frames, total_frames), dtype=int)
        frame_results = []
        
        for frame_idx in frame_indices:
            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
            ret, frame = cap.read()
            if ret:
                # Convert BGR to RGB
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                result = self.detect_image(frame_rgb)
                frame_results.append(result)
        
        cap.release()
        
        # Aggregate results
        if not frame_results:
            raise ValueError("No frames could be read from video")
        
        # Calculate video-level metrics
        confidences = [r['confidence'] for r in frame_results]
        fake_flags = [r['is_fake'] for r in frame_results]
        
        final_result = {
            'is_fake': np.mean(fake_flags) > 0.5,
            'confidence': np.mean(confidences),
            'duration': total_frames / fps,
            'frames_analyzed': len(frame_results),
            'resolution': f"{width}x{height}",
            'fps': fps,
            'frame_results': frame_results,
            'processing_time': time.time() - start_time,
            'fake_segments': self.identify_fake_segments(frame_results, frame_indices, fps)
        }
        
        return final_result
    
    def analyze_frequency(self, image):
        """Analyze frequency domain"""
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        else:
            gray = image
        
        # Fourier Transform
        f = np.fft.fft2(gray)
        fshift = np.fft.fftshift(f)
        magnitude = np.log(np.abs(fshift) + 1)
        
        # Analyze frequency patterns
        height, width = magnitude.shape
        center_h, center_w = height // 2, width // 2
        
        # Check for grid-like patterns common in GANs
        low_freq = magnitude[center_h-20:center_h+20, center_w-20:center_w+20]
        high_freq = np.copy(magnitude)
        high_freq[center_h-20:center_h+20, center_w-20:center_w+20] = 0
        
        low_energy = np.mean(low_freq)
        high_energy = np.mean(high_freq)
        
        # Deepfakes often have different frequency distributions
        score = min(high_energy / (low_energy + 1e-10) * 0.5, 1.0)
        
        return score
    
    def analyze_faces(self, image):
        """Analyze faces in image"""
        # Detect faces
        boxes, probs = self.face_detector.detect(image)
        
        if boxes is None:
            return {'confidence': 0.0, 'num_faces': 0}
        
        num_faces = len(boxes)
        face_scores = []
        
        for i, box in enumerate(boxes):
            if probs[i] < 0.9:
                continue
            
            # Extract face
            x1, y1, x2, y2 = map(int, box)
            face = image[y1:y2, x1:x2]
            
            if face.size == 0:
                continue
            
            # Analyze face artifacts
            score = self.analyze_face_artifacts(face)
            face_scores.append(score)
        
        if not face_scores:
            return {'confidence': 0.0, 'num_faces': num_faces}
        
        return {
            'confidence': np.mean(face_scores),
            'num_faces': num_faces
        }
    
    def analyze_face_artifacts(self, face_img):
        """Analyze artifacts in face image"""
        # Check for unnatural symmetry
        if face_img.shape[1] > 10:  # Ensure face is wide enough
            left_half = face_img[:, :face_img.shape[1]//2]
            right_half = face_img[:, face_img.shape[1]//2:]
            right_half_flipped = np.fliplr(right_half)
            
            # Resize to match
            min_height = min(left_half.shape[0], right_half_flipped.shape[0])
            min_width = min(left_half.shape[1], right_half_flipped.shape[1])
            
            left_cropped = left_half[:min_height, :min_width]
            right_cropped = right_half_flipped[:min_height, :min_width]
            
            # Calculate symmetry
            if left_cropped.size > 0 and right_cropped.size > 0:
                symmetry_error = np.mean(np.abs(left_cropped - right_cropped))
                symmetry_score = min(symmetry_error / 10.0, 1.0)
            else:
                symmetry_score = 0.5
        else:
            symmetry_score = 0.5
        
        # Check for unnatural edges
        gray = cv2.cvtColor(face_img, cv2.COLOR_RGB2GRAY)
        edges = cv2.Canny(gray, 100, 200)
        edge_density = np.sum(edges) / edges.size
        
        # Combine scores
        final_score = (symmetry_score * 0.6 + edge_density * 0.4)
        
        return final_score
    
    def predict_with_model(self, image, model):
        """Predict using a specific model"""
        # Preprocess image
        transform = self.get_transform()
        
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        
        input_tensor = transform(image).unsqueeze(0).to(self.device)
        
        with torch.no_grad():
            output = model(input_tensor)
            probabilities = torch.softmax(output, dim=1)
            fake_prob = probabilities[0][1].item()
        
        return fake_prob
    
    def get_transform(self):
        """Get image transformation pipeline"""
        from torchvision import transforms
        
        return transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])
    
    def assess_quality(self, image):
        """Assess image quality"""
        # Simple quality metrics
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        else:
            gray = image
        
        # Calculate sharpness (variance of Laplacian)
        laplacian_var = cv2.Laplacian(gray, cv2.CV_64F).var()
        sharpness_score = min(laplacian_var / 1000.0, 1.0)
        
        # Calculate contrast
        contrast_score = np.std(gray) / 255.0
        
        return (sharpness_score + contrast_score) / 2
    
    def identify_fake_segments(self, frame_results, frame_indices, fps):
        """Identify segments in video that are likely deepfakes"""
        if not frame_results:
            return []
        
        segments = []
        current_segment = None
        
        for i, result in enumerate(frame_results):
            if result['is_fake']:
                if current_segment is None:
                    current_segment = {
                        'start': frame_indices[i] / fps,
                        'end': frame_indices[i] / fps,
                        'confidence': [result['confidence']]
                    }
                else:
                    current_segment['end'] = frame_indices[i] / fps
                    current_segment['confidence'].append(result['confidence'])
            else:
                if current_segment is not None:
                    current_segment['confidence'] = np.mean(current_segment['confidence'])
                    segments.append(current_segment)
                    current_segment = None
        
        # Add last segment if exists
        if current_segment is not None:
            current_segment['confidence'] = np.mean(current_segment['confidence'])
            segments.append(current_segment)
        
        return segments
    
    def visualize_result(self, image, result):
        """Create visualization of detection result"""
        # Convert to BGR for OpenCV
        if isinstance(image, Image.Image):
            image = np.array(image)
        
        if len(image.shape) == 3 and image.shape[2] == 3:
            vis = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        else:
            vis = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
        
        # Add result text
        text = "REAL" if not result['is_fake'] else "DEEPFAKE"
        color = (0, 255, 0) if not result['is_fake'] else (0, 0, 255)
        
        # Add text background
        text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 2, 3)[0]
        cv2.rectangle(vis, (10, 10), (10 + text_size[0] + 20, 10 + text_size[1] + 20), (0, 0, 0), -1)
        
        # Add text
        cv2.putText(vis, text, (20, 20 + text_size[1]), 
                   cv2.FONT_HERSHEY_SIMPLEX, 2, color, 3)
        
        # Add confidence
        conf_text = f"Confidence: {result['confidence']:.2%}"
        cv2.putText(vis, conf_text, (20, 80), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
        
        # Convert back to RGB
        vis = cv2.cvtColor(vis, cv2.COLOR_BGR2RGB)
        
        return vis
    
    def detect_file(self, file_path):
        """Detect deepfake in file (auto-detect type)"""
        if file_path.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
            # Image file
            image = Image.open(file_path)
            result = self.detect_image(image)
            result['type'] = 'image'
        elif file_path.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
            # Video file
            result = self.detect_video(file_path)
            result['type'] = 'video'
        else:
            raise ValueError(f"Unsupported file type: {file_path}")
        
        result['filename'] = file_path
        return result