61 lines
2.0 KiB
Python
61 lines
2.0 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import segmentation_models_pytorch as smp
|
|
|
|
|
|
class CustomLoss(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.dice_loss = smp.losses.DiceLoss(mode="multiclass")
|
|
self.ce_loss = nn.CrossEntropyLoss()
|
|
|
|
def forward(self, y_pred, y_true):
|
|
return self.dice_loss(y_pred, y_true) + self.ce_loss(y_pred, y_true)
|
|
|
|
|
|
# def dice(pred, target, eps=1.0):
|
|
# """
|
|
# 计算每个类别的 Dice 系数，并返回平均值。
|
|
# pred: 预测的类别（预测的最大类别索引）。
|
|
# target: 真实标签类别（每个像素的类别索引）。
|
|
# """
|
|
# num_classes = pred.shape[1]
|
|
# pred = torch.argmax(pred, dim=1)
|
|
|
|
# dice_list = []
|
|
# for c in range(1, num_classes):
|
|
# # 对于每个类别，计算该类别的 Dice 系数
|
|
# pred_c = (pred == c)
|
|
# target_c = (target == c)
|
|
|
|
# intersection = torch.sum(pred_c & target_c)
|
|
|
|
# # 计算 Dice 系数
|
|
# dice_list.append((2 * intersection) / (torch.sum(pred_c) + torch.sum(target_c) + eps))
|
|
|
|
# return torch.mean(torch.tensor(dice_list)) # 返回所有类别的平均 Dice 系数
|
|
|
|
def iou(pred, target, eps=1.0):
|
|
"""
|
|
计算每个类别的 IoU 系数，并返回平均值。
|
|
pred: 预测的类别（预测的最大类别索引）。
|
|
target: 真实标签类别（每个像素的类别索引）。
|
|
"""
|
|
num_classes = pred.shape[1]
|
|
pred = torch.argmax(pred, dim=1)
|
|
|
|
iou_list = []
|
|
for c in range(1, num_classes):
|
|
# 对于每个类别，计算该类别的 IoU
|
|
pred_c = (pred == c)
|
|
target_c = (target == c)
|
|
|
|
intersection = torch.sum(pred_c & target_c)
|
|
union = torch.sum(pred_c | target_c)
|
|
|
|
# 计算 IoU
|
|
if union != 0:
|
|
iou_list.append(intersection / (union + eps))
|
|
|
|
return torch.mean(torch.tensor(iou_list)) # 返回所有类别的平均 IoU
|