医学图像分割任务的测试代码

测试集进行测试

import os
import torch
import numpy as np
from torch.utils.data import DataLoader
from sklearn.metrics import (precision_score,recall_score,f1_score,roc_curve,auc,confusion_matrix,
)
import matplotlib.pyplot as plt
from utils import NiiDataset
from model.UNet import UNet# 加载最佳模型
best_unet_model = r"D:\PytnonProject\Segment\best_unet_model.pth"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNet(in_channels=1, out_channels=1).to(device)
model.load_state_dict(torch.load(best_unet_model))
model.eval()  # 设置为评估模式# 定义测试数据集
test_image_paths = [r"D:\Data\DegmentData\OriginalNii\DCE\CAO_ZHAN_GUO.nii",  # 替换为实际的测试图像路径r"D:\Data\DegmentData\OriginalNii\DCE\CHAI_GUI_LAN.nii",# 添加其他测试数据路径...
]test_mask_paths = [r"D:\Data\DegmentData\ROI\CAO_ZHAN_GUO-label.nii",  # 替换为实际的测试掩码路径r"D:\Data\DegmentData\ROI\CHAI_GUI_LAN-label.nii",# 添加其他测试掩码路径...
]# 创建测试数据集和数据加载器
test_dataset = NiiDataset(test_image_paths, test_mask_paths)
test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False)  # 批量大小为 1# 定义评估指标
def dice_coefficient(y_true, y_pred):intersection = np.sum(y_true * y_pred)denominator = np.sum(y_true) + np.sum(y_pred)if denominator == 0:return 1.0  # 如果分母为零，返回 1（表示完全匹配）return (2.0 * intersection) / denominatordef iou(y_true, y_pred):intersection = np.sum(y_true * y_pred)union = np.sum(y_true) + np.sum(y_pred) - intersectionif union == 0:return 1.0  # 如果分母为零，返回 1（表示完全匹配）return intersection / union# 初始化指标
dice_scores = []
iou_scores = []
precisions = []
recalls = []
f1_scores = []
sensitivities = []
specificities = []
auc_scores = []# 用于 ROC 曲线的数据
all_masks = []
all_predictions = []
all_probabilities = []# 测试过程
with torch.no_grad():  # 禁用梯度计算for images, masks in test_dataloader:images = images.to(device)  # 形状: [1, 1, 480, 480]masks = masks.to(device)    # 形状: [1, 1, 480, 480]# 前向传播outputs = model(images)  # 形状: [1, 1, 480, 480]# 将输出转换为概率和二进制掩码probabilities = outputs.cpu().numpy().flatten()  # 形状: [480 * 480]predictions = (outputs > 0.5).float().cpu().numpy().flatten()  # 形状: [480 * 480]masks = masks.cpu().numpy().flatten()  # 形状: [480 * 480]# 保存用于 ROC 曲线的数据all_masks.extend(masks)all_predictions.extend(predictions)all_probabilities.extend(probabilities)# 检查 masks 和 predictions 是否只包含一个类别if np.all(masks == 0) and np.all(predictions == 0):# 如果 masks 和 predictions 都为全 0，则跳过该样本continue# 计算指标dice = dice_coefficient(masks, predictions)iou_score = iou(masks, predictions)precision = precision_score(masks, predictions, zero_division=0)recall = recall_score(masks, predictions, zero_division=0)f1 = f1_score(masks, predictions, zero_division=0)# 计算混淆矩阵cm = confusion_matrix(masks, predictions)if cm.size == 1:# 如果混淆矩阵只有一个值（全 0 或全 1）if np.all(masks == 0):tn, fp, fn, tp = cm[0, 0], 0, 0, 0else:tn, fp, fn, tp = 0, 0, 0, cm[0, 0]else:tn, fp, fn, tp = cm.ravel()sensitivity = tp / (tp + fn) if (tp + fn) > 0 else 0  # 灵敏度specificity = tn / (tn + fp) if (tn + fp) > 0 else 0  # 特异度# 保存指标dice_scores.append(dice)iou_scores.append(iou_score)precisions.append(precision)recalls.append(recall)f1_scores.append(f1)sensitivities.append(sensitivity)specificities.append(specificity)# 计算 AUC
fpr, tpr, thresholds = roc_curve(all_masks, all_probabilities)
roc_auc = auc(fpr, tpr)
auc_scores.append(roc_auc)# 打印平均指标
print(f"Average Dice Coefficient: {np.mean(dice_scores):.4f}")
print(f"Average IoU: {np.mean(iou_scores):.4f}")
print(f"Average Precision: {np.mean(precisions):.4f}")
print(f"Average Recall: {np.mean(recalls):.4f}")
print(f"Average F1 Score: {np.mean(f1_scores):.4f}")
print(f"Average Sensitivity: {np.mean(sensitivities):.4f}")
print(f"Average Specificity: {np.mean(specificities):.4f}")
print(f"Average AUC: {np.mean(auc_scores):.4f}")# 创建 checkpoint 文件夹（如果不存在）
checkpoint_dir = "checkpoint"
os.makedirs(checkpoint_dir, exist_ok=True)# 保存 ROC 曲线的数据到 checkpoint 文件夹下的 .npz 文件
roc_data_path = os.path.join(checkpoint_dir, "roc_data.npz")
np.savez(roc_data_path, fpr=fpr, tpr=tpr, thresholds=thresholds, roc_auc=roc_auc)
print(f"ROC 曲线的数据已保存到文件: {roc_data_path}")# 绘制 ROC 曲线
plt.figure()
plt.plot(fpr, tpr, color="darkorange", lw=2, label=f"ROC curve (AUC = {roc_auc:.2f})")
plt.plot([0, 1], [0, 1], color="navy", lw=2, linestyle="--")
plt.xlabel("False Positive Rate")
plt.ylabel("True Positive Rate")
plt.title("Receiver Operating Characteristic (ROC) Curve")
plt.legend(loc="lower right")
plt.show()