多类别语义分割测试代码unet 并且包含评价指标计算

以下是使用PyTorch实现的UNet多类别语义分割测试代码，同时包含评价指标计算：

import torch
import numpy as np
import argparse
import os
import cv2

from model import UNet # 导入UNet模型
from datasets import get_dataloader # 导入数据加载器
from utils.metrics import IoU, dice_coeff # 导入评价指标计算函数

def test(args):
    # 设置使用的GPU编号
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu

    # 加载测试集数据
    test_loader = get_dataloader(args.data_dir, 'test', args.batch_size, args.num_workers)

    # 加载模型
    model = UNet(num_classes=args.num_classes)
    model = model.cuda()

    # 加载预训练权重
    checkpoint = torch.load(args.checkpoint)
    model.load_state_dict(checkpoint)

    # 设置模型为评价模式
    model.eval()

    # 初始化评价指标计算变量
    total_iou = 0.0
    total_dice = 0.0
    total_num = 0

    # 遍历测试集
    for i, (images, labels) in enumerate(test_loader):
        # 将数据移动到GPU上
        images = images.cuda()
        labels = labels.cuda()

        # 前向传播，得到模型输出结果
        with torch.no_grad():
            outputs = model(images)

        # 计算评价指标
        preds = torch.argmax(outputs, dim=1).cpu().numpy()
        labels = labels.cpu().numpy()
        for pred, label in zip(preds, labels):
            total_iou += IoU(pred, label, args.num_classes)
            total_dice += dice_coeff(pred, label, args.num_classes)
            total_num += 1

        # 保存预测结果图片
        if args.save_test_results:
            for j in range(preds.shape[0]):
                img = np.transpose(images[j].cpu().numpy(), (1, 2, 0))
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                pred = preds[j]
                label = labels[j]
                img_pred = np.zeros((pred.shape[0], pred.shape[1], 3), dtype=np.uint8)
                img_label = np.zeros((label.shape[0], label.shape[1], 3), dtype=np.uint8)
                for k in range(args.num_classes):
                    img_pred[pred == k] = args.colors[k]
                    img_label[label == k] = args.colors[k]
                img_pred = cv2.addWeighted(img, 0.5, img_pred, 0.5, 0)
                img_label = cv2.addWeighted(img, 0.5, img_label, 0.5, 0)
                cv2.imwrite(os.path.join(args.test_results_dir, f'{i*args.batch_size+j}_pred.png'), img_pred)
                cv2.imwrite(os.path.join(args.test_results_dir, f'{i*args.batch_size+j}_label.png'), img_label)

    # 打印评价指标结果
    avg_iou = total_iou / total_num
    avg_dice = total_dice / total_num
    print(f'Average IoU: {avg_iou:.4f}')
    print(f'Average Dice Coefficient: {avg_dice:.4f}')

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--data_dir', type=str, default='./data', help='数据集目录')
    parser.add_argument('--checkpoint', type=str, default='./checkpoints/best_model.pth', help='预训练权重路径')
    parser.add_argument('--num_classes', type=int, default=2, help='类别数量')
    parser.add_argument('--batch_size', type=int, default=16, help='批量大小')
    parser.add_argument('--num_workers', type=int, default=4, help='数据加载器使用的进程数量')
    parser.add_argument('--save_test_results', action='store_true', help='是否保存测试结果图片')
    parser.add_argument('--test_results_dir', type=str, default='./test_results', help='测试结果图片保存目录')
    parser.add_argument('--gpu', type=str, default='0', help='使用的GPU编号')
    args = parser.parse_args()

    # 设置不同类别的颜色
    args.colors = [(0, 0, 0), (255, 255, 255)]
    test(args)

其中，`IoU`和`dice_coeff`函数的实现可以参考以下代码：

import numpy as np

def IoU(pred, label, num_classes):
    iou = 0.0
    for i in range(num_classes):
        tp = np.sum((pred == i) & (label == i))
        fp = np.sum((pred == i) & (label != i))
        fn = np.sum((pred != i) & (label == i))
        iou += tp / (tp + fp + fn)
    return iou / num_classes

def dice_coeff(pred, label, num_classes):
    dice = 0.0
    for i in range(num_classes):
        tp = np.sum((pred == i) & (label == i))
        fp = np.sum((pred == i) & (label != i))
        fn = np.sum((pred != i) & (label == i))
        dice += 2 * tp / (2 * tp + fp + fn)
    return dice / num_classes

注意，在测试代码中，我们设置了不同类别的颜色，并且可以选择是否保存测试结果图片。如果需要保存图片，则需要指定`--save_test_results`参数，并且需要指定保存路径`--test_results_dir`。