Update app.py

app.py

@@ -1,83 +1,95 @@
-import gradio as gr
+import os
+import cv2
 import numpy as np
+import gradio as gr
 import tensorflow as tf
-import cv2
 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
+from huggingface_hub import hf_hub_download
 from mtcnn import MTCNN
-import os
-import warnings
-warnings.filterwarnings("ignore")
+import matplotlib.pyplot as plt
 
-# Force TF to suppress log-level warnings
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+# Download models
+xcp_path = hf_hub_download(repo_id="Zeyadd-Mostaffa/deepfake-image-detector", filename="xception_model.h5")
+eff_path = hf_hub_download(repo_id="Zeyadd-Mostaffa/deepfake-image-detector", filename="efficientnet_model.h5")
 
-# Load models from local (downloaded from HF first in app setup)
-xcp_model = load_model("xception_model.h5")
-eff_model = load_model("efficientnet_model.h5")
+# Load models
+xcp_model = load_model(xcp_path)
+eff_model = load_model(eff_path)
 
-# Grad-CAM for Xception
-def grad_cam(model, img_array, size, preprocess_fn):
-    img = cv2.resize(img_array, size)
-    input_tensor = preprocess_fn(np.expand_dims(img, axis=0).astype(np.float32))
-    input_tensor = tf.convert_to_tensor(input_tensor)
+# Face detector
+detector = MTCNN()
+
+def extract_face(image):
+    faces = detector.detect_faces(image)
+    if not faces:
+        return None
+    x, y, w, h = faces[0]['box']
+    x, y = max(0, x), max(0, y)
+    return image[y:y+h, x:x+w]
+
+def grad_cam(model, image, size, preprocess_fn):
+    img = cv2.resize(image, size)
+    input_tensor = preprocess_fn(img.astype(np.float32))[np.newaxis, ...]
+    
+    grad_model = tf.keras.models.Model(
+        [model.inputs], [model.get_layer(index=-1).output, model.output]
+    )
 
     with tf.GradientTape() as tape:
-        conv_layer = model.get_layer(index=-5).output
-        grad_model = tf.keras.models.Model([model.inputs], [conv_layer, model.output])
         conv_outputs, predictions = grad_model(input_tensor)
-        loss = predictions[:, 0]
+        loss = predictions[:, 0]  # Assuming binary classification
+
+    grads = tape.gradient(loss, conv_outputs)[0]
+    conv_outputs = conv_outputs[0]
 
-    grads = tape.gradient(loss, conv_outputs)
-    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
-    cam = tf.reduce_sum(tf.multiply(pooled_grads, conv_outputs), axis=-1).numpy()[0]
+    weights = tf.reduce_mean(grads, axis=(0, 1))
+    cam = np.zeros(conv_outputs.shape[:2], dtype=np.float32)
+    for i, w in enumerate(weights):
+        cam += w * conv_outputs[:, :, i]
 
     cam = np.maximum(cam, 0)
     cam = cam / (cam.max() + 1e-8)
-    cam = (cam * 255).astype(np.uint8)
-    cam = cam.numpy() if hasattr(cam, 'numpy') else cam
-    cam = cv2.resize(cam, size)
-    heatmap = cv2.applyColorMap(cam, cv2.COLORMAP_JET)
-    superimposed_img = cv2.addWeighted(cv2.cvtColor(img, cv2.COLOR_RGB2BGR), 0.6, heatmap, 0.4, 0)
+    heatmap = cv2.resize(cam, size)
+    heatmap = np.uint8(255 * heatmap)
+    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
+    superimposed_img = cv2.addWeighted(cv2.resize(img, size), 0.6, heatmap, 0.4, 0)
     return superimposed_img
 
-# Face detector
-detector = MTCNN()
-
-def detect_face(image):
-    faces = detector.detect_faces(image)
-    if not faces:
-        raise ValueError("No face detected.")
-    x, y, w, h = faces[0]['box']
-    return image[y:y+h, x:x+w]
-
 def predict(image):
-    try:
-        face = detect_face(image)
-
-        xcp_img = cv2.resize(face, (299, 299))
-        eff_img = cv2.resize(face, (224, 224))
+    face = extract_face(image)
+    if face is None:
+        return "No face detected", None
 
-        xcp_input = np.expand_dims(xcp_pre(xcp_img.astype(np.float32)), axis=0)
-        eff_input = np.expand_dims(eff_pre(eff_img.astype(np.float32)), axis=0)
+    # Xception
+    xcp_img = cv2.resize(face, (299, 299))
+    xcp_tensor = xcp_pre(xcp_img.astype(np.float32))[np.newaxis, ...]
+    xcp_pred = xcp_model.predict(xcp_tensor)[0][0]
 
-        xcp_pred = xcp_model.predict(xcp_input)[0][0]
-        eff_pred = eff_model.predict(eff_input)[0][0]
+    # EfficientNet
+    eff_img = cv2.resize(face, (224, 224))
+    eff_tensor = eff_pre(eff_img.astype(np.float32))[np.newaxis, ...]
+    eff_pred = eff_model.predict(eff_tensor)[0][0]
 
-        ensemble_pred = (xcp_pred + eff_pred) / 2
-        label = "Fake" if ensemble_pred > 0.5 else "Real"
+    # Ensemble average
+    avg_pred = (xcp_pred + eff_pred) / 2
+    label = "Fake" if avg_pred > 0.5 else "Real"
 
-        cam_img = grad_cam(xcp_model, face, (299, 299), xcp_pre)
+    # Grad-CAM on Xception
+    cam_img = grad_cam(xcp_model, face, (299, 299), xcp_pre)
 
-        return label, cam_img
-    except Exception as e:
-        return "خطأ", "خطأ"
+    return label, cam_img
 
-gr.Interface(
+interface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="numpy", label="Upload Face Image"),
-    outputs=[gr.Label(label="Prediction"), gr.Image(label="Grad-CAM Explanation")],
+    inputs=gr.Image(type="numpy", label="image"),
+    outputs=[
+        gr.Label(label="Prediction"),
+        gr.Image(type="numpy", label="Grad-CAM")
+    ],
     title="Deepfake Image Detector (with Grad-CAM)",
     description="Upload an image. We detect the face, classify using an ensemble (Xception + EfficientNetB4), and explain the prediction using Grad-CAM on Xception."
-).launch()
+)
+
+interface.launch()