Update gradcam_utils.py

gradcam_utils.py

@@ -1,187 +1,255 @@
-"""
-Grad-CAM Implementation for Chest X-Ray Classification
-========================================================
-
-Visualizes which regions of the X-ray the model focuses on when making predictions.
-
-Reference: Selvaraju et al. (2017) - Grad-CAM: Visual Explanations from Deep Networks
-"""
-
-import tensorflow as tf
-import numpy as np
-import cv2
-from PIL import Image
-
-
-def make_gradcam_heatmap(img_array, model, last_conv_layer_name, pred_index=None):
-    """
-    Generate Grad-CAM heatmap for a given image and prediction.
-    
-    Args:
-        img_array: Preprocessed image (batch_size, height, width, channels)
-        model: Trained Keras model
-        last_conv_layer_name: Name of last convolutional layer
-        pred_index: Target class index (if None, uses predicted class)
-    
-    Returns:
-        heatmap: Normalized heatmap (0-1 range)
-    """
-    # Create a model that maps the input image to the activations of the last conv layer
-    # as well as the output predictions
-    grad_model = tf.keras.models.Model(
-        [model.inputs],
-        [model.get_layer(last_conv_layer_name).output, model.output]
-    )
-    
-    # Compute the gradient of the top predicted class for our input image
-    # with respect to the activations of the last conv layer
-    with tf.GradientTape() as tape:
-        last_conv_layer_output, preds = grad_model(img_array)
-        if pred_index is None:
-            pred_index = tf.argmax(preds[0])
-        class_channel = preds[:, pred_index]
-    
-    # Gradient of the output neuron with regard to the output feature map of the last conv layer
-    grads = tape.gradient(class_channel, last_conv_layer_output)
-    
-    # Vector where each entry is the mean intensity of the gradient over a specific feature map channel
-    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
-    
-    # Multiply each channel in the feature map array by "how important this channel is"
-    last_conv_layer_output = last_conv_layer_output[0]
-    heatmap = last_conv_layer_output @ pooled_grads[..., tf.newaxis]
-    heatmap = tf.squeeze(heatmap)
-    
-    # Normalize the heatmap between 0 & 1 for visualization
-    heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
-    return heatmap.numpy()
-
-
-def overlay_heatmap_on_image(img, heatmap, alpha=0.4, colormap=cv2.COLORMAP_JET):
-    """
-    Overlay Grad-CAM heatmap on original image.
-    
-    Args:
-        img: Original PIL Image or numpy array
-        heatmap: Grad-CAM heatmap (0-1 range)
-        alpha: Transparency of heatmap overlay (0-1)
-        colormap: OpenCV colormap (default: JET - red=hot, blue=cold)
-    
-    Returns:
-        superimposed_img: PIL Image with heatmap overlay
-    """
-    # Convert PIL to numpy if needed
-    if isinstance(img, Image.Image):
-        img = np.array(img)
-    
-    # Resize heatmap to match image size
-    heatmap_resized = cv2.resize(heatmap, (img.shape[1], img.shape[0]))
-    
-    # Convert heatmap to RGB
-    heatmap_colored = np.uint8(255 * heatmap_resized)
-    heatmap_colored = cv2.applyColorMap(heatmap_colored, colormap)
-    heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)
-    
-    # Superimpose the heatmap on original image
-    superimposed_img = heatmap_colored * alpha + img * (1 - alpha)
-    superimposed_img = np.uint8(superimposed_img)
-    
-    return Image.fromarray(superimposed_img)
-
-
-def generate_gradcam_for_disease(image, model, disease_name, label_encoder, 
-                                  last_conv_layer_name='top_conv', img_size=224):
-    """
-    Generate Grad-CAM visualization for a specific disease prediction.
-    
-    Args:
-        image: PIL Image
-        model: Trained model
-        disease_name: Name of disease to visualize
-        label_encoder: Disease name -> index mapping
-        last_conv_layer_name: Name of last conv layer in EfficientNetB0
-        img_size: Input image size
-    
-    Returns:
-        overlaid_image: PIL Image with Grad-CAM overlay
-        heatmap: Raw heatmap array
-    """
-    # Preprocess image
-    img_resized = image.convert('RGB').resize((img_size, img_size))
-    img_array = np.array(img_resized) / 255.0
-    img_array = np.expand_dims(img_array, axis=0).astype(np.float32)
-    
-    # Get disease index
-    disease_idx = label_encoder[disease_name]
-    
-    # Generate heatmap
-    heatmap = make_gradcam_heatmap(img_array, model, last_conv_layer_name, disease_idx)
-    
-    # Overlay on original image
-    overlaid_image = overlay_heatmap_on_image(img_resized, heatmap, alpha=0.4)
-    
-    return overlaid_image, heatmap
-
-
-def generate_gradcam_for_top_predictions(image, model, predictions, label_encoder, 
-                                          top_k=3, last_conv_layer_name='top_conv'):
-    """
-    Generate Grad-CAM for top K predicted diseases.
-    
-    Args:
-        image: PIL Image
-        model: Trained model
-        predictions: List of prediction dicts from main app
-        label_encoder: Disease name -> index mapping
-        top_k: Number of top predictions to visualize
-        last_conv_layer_name: Name of last conv layer
-    
-    Returns:
-        gradcam_images: List of (disease_name, overlaid_image, probability) tuples
-    """
-    gradcam_images = []
-    
-    # Sort predictions by probability
-    sorted_preds = sorted(predictions, key=lambda x: x['probability'], reverse=True)[:top_k]
-    
-    for pred in sorted_preds:
-        disease_name = pred['disease']
-        probability = pred['probability']
-        
-        # Generate Grad-CAM
-        overlaid_img, _ = generate_gradcam_for_disease(
-            image, model, disease_name, label_encoder, last_conv_layer_name
-        )
-        
-        gradcam_images.append((disease_name, overlaid_img, probability))
-    
-    return gradcam_images
-
-
-def get_last_conv_layer_name(model):
-    """
-    Automatically find the last convolutional layer in the model.
-    
-    For EfficientNetB0, it's typically 'top_conv' or the last Conv2D layer.
-    
-    Args:
-        model: Keras model
-    
-    Returns:
-        layer_name: Name of last conv layer
-    """
-    # Try common names first
-    common_names = ['top_conv', 'block7a_project_conv', 'conv_head']
-    for name in common_names:
-        try:
-            model.get_layer(name)
-            return name
-        except:
-            pass
-    
-    # Search backwards for Conv2D layer
-    for layer in reversed(model.layers):
-        if isinstance(layer, tf.keras.layers.Conv2D):
-            return layer.name
-    
-    raise ValueError("No convolutional layer found in model!")
+"""
+Improved Grad-CAM Implementation for Medical Images
+====================================================
+Fixed version with better visualization and noise reduction
+"""
+
+import tensorflow as tf
+import numpy as np
+import cv2
+from PIL import Image
+
+
+def make_gradcam_heatmap(img_array, model, last_conv_layer_name, pred_index=None):
+    """
+    Generate improved Grad-CAM heatmap with noise reduction.
+    
+    Args:
+        img_array: Preprocessed image (batch_size, height, width, channels)
+        model: Trained Keras model
+        last_conv_layer_name: Name of last convolutional layer
+        pred_index: Target class index (if None, uses predicted class)
+    
+    Returns:
+        heatmap: Normalized heatmap (0-1 range)
+    """
+    # Create gradient model
+    grad_model = tf.keras.models.Model(
+        [model.inputs],
+        [model.get_layer(last_conv_layer_name).output, model.output]
+    )
+    
+    with tf.GradientTape() as tape:
+        conv_outputs, predictions = grad_model(img_array)
+        
+        if pred_index is None:
+            pred_index = tf.argmax(predictions[0])
+        
+        # Get the score for target class
+        class_channel = predictions[:, pred_index]
+    
+    # Compute gradients
+    grads = tape.gradient(class_channel, conv_outputs)
+    
+    # Global average pooling of gradients
+    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
+    
+    # Weight the channels by importance
+    conv_outputs = conv_outputs[0]
+    pooled_grads = pooled_grads.numpy()
+    conv_outputs = conv_outputs.numpy()
+    
+    # Multiply each channel by its importance
+    for i in range(pooled_grads.shape[-1]):
+        conv_outputs[:, :, i] *= pooled_grads[i]
+    
+    # Average over all channels to get the heatmap
+    heatmap = np.mean(conv_outputs, axis=-1)
+    
+    # Apply ReLU to heatmap (only positive influence)
+    heatmap = np.maximum(heatmap, 0)
+    
+    # Normalize between 0 and 1
+    if heatmap.max() > 0:
+        heatmap = heatmap / heatmap.max()
+    
+    # Apply slight gaussian blur to reduce noise
+    heatmap = cv2.GaussianBlur(heatmap, (3, 3), 0)
+    
+    return heatmap
+
+
+def overlay_heatmap_on_image(img, heatmap, alpha=0.5, colormap=cv2.COLORMAP_JET):
+    """
+    Overlay Grad-CAM heatmap on original image with better contrast.
+    
+    Args:
+        img: Original PIL Image or numpy array
+        heatmap: Grad-CAM heatmap (0-1 range)
+        alpha: Transparency of heatmap overlay (default: 0.5 for better visibility)
+        colormap: OpenCV colormap (JET: red=important, blue=not important)
+    
+    Returns:
+        superimposed_img: PIL Image with heatmap overlay
+    """
+    # Convert PIL to numpy if needed
+    if isinstance(img, Image.Image):
+        img = np.array(img)
+    
+    # Ensure image is RGB
+    if len(img.shape) == 2:
+        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
+    
+    # Resize heatmap to match image size
+    heatmap_resized = cv2.resize(heatmap, (img.shape[1], img.shape[0]))
+    
+    # Apply threshold to remove very weak activations (noise reduction)
+    threshold = 0.2  # Only show activations above 20%
+    heatmap_resized[heatmap_resized < threshold] = 0
+    
+    # Convert heatmap to RGB colormap
+    heatmap_colored = np.uint8(255 * heatmap_resized)
+    heatmap_colored = cv2.applyColorMap(heatmap_colored, colormap)
+    heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)
+    
+    # Normalize original image to 0-255 range
+    if img.max() <= 1.0:
+        img = np.uint8(255 * img)
+    
+    # Create mask for non-zero heatmap areas
+    mask = heatmap_resized > 0
+    
+    # Create output image
+    superimposed_img = img.copy().astype(float)
+    
+    # Only apply heatmap where mask is True
+    superimposed_img[mask] = (
+        heatmap_colored[mask] * alpha + img[mask] * (1 - alpha)
+    )
+    
+    superimposed_img = np.uint8(np.clip(superimposed_img, 0, 255))
+    
+    return Image.fromarray(superimposed_img)
+
+
+def get_last_conv_layer_name(model):
+    """
+    Find the last convolutional layer in EfficientNetB0.
+    
+    Args:
+        model: Keras model
+    
+    Returns:
+        layer_name: Name of last conv layer
+    """
+    # EfficientNetB0 specific layer names (in order of preference)
+    efficientnet_layers = [
+        'top_conv',
+        'block7a_project_conv',
+        'block6d_project_conv',
+        'block6c_project_conv',
+        'conv_head'
+    ]
+    
+    # Try EfficientNet specific layers first
+    for layer_name in efficientnet_layers:
+        try:
+            layer = model.get_layer(layer_name)
+            print(f"✅ Found Grad-CAM layer: {layer_name}")
+            return layer_name
+        except:
+            continue
+    
+    # Fallback: search for last Conv2D layer
+    for layer in reversed(model.layers):
+        if isinstance(layer, tf.keras.layers.Conv2D):
+            print(f"✅ Using fallback Conv2D layer: {layer.name}")
+            return layer.name
+    
+    # Last resort: search in nested models
+    for layer in reversed(model.layers):
+        if hasattr(layer, 'layers'):
+            for sublayer in reversed(layer.layers):
+                if isinstance(sublayer, tf.keras.layers.Conv2D):
+                    print(f"✅ Using nested Conv2D layer: {sublayer.name}")
+                    return sublayer.name
+    
+    raise ValueError("❌ No convolutional layer found in model!")
+
+
+def create_gradcam_comparison(original_img, heatmap, predictions, disease_name):
+    """
+    Create a side-by-side comparison with original, heatmap, and overlay.
+    
+    Args:
+        original_img: Original PIL Image
+        heatmap: Grad-CAM heatmap
+        predictions: Model predictions
+        disease_name: Name of disease being visualized
+    
+    Returns:
+        comparison_img: PIL Image with 3-panel comparison
+    """
+    # Convert original to numpy
+    if isinstance(original_img, Image.Image):
+        original_np = np.array(original_img)
+    else:
+        original_np = original_img
+    
+    # Resize heatmap
+    heatmap_resized = cv2.resize(heatmap, (original_np.shape[1], original_np.shape[0]))
+    
+    # Create colored heatmap
+    heatmap_colored = np.uint8(255 * heatmap_resized)
+    heatmap_colored = cv2.applyColorMap(heatmap_colored, cv2.COLORMAP_JET)
+    heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)
+    
+    # Create overlay
+    overlay = overlay_heatmap_on_image(original_img, heatmap, alpha=0.5)
+    overlay_np = np.array(overlay)
+    
+    # Ensure all images are same size and RGB
+    if len(original_np.shape) == 2:
+        original_np = cv2.cvtColor(original_np, cv2.COLOR_GRAY2RGB)
+    
+    # Stack horizontally
+    comparison = np.hstack([original_np, heatmap_colored, overlay_np])
+    
+    return Image.fromarray(comparison)
+
+
+def generate_multi_disease_gradcam(image, model, predictions, all_diseases, 
+                                    last_conv_layer_name, top_k=3, img_size=224):
+    """
+    Generate Grad-CAM visualizations for multiple diseases.
+    
+    Args:
+        image: Input PIL Image or numpy array
+        model: Trained model
+        predictions: Prediction probabilities for all diseases
+        all_diseases: List of disease names
+        last_conv_layer_name: Name of last conv layer
+        top_k: Number of top predictions to visualize
+        img_size: Image size for model input
+    
+    Returns:
+        gradcam_results: List of (disease_name, probability, gradcam_image) tuples
+    """
+    # Preprocess image
+    if isinstance(image, np.ndarray):
+        img_pil = Image.fromarray(image.astype('uint8'))
+    else:
+        img_pil = image
+    
+    img_resized = img_pil.convert('RGB').resize((img_size, img_size))
+    img_array = np.array(img_resized) / 255.0
+    img_array = np.expand_dims(img_array, axis=0).astype(np.float32)
+    
+    # Get top K diseases
+    top_indices = np.argsort(predictions)[::-1][:top_k]
+    
+    results = []
+    
+    for idx in top_indices:
+        disease_name = all_diseases[idx]
+        probability = float(predictions[idx])
+        
+        # Generate heatmap
+        heatmap = make_gradcam_heatmap(img_array, model, last_conv_layer_name, idx)
+        
+        # Create overlay
+        gradcam_img = overlay_heatmap_on_image(img_resized, heatmap, alpha=0.5)
+        
+        results.append((disease_name, probability, gradcam_img))
+    
+    return results