Update app.py

app.py

@@ -1,88 +1,84 @@
-import os
-import cv2
-import numpy as np
-import gradio as gr
-import tensorflow as tf
-from mtcnn import MTCNN
-from huggingface_hub import hf_hub_download
-from tensorflow.keras.models import load_model
-from tensorflow.keras.applications.xception import preprocess_input as xcp_pre
-from tensorflow.keras.applications.efficientnet import preprocess_input as eff_pre
-
-# ---------------------------------------------------------
-# Load models from Hugging Face Hub
-# ---------------------------------------------------------
-xcp_path = hf_hub_download(repo_id="Zeyadd-Mostaffa/deepfake-image-detector_final", filename="xception_model.h5")
-eff_path = hf_hub_download(repo_id="Zeyadd-Mostaffa/deepfake-image-detector_final", filename="efficientnet_model.h5")
-xcp_model = load_model(xcp_path)
-eff_model = load_model(eff_path)
-
-# ---------------------------------------------------------
-# Face Detection
-# ---------------------------------------------------------
-detector = MTCNN()
-
-def extract_faces(image):
-    faces = detector.detect_faces(image)
-    if not faces:
-        return []
-    results = []
-    for i, face in enumerate(faces):
-        x, y, w, h = face['box']
-        x, y = max(0, x), max(0, y)
-        cropped = image[y:y+h, x:x+w]
-        if cropped.shape[0] >= 60 and cropped.shape[1] >= 60:
-            results.append((cropped, (x, y, w, h)))
-    return results
-
-# ---------------------------------------------------------
-# Inference Function
-# ---------------------------------------------------------
-def predict_faces(image):
-    faces = extract_faces(image)
-    if not faces:
-        return "No faces detected", None
-
-    annotated = image.copy()
-    results = []
-
-    for i, (face, (x, y, w, h)) in enumerate(faces):
-        # Preprocess
-        xcp_img = cv2.resize(face, (299, 299))
-        eff_img = cv2.resize(face, (224, 224))
-
-        xcp_tensor = xcp_pre(xcp_img.astype(np.float32))[np.newaxis, ...]
-        eff_tensor = eff_pre(eff_img.astype(np.float32))[np.newaxis, ...]
-
-        # Predict
-        xcp_pred = xcp_model.predict(xcp_tensor, verbose=0).flatten()[0]
-        eff_pred = eff_model.predict(eff_tensor, verbose=0).flatten()[0]
-        avg_pred = (xcp_pred + eff_pred) / 2
-
-        label = "Real" if avg_pred > 0.5 else "Fake"
-        confidence = f"{avg_pred:.2f}"
-
-        # Annotate
-        color = (0, 255, 0) if label == "Real" else (0, 0, 255)
-        cv2.rectangle(annotated, (x, y), (x + w, y + h), color, 2)
-        cv2.putText(annotated, f"{label} ({confidence})", (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)
-
-        results.append(f"Face {i+1}: {label} (Avg: {avg_pred:.3f}, XCP: {xcp_pred:.3f}, EFF: {eff_pred:.3f})")
-
-    return "\n".join(results), annotated
-
-# ---------------------------------------------------------
-# Gradio Interface
-# ---------------------------------------------------------
-interface = gr.Interface(
-    fn=predict_faces,
-    inputs=gr.Image(type="numpy", label="Upload Image"),
-    outputs=[
-        gr.Textbox(label="Predictions"),
-        gr.Image(type="numpy", label="Annotated Image")
-    ],
-    title="Deepfake Detector (Multi-Face Ensemble)",
-    description="This model detects all faces in an image and classifies each one as real or fake using Xception and EfficientNetB4 ensemble."
-)
-
-interface.launch()
+import cv2
+import numpy as np
+import gradio as gr
+from mtcnn import MTCNN
+from tensorflow.keras.models import load_model
+from tensorflow.keras.applications.xception import preprocess_input as xcp_pre
+from tensorflow.keras.applications.efficientnet import preprocess_input as eff_pre
+
+# Load models
+xcp_model = load_model("xception_model.h5")
+eff_model = load_model("efficientnet_model.h5")
+
+# Load face detector
+detector = MTCNN()
+
+def expand_box(x, y, w, h, scale=1.5, img_shape=None):
+    """Expand face bounding box with margin."""
+    cx, cy = x + w // 2, y + h // 2
+    new_w, new_h = int(w * scale), int(h * scale)
+    x1 = max(0, cx - new_w // 2)
+    y1 = max(0, cy - new_h // 2)
+    x2 = min(img_shape[1], cx + new_w // 2)
+    y2 = min(img_shape[0], cy + new_h // 2)
+    return x1, y1, x2, y2
+
+def predict(image):
+    faces = detector.detect_faces(image)
+    if not faces:
+        return "No faces detected", image
+
+    results = []
+    annotated = image.copy()
+
+    for i, face in enumerate(faces):
+        x, y, w, h = face['box']
+        x, y, w, h = max(0, x), max(0, y), w, h
+        x1, y1, x2, y2 = expand_box(x, y, w, h, scale=1.6, img_shape=image.shape)
+
+        face_crop = image[y1:y2, x1:x2]
+
+        # Preprocess for each model
+        xcp_img = cv2.resize(face_crop, (299, 299))
+        eff_img = cv2.resize(face_crop, (224, 224))
+
+        xcp_tensor = xcp_pre(xcp_img.astype(np.float32))[np.newaxis, ...]
+        eff_tensor = eff_pre(eff_img.astype(np.float32))[np.newaxis, ...]
+
+        xcp_pred = xcp_model.predict(xcp_tensor, verbose=0).flatten()[0]
+        eff_pred = eff_model.predict(eff_tensor, verbose=0).flatten()[0]
+        avg_pred = (xcp_pred + eff_pred) / 2
+        label = "Real" if avg_pred > 0.5 else "Fake"
+
+        results.append(
+            f"Face {i+1}: {label} (Avg: {avg_pred:.3f}, XCP: {xcp_pred:.3f}, EFF: {eff_pred:.3f})"
+        )
+
+        # Draw
+        color = (0, 255, 0) if label == "Real" else (255, 0, 0)
+        cv2.rectangle(annotated, (x1, y1), (x2, y2), color, 2)
+        cv2.putText(
+            annotated,
+            f"{label} ({avg_pred:.2f})",
+            (x1, y1 - 10),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.6,
+            color,
+            2,
+        )
+
+    return "\n".join(results), annotated
+
+# Gradio Interface
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="numpy", label="Upload Image"),
+    outputs=[
+        gr.Textbox(label="Predictions"),
+        gr.Image(type="numpy", label="Annotated Image"),
+    ],
+    title="Deepfake Detector (Multi-Face Ensemble)",
+    description="Detects all faces in an image and classifies each one as real or fake using Xception and EfficientNetB4 ensemble.",
+)
+
+interface.launch()