Update app.py

app.py

@@ -16,23 +16,33 @@ except Exception as e:
     model = tf.keras.Model(inputs, outputs)
 
 # ==============================================================================
-# --- Grad-CAM Heatmap Generation Functions ---
+# --- Grad-CAM Heatmap Generation Functions (with final fix) ---
 # ==============================================================================
-
 def get_last_conv_layer_name(model):
+    """Finds the name of the last convolutional layer in the model."""
     for layer in reversed(model.layers):
         if len(layer.output.shape) == 4:
             return layer.name
     raise ValueError("Could not find a conv layer in the model")
 
 def make_gradcam_heatmap(img_array, model, last_conv_layer_name):
+    """Generates the Grad-CAM heatmap."""
     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])
+        last_conv_layer_output, preds = grad_model([img_array], training=False)
         class_channel = preds[0]
+        
     grads = tape.gradient(class_channel, last_conv_layer_output)
+
+    # <-- FIX: Add a safety check in case the gradient does not exist.
+    if grads is None:
+        print("Warning: Gradient is None. Cannot compute heatmap. Returning a blank map.")
+        # Return a blank (black) map of the same size as the feature map.
+        h, w = last_conv_layer_output.shape[1:3]
+        return np.zeros((h, w), dtype=np.float32)
+
     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]
@@ -41,6 +51,7 @@ def make_gradcam_heatmap(img_array, model, last_conv_layer_name):
     return heatmap.numpy()
 
 def superimpose_gradcam(original_img_pil, heatmap):
+    """Overlays the heatmap on the original image."""
     original_img = np.array(original_img_pil)
     heatmap = cv2.resize(heatmap, (original_img.shape[1], original_img.shape[0]))
     heatmap = np.uint8(255 * heatmap)
@@ -51,11 +62,14 @@ def superimpose_gradcam(original_img_pil, heatmap):
 # ==============================================================================
 # --- Main Prediction Function ---
 # ==============================================================================
-
 last_conv_layer_name = get_last_conv_layer_name(model)
 
 def predict_and_visualize(img):
+    """Performs prediction and generates the Grad-CAM heatmap."""
     try:
+        if img is None: # Handle case where user clears the image
+            return None, None
+
         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)
@@ -69,7 +83,7 @@ def predict_and_visualize(img):
         superimposed_image = superimpose_gradcam(img, heatmap)
         
         return labels, superimposed_image
-
+        
     except Exception as e:
         print("--- GRADIO APP ERROR ---")
         traceback.print_exc()
@@ -78,9 +92,8 @@ def predict_and_visualize(img):
         return {error_msg: 0.0}, None
 
 # ==============================================================================
-# --- Gradio Interface with Improved Design ---
+# --- Gradio Interface ---
 # ==============================================================================
-
 gr.Interface(
     fn=predict_and_visualize,
     inputs=gr.Image(type="pil", label="📷 Upload a Face Image"),
@@ -90,15 +103,8 @@ gr.Interface(
     ],
     title="✨ Deepfake Image Detector with Visual Explanation ✨",
     description="""
-    **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 the prediction and heatmap.  
-    3. Observe the confidence bars and heatmap for model explanation.
+    **Detect whether an uploaded image is Real or AI-Generated (Deepfake).**
+    The confidence bars show the model's certainty, and the heatmap highlights the regions the model focused on (red = most important).
     """,
-    theme="default"  # you can later try 'soft', 'grass', or 'peach'
-).launch()
+    theme="default"
+).launch()