Update app.py

app.py

@@ -6,45 +6,36 @@ from PIL import Image
 import cv2
 import traceback
 
-# =========================
-# --- Load the model ---
-# =========================
+# Load the trained model
 try:
     model = tf.keras.models.load_model("./model/deepfake_mobilenet_model.h5")
 except Exception as e:
     print(f"Error loading model. Make sure the path is correct. Error: {e}")
-    # Fallback dummy model
     inputs = tf.keras.Input(shape=(224, 224, 3))
-    x = tf.keras.layers.GlobalAveragePooling2D()(inputs)
-    outputs = tf.keras.layers.Dense(1, activation="sigmoid")(x)
+    outputs = tf.keras.layers.Dense(1, activation="sigmoid")(tf.keras.layers.GlobalAveragePooling2D()(inputs))
     model = tf.keras.Model(inputs, outputs)
 
-# =========================
-# --- Grad-CAM helpers ---
-# =========================
+# ==============================================================================
+# --- Grad-CAM Heatmap Generation Functions ---
+# ==============================================================================
 
 def get_last_conv_layer_name(model):
     for layer in reversed(model.layers):
-        if hasattr(layer.output, "shape") and len(layer.output.shape) == 4:
+        if len(layer.output.shape) == 4:
             return layer.name
-    raise ValueError("No convolutional layer found in model.")
+    raise ValueError("Could not find a conv layer in the model")
 
 def make_gradcam_heatmap(img_array, model, last_conv_layer_name):
     grad_model = tf.keras.models.Model(
         model.inputs, [model.get_layer(last_conv_layer_name).output, model.output]
     )
-
     with tf.GradientTape() as tape:
-        last_conv_output, preds = grad_model(img_array)
-        
-        # ✅ Robustly handle list/tensor output
-        preds = preds[0] if isinstance(preds, (list, tuple)) else preds
-        class_channel = preds[:, 0] if len(preds.shape) > 1 else preds
-
-    grads = tape.gradient(class_channel, last_conv_output)
+        last_conv_layer_output, preds = grad_model([img_array])
+        class_channel = preds[0][:, 0]
+    grads = tape.gradient(class_channel, last_conv_layer_output)
     pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
-    last_conv_output = last_conv_output[0]
-    heatmap = last_conv_output @ pooled_grads[..., tf.newaxis]
+    last_conv_layer_output = last_conv_layer_output[0]
+    heatmap = last_conv_layer_output @ pooled_grads[..., tf.newaxis]
     heatmap = tf.squeeze(heatmap)
     heatmap = tf.maximum(heatmap, 0) / (tf.math.reduce_max(heatmap) + 1e-8)
     return heatmap.numpy()
@@ -57,54 +48,38 @@ def superimpose_gradcam(original_img_pil, heatmap):
     superimposed_img = cv2.addWeighted(original_img, 0.6, heatmap, 0.4, 0)
     return Image.fromarray(superimposed_img)
 
-# =========================
-# --- Main prediction ---
-# =========================
+# ==============================================================================
+# --- Main Prediction Function ---
+# ==============================================================================
 
 last_conv_layer_name = get_last_conv_layer_name(model)
 
 def predict_and_visualize(img):
     try:
-        # --- Preprocess Image ---
         img_resized = img.resize((224, 224))
         img_array = image.img_to_array(img_resized) / 255.0
-        img_array_exp = np.expand_dims(img_array, axis=0)
-
-        # --- Model Prediction ---
-        prediction = model.predict(img_array_exp, verbose=0)
+        img_array_expanded = np.expand_dims(img_array, axis=0)
         
-        # Robust scalar extraction
-        if isinstance(prediction, (list, tuple)):
-            prediction = np.array(prediction[0])
-        if isinstance(prediction, np.ndarray):
-            prediction = prediction.item()
-        prediction = float(prediction)
-
-        # Confidence bars
-        real_conf = max(0.0, min(1.0, prediction))  # clamp between 0 and 1
-        fake_conf = 1.0 - real_conf
+        prediction = model.predict(img_array_expanded, verbose=0)[0][0]
+        real_conf = float(prediction)
+        fake_conf = float(1 - prediction)
         labels = {"Real Image": real_conf, "Fake Image": fake_conf}
-
-        # --- Grad-CAM (optional, fail-safe) ---
-        try:
-            heatmap = make_gradcam_heatmap(img_array_exp, model, last_conv_layer_name)
-            superimposed_img = superimpose_gradcam(img, heatmap)
-        except Exception as e:
-            print("Grad-CAM failed:", e)
-            superimposed_img = img  # fallback to original image
-
-        return labels, superimposed_img
+        
+        heatmap = make_gradcam_heatmap(img_array_expanded, model, last_conv_layer_name)
+        superimposed_image = superimpose_gradcam(img, heatmap)
+        
+        return labels, superimposed_image
 
     except Exception as e:
-        print("--- PREDICTION ERROR ---")
+        print("--- GRADIO APP ERROR ---")
         traceback.print_exc()
-        # Always return a valid dict for gr.Label
-        return {"Real Image": 0.0, "Fake Image": 0.0}, img
+        print("------------------------")
+        error_msg = f"Error: {e}"
+        return {error_msg: 0.0}, None
 
-
-# =========================
-# --- Gradio Interface ---
-# =========================
+# ==============================================================================
+# --- Gradio Interface with Improved Design ---
+# ==============================================================================
 
 gr.Interface(
     fn=predict_and_visualize,
@@ -113,16 +88,17 @@ gr.Interface(
         gr.Label(num_top_classes=2, label="🧠 Model Prediction"),
         gr.Image(label="🔥 Grad-CAM Heatmap Overlay")
     ],
-    title="✨ Deepfake Image Detector with Grad-CAM ✨",
+    title="✨ Deepfake Image Detector with Visual Explanation ✨",
     description="""
-**Detect Real vs AI-Generated (Deepfake) Images.**  
-The confidence bars show the model's certainty for **Real** and **Fake**.  
-The Grad-CAM heatmap highlights areas most important to the model (red = most important).  
+    **Detect whether an uploaded image is Real or AI-Generated (Deepfake).**  
+    The confidence bars show the model's certainty for both **Real** and **Fake** categories.  
+
+    Below, the **Grad-CAM heatmap** highlights the regions the model focused on (red = most important).  
 
-⚡ **Instructions:**  
-1. Upload a face image (JPEG/PNG).  
-2. Wait a few seconds for prediction & heatmap.  
-3. Observe confidence bars and Grad-CAM overlay.
-""",
-    theme="default"
+    ⚡ **Instructions:**  
+    1. Upload a face image (JPEG/PNG).  
+    2. Wait a few seconds for the prediction and heatmap.  
+    3. Observe the confidence bars and heatmap for model explanation.
+    """,
+    theme="default"  # you can later try 'soft', 'grass', or 'peach'
 ).launch()