Update gradcam_utils.py

gradcam_utils.py

@@ -77,68 +77,103 @@
 
 
 
+# import numpy as np
+# import cv2
+# import tensorflow as tf
+# from tensorflow.keras.preprocessing.image import img_to_array, load_img
+# from matplotlib.colors import LinearSegmentedColormap
+
+# def preprocess_image(img_path, target_size):
+#     img = load_img(img_path, target_size=target_size)
+#     img = img_to_array(img)
+#     img = np.expand_dims(img, axis=0)
+#     img = img / 255.0
+#     return img
+
+# def make_gradcam_heatmap(model, img_tensor):
+#     grad_model = tf.keras.models.Model([model.input], [model.output])
+
+#     with tf.GradientTape() as tape:
+#         conv_outputs = model(img_tensor)
+#         loss = conv_outputs[:, 1]  # class index 1 = pneumonia
+
+#     grads = tape.gradient(loss, conv_outputs)
+#     pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
+
+#     conv_outputs = conv_outputs[0]
+#     heatmap = conv_outputs @ pooled_grads[..., tf.newaxis]
+#     heatmap = tf.squeeze(heatmap)
+#     heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
+#     heatmap = tf.where(tf.math.is_nan(heatmap), tf.zeros_like(heatmap), heatmap)
+#     return heatmap.numpy()
+
+# def create_custom_colormap():
+#     colors = ['blue', 'green', 'yellow', 'red']
+#     cmap = LinearSegmentedColormap.from_list('custom', colors, N=256)
+#     return cmap
+
+# def apply_custom_colormap(heatmap, cmap):
+#     colored_heatmap = cmap(heatmap)
+#     return np.uint8(colored_heatmap * 255)
+
+# def enhance_heatmap(heatmap, gamma=0.7, percentile=99):
+#     heatmap = np.power(heatmap, gamma)
+#     heatmap = heatmap / np.percentile(heatmap, percentile)
+#     return np.clip(heatmap, 0, 1)
+
+# def generate_and_merge_heatmaps(img_path, vgg_model, efficientnet_model, densenet_model, img_size=(224, 224)):
+#     img_tensor = preprocess_image(img_path, img_size)
+
+#     vgg_heatmap = make_gradcam_heatmap(vgg_model, img_tensor)
+#     efficientnet_heatmap = make_gradcam_heatmap(efficientnet_model, img_tensor)
+#     densenet_heatmap = make_gradcam_heatmap(densenet_model, img_tensor)
+
+#     vgg_heatmap = cv2.resize(vgg_heatmap, img_size)
+#     efficientnet_heatmap = cv2.resize(efficientnet_heatmap, img_size)
+#     densenet_heatmap = cv2.resize(densenet_heatmap, img_size)
+
+#     merged = (vgg_heatmap + efficientnet_heatmap + densenet_heatmap) / 3.0
+#     enhanced = enhance_heatmap(merged)
+#     colored = apply_custom_colormap(enhanced, create_custom_colormap())
+
+#     original = cv2.imread(img_path)
+#     original = cv2.resize(original, img_size)
+#     original = cv2.cvtColor(original, cv2.COLOR_BGR2RGB)
+
+#     superimposed_img = cv2.addWeighted(original, 0.6, colored[:, :, :3], 0.4, 0)
+#     return superimposed_img
+
+
+
+
+
+from tf_explain.core.grad_cam import GradCAM
 import numpy as np
-import cv2
-import tensorflow as tf
-from tensorflow.keras.preprocessing.image import img_to_array, load_img
-from matplotlib.colors import LinearSegmentedColormap
-
-def preprocess_image(img_path, target_size):
-    img = load_img(img_path, target_size=target_size)
-    img = img_to_array(img)
-    img = np.expand_dims(img, axis=0)
-    img = img / 255.0
-    return img
-
-def make_gradcam_heatmap(model, img_tensor):
-    grad_model = tf.keras.models.Model([model.input], [model.output])
-
-    with tf.GradientTape() as tape:
-        conv_outputs = model(img_tensor)
-        loss = conv_outputs[:, 1]  # class index 1 = pneumonia
-
-    grads = tape.gradient(loss, conv_outputs)
-    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
-
-    conv_outputs = conv_outputs[0]
-    heatmap = conv_outputs @ pooled_grads[..., tf.newaxis]
-    heatmap = tf.squeeze(heatmap)
-    heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
-    heatmap = tf.where(tf.math.is_nan(heatmap), tf.zeros_like(heatmap), heatmap)
-    return heatmap.numpy()
-
-def create_custom_colormap():
-    colors = ['blue', 'green', 'yellow', 'red']
-    cmap = LinearSegmentedColormap.from_list('custom', colors, N=256)
-    return cmap
-
-def apply_custom_colormap(heatmap, cmap):
-    colored_heatmap = cmap(heatmap)
-    return np.uint8(colored_heatmap * 255)
-
-def enhance_heatmap(heatmap, gamma=0.7, percentile=99):
-    heatmap = np.power(heatmap, gamma)
-    heatmap = heatmap / np.percentile(heatmap, percentile)
-    return np.clip(heatmap, 0, 1)
-
-def generate_and_merge_heatmaps(img_path, vgg_model, efficientnet_model, densenet_model, img_size=(224, 224)):
-    img_tensor = preprocess_image(img_path, img_size)
-
-    vgg_heatmap = make_gradcam_heatmap(vgg_model, img_tensor)
-    efficientnet_heatmap = make_gradcam_heatmap(efficientnet_model, img_tensor)
-    densenet_heatmap = make_gradcam_heatmap(densenet_model, img_tensor)
-
-    vgg_heatmap = cv2.resize(vgg_heatmap, img_size)
-    efficientnet_heatmap = cv2.resize(efficientnet_heatmap, img_size)
-    densenet_heatmap = cv2.resize(densenet_heatmap, img_size)
-
-    merged = (vgg_heatmap + efficientnet_heatmap + densenet_heatmap) / 3.0
-    enhanced = enhance_heatmap(merged)
-    colored = apply_custom_colormap(enhanced, create_custom_colormap())
-
-    original = cv2.imread(img_path)
-    original = cv2.resize(original, img_size)
-    original = cv2.cvtColor(original, cv2.COLOR_BGR2RGB)
-
-    superimposed_img = cv2.addWeighted(original, 0.6, colored[:, :, :3], 0.4, 0)
-    return superimposed_img
+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
+    img_array = np.array(image_pil.resize((224, 224))) / 255.0
+    img_array = np.expand_dims(img_array, axis=0)
+
+    # Generate Grad-CAM explanation
+    explainer = GradCAM()
+    explanation = explainer.explain(
+        validation_data=(img_array, None),
+        model=model,
+        class_index=class_index
+    )
+
+    return Image.fromarray(explanation)
+