Update app.py

app.py

@@ -10,68 +10,7 @@ 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
+# [把你原来代码中的AttentionBlock和AdvancedFaceDetectionModel类定义放在这里]
 class AttentionBlock(nn.Module):
     def __init__(self, in_features):
         super(AttentionBlock, self).__init__()
@@ -90,7 +29,7 @@ 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)
+        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():
@@ -169,12 +108,72 @@ class AdvancedFaceDetectionModel(nn.Module):
         output = self.fusion(combined_features)
         return output.squeeze(1)
 
-# Initialize model and transform
+# 特征提取函数
+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(),
@@ -183,64 +182,61 @@ transform = transforms.Compose([
 
 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)
+    if image is None:
+        return "请上传图片"
     
-    # Convert to RGB if needed
-    if image.mode != 'RGB':
-        image = image.convert('RGB')
+    # 转换图片格式
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
     
-    # 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"
+        prediction = torch.sigmoid(outputs).item()
         
-        return prediction, f"Confidence: {probability:.2%}"
+        # 返回预测结果
+        if prediction > 0.5:
+            return "真实人脸图片"
+        else:
+            return "虚假人脸图片"
 
-# Create Gradio interface
+# 创建Gradio界面
 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.",
+    outputs=gr.Text(label="预测结果"),
+    title="人脸真伪检测",
+    description="上传一张人脸图片，模型将判断是真实人脸还是虚假人脸。",
     examples=[
-        ["example1.jpg"],
-        ["example2.jpg"]
-    ] if os.path.exists("example1.jpg") else None,
+        # 这里可以添加示例图片路径
+    ]
 )
 
-# Launch the app
-if __name__ == "__main__":
-    iface.launch()
+# 启动应用
+iface.launch()