ncuiew7/Advanced_Bone-Fracture_Detection_and_Localized_Attention_Mapping

1import torch2import numpy as np3import cv24import random5from PIL import Image6import matplotlib.pyplot as plt7 8def get_gradcam(model, image_tensor, target_layer, task_index=1):9    """Generates Grad-CAM heatmap for a specific task in the multi-label model."""10    gradients = []11    activations = []12 13    def save_gradient(module, grad_input, grad_output):14        gradients.append(grad_output[0])15 16    def save_activation(module, input, output):17        activations.append(output)18 19    # Register hooks20    handle_act = target_layer.register_forward_hook(save_activation)21    handle_grad = target_layer.register_full_backward_hook(save_gradient)22 23    model.eval()24    logits = model(image_tensor).logits25    target_logit = logits[:, task_index]26 27    model.zero_grad()28    target_logit.backward()29 30    # Capture and detach data31    grad = gradients[0].cpu().detach().numpy()[0]32    act = activations[0].cpu().detach().numpy()[0]33 34    # Reshape from sequence (N, C) to spatial grid (H, W, C)35    grid_size = int(np.sqrt(act.shape[0]))36    act = act.reshape(grid_size, grid_size, -1)37    grad = grad.reshape(grid_size, grid_size, -1)38 39    # Calculate weighted sum of activations40    weights = np.mean(grad, axis=(0, 1))41    cam = np.dot(act, weights)42 43    # ReLU and Normalize44    cam = np.maximum(cam, 0)45    cam = cv2.resize(cam, (224, 224))46    cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)47 48    handle_act.remove()49    handle_grad.remove()50 51    return cam52 53# --- Execution ---54fracture_samples = [i for i, x in enumerate(binary_multitask_ds['train']) if 'positive' in x['image'].filename.lower()]55sidx = random.choice(fracture_samples)56sample = binary_multitask_ds['train'][sidx]57img = sample['image'].convert('RGB')58 59inputs = processor(img, return_tensors="pt").to(model.device)60inputs['pixel_values'].requires_grad = True61 62# Correct layer name for Swin: 'layernorm_before' or 'layernorm_after'63target_layer = model.swin.encoder.layers[-1].blocks[-1].layernorm_before64 65heatmap = get_gradcam(model, inputs['pixel_values'], target_layer)66 67# Overlay and Plot68img_np = np.array(img.resize((224, 224)))69heatmap_color = cv2.applyColorMap(np.uint8(255 * heatmap), cv2.COLORMAP_JET)70overlay = cv2.addWeighted(img_np, 0.6, heatmap_color, 0.4, 0)71 72plt.figure(figsize=(10, 5))73plt.subplot(1, 2, 1)74plt.imshow(img_np)75plt.title("Original X-Ray")76plt.axis('off')77 78plt.subplot(1, 2, 2)79plt.imshow(overlay)80plt.title("Fracture Evidence Heatmap (Grad-CAM)")81plt.axis('off')82plt.show()