在原始图像上可视化分割结果

1.预测分割图

训练的得分最高的模型参数文件为best_score_ISIC2018_checkpoint.pth.tar
定义一个转换图像，通过调用模型文件和训练的参数文件进行对单个输入图像进行推理预测。并将分割结果保存显示。

import os

import torch
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
from matplotlib import pyplot as plt
import matplotlib.pyplot as plt
import transform
from Models.networks.msca_net import msca_net

# 定义图像转换函数
def transform_image(image_path):
image = Image.open(image_path).convert('RGB')
transform = transforms.Compose([
transforms.Resize((224, 320)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
image = transform(image).unsqueeze(0) # 扩展一个维度即batch
return image

# 加载checkpoint文件
checkpoint = torch.load('../best_score_ISIC2018_checkpoint.pth.tar')
model = ca_net()
model.load_state_dict(checkpoint['state_dict']) # 这个函数的作用是将 checkpoint['state_dict'] 中保存的模型参数加载到 model 中。

# 进行推理
model.eval()
image = transform_image('../image/ISIC_0000003.jpg')

with torch.no_grad():
output = model(image)

# 处理输出
# output = torch.sigmoid(output)
output = output.squeeze(0).permute(1, 2, 0).numpy()

# 可视化结果
# 将灰度图像转换为二值图像
threshold = 0.5
binary_output = np.where(output[:,:,0] >= threshold, 1, 0)

# 显示二值图像
plt.imshow(binary_output, cmap='binary')
plt.show()

# 保存二值图像
save_dir = 'result'
if not os.path.exists(save_dir):
os.makedirs(save_dir)
save_path = os.path.join(save_dir, 'ISIC_0000003.png')
plt.imsave(save_path, binary_output, cmap='binary')

2.可视化预测分割

将推理预测的分割掩码与原始图像融合，从而在原始图像中可视化预测分割，用红色表示预测分割

import cv2
import numpy as np

# 加载原始图像和分割掩码
img = cv2.imread('../result_image.jpg')
mask = cv2.imread('../ISIC_0000003_segmentation.png', cv2.IMREAD_GRAYSCALE)

# 使用Canny边缘检测算法检测分割掩码的边缘
edges = cv2.Canny(mask, 100, 200)

# 将检测到的边缘绘制为红色线条
result = img.copy()
result[edges != 0] = [0, 255, 0]

# 保存结果
cv2.imwrite('../result_image1.jpg', result)
# 显示结果
cv2.imshow('Result', result)
cv2.waitKey(0)
cv2.destroyAllWindows()

3.可视化Ground truth

将Ground truth的分割掩码与原始图像融合，从而在原始图像中可视化Ground truth的分割掩码，用绿色表示分割。

import cv2
import numpy as np

# 加载原始图像和Ground truth掩码
img = cv2.imread('../result_image.jpg')
gt_mask = cv2.imread('../ISIC_0000003_segmentation.png', cv2.IMREAD_GRAYSCALE)

# 将Ground truth掩码转换为二值图像
gt_mask_bin = cv2.threshold(gt_mask, 127, 255, cv2.THRESH_BINARY)[1]

# 使用Canny边缘检测算法检测Ground truth掩码的边缘
gt_edges = cv2.Canny(gt_mask_bin, 100, 200)

# 将检测到的Ground truth边缘绘制为绿色线条
result = img.copy()
result[gt_edges != 0] = [0, 255, 0]

cv2.imwrite('../result_image1.jpg', result)
# 显示结果
cv2.imshow('Result', result)
cv2.waitKey(0)
cv2.destroyAllWindows()