Update app.py

app.py

@@ -11,13 +11,15 @@ 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 for deployment debugging
     inputs = tf.keras.Input(shape=(224, 224, 3))
     outputs = tf.keras.layers.Dense(1, activation="sigmoid")(tf.keras.layers.GlobalAveragePooling2D()(inputs))
     model = tf.keras.Model(inputs, outputs)
 
 # ==============================================================================
-# --- Grad-CAM Heatmap Generation Functions (with final fix) ---
+# --- Grad-CAM Heatmap Generation Functions (Improved Logic) ---
 # ==============================================================================
+
 def get_last_conv_layer_name(model):
     """Finds the name of the last convolutional layer in the model."""
     for layer in reversed(model.layers):
@@ -26,20 +28,22 @@ def get_last_conv_layer_name(model):
     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."""
+    """Generates the Grad-CAM heatmap with robust logic."""
     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], training=False)
-        class_channel = preds[0]
+        
+        # <-- IMPROVEMENT: Use a robust method to get the prediction value.
+        # This works correctly regardless of tensor shape (e.g., (1,1) vs (1,))
+        # and fixes the root cause of previous errors and poor quality.
+        class_channel = tf.reduce_mean(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.
+        print("Warning: Gradient is None. Cannot compute heatmap.")
         h, w = last_conv_layer_output.shape[1:3]
         return np.zeros((h, w), dtype=np.float32)
 
@@ -51,12 +55,21 @@ 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."""
+    """Overlays the heatmap on the original image with visual improvements."""
     original_img = np.array(original_img_pil)
+    
     heatmap = cv2.resize(heatmap, (original_img.shape[1], original_img.shape[0]))
+    
+    # <-- IMPROVEMENT: Add a blur to make the heatmap smoother.
+    heatmap = cv2.GaussianBlur(heatmap, (15, 15), 0)
+    
     heatmap = np.uint8(255 * heatmap)
-    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
+    
+    # <-- IMPROVEMENT: Use a better color map.
+    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_VIRIDIS)
+    
     superimposed_img = cv2.addWeighted(original_img, 0.6, heatmap, 0.4, 0)
+    
     return Image.fromarray(superimposed_img)
 
 # ==============================================================================
@@ -65,15 +78,16 @@ def superimpose_gradcam(original_img_pil, heatmap):
 last_conv_layer_name = get_last_conv_layer_name(model)
 
 def predict_and_visualize(img):
-    """Performs prediction and generates the Grad-CAM heatmap."""
+    """Performs prediction and generates the improved Grad-CAM heatmap."""
     try:
-        if img is None: # Handle case where user clears the image
+        if img is None:
             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)
         
+        # The prediction logic remains unchanged and should work as before
         prediction = model.predict(img_array_expanded, verbose=0)[0][0]
         real_conf = float(prediction)
         fake_conf = float(1 - prediction)
@@ -104,7 +118,7 @@ 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, and the heatmap highlights the regions the model focused on (red = most important).
+    The confidence bars show the model's certainty, and the heatmap highlights the regions the model focused on.
     """,
     theme="default"
 ).launch()