Merge pull request #23 from HiLab-git/dev

Dev

Author: taigw

README.md

@@ -35,7 +35,11 @@ Run the following command to install the current released version of PyMIC:
 pip install PYMIC
 ```
 
-Alternatively, you can download the source code and add the path of pymic to the `PYTHONPATH` environment variable. 
+Alternatively, you can download the source code for the latest version. Run the following command to compile and install:
+
+```bash
+python setup.py install
+``` 
 
 ## Examples
 [PyMIC_examples][examples] provides some examples of starting to use PyMIC. For beginners, you only need to simply change the configuration files to select different datasets, networks and training methods for running the code. For advanced users, you can develop your own modules based on this package. You can find both types of examples 

pymic/io/image_read_write.py

@@ -79,7 +79,10 @@ def save_array_as_nifty_volume(data, image_name, reference_name = None):
     img = sitk.GetImageFromArray(data)
     if(reference_name is not None):
         img_ref = sitk.ReadImage(reference_name)
-        img.CopyInformation(img_ref)
+        #img.CopyInformation(img_ref)
+        img.SetSpacing(img_ref.GetSpacing())
+        img.SetOrigin(img_ref.GetOrigin())
+        img.SetDirection(img_ref.GetDirection())
     sitk.WriteImage(img, image_name)
 
 def save_array_as_rgb_image(data, image_name):

pymic/loss/cls/ce.py

@@ -37,7 +37,8 @@ def __init__(self, params):
     def forward(self, loss_input_dict):
         predict = loss_input_dict['prediction']
         labels  = loss_input_dict['ground_truth']
-        predict = nn.Sigmoid()(predict)
+        # for numeric stability
+        predict = nn.Sigmoid()(predict) * 0.999 + 5e-4
         loss = - labels * torch.log(predict) - (1 - labels) * torch.log( 1 - predict)
         loss = loss.mean()
         return loss

pymic/loss/loss_dict_seg.py

@@ -3,15 +3,15 @@
 import torch.nn as nn 
 from pymic.loss.seg.ce import CrossEntropyLoss, GeneralizedCrossEntropyLoss
 from pymic.loss.seg.dice import DiceLoss, MultiScaleDiceLoss
-from pymic.loss.seg.dice import DiceWithCrossEntropyLoss, NoiseRobustDiceLoss
+from pymic.loss.seg.dice import FocalDiceLoss, NoiseRobustDiceLoss
 from pymic.loss.seg.exp_log import ExpLogLoss
 from pymic.loss.seg.mse import MSELoss, MAELoss
 
 SegLossDict = {'CrossEntropyLoss': CrossEntropyLoss,
     'GeneralizedCrossEntropyLoss': GeneralizedCrossEntropyLoss,
     'DiceLoss': DiceLoss,
     'MultiScaleDiceLoss': MultiScaleDiceLoss,
-    'DiceWithCrossEntropyLoss': DiceWithCrossEntropyLoss,
+    'FocalDiceLoss': FocalDiceLoss,
     'NoiseRobustDiceLoss': NoiseRobustDiceLoss,
     'ExpLogLoss': ExpLogLoss,
     'MSELoss': MSELoss,

pymic/loss/seg/ce.py

@@ -8,8 +8,8 @@
 class CrossEntropyLoss(nn.Module):
     def __init__(self, params):
         super(CrossEntropyLoss, self).__init__()
-        self.enable_pix_weight = params['CrossEntropyLoss_Enable_Pixel_Weight'.lower()]
-        self.enable_cls_weight = params['CrossEntropyLoss_Enable_Class_Weight'.lower()]
+        self.enable_pix_weight = params.get('CrossEntropyLoss_Enable_Pixel_Weight'.lower(), False)
+        self.enable_cls_weight = params.get('CrossEntropyLoss_Enable_Class_Weight'.lower(), False)
 
     def forward(self, loss_input_dict):
         predict = loss_input_dict['prediction']
@@ -25,6 +25,8 @@ def forward(self, loss_input_dict):
         predict = reshape_tensor_to_2D(predict)
         soft_y  = reshape_tensor_to_2D(soft_y)
 
+        # for numeric stability
+        predict = predict * 0.999 + 5e-4
         ce = - soft_y* torch.log(predict)
         if(self.enable_cls_weight):
             if(cls_w is None):

pymic/loss/seg/combined.py

@@ -22,5 +22,5 @@ def __init__(self, params, loss_dict):
     def forward(self, loss_input_dict):
         loss_value = 0.0
         for i in range(len(self.loss_list)):
-            loss_value = self.loss_weight[i] + self.loss_list[i](loss_input_dict)
-        return loss_value 
+            loss_value += self.loss_weight[i]*self.loss_list[i](loss_input_dict)
+        return loss_value

pymic/loss/seg/dice.py

@@ -40,24 +40,46 @@ def forward(self, loss_input_dict):
         dice_loss  = 1.0 - average_dice
         return dice_loss
 
-class DiceWithCrossEntropyLoss(nn.Module):
-    def __init__(self, params):
-        super(DiceWithCrossEntropyLoss, self).__init__()
-        self.enable_pix_weight = params['DiceWithCrossEntropyLoss_Enable_Pixel_Weight'.lower()]
-        self.enable_cls_weight = params['DiceWithCrossEntropyLoss_Enable_Class_Weight'.lower()]
-        self.ce_weight = params['DiceWithCrossEntropyLoss_CE_Weight'.lower()]
-        dice_params = {'DiceLoss_Enable_Pixel_Weight'.lower(): self.enable_pix_weight,
-                       'DiceLoss_Enable_Class_Weight'.lower(): self.enable_cls_weight}
-        ce_params   = {'CrossEntropyLoss_Enable_Pixel_Weight'.lower(): self.enable_pix_weight,
-                       'CrossEntropyLoss_Enable_Class_Weight'.lower(): self.enable_cls_weight}
-        self.dice_loss = DiceLoss(dice_params)
-        self.ce_loss   = CrossEntropyLoss(ce_params)
+class FocalDiceLoss(nn.Module):
+    """
+    focal Dice according to the following paper:
+    Pei Wang and Albert C. S. Chung, Focal Dice Loss and Image Dilation for 
+    Brain Tumor Segmentation, 2018
+    """
+    def __init__(self, params = None):
+        super(FocalDiceLoss, self).__init__()
+        self.beta = params['FocalDiceLoss_beta'.lower()] #beta should be >=1.0
 
     def forward(self, loss_input_dict):
-        loss1 = self.dice_loss(loss_input_dict)
-        loss2 = self.ce_loss(loss_input_dict)
-        loss = loss1 + self.ce_weight * loss2
-        return loss 
+        predict = loss_input_dict['prediction']
+        soft_y  = loss_input_dict['ground_truth']
+        img_w   = loss_input_dict['image_weight']
+        pix_w   = loss_input_dict['pixel_weight']
+        cls_w   = loss_input_dict['class_weight']
+        softmax = loss_input_dict['softmax']
+
+        if(isinstance(predict, (list, tuple))):
+            predict = predict[0]
+        tensor_dim = len(predict.size())
+        if(softmax):
+            predict = nn.Softmax(dim = 1)(predict)
+        predict = reshape_tensor_to_2D(predict)
+        soft_y  = reshape_tensor_to_2D(soft_y) 
+
+        # combien pixel weight and image weight
+        if(tensor_dim == 5):
+            img_w = img_w[:, None, None, None, None]
+        else:
+            img_w = img_w[:, None, None, None]
+        pix_w = pix_w * img_w
+        pix_w = reshape_tensor_to_2D(pix_w)
+        dice_score = get_classwise_dice(predict, soft_y, pix_w)
+
+        dice_score = torch.pow(dice_score, 1.0 / self.beta)
+        weighted_dice = dice_score * cls_w
+        average_dice  =  weighted_dice.sum() / cls_w.sum()
+        dice_loss  = 1.0 - average_dice
+        return dice_loss
 
 class MultiScaleDiceLoss(nn.Module):
     def __init__(self, params):

pymic/net/cls/torch_pretrained_net.py

@@ -73,3 +73,30 @@ def get_parameters_to_update(self):
         else:
             raise(ValueError("update_layers can only be 0 (all layers) " +
                 "or -1 (the last layer)"))
+
+class MobileNetV2(nn.Module):
+    def __init__(self, params):
+        super(MobileNetV2, self).__init__()
+        self.params = params
+        net_name = params['net_type']
+        cls_num  = params['class_num']
+        in_chns  = params['input_chns']
+        self.pretrain = params['pretrain']
+        self.update_layers = params.get('update_layers', 0)
+        self.net = models.mobilenet_v2(pretrained = self.pretrain)
+        
+        # replace the last layer 
+        num_ftrs = self.net.last_channel
+        self.net.classifier[-1] = nn.Linear(num_ftrs, cls_num)
+    
+    def forward(self, x):
+        return self.net(x)
+    
+    def get_parameters_to_update(self):
+        if(self.pretrain == False or self.update_layers == 0):
+            return self.net.parameters()
+        elif(self.update_layers == -1):
+            return self.net.classifier[-1].parameters()
+        else:
+            raise(ValueError("update_layers can only be 0 (all layers) " +
+                "or -1 (the last layer)"))

pymic/net/net_dict_cls.py

@@ -4,5 +4,6 @@
 
 TorchClsNetDict = {
     'resnet18':   ResNet18,
-    'vgg16':      VGG16
+    'vgg16':      VGG16,
+    'mobilenetv2':MobileNetV2
 }

pymic/net_run/agent_cls.py

@@ -278,7 +278,7 @@ def infer(self):
         checkpoint = torch.load(checkpoint_name, map_location = device)
         self.net.load_state_dict(checkpoint['model_state_dict'])
 
-        if(self.config['testing']['evaluation_mode'] == True):
+        if(self.config['testing'].get('evaluation_mode', True)):
             self.net.eval()
 
         output_csv   = self.config['testing']['output_csv']