File size: 2,289 Bytes
226675b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import torch
import torch.nn as nn
class dice_loss(nn.Module):
def __init__(self, batch=True):
super(dice_loss, self).__init__()
# batch equal to True means views all batch images as an entity and calculate loss
# batch equal to False means calculate loss of every single image in batch and get their mean
self.batch = batch
def soft_dice_coeff(self, y_pred, y_true):
smooth = 0.00001
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)
return score.mean()
def soft_dice_loss(self, y_pred, y_true):
loss = 1 - self.soft_dice_coeff(y_pred, y_true)
return loss
def __call__(self, y_pred, y_true):
return self.soft_dice_loss(y_pred.to(dtype=torch.float32), y_true)
class dice_focal_loss(nn.Module):
def __init__(self):
super(dice_focal_loss, self).__init__()
self.focal_loss = nn.BCEWithLogitsLoss()
self.binnary_dice = dice_loss()
def __call__(self, scores, labels):
diceloss = self.binnary_dice(torch.sigmoid(scores.clone()), labels)
foclaloss = self.focal_loss(scores.clone(), labels)
return diceloss, foclaloss
class FCCDN_loss_without_seg(nn.Module):
def __init__(self):
super(FCCDN_loss_without_seg, self).__init__()
def __call__(self, scores, labels):
# scores = change_pred
# labels = binary_cd_labels
scores = scores.squeeze(1) if len(scores.shape) > 3 else scores
labels = labels.squeeze(1) if len(labels.shape) > 3 else labels
# if len(scores.shape) > 3:
# scores = scores.squeeze(1)
# if len(labels.shape) > 3:
# labels = labels.squeeze(1)
""" for binary change detection task"""
criterion_change = dice_focal_loss()
# change loss
diceloss, foclaloss = criterion_change(scores, labels.float())
loss_change = diceloss + foclaloss
return loss_change.mean()
|