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

# [把你原来代码中的AttentionBlock和AdvancedFaceDetectionModel类定义放在这里]
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=True, 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)

# 特征提取函数
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)

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)

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

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)

# 加载模型
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):
    """
    处理上传的图片并返回预测结果
    """
    if image is None:
        return "请上传图片"
    
    # 转换图片格式
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    
    # 应用转换
    image_tensor = transform(image).unsqueeze(0)
    
    # 准备特征提取用的图像
    np_image = image_tensor.cpu().numpy().squeeze(0).transpose(1, 2, 0)
    np_image = np.clip(np_image, 0, 1)
    
    # 转换为灰度图
    gray_image = cv2.cvtColor((np_image * 255).astype(np.uint8), cv2.COLOR_RGB2GRAY)
    gray_image = gray_image.astype(np.float32) / 255.0
    
    # 提取特征
    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)
    
    # 转换为张量并移到设备
    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)
        
        # 模型预测
        outputs = model(image_tensor, glcm_features, spectrum_features, edge_features, lbp_features)
        prediction = torch.sigmoid(outputs).item()
        
        # 返回预测结果
        if prediction > 0.5:
            return "真实人脸图片"
        else:
            return "虚假人脸图片"

# 创建Gradio界面
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Text(label="预测结果"),
    title="人脸真伪检测",
    description="上传一张人脸图片，模型将判断是真实人脸还是虚假人脸。",
    examples=[
        # 这里可以添加示例图片路径
    ]
)

# 启动应用
iface.launch()