import torch
import torch.nn as nn
from torch.autograd import Variable as V
import cv2
import numpy as np
class dice_bce_loss(nn.Module):
def __init__(self, batch=True):
super(dice_bce_loss, self).__init__()
self.batch = batch
self.bce_loss = nn.BCELoss()
def soft_dice_coeff(self, y_true, y_pred):
smooth = 0.0 # may change
if self.batch:
i = torch.sum(y_true)
j = torch.sum(y_pred)
intersection = torch.sum(y_true * y_pred)
else:
i = y_true.sum(1).sum(1).sum(1)
j = y_pred.sum(1).sum(1).sum(1)
intersection = (y_true * y_pred).sum(1).sum(1).sum(1)
score = (2. * intersection + smooth) / (i + j + smooth)
# score = (intersection + smooth) / (i + j - intersection + smooth)#iou
return score.mean()
def soft_dice_loss(self, y_true, y_pred):
loss = 1 - self.soft_dice_coeff(y_true, y_pred)
return loss
def forward(self, y_true, y_pred):
a = self.bce_loss(y_pred, y_true)
b = self.soft_dice_loss(y_true, y_pred)
return a + b
if __name__ == "__main__":
# logit.shape:[2,13,320,640]
logit = np.random.random((2, 2, 320, 640))
target1 = np.random.randint(0, 2, size=(320, 640))
target1 = target1[np.newaxis, :, :]
target2 = np.random.randint(0, 2, size=(320, 640))
target2 = target2[np.newaxis, :, :]
# target.shape:[2,320,640]
target = np.vstack([target1, target2])
# numpy --> tensor
logit = torch.tensor(logit)
logit = logit.argmax(dim=1).to(torch.float32)
target = torch.tensor(target).to(torch.float32)
# loss forword
F = dice_bce_loss()
loss = F.forward(target, logit)
print(loss)
dice_loss+BCE_loss 源码
最新推荐文章于 2023-12-28 17:43:58 发布