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

测试集进行测试

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",
    # 添加其他测试掩码路径...
]

# 创建测试数据集