Update gradcam_utils.py

gradcam_utils.py

@@ -151,28 +151,23 @@ from tf_explain.core.grad_cam import GradCAM
 import numpy as np
 from PIL import Image
 
-def generate_heatmap_tf_explain(image_pil, model, class_index):
-    """
-    Generates a Grad-CAM heatmap using tf-explain and overlays it on the original image.
-
-    Parameters:
-        image_pil (PIL.Image): Input chest X-ray image.
-        model (tf.keras.Model): CNN model for explanation (e.g. VGG19).
-        class_index (int): Index of the predicted class (0 or 1).
-
-    Returns:
-        heatmap_image (PIL.Image): Heatmap image overlaid on original image.
-    """
-    # Resize and preprocess image
+def generate_heatmap_tf_explain(image_pil, model, class_index, layer_name="block5_conv4"):
+    from tf_explain.core.grad_cam import GradCAM
+
+    # Preprocess image
     img_array = np.array(image_pil.resize((224, 224))) / 255.0
     img_array = np.expand_dims(img_array, axis=0)
 
-    # Generate Grad-CAM explanation
+    # Reconstruct model to include target layer
+    from tensorflow.keras.models import Model
+    model_for_explanation = Model(inputs=model.input, outputs=model.output)
+
     explainer = GradCAM()
     explanation = explainer.explain(
         validation_data=(img_array, None),
-        model=model,
-        class_index=class_index
+        model=model_for_explanation,
+        class_index=class_index,
+        layer_name=layer_name
     )
 
     return Image.fromarray(explanation)