app.py

import os
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
import cv2
from scipy.fftpack import fft2, fftshift
from skimage.feature import graycomatrix, graycoprops, local_binary_pattern
import timm
import gradio as gr

# GLCM feature extraction
def extract_glcm_features(image):
    image_uint8 = (image * 255).astype(np.uint8)
    image_uint8 = image_uint8 // 4

    glcm = graycomatrix(
        image_uint8,
        distances=[1],
        angles=[0, np.pi / 4, np.pi / 2, 3 * np.pi / 4],
        levels=64,
        symmetric=True,
        normed=True
    )

    contrast = graycoprops(glcm, 'contrast').flatten()
    dissimilarity = graycoprops(glcm, 'dissimilarity').flatten()
    homogeneity = graycoprops(glcm, 'homogeneity').flatten()
    energy = graycoprops(glcm, 'energy').flatten()
    correlation = graycoprops(glcm, 'correlation').flatten()

    features = np.hstack([contrast, dissimilarity, homogeneity, energy, correlation])
    return features.astype(np.float32)

# Spectrum analysis
def analyze_spectrum(image, target_spectrum_length=181):
    f = fft2(image)
    fshift = fftshift(f)
    magnitude_spectrum = 20 * np.log(np.abs(fshift) + 1e-8)

    center = np.array(magnitude_spectrum.shape) // 2
    y, x = np.indices(magnitude_spectrum.shape)
    r = np.sqrt((x - center[1])**2 + (y - center[0])**2).astype(int)

    radial_mean = np.bincount(r.ravel(), magnitude_spectrum.ravel()) / np.bincount(r.ravel())

    if len(radial_mean) < target_spectrum_length:
        radial_mean = np.pad(radial_mean, (0, target_spectrum_length - len(radial_mean)), 'constant')
    else:
        radial_mean = radial_mean[:target_spectrum_length]

    return radial_mean.astype(np.float32)

# Edge feature extraction
def extract_edge_features(image):
    image_uint8 = (image * 255).astype(np.uint8)
    edges = cv2.Canny(image_uint8, 100, 200)
    edges_resized = cv2.resize(edges, (64, 64), interpolation=cv2.INTER_AREA)
    return edges_resized.astype(np.float32) / 255.0

# LBP feature extraction
def extract_lbp_features(image):
    radius = 1
    n_points = 8 * radius
    METHOD = 'uniform'

    lbp = local_binary_pattern(image, n_points, radius, METHOD)
    n_bins = n_points + 2
    hist, _ = np.histogram(lbp.ravel(), bins=n_bins, range=(0, n_bins), density=True)

    return hist.astype(np.float32)

# Model architecture
class AttentionBlock(nn.Module):
    def __init__(self, in_features):
        super(AttentionBlock, self).__init__()
        self.attention = nn.Sequential(
            nn.Linear(in_features, max(in_features // 8, 1)),
            nn.ReLU(),
            nn.Linear(max(in_features // 8, 1), in_features),
            nn.Sigmoid()
        )

    def forward(self, x):
        attention_weights = self.attention(x)
        return x * attention_weights

class AdvancedFaceDetectionModel(nn.Module):
    def __init__(self, spectrum_length=181, lbp_n_bins=10):
        super(AdvancedFaceDetectionModel, self).__init__()

        self.efficientnet = timm.create_model('tf_efficientnetv2_b2', pretrained=False, num_classes=0)
        for param in self.efficientnet.conv_stem.parameters():
            param.requires_grad = False
        for param in self.efficientnet.bn1.parameters():
            param.requires_grad = False

        self.glcm_fc = nn.Sequential(
            nn.Linear(20, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Dropout(0.5)
        )

        self.spectrum_conv = nn.Sequential(
            nn.Conv1d(1, 64, kernel_size=3, padding=1),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.AdaptiveAvgPool1d(1)
        )

        self.edge_conv = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((8, 8))
        )

        self.lbp_fc = nn.Sequential(
            nn.Linear(lbp_n_bins, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Dropout(0.5)
        )

        image_feature_size = self.efficientnet.num_features
        self.image_attention = AttentionBlock(image_feature_size)
        self.glcm_attention = AttentionBlock(64)
        self.spectrum_attention = AttentionBlock(64)
        self.edge_attention = AttentionBlock(32 * 8 * 8)
        self.lbp_attention = AttentionBlock(64)

        total_features = image_feature_size + 64 + 64 + (32 * 8 * 8) + 64
        self.fusion = nn.Sequential(
            nn.Linear(total_features, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 1)
        )

    def forward(self, image, glcm_features, spectrum_features, edge_features, lbp_features):
        image_features = self.efficientnet(image)
        image_features = self.image_attention(image_features)

        glcm_features = self.glcm_fc(glcm_features)
        glcm_features = self.glcm_attention(glcm_features)

        spectrum_features = self.spectrum_conv(spectrum_features.unsqueeze(1))
        spectrum_features = spectrum_features.squeeze(2)
        spectrum_features = self.spectrum_attention(spectrum_features)

        edge_features = self.edge_conv(edge_features.unsqueeze(1))
        edge_features = edge_features.view(edge_features.size(0), -1)
        edge_features = self.edge_attention(edge_features)

        lbp_features = self.lbp_fc(lbp_features)
        lbp_features = self.lbp_attention(lbp_features)

        combined_features = torch.cat(
            (image_features, glcm_features, spectrum_features, edge_features, lbp_features), dim=1
        )

        output = self.fusion(combined_features)
        return output.squeeze(1)

# Initialize model and transform
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = AdvancedFaceDetectionModel(spectrum_length=181, lbp_n_bins=10).to(device)
model.load_state_dict(torch.load('best_model.pth', map_location=device))
model.eval()

transform = 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 predict_image(image):
    """
    Process a single image and return prediction
    """
    # Convert to PIL Image if needed
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    
    # Convert to RGB if needed
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    # Apply transformations
    image_tensor = transform(image).unsqueeze(0)

    # Convert to NumPy array for feature extraction
    np_image = image_tensor.cpu().numpy().squeeze(0).transpose(1, 2, 0)
    np_image = np.clip(np_image, 0, 1)

    # Convert to grayscale
    gray_image = cv2.cvtColor((np_image * 255).astype(np.uint8), cv2.COLOR_RGB2GRAY)
    gray_image = gray_image.astype(np.float32) / 255.0

    # Extract features
    glcm_features = extract_glcm_features(gray_image)
    spectrum_features = analyze_spectrum(gray_image)
    edge_features = extract_edge_features(gray_image)
    lbp_features = extract_lbp_features(gray_image)

    # Move everything to device
    with torch.no_grad():
        image_tensor = image_tensor.to(device)
        glcm_features = torch.from_numpy(glcm_features).unsqueeze(0).to(device)
        spectrum_features = torch.from_numpy(spectrum_features).unsqueeze(0).to(device)
        edge_features = torch.from_numpy(edge_features).unsqueeze(0).to(device)
        lbp_features = torch.from_numpy(lbp_features).unsqueeze(0).to(device)

        # Forward pass
        outputs = model(image_tensor, glcm_features, spectrum_features, edge_features, lbp_features)
        probability = torch.sigmoid(outputs).item()
        prediction = "Real Face" if probability > 0.5 else "Fake Face"
        
        return prediction, f"Confidence: {probability:.2%}"

# Create Gradio interface
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Label(label="Prediction"),
        gr.Label(label="Confidence")
    ],
    title="Face Authentication System",
    description="Upload an image to determine if it contains a real or fake face.",
    examples=[
        ["example1.jpg"],
        ["example2.jpg"]
    ] if os.path.exists("example1.jpg") else None,
)

# Launch the app
if __name__ == "__main__":
    iface.launch()