add dice focal

Author: Hsuxu

seg_loss/dice_loss.py

@@ -3,6 +3,8 @@
 import torch.nn.functional as F
 import numpy as np
 
+from .focal_loss import BinaryFocalLoss
+
 
 def make_one_hot(input, num_classes=None):
     """Convert class index tensor to one hot encoding tensor.
@@ -207,6 +209,19 @@ def forward(self, output, target):
         return loss
 
 
+class Binary_Focal_Dice(nn.Module):
+    def __init__(self, **kwargs):
+        super(Binary_Focal_Dice, self).__init__()
+        self.dice = BinaryDiceLoss(**kwargs)
+        self.focal = BinaryFocalLoss(**kwargs)
+
+    def forward(self, logits, target):
+        dice_loss = self.dice(logits, target)
+        focal_loss = self.focal(logits, target)
+        loss = dice_loss + focal_loss
+        return loss, (dice_loss.detach(), focal_loss.detach())
+
+
 def test():
     input = torch.rand((3, 1, 32, 32, 32))
     model = nn.Conv3d(1, 4, 3, padding=1)