Update app.py

app.py

@@ -6,57 +6,50 @@ from PIL import Image
 import cv2
 import traceback
 
-# ==============================================================================
-# --- Load the trained model ---
-# ==============================================================================
+# =========================
+# --- Load the model ---
+# =========================
 try:
-    model = tf.keras.models.load_model("./deepfake_mobilenet_model.h5")
+    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: create dummy model (won't produce Grad-CAM)
+    # Fallback dummy model
     inputs = tf.keras.Input(shape=(224, 224, 3))
-    outputs = tf.keras.layers.Dense(1, activation="sigmoid")(
-        tf.keras.layers.GlobalAveragePooling2D()(inputs)
-    )
+    x = tf.keras.layers.GlobalAveragePooling2D()(inputs)
+    outputs = tf.keras.layers.Dense(1, activation="sigmoid")(x)
     model = tf.keras.Model(inputs, outputs)
 
-# ==============================================================================
-# --- Grad-CAM Heatmap Generation Functions ---
-# ==============================================================================
+# =========================
+# --- Grad-CAM helpers ---
+# =========================
+
 def get_last_conv_layer_name(model):
     for layer in reversed(model.layers):
-        if len(layer.output.shape) == 4:  # conv layer has 4D output
-            print(f"Using conv layer for Grad-CAM: {layer.name}")
+        if hasattr(layer.output, "shape") and len(layer.output.shape) == 4:
             return layer.name
-    print("No convolutional layer found; Grad-CAM will be disabled.")
-    return None
-
+    raise ValueError("No convolutional layer found in model.")
 
 def make_gradcam_heatmap(img_array, model, last_conv_layer_name):
-    if last_conv_layer_name is None:
-        return None
     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_layer_output, preds = grad_model([img_array])
-        # Robustly extract scalar
-        if len(preds.shape) == 2:
-            class_channel = preds[:, 0]
-        else:
-            class_channel = preds
-    grads = tape.gradient(class_channel, last_conv_layer_output)
+        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)
     pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
-    last_conv_layer_output = last_conv_layer_output[0]
-    heatmap = last_conv_layer_output @ pooled_grads[..., tf.newaxis]
+    last_conv_output = last_conv_output[0]
+    heatmap = last_conv_output @ pooled_grads[..., tf.newaxis]
     heatmap = tf.squeeze(heatmap)
     heatmap = tf.maximum(heatmap, 0) / (tf.math.reduce_max(heatmap) + 1e-8)
     return heatmap.numpy()
 
-
 def superimpose_gradcam(original_img_pil, heatmap):
-    if heatmap is None:
-        return None
     original_img = np.array(original_img_pil)
     heatmap = cv2.resize(heatmap, (original_img.shape[1], original_img.shape[0]))
     heatmap = np.uint8(255 * heatmap)
@@ -64,69 +57,62 @@ 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 ---
+# =========================
 
 last_conv_layer_name = get_last_conv_layer_name(model)
 
-# ==============================================================================
-# --- Main Prediction Function ---
-# ==============================================================================
 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_expanded = np.expand_dims(img_array, axis=0)
-
-        # --- Model Prediction ---
-        prediction = model.predict(img_array_expanded, verbose=0)
+        img_array_exp = np.expand_dims(img_array, axis=0)
 
-        # Convert prediction to scalar safely
+        # --- Model prediction ---
+        prediction = model.predict(img_array_exp, verbose=0)
+        
+        # ✅ Ensure scalar prediction
         if isinstance(prediction, (list, tuple)):
-            prediction = prediction[0]
-        prediction = np.array(prediction)
-        if prediction.ndim > 0:
-            prediction = prediction.item()
+            prediction = np.array(prediction[0])
+        prediction = prediction.item() if isinstance(prediction, np.ndarray) else float(prediction)
 
-        # Confidence bars
         real_conf = float(prediction)
         fake_conf = float(1 - prediction)
         labels = {"Real Image": real_conf, "Fake Image": fake_conf}
 
-        # --- Grad-CAM Overlay ---
-        heatmap = make_gradcam_heatmap(img_array_expanded, model, last_conv_layer_name)
-        superimposed_image = superimpose_gradcam(img, heatmap)
+        # --- Grad-CAM ---
+        heatmap = make_gradcam_heatmap(img_array_exp, model, last_conv_layer_name)
+        superimposed_img = superimpose_gradcam(img, heatmap)
 
-        return labels, superimposed_image
+        return labels, superimposed_img
 
     except Exception as e:
-        print("--- GRADIO APP ERROR ---")
+        print("--- ERROR ---")
         traceback.print_exc()
-        print("------------------------")
-        error_msg = f"Error: {e}"
-        return {error_msg: 0.0}, None
-
+        return {"Error": 0.0}, None
 
-# ==============================================================================
+# =========================
 # --- Gradio Interface ---
-# ==============================================================================
+# =========================
+
 gr.Interface(
     fn=predict_and_visualize,
     inputs=gr.Image(type="pil", label="📷 Upload a Face Image"),
     outputs=[
         gr.Label(num_top_classes=2, label="🧠 Model Prediction"),
-        gr.Image(label="🔥 Grad-CAM Heatmap Overlay (if available)")
+        gr.Image(label="🔥 Grad-CAM Heatmap Overlay")
     ],
-    title="✨ Deepfake Image Detector with Visual Explanation ✨",
+    title="✨ Deepfake Image Detector with Grad-CAM ✨",
     description="""
-**Detect whether an uploaded image is Real or AI-Generated (Deepfake).**  
-Confidence bars show certainty for **Real** and **Fake** categories.  
-
-The **Grad-CAM heatmap** highlights regions the model focused on (red = most important).  
+**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).  
 
 ⚡ **Instructions:**  
 1. Upload a face image (JPEG/PNG).  
-2. Wait a few seconds for prediction and heatmap.  
-3. Observe the confidence bars and Grad-CAM overlay (if available).
+2. Wait a few seconds for prediction & heatmap.  
+3. Observe confidence bars and Grad-CAM overlay.
 """,
     theme="default"
 ).launch()