update jsrt2

Author: taigw

segmentation/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 define a focal dice loss. 
+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 named as MyFocalDiceLoss. We use it `MyFocalDiceLoss + CrossEntropyLoss` to train the custermized network.
 
-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.
+We also write a customized main function in `jsrt_net_run.py` so that we can combine SegmentationAgent 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.
@@ -37,4 +37,4 @@ python net_run_jsrt.py test config/train_test.cfg
 pymic_evaluate_seg config/evaluation.cfg
 ```
 
-The obtained dice score by default setting should be close to 93.90%. 
+The obtained dice score by default setting should be close to 94.35%. 

segmentation/JSRT2/config/train_test.cfg

@@ -56,12 +56,16 @@ bilinear     = True
 # list of gpus
 gpus          = [0]
 
-loss_type     = MyFocalDiceLoss
+loss_type     = [MyFocalDiceLoss, CrossEntropyLoss]
+loss_weight   = [1.0, 1.0]
 MyFocalDiceLoss_Enable_Pixel_Weight = False
 MyFocalDiceLoss_Enable_Class_Weight = True
 MyFocalDiceLoss_beta = 1.5
 class_weight = [0.2, 1.0]
 
+CrossEntropyLoss_Enable_Pixel_Weight = False
+CrossEntropyLoss_Enable_Class_Weight = False
+
 # for optimizers
 optimizer     = Adam
 learning_rate = 1e-3

segmentation/JSRT2/net_run_jsrt.py

@@ -4,24 +4,27 @@
 import sys
 from pymic.util.parse_config import parse_config
 from pymic.net_run.agent_seg import  SegmentationAgent
+from pymic.loss.loss_dict_seg import SegLossDict
 from my_net2d import MyUNet2D 
 from my_loss  import MyFocalDiceLoss
 
+loss_dict = {'MyFocalDiceLoss':MyFocalDiceLoss}
+loss_dict.update(SegLossDict)
+
 def main():
     if(len(sys.argv) < 3):
         print('Number of arguments should be 3. e.g.')
-        print('    python train_infer.py train config.cfg')
+        print('    python net_run_jsrt.py train config.cfg')
         exit()
     stage    = str(sys.argv[1])
     cfg_file = str(sys.argv[2])
     config   = parse_config(cfg_file)
 
+    agent  = SegmentationAgent(config, stage)
     # use custormized CNN and loss function
     mynet  = MyUNet2D(config['network'])
-    myloss = MyFocalDiceLoss(config['training'])
-    agent  = SegmentationAgent(config, stage)
     agent.set_network(mynet)
-    agent.set_loss(myloss)
+    agent.set_loss_dict(loss_dict)
     agent.run()
 
 if __name__ == "__main__":