code format

Author: Hsuxu

DiceLoss/__init__.py

@@ -0,0 +1,4 @@
+from .dice_loss import *
+from focal_loss import *
+from lovasz_loss import *
+from tverskyloss import *

FocalLoss/README.md

@@ -40,12 +40,12 @@ class BinaryDiceLoss(nn.Module):
         Exception if unexpected reduction
     """
 
-    def __init__(self, ignore_index=None, reduction='mean',**kwargs):
+    def __init__(self, ignore_index=None, reduction='mean', **kwargs):
         super(BinaryDiceLoss, self).__init__()
-        self.smooth = 1   # suggest set a large number when target area is large,like '10|100'
+        self.smooth = 1  # suggest set a large number when target area is large,like '10|100'
         self.ignore_index = ignore_index
         self.reduction = reduction
-        self.batch_dice = False # treat a large map when True
+        self.batch_dice = False  # treat a large map when True
         if 'batch_loss' in kwargs.keys():
             self.batch_dice = kwargs['batch_loss']
 
@@ -59,7 +59,7 @@ def forward(self, output, target, use_sigmoid=True):
             output = output.mul(validmask)  # can not use inplace for bp
             target = target.float().mul(validmask)
 
-        dim0= output.shape[0]
+        dim0 = output.shape[0]
         if self.batch_dice:
             dim0 = 1
 
@@ -212,7 +212,7 @@ def test():
     model = nn.Conv3d(1, 4, 3, padding=1)
     target = torch.randint(0, 4, (3, 1, 32, 32, 32)).float()
     target = make_one_hot(target, num_classes=4)
-    criterion = DiceLoss(ignore_index=[2,3], reduction='mean')
+    criterion = DiceLoss(ignore_index=[2, 3], reduction='mean')
     loss = criterion(model(input), target)
     loss.backward()
     print(loss.item())

FocalLoss/__init__.py

@@ -1,14 +1,11 @@
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
-from torch.autograd import Function
 
-from torch.autograd import Variable
+from torch.autograd import Function
 
 
 class FocalBinaryTverskyLoss(Function):
 
-
     @staticmethod
     def forward(ctx, input, target):
         _alpha = 0.5
@@ -142,3 +139,46 @@ def forward(self, output, target, mask=None):
         else:
             loss = torch.mean(loss)
         return loss
+
+
+class MultiTverskyLoss(nn.Module):
+    """
+    Tversky Loss for segmentation adaptive with multi class segmentation
+    """
+
+    def __init__(self, alpha=0.5, beta=0.5, gamma=1.0, weights=None):
+        """
+        :param alpha (Tensor, float, optional): controls the penalty for false positives.
+        :param beta (Tensor, float, optional): controls the penalty for false negative.
+        :param gamma (Tensor, float, optional): focal coefficient
+        :param weights (Tensor, optional): a manual rescaling weight given to each
+            class. If given, it has to be a Tensor of size `C`
+
+        """
+        super(MultiTverskyLoss, self).__init__()
+        self.alpha = alpha
+        self.beta = beta
+        self.gamma = gamma
+        self.weights = weights
+
+    def forward(self, inputs, targets):
+
+        num_class = inputs.size(1)
+        weight_losses = 0.0
+        if self.weights is not None:
+            assert len(self.weights) == num_class, 'number of classes should be equal to length of weights '
+            weights = self.weights
+        else:
+            weights = [1.0 / num_class] * num_class
+        input_slices = torch.split(inputs, [1] * num_class, dim=1)
+        for idx in range(num_class):
+            input_idx = input_slices[idx]
+            input_idx = torch.cat((1 - input_idx, input_idx), dim=1)
+            target_idx = (targets == idx) * 1
+            loss_func = FocalBinaryTverskyLoss(self.alpha, self.beta, self.gamma)
+            loss_idx = loss_func.apply(input_idx, target_idx)
+            weight_losses += loss_idx * weights[idx]
+        # loss = torch.Tensor(weight_losses)
+        # loss = loss.to(inputs.device)
+        # loss = torch.sum(loss)
+        return weight_losses