File size: 2,957 Bytes
011f070 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | from pathlib import Path
import torch
import torch.nn as nn
import torchio as tio
import matplotlib.pyplot as plt
import numpy as np
from auto_detect_breast_mri.config import resolve_path
from auto_detect_breast_mri.models.resnets import model_names
from auto_detect_breast_mri.data.breast_mri_dataset import BreastMRISubjects
# Paths come from the site config (see config.example.yaml).
path_base = resolve_path(None, "data_root", "root folder of the NIfTI data")
split_root = resolve_path(None, "split_root", "folder holding the split files")
model_root = resolve_path(None, "model_root", "folder holding the checkpoints", required=False) \
or resolve_path(None, "output_root", "folder holding the checkpoints")
batch_size = 1
pre_image_shape = (32, 512, 512)
heatmap_size = (32, 256, 256)
# Get Model
model_name = 'resnet18_abrv'
protocol = ['Dyn_0', 'Sub_1']
model_path = str(Path(model_root) / "malignity_{}.pth".format(model_name))
model = model_names.get(model_name)
checkpoint = torch.load(model_path)
model.load_state_dict(checkpoint)
model.eval()
# Register forward hook
target_layer = model.layer4[-1] # Choose the target layer
activation = {}
def get_activation(name):
def hook(model, input, output):
activation[name] = output.detach()
return hook
target_layer.register_forward_hook(get_activation('target_layer'))
# forward pass
data_set = BreastMRISubjects(path_base, str(Path(split_root) / "fold0" / "stratified_evaluation_set-f0-0.csv"),
protocol=protocol)
data_loader = tio.SubjectsLoader(data_set, batch_size=batch_size, shuffle=False, num_workers=16)
for batch_id, batch in enumerate(data_loader):
data = batch['image']['data']
ground_truth = batch['label']
path = batch['path']
output = model(data)
# Calculate gradients
model.zero_grad()
one_hot = torch.zeros_like(output)
one_hot[0, output.argmax()] = 1.0
output.backward(gradient=one_hot, retain_graph=True)
# Weighted sum of feature maps
grads = activation['target_layer'].grad.mean(dim=(2, 3), keepdim=True)
heatmap = torch.sum(grads * activation['target_layer'], dim=1, keepdim=True)
heatmap = nn.functional.relu(heatmap)
# Resize the heatmap to match the original image size
heatmap = nn.functional.interpolate(heatmap, size=heatmap_size, mode="bilinear", align_corners=False)
heatmap = heatmap / heatmap.max() # Normalize
heatmap = heatmap.squeeze().numpy() # transform to numpy array
# Display the original image and Grad-CAM heatmap
plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.imshow(np.transpose(data[0, 0, 16].squeeze().numpy(), (1, 2, 0)))
plt.title('Original Image')
plt.subplot(1, 2, 2)
plt.imshow(heatmap[0, 0, 16], cmap='viridis') # TODO maybe only 0, 16 ?
plt.title('Grad-CAM Heatmap')
plt.show()
|