update NetRunAgent

Author: taigw

examples/JSRT2/README.md

@@ -5,9 +5,9 @@
 
 In this example, we show how to use a customized CNN and a customized loss function to segment the heart from X-Ray images. The configurations are the same as those in the `JSRT` example except the network structure and loss function. 
 
-The customized CNN is detailed in `my_net2d.py`, which is a modification of the 2D UNet. In this new network, we use a residual connection in each block. The customized loss is detailed in `my_loss.py`, where we combine Dice loss and MAE loss as our new loss function.
+The customized CNN is detailed in `my_net2d.py`, which is a modification of the 2D UNet. In this new network, we use a residual connection in each block. The customized loss is detailed in `my_loss.py`, where we define a focal dice loss. 
 
-We also write a customized main function in `jsrt_train_infer.py` so that we can combine TrainInferAgent from PyMIC with our customized CNN and loss function.
+We also write a customized main function in `jsrt_net_run.py` so that we can combine NetRunAgent from PyMIC with our customized CNN and loss function.
 
 ## Data and preprocessing
 1. Data preprocessing is the same as that in the the `JSRT` example. Please follow that example for details.

examples/JSRT2/config/train_test.cfg

@@ -53,7 +53,11 @@ bilinear     = True
 device_name = cuda:0
 
 batch_size    = 4
-loss_function = my_loss
+loss_type     = MyFocalDiceLoss
+MyFocalDiceLoss_Enable_Pixel_Weight = False
+MyFocalDiceLoss_Enable_Class_Weight = True
+MyFocalDiceLoss_beta = 1.5
+class_weight = [0.2, 1.0]
 
 # for optimizers
 optimizer     = Adam

examples/JSRT2/jsrt_net_run.py

@@ -2,11 +2,12 @@
 from __future__ import print_function, division
 
 import sys
-from pymic.net_run.net_run import  TrainInferAgent
-from pymic.net_run.net_factory import net_dict
 from pymic.util.parse_config import parse_config
+from pymic.net_run.net_run_agent import  NetRunAgent
+from pymic.net.net_factory import net_dict
+from pymic.loss.loss_factory import loss_dict
 from my_net2d import MyUNet2D 
-from my_loss  import MySegmentationLossCalculator
+from my_loss  import MyFocalDiceLoss
 
 my_net_dict = {
     "MyUNet2D": MyUNet2D
@@ -22,6 +23,20 @@ def get_network(params):
         raise ValueError("Undefined network: {0:}".format(net_type))
     return net 
 
+my_loss_dict = {
+    "MyFocalDiceLoss": MyFocalDiceLoss
+}
+
+def get_loss(params):
+    loss_type = params["loss_type"]
+    if(loss_type in my_loss_dict):
+        loss_obj = my_loss_dict[loss_type](params)
+    elif(loss_type in net_dict):
+        loss_obj = loss_dict[loss_type](params)
+    else:
+        raise ValueError("Undefined loss: {0:}".format(loss_type))
+    return loss_obj
+
 def main():
     if(len(sys.argv) < 3):
         print('Number of arguments should be 3. e.g.')
@@ -33,10 +48,10 @@ def main():
 
     # use custormized CNN and loss function
     net      = get_network(config['network'])
-    loss_cal = MySegmentationLossCalculator(config['training'])
-    agent  = TrainInferAgent(config, stage)
+    loss_obj = get_loss(config['training'])
+    agent  = NetRunAgent(config, stage)
     agent.set_network(net)
-    agent.set_loss_calculater(loss_cal)
+    agent.set_loss_calculater(loss_obj)
     agent.run()
 
 if __name__ == "__main__":

examples/JSRT2/my_loss.py

@@ -0,0 +1,50 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+
+import torch
+import torch.nn as nn
+import numpy as np 
+from pymic.loss.util import reshape_tensor_to_2D, get_classwise_dice
+
+class MyFocalDiceLoss(nn.Module):
+    """
+    Focal Dice loss proposed in the following paper:
+       P. Wang et al. Focal dice loss and image dilatin for brain tumor segmentation.
+       in Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical
+       Decision Support, 2018.
+    """
+    def __init__(self, params):
+        super(MyFocalDiceLoss, self).__init__()
+        self.enable_pix_weight  = params['MyFocalDiceLoss_Enable_Pixel_Weight'.lower()]
+        self.enable_cls_weight  = params['MyFocalDiceLoss_Enable_Class_Weight'.lower()]
+        self.beta = params['MyFocalDiceLoss_beta'.lower()]
+        assert(self.beta >= 1.0)
+
+    def forward(self, loss_input_dict):
+        predict = loss_input_dict['prediction']
+        soft_y  = loss_input_dict['ground_truth']
+        pix_w   = loss_input_dict['pixel_weight']
+        cls_w   = loss_input_dict['class_weight']
+        softmax = loss_input_dict['softmax']
+
+        if(softmax):
+            predict = nn.Softmax(dim = 1)(predict)
+        predict = reshape_tensor_to_2D(predict)
+        soft_y  = reshape_tensor_to_2D(soft_y) 
+
+        if(self.enable_pix_weight):
+            if(pix_w is None):
+                raise ValueError("Pixel weight is enabled but not defined")
+            pix_w = reshape_tensor_to_2D(pix_w)
+        dice_score = get_classwise_dice(predict, soft_y, pix_w)
+        dice_score = 0.01 + dice_score * 0.98
+        dice_loss  = 1.0 - torch.pow(dice_score, 1.0 / self.beta)
+
+        if(self.enable_cls_weight):
+            if(cls_w is None):
+                raise ValueError("Class weight is enabled but not defined")
+            weighted_loss = dice_loss * cls_w
+            avg_loss =  weighted_loss.sum() / cls_w.sum()
+        else:
+            avg_loss = torch.mean(dice_loss)   
+        return avg_loss

examples/JSRT2/my_net2d.py

@@ -107,6 +107,7 @@ def forward(self, x):
             new_shape = [N, D] + list(output.shape)[1:]
             output = torch.reshape(output, new_shape)
             output = torch.transpose(output, 1, 2)
+
         return output
 
 if __name__ == "__main__":
@@ -122,7 +123,7 @@ def forward(self, x):
     xt = torch.from_numpy(x)
     xt = torch.tensor(xt)
 
-    y = Net(xt)
+    y, y2 = Net(xt)
     print(len(y.size()))
     y = y.detach().numpy()
     print(y.shape)

pymic/loss/exp_log.py

@@ -0,0 +1,47 @@
+
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+
+import torch
+import torch.nn as nn
+from pymic.loss.util import reshape_tensor_to_2D, get_classwise_dice
+
+class ExpLogLoss(nn.Module):
+    """
+    The exponential logarithmic loss in this paper: 
+        K. Wong et al.: 3D Segmentation with Exponential Logarithmic Loss for Highly 
+        Unbalanced Object Sizes. MICCAI 2018.
+    """
+    def __init__(self, params):
+        super(ExpLogLoss, self).__init__()
+        self.w_dice = params['ExpLogLoss_w_dice'.lower()]
+        self.gamma  = params['ExpLogLoss_gamma'.lower()]
+
+    def forward(self, loss_input_dict):
+        predict = loss_input_dict['prediction']
+        soft_y  = loss_input_dict['ground_truth']
+        softmax = loss_input_dict['softmax']
+
+        if(softmax):
+            predict = nn.Softmax(dim = 1)(predict)
+        predict = reshape_tensor_to_2D(predict)
+        soft_y  = reshape_tensor_to_2D(soft_y)
+
+
+        dice_score = get_classwise_dice(predict, soft_y)
+        dice_score = 0.01 + dice_score * 0.98
+        exp_dice   = -torch.log(dice_score)
+        exp_dice   = torch.pow(exp_dice, self.gamma)
+        exp_dice   = torch.mean(exp_dice)
+
+        predict= 0.01 + predict * 0.98
+        wc     = torch.mean(soft_y, dim = 0)
+        wc     = 1.0 / (wc + 0.1)
+        wc     = torch.pow(wc, 0.5)
+        ce     = - torch.log(predict)
+        exp_ce = wc * torch.pow(ce, self.gamma)
+        exp_ce = torch.sum(soft_y * exp_ce, dim = 1)
+        exp_ce = torch.mean(exp_ce)
+
+        loss = exp_dice * self.w_dice + exp_ce * (1.0 - self.w_dice)
+        return loss

pymic/net/__init__.py

@@ -10,7 +10,7 @@
 import torch
 import torch.nn as nn
 import numpy as np 
-from pymic.net2d.squeeze_and_excitation import *
+from pymic.net.net2d.squeeze_and_excitation import *
 
 class ConvScSEBlock(nn.Module):
     """two convolution layers with batch norm and leaky relu"""