ai_detector.py

"""
Enhanced AI Image Detector - Hugging Face Version
This model detects whether an image is real or AI-generated using a trained PyTorch model.
"""

import os
import cv2
import numpy as np
from PIL import Image
import json
from pytorch_model import PyTorchAIDetector

class EnhancedAIDetector:
    """
    Enhanced detector for AI-generated images using a trained PyTorch model.
    """
    
    def __init__(self, model_path='best_model_improved.pth'):
        """
        Initialize the enhanced detector with a trained PyTorch model.
        
        Args:
            model_path: Path to the trained PyTorch model
        """
        # Initialize the PyTorch model
        self.model = PyTorchAIDetector(model_path)
    
    def analyze_image(self, image_path):
        """
        Analyze an image to detect if it's AI-generated using the PyTorch model.
        
        Args:
            image_path: Path to the image.
            
        Returns:
            Dictionary with analysis results.
        """
        # Use the PyTorch model to analyze the image
        try:
            # Get the model prediction
            result = self.model.analyze_image(image_path)
            
            # Enhance the result with additional metadata
            result.update({
                "model_name": "Enhanced AI Image Detector",
                "model_version": "1.0.0",
                "confidence": float(result["overall_score"] if result["is_ai_generated"] else 1.0 - result["overall_score"])
            })
            
            return result
            
        except Exception as e:
            raise ValueError(f"Failed to analyze image: {str(e)}")
    
    def _analyze_noise_patterns(self, cv_img):
        """
        Analyze noise patterns in the image.
        """
        # Convert to grayscale
        gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
        
        # Apply ELA (Error Level Analysis)
        # Save image at 90% quality
        temp_path = "temp_ela.jpg"
        cv2.imwrite(temp_path, cv_img, [cv2.IMWRITE_JPEG_QUALITY, 90])
        
        # Reload the image
        compressed_img = cv2.imread(temp_path)
        os.remove(temp_path)
        
        # Calculate difference between original and compressed
        if compressed_img is not None:
            diff = cv2.absdiff(cv_img, compressed_img)
            diff_gray = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
            
            # Calculate difference statistics
            mean_diff = np.mean(diff_gray)
            std_diff = np.std(diff_gray)
            
            # Normalize result to [0, 1] range
            # Higher values indicate higher likelihood of AI-generated
            ela_score = min(mean_diff / 8.0, 1.0)
        else:
            ela_score = 0.5  # Neutral value if ELA analysis fails
        
        # Analyze noise using Laplacian filter
        laplacian = cv2.Laplacian(gray, cv2.CV_64F)
        laplacian_std = np.std(laplacian)
        
        # Normalize result to [0, 1] range
        # Lower values indicate higher likelihood of AI-generated
        # (AI-generated images tend to be smoother)
        laplacian_score = 1.0 - min(laplacian_std / 15.0, 1.0)
        
        # Analyze noise in frequency domain
        # Apply Fourier transform
        f_transform = np.fft.fft2(gray)
        f_shift = np.fft.fftshift(f_transform)
        magnitude_spectrum = 20 * np.log(np.abs(f_shift) + 1)
        
        # Calculate ratio of high to low frequencies
        h, w = gray.shape
        center_y, center_x = h // 2, w // 2
        radius = min(center_y, center_x) // 4
        
        # Create mask for low frequencies
        y, x = np.ogrid[:h, :w]
        low_freq_mask = ((y - center_y) ** 2 + (x - center_x) ** 2) <= radius ** 2
        
        # Calculate average of low and high frequencies
        low_freq_mean = np.mean(magnitude_spectrum[low_freq_mask])
        high_freq_mean = np.mean(magnitude_spectrum[~low_freq_mask])
        
        # Calculate ratio
        freq_ratio = high_freq_mean / low_freq_mean if low_freq_mean > 0 else 0
        
        # Normalize result to [0, 1] range
        # Lower values indicate higher likelihood of AI-generated
        freq_score = 1.0 - min(freq_ratio / 0.4, 1.0)
        
        # Combine results with optimized weights
        noise_score = 0.4 * ela_score + 0.35 * laplacian_score + 0.25 * freq_score
        
        return noise_score
    
    def _analyze_texture(self, cv_img):
        """
        Analyze image texture.
        """
        # Convert to grayscale
        gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
        
        # Calculate Gray Level Co-occurrence Matrix (GLCM)
        distances = [1]
        angles = [0, np.pi/4, np.pi/2, 3*np.pi/4]
        
        # Use skimage to calculate GLCM
        try:
            glcm = feature.graycomatrix(gray, distances, angles, 256, symmetric=True, normed=True)
            
            # Extract GLCM features
            contrast = feature.graycoprops(glcm, 'contrast')[0, 0]
            dissimilarity = feature.graycoprops(glcm, 'dissimilarity')[0, 0]
            homogeneity = feature.graycoprops(glcm, 'homogeneity')[0, 0]
            energy = feature.graycoprops(glcm, 'energy')[0, 0]
            correlation = feature.graycoprops(glcm, 'correlation')[0, 0]
            
            # Normalize values
            contrast_norm = min(contrast / 80.0, 1.0)
            dissimilarity_norm = min(dissimilarity / 8.0, 1.0)
            homogeneity_norm = homogeneity
            energy_norm = energy
            correlation_norm = correlation
            
            # AI-generated images tend to be more homogeneous and less contrasted
            texture_score = (
                0.15 * (1.0 - contrast_norm) +
                0.15 * (1.0 - dissimilarity_norm) +
                0.25 * homogeneity_norm +
                0.25 * energy_norm +
                0.20 * correlation_norm
            )
        except Exception:
            # Use alternative analysis if GLCM fails
            # Calculate image variance
            variance = np.var(gray)
            variance_norm = min(variance / 800.0, 1.0)
            
            # Calculate entropy
            hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
            hist = hist / hist.sum()
            entropy = -np.sum(hist * np.log2(hist + 1e-10))
            entropy_norm = min(entropy / 7.5, 1.0)
            
            # AI-generated images tend to have less variance and entropy
            texture_score = 0.6 * (1.0 - variance_norm) + 0.4 * (1.0 - entropy_norm)
        
        return texture_score
    
    def _analyze_color_coherence(self, cv_img):
        """
        Analyze color coherence in the image.
        """
        # Split image into BGR channels
        b, g, r = cv2.split(cv_img)
        
        # Calculate standard deviation for each channel
        std_b = np.std(b)
        std_g = np.std(g)
        std_r = np.std(r)
        
        # Calculate average standard deviation
        avg_std = (std_b + std_g + std_r) / 3.0
        
        # Normalize result to [0, 1] range
        # Lower values indicate higher likelihood of AI-generated
        std_score = 1.0 - min(avg_std / 45.0, 1.0)
        
        # Analyze color distribution
        hist_b = cv2.calcHist([b], [0], None, [64], [0, 256])
        hist_g = cv2.calcHist([g], [0], None, [64], [0, 256])
        hist_r = cv2.calcHist([r], [0], None, [64], [0, 256])
        
        # Normalize histograms
        hist_b = hist_b / hist_b.sum()
        hist_g = hist_g / hist_g.sum()
        hist_r = hist_r / hist_r.sum()
        
        # Calculate histogram entropy
        entropy_b = -np.sum(hist_b * np.log2(hist_b + 1e-10))
        entropy_g = -np.sum(hist_g * np.log2(hist_g + 1e-10))
        entropy_r = -np.sum(hist_r * np.log2(hist_r + 1e-10))
        
        # Calculate average entropy
        avg_entropy = (entropy_b + entropy_g + entropy_r) / 3.0
        
        # Normalize result to [0, 1] range
        # Lower values indicate higher likelihood of AI-generated
        entropy_score = 1.0 - min(avg_entropy / 5.8, 1.0)
        
        # Analyze channel relationships
        correlation_bg = np.corrcoef(b.flatten(), g.flatten())[0, 1]
        correlation_br = np.corrcoef(b.flatten(), r.flatten())[0, 1]
        correlation_gr = np.corrcoef(g.flatten(), r.flatten())[0, 1]
        
        # Calculate average correlation
        avg_correlation = (abs(correlation_bg) + abs(correlation_br) + abs(correlation_gr)) / 3.0
        
        # Normalize result to [0, 1] range
        # Higher values indicate higher likelihood of AI-generated
        correlation_score = min(avg_correlation, 1.0)
        
        # Combine results with optimized weights
        color_score = 0.25 * std_score + 0.25 * entropy_score + 0.5 * correlation_score
        
        return color_score
    
    def _analyze_edges(self, cv_img):
        """
        Analyze edges in the image.
        """
        # Convert to grayscale
        gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
        
        # Detect edges using Canny
        edges = cv2.Canny(gray, 100, 200)
        
        # Calculate ratio of edge pixels
        edge_ratio = np.count_nonzero(edges) / edges.size
        
        # Normalize result to [0, 1] range
        # Lower values indicate higher likelihood of AI-generated
        edge_ratio_score = 1.0 - min(edge_ratio / 0.08, 1.0)
        
        # Analyze edge consistency
        # Apply Sobel filter
        sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
        sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
        
        # Calculate gradient direction
        gradient_direction = np.arctan2(sobely, sobelx)
        
        # Calculate histogram of gradient directions
        hist, _ = np.histogram(gradient_direction, bins=36, range=(-np.pi, np.pi))
        hist = hist / hist.sum()
        
        # Calculate histogram entropy
        entropy = -np.sum(hist * np.log2(hist + 1e-10))
        
        # Normalize result to [0, 1] range
        # Lower values indicate higher likelihood of AI-generated
        entropy_score = 1.0 - min(entropy / 5.0, 1.0)
        
        # Additional analysis: edge consistency
        # Calculate gradient magnitude
        gradient_magnitude = np.sqrt(sobelx**2 + sobely**2)
        
        # Calculate standard deviation of gradient magnitude
        gradient_std = np.std(gradient_magnitude)
        
        # Normalize result to [0, 1] range
        # Lower values indicate higher likelihood of AI-generated
        gradient_score = 1.0 - min(gradient_std / 30.0, 1.0)
        
        # Combine results with optimized weights
        edge_score = 0.4 * edge_ratio_score + 0.3 * entropy_score + 0.3 * gradient_score
        
        return edge_score
    
    def _analyze_faces(self, cv_img):
        """
        Analyze faces in the image.
        """
        # Convert to grayscale
        gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
        
        # Detect faces
        faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)
        
        if len(faces) == 0:
            return None
        
        face_scores = []
        
        for (x, y, w, h) in faces:
            # Extract face region
            face_roi = gray[y:y+h, x:x+w]
            
            # Detect eyes
            eyes = self.eye_cascade.detectMultiScale(face_roi)
            
            # Analyze face symmetry
            symmetry_score = self._analyze_face_symmetry(face_roi)
            
            # Analyze face texture
            texture_score = self._analyze_face_texture(face_roi)
            
            # Analyze face proportions
            proportion_score = self._analyze_face_proportions(face_roi, eyes)
            
            # Additional analysis: face details
            detail_score = self._analyze_face_details(face_roi)
            
            # Combine results with optimized weights
            face_score = 0.3 * symmetry_score + 0.25 * texture_score + 0.25 * proportion_score + 0.2 * detail_score
            face_scores.append(face_score)
        
        # Calculate average face score
        avg_face_score = np.mean(face_scores) if face_scores else 0.5
        
        return avg_face_score
    
    def _analyze_face_symmetry(self, face_roi):
        """
        Analyze face symmetry.
        """
        # Split face into halves
        h, w = face_roi.shape
        midpoint = w // 2
        
        left_half = face_roi[:, :midpoint]
        right_half = face_roi[:, midpoint:]
        
        # Flip right half
        right_half_flipped = cv2.flip(right_half, 1)
        
        # Resize flipped right half to match left half
        if left_half.shape[1] != right_half_flipped.shape[1]:
            right_half_flipped = cv2.resize(right_half_flipped, (left_half.shape[1], left_half.shape[0]))
        
        # Calculate difference between halves
        diff = cv2.absdiff(left_half, right_half_flipped)
        
        # Calculate average difference
        mean_diff = np.mean(diff)
        
        # Normalize result to [0, 1] range
        # Lower values indicate more symmetry
        # AI-generated faces tend to be more symmetrical
        symmetry_score = 1.0 - min(mean_diff / 25.0, 1.0)
        
        return symmetry_score
    
    def _analyze_face_texture(self, face_roi):
        """
        Analyze face texture.
        """
        # Apply Laplacian filter to detect details
        laplacian = cv2.Laplacian(face_roi, cv2.CV_64F)
        laplacian_std = np.std(laplacian)
        
        # Normalize result to [0, 1] range
        # Lower values indicate fewer details
        # AI-generated faces tend to have fewer details
        texture_score = 1.0 - min(laplacian_std / 15.0, 1.0)
        
        return texture_score
    
    def _analyze_face_proportions(self, face_roi, eyes):
        """
        Analyze face proportions.
        """
        h, w = face_roi.shape
        
        if len(eyes) >= 2:
            # Sort eyes from left to right
            eyes = sorted(eyes, key=lambda e: e[0])
            
            # Calculate distance between eyes
            eye1_center = (eyes[0][0] + eyes[0][2] // 2, eyes[0][1] + eyes[0][3] // 2)
            eye2_center = (eyes[1][0] + eyes[1][2] // 2, eyes[1][1] + eyes[1][3] // 2)
            
            eye_distance = np.sqrt((eye2_center[0] - eye1_center[0]) ** 2 + (eye2_center[1] - eye1_center[1]) ** 2)
            
            # Calculate ratio of eye distance to face width
            eye_ratio = eye_distance / w
            
            # Analyze ratio
            # In natural faces, eye distance is about 0.3 of face width
            proportion_score = 1.0 - min(abs(eye_ratio - 0.3) / 0.2, 1.0)
        else:
            # If two eyes are not found, use neutral value
            proportion_score = 0.5
        
        return proportion_score
    
    def _analyze_face_details(self, face_roi):
        """
        Analyze face details.
        """
        # Apply Gabor filter to detect fine details
        # Create Gabor filter
        ksize = 31
        sigma = 4.0
        theta = 0
        lambd = 10.0
        gamma = 0.5
        
        # Apply Gabor filter at different angles
        angles = [0, np.pi/4, np.pi/2, 3*np.pi/4]
        responses = []
        
        for theta in angles:
            kernel = cv2.getGaborKernel((ksize, ksize), sigma, theta, lambd, gamma, 0, ktype=cv2.CV_32F)
            filtered = cv2.filter2D(face_roi, cv2.CV_8UC3, kernel)
            responses.append(filtered)
        
        # Combine responses
        combined_response = np.zeros_like(face_roi)
        for response in responses:
            combined_response = cv2.add(combined_response, response)
        
        # Calculate standard deviation of combined response
        response_std = np.std(combined_response)
        
        # Normalize result to [0, 1] range
        # Lower values indicate fewer details
        # AI-generated faces tend to have fewer details
        detail_score = 1.0 - min(response_std / 25.0, 1.0)
        
        return detail_score
    
    def _analyze_additional_features(self, cv_img):
        """
        Analyze additional features for borderline images.
        """
        # Convert to grayscale
        gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
        
        # Analyze repetition in the image
        # Apply Fourier transform
        f_transform = np.fft.fft2(gray)
        f_shift = np.fft.fftshift(f_transform)
        magnitude_spectrum = np.log(np.abs(f_shift) + 1)
        
        # Calculate standard deviation of spectrum
        spectrum_std = np.std(magnitude_spectrum)
        
        # Normalize result to [0, 1] range
        # Lower values indicate higher likelihood of AI-generated
        spectrum_score = 1.0 - min(spectrum_std / 2.0, 1.0)
        
        # Analyze color gradients
        # Convert to HSV space
        hsv = cv2.cvtColor(cv_img, cv2.COLOR_BGR2HSV)
        h, s, v = cv2.split(hsv)
        
        # Calculate histogram of saturation channel
        hist_s = cv2.calcHist([s], [0], None, [64], [0, 256])
        hist_s = hist_s / hist_s.sum()
        
        # Calculate histogram entropy
        entropy_s = -np.sum(hist_s * np.log2(hist_s + 1e-10))
        
        # Normalize result to [0, 1] range
        # Lower values indicate higher likelihood of AI-generated
        entropy_score = 1.0 - min(entropy_s / 5.0, 1.0)
        
        # Combine results
        additional_score = 0.6 * spectrum_score + 0.4 * entropy_score
        
        return additional_score