AI_MRI / scripts /validation /generate_heatmaps.py

upload job examples and scripts

011f070 verified 28 days ago

2.96 kB

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()

	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()