Calculation of evaluation indicators for each class

[RY-97]

def validation(epoch_iterator_val, model, args): # 设置模型为验证模式 model.eval()

# 初始化度量指标
dice_metric = DiceMetric(include_background=False, reduction="mean", get_not_nans=False)
dice_metric_batch = DiceMetric(include_background=False, reduction="mean_batch", get_not_nans=False)
hd_metric = HausdorffDistanceMetric(include_background=False, reduction="mean", percentile=95)
hd_metric_batch = HausdorffDistanceMetric(include_background=False, reduction="mean_batch", percentile=95)

# 后处理函数
post_label = AsDiscrete(to_onehot=args.num_classes)
post_pred = AsDiscrete(argmax=True, to_onehot=args.num_classes)

# 初始化度量值列表
metric_values_dice = []
metric_values_hd = []
metric_values_dice_per_class = [[] for _ in range(args.num_classes-1)]
metric_values_hd_per_class = [[] for _ in range(args.num_classes-1)]

with torch.no_grad():
    for step, batch in enumerate(epoch_iterator_val):
        val_inputs, val_labels = batch["image"].cuda(), batch["label"].cuda()

        # 使用滑动窗口推理
        val_outputs = sliding_window_inference(val_inputs, (96, 96, 96), 2, model)

        # 解码并转换标签和预测
        val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]
        val_labels = [post_label(i) for i in decollate_batch(val_labels)]

        # 计算每个类别的Dice和HD
        dice_metric(y_pred=val_outputs, y=val_labels)
        dice_metric_batch(y_pred=val_outputs, y=val_labels)

        hd_metric(y_pred=val_outputs, y=val_labels)
        hd_metric_batch(y_pred=val_outputs, y=val_labels)

        # 获取总体 Dice 和 HD 度量
        metric_dice = dice_metric.aggregate().item()
        metric_values_dice.append(metric_dice)
        metric_hd = hd_metric.aggregate().item()
        metric_values_hd.append(metric_hd)

        # 获取每个类别的度量结果
        dice_metric_per_class = dice_metric_batch.aggregate()
        hd_metric_per_class = hd_metric_batch.aggregate()

        # 提取每个类别的 Dice 和 HD
        for i in range(0, args.num_classes - 1):
            metric_dice_i = dice_metric_per_class[i].item()  # 修正索引
            metric_values_dice_per_class[i].append(metric_dice_i)

            metric_hd_i = hd_metric_per_class[i].item()  # 修正索引
            metric_values_hd_per_class[i].append(metric_hd_i)

        # 重置度量器以避免累计误差
        dice_metric.reset()
        dice_metric_batch.reset()
        hd_metric.reset()
        hd_metric_batch.reset()

    # 计算总体平均 Dice 和 HD
    avg_dice = np.mean(metric_values_dice)
    avg_hd = np.mean(metric_values_hd)

    print("\n" + "--" * 40 + "\n")

    # 输出总体平均 Dice 和 HD
    print(f"Overall Average - Dice: {avg_dice:.4f}, HD: {avg_hd:.4f}")

    # 输出每个类别的 Dice 和 HD 度量
    for i in range(1, args.num_classes):  # 从 1 开始跳过背景类
        avg_dice_per_class = np.mean(metric_values_dice_per_class[i - 1])
        avg_hd_per_class = np.mean(metric_values_hd_per_class[i - 1])
        print(f"Class {i} - Dice: {avg_dice_per_class:.4f}, HD: {avg_hd_per_class:.4f}")

    print("\n" + "--" * 40 + "\n")

    # 返回总的平均和每个类别的平均度量值
    return {
        "overall_dice": avg_dice,
        "overall_hd": avg_hd,
        "dice_per_class": [np.mean(metric_values_dice_per_class[i]) for i in range(args.num_classes - 1)],
        "hd_per_class": [np.mean(metric_values_hd_per_class[i]) for i in range(args.num_classes - 1)]
    }

I am using MONAI for research. My task is multi-category medical image segmentation. Now I need to obtain the evaluation index of each category. Is my code correct?