Update app.py

app.py

@@ -17,7 +17,7 @@ except Exception as e:
     model = tf.keras.Model(inputs, outputs)
 
 # ==============================================================================
-# --- Grad-CAM Heatmap Generation Functions (Your Original Structure) ---
+# --- Grad-CAM Heatmap Generation Functions (with safety check restored) ---
 # ==============================================================================
 
 def get_last_conv_layer_name(model):
@@ -28,25 +28,28 @@ 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 using your original logic, but without crashing."""
+    """Generates the Grad-CAM heatmap with a stability check."""
     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)
         
-        # <-- THE ONLY FIX: A robust check that preserves your original logic.
-        # This handles tensor shape variations without changing the code's structure.
         preds_tensor = preds[0]
         if tf.rank(preds_tensor) > 1:
-            # If tensor is 2D (e.g., shape=(1,1)), slice it like your original code.
             class_channel = preds_tensor[:, 0]
         else:
-            # If tensor is 1D (e.g., shape=(1,)), use it directly.
             class_channel = preds_tensor
             
     grads = tape.gradient(class_channel, last_conv_layer_output)
     
+    # <-- FIX: The critical safety check is restored here.
+    # This prevents the "None to a Tensor" crash without changing the logic.
+    if grads is None:
+        print("Warning: Gradient is None. Cannot compute heatmap for this image. Returning a blank 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]
@@ -55,13 +58,11 @@ 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 (Restored to your original style)."""
+    """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)
-    # Restored to COLORMAP_JET as in your code.
     heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) 
-    # Removed the blur and restored your original transparency.
     superimposed_img = cv2.addWeighted(original_img, 0.6, heatmap, 0.4, 0) 
     return Image.fromarray(superimposed_img)