grad-cam医学图像分割可视化代码

以下是利用Grad-CAM算法进行医学图像分割可视化的示例代码。代码使用的是PyTorch框架。

import torch
import torch.nn.functional as F
import cv2
import numpy as np
import matplotlib.pyplot as plt

# 定义Grad-CAM函数
def grad_cam(model, x, layer_name):
    # 获取目标层的输出特征图和梯度
    features_blobs = []
    def hook_feature(module, input, output):
        features_blobs.append(output.data.cpu().numpy())
    net._modules.get(layer_name).register_forward_hook(hook_feature)
    grad_blobs = []
    def hook_grad(module, grad_in, grad_out):
        grad_blobs.append(grad_out[0].cpu().numpy())
    net._modules.get(layer_name).register_backward_hook(hook_grad)

    # 运行模型并计算梯度
    output = model(x)
    one_hot = np.zeros((1, output.size()[-1]), dtype=np.float32)
    one_hot[0][np.argmax(output.data.cpu().numpy())] = 1
    one_hot = torch.from_numpy(one_hot)
    one_hot.requires_grad = True
    one_hot = torch.sum(one_hot * output)

    model.zero_grad()
    one_hot.backward()
    grad = grad_blobs[0]
    feature = features_blobs[0]
    weights = np.mean(grad, axis=(2, 3))[0, :]
    cam = np.zeros(feature.shape[2:], dtype=np.float32)
    for i, w in enumerate(weights):
        cam += w * feature[0, i, :, :]

    cam = cv2.resize(cam, x.shape[2:])
    cam = np.maximum(cam, 0)
    heatmap = cam / np.max(cam)

    return heatmap

# 读取图像
img = cv2.imread('test.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
plt.imshow(img)
plt.show()

# 转换图像为Tensor，并进行预处理
img = cv2.resize(img, (224, 224))
img = np.float32(img) / 255
img = np.transpose(img, (2, 0, 1))
img = torch.from_numpy(img).unsqueeze(0)

# 加载模型
model = torch.load('model.pth')
model.eval()

# 应用Grad-CAM算法
heatmap = grad_cam(model, img, 'features')

# 可视化结果
heatmap = cv2.resize(heatmap, (img.shape[3], img.shape[2]))
heatmap = np.uint8(heatmap * 255)
heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
result = cv2.addWeighted(img.squeeze().numpy().transpose(1, 2, 0), 0.5, heatmap, 0.5, 0)
plt.imshow(result)
plt.show()

其中，`model.pth`为已经训练好的模型文件，`test.jpg`为待处理的图像文件。需要根据具体情况进行修改。