Merge pull request #33 from HiLab-git/dev

Dev

Author: taigw

README.md

@@ -4,8 +4,8 @@ PyMIC is a pytorch-based toolkit for medical image computing with annotation-eff
 
 Currently PyMIC supports 2D/3D medical image classification and segmentation, and it is still under development. If you use this toolkit, please cite the following paper:
 
-* G. Wang, X. Luo, R. Gu, S. Yang, Y. Qu, S. Zhai, Q. Zhao, K. Li, S. Zhang. (2022). 
-[PyMIC: A deep learning toolkit for annotation-efficient medical image segmentation.][arxiv2022] arXiv, 2208.09350.
+* G. Wang, X. Luo, R. Gu, S. Yang, Y. Qu, S. Zhai, Q. Zhao, K. Li, S. Zhang. (2023). 
+[PyMIC: A deep learning toolkit for annotation-efficient medical image segmentation.][arxiv2022] Computer Methods and Programs in Biomedicine (CMPB). February 2023, 107398.
 
 [arxiv2022]:http://arxiv.org/abs/2208.09350
 
@@ -14,11 +14,11 @@ BibTeX entry:
     @article{Wang2022pymic,
     author = {Guotai Wang and Xiangde Luo and Ran Gu and Shuojue Yang and Yijie Qu and Shuwei Zhai and Qianfei Zhao and Kang Li and Shaoting Zhang},
     title = {{PyMIC: A deep learning toolkit for annotation-efficient medical image segmentation}},
-    year = {2022},
+    year = {2023},
     url = {http://arxiv.org/abs/2208.09350},
-    journal = {arXiv},
-    volume = {2208.09350},
-    pages = {1-10},
+    journal = {Computer Methods and Programs in Biomedicine},
+    volume = {February},
+    pages = {107398},
     }
 
 # Features

pymic/loss/seg/deep_sup.py

@@ -2,8 +2,43 @@
 from __future__ import print_function, division
 
 import torch.nn as nn
+from torch.nn.functional import interpolate
 from pymic.loss.seg.abstract import AbstractSegLoss
 
+def match_prediction_and_gt_shape(pred, gt, mode = 0):
+    pred_shape = list(pred.shape)
+    gt_shape   = list(gt.shape)
+    dim = len(pred_shape) - 2
+    shape_match = False 
+    if(dim == 2):
+        if(pred_shape[-1] == gt_shape[-1] and pred_shape[-2] == gt_shape[-2]):
+            shape_match = True
+    else:
+        if(pred_shape[-1] == gt_shape[-1] and pred_shape[-2] == gt_shape[-2]
+           and pred_shape[-3] == gt_shape[-3]):
+            shape_match = True
+    if(shape_match):
+        return pred, gt 
+    
+    interp_mode = 'bilinear' if dim == 2 else 'trilinear'
+    if(mode == 0):
+        pred_new = interpolate(pred, gt_shape[2:], mode = interp_mode)
+        gt_new   = gt  
+    elif(mode == 1):
+        pred_new = pred
+        gt_new   = interpolate(gt, pred_shape[2:], mode = interp_mode)
+    elif(mode == 2):
+        pred_new = pred
+        if(dim == 2):
+            avg_pool = nn.AdaptiveAvgPool2d(pred_shape[-2:])
+        else:
+            avg_pool = nn.AdaptiveAvgPool3d(pred_shape[-3:])
+        gt_new = avg_pool(gt)
+    else:
+        raise ValueError("mode shoud be 0, 1 or 2, but {0:} was given".format(mode))
+    return pred_new, gt_new
+            
+
 class DeepSuperviseLoss(AbstractSegLoss):
     '''
     Combine deep supervision with a basic loss function.  
@@ -12,28 +47,36 @@ class DeepSuperviseLoss(AbstractSegLoss):
 
     :param `loss_softmax`: (optional, bool) 
         Apply softmax to the prediction of network or not. Default is True.
-    :param `deep_suervise_weight`: (list) A list of weight for each deep supervision scale. \n
     :param `base_loss`: (nn.Module) The basic function used for each scale.
+    :param `deep_supervise_weight`: (list) A list of weight for each deep supervision scale. 
+    :param `deep_supervise_model`: (int) Mode for deep supervision when the prediction
+        has a smaller shape than the ground truth. 0: upsample the prediction to the size 
+        of the ground truth. 1: downsample the ground truth to the size of the prediction
+        via interpolation. 2: downsample the ground truth via adaptive average pooling.
 
     '''
     def __init__(self, params):
         super(DeepSuperviseLoss, self).__init__(params)
-        self.deep_sup_weight = params.get('deep_suervise_weight', None)
-        self.base_loss = params['base_loss']
+        self.base_loss       = params['base_loss']
+        self.deep_sup_weight = params.get('deep_supervise_weight', None)
+        self.deep_sup_mode   = params.get('deep_supervise_mode', 0)
 
     def forward(self, loss_input_dict):
-        predict = loss_input_dict['prediction']
-        if(not isinstance(predict, (list,tuple))):
+        pred = loss_input_dict['prediction']
+        gt   = loss_input_dict['ground_truth']
+        if(not isinstance(pred, (list,tuple))):
             raise ValueError("""For deep supervision, the prediction should
                 be a list or a tuple""")
-        predict_num = len(predict)
+        pred_num = len(pred)
         if(self.deep_sup_weight is None):
-            self.deep_sup_weight = [1.0] * predict_num
+            self.deep_sup_weight = [1.0] * pred_num
         else:
-            assert(predict_num == len(self.deep_sup_weight))
+            assert(pred_num == len(self.deep_sup_weight))
         loss_sum, weight_sum  = 0.0, 0.0
-        for i in range(predict_num):
-            loss_input_dict['prediction'] =  predict[i]
+        for i in range(pred_num):
+            pred_i, gt_i = match_prediction_and_gt_shape(pred[i], gt, self.deep_sup_mode)
+            loss_input_dict['prediction']   = pred_i
+            loss_input_dict['ground_truth'] = gt_i
             temp_loss   = self.base_loss(loss_input_dict)
             loss_sum   += temp_loss * self.deep_sup_weight[i]
             weight_sum += self.deep_sup_weight[i]

pymic/net/net2d/unet2d.py

@@ -240,18 +240,14 @@ def forward(self, x):
         x_d0 = self.up4(x_d1, x0)
         output = self.out_conv(x_d0)
         if(self.deep_sup):
-            out_shape = list(output.shape)[2:]
             output1 = self.out_conv1(x_d1)
-            output1 = interpolate(output1, out_shape, mode = 'bilinear')
             output2 = self.out_conv2(x_d2)
-            output2 = interpolate(output2, out_shape, mode = 'bilinear')
             output3 = self.out_conv3(x_d3)
-            output3 = interpolate(output3, out_shape, mode = 'bilinear')
             output = [output, output1, output2, output3]
 
             if(len(x_shape) == 5):
-                new_shape = [N, D] + list(output[0].shape)[1:]
                 for i in range(len(output)):
+                    new_shape = [N, D] + list(output[i].shape)[1:]
                     output[i] = torch.transpose(torch.reshape(output[i], new_shape), 1, 2)
         elif(len(x_shape) == 5):
             new_shape = [N, D] + list(output.shape)[1:]

pymic/net/net3d/unet3d.py

@@ -77,6 +77,93 @@ def forward(self, x1, x2):
         x = torch.cat([x2, x1], dim=1)
         return self.conv(x)
 
+class Encoder(nn.Module):
+    """
+    Encoder of 3D UNet.
+
+    Parameters are given in the `params` dictionary, and should include the
+    following fields:
+
+    :param in_chns: (int) Input channel number.
+    :param feature_chns: (list) Feature channel for each resolution level. 
+      The length should be 4 or 5, such as [16, 32, 64, 128, 256].
+    :param dropout: (list) The dropout ratio for each resolution level. 
+      The length should be the same as that of `feature_chns`.
+    """
+    def __init__(self, params):
+        super(Encoder, self).__init__()
+        self.params    = params
+        self.in_chns   = self.params['in_chns']
+        self.ft_chns   = self.params['feature_chns']
+        self.dropout   = self.params['dropout']
+        assert(len(self.ft_chns) == 5 or len(self.ft_chns) == 4)
+
+        self.in_conv= ConvBlock(self.in_chns, self.ft_chns[0], self.dropout[0])
+        self.down1  = DownBlock(self.ft_chns[0], self.ft_chns[1], self.dropout[1])
+        self.down2  = DownBlock(self.ft_chns[1], self.ft_chns[2], self.dropout[2])
+        self.down3  = DownBlock(self.ft_chns[2], self.ft_chns[3], self.dropout[3])
+        if(len(self.ft_chns) == 5):
+            self.down4  = DownBlock(self.ft_chns[3], self.ft_chns[4], self.dropout[4])
+
+    def forward(self, x):
+        x0 = self.in_conv(x)
+        x1 = self.down1(x0)
+        x2 = self.down2(x1)
+        x3 = self.down3(x2)
+        output = [x0, x1, x2, x3]
+        if(len(self.ft_chns) == 5):
+          x4 = self.down4(x3)
+          output.append(x4)
+        return output
+
+class Decoder(nn.Module):
+    """
+    Decoder of 3D UNet.
+
+    Parameters are given in the `params` dictionary, and should include the
+    following fields:
+
+    :param in_chns: (int) Input channel number.
+    :param feature_chns: (list) Feature channel for each resolution level. 
+      The length should be 4 or 5, such as [16, 32, 64, 128, 256].
+    :param dropout: (list) The dropout ratio for each resolution level. 
+      The length should be the same as that of `feature_chns`.
+    :param class_num: (int) The class number for segmentation task. 
+    :param trilinear: (bool) Using bilinear for up-sampling or not. 
+        If False, deconvolution will be used for up-sampling.
+    """
+    def __init__(self, params):
+        super(Decoder, self).__init__()
+        self.params    = params
+        self.in_chns   = self.params['in_chns']
+        self.ft_chns   = self.params['feature_chns']
+        self.dropout   = self.params['dropout']
+        self.n_class   = self.params['class_num']
+        self.trilinear = self.params['trilinear']
+
+        assert(len(self.ft_chns) == 5 or len(self.ft_chns) == 4)
+
+        if(len(self.ft_chns) == 5):
+            self.up1 = UpBlock(self.ft_chns[4], self.ft_chns[3], self.ft_chns[3], self.dropout[3], self.bilinear) 
+        self.up2 = UpBlock(self.ft_chns[3], self.ft_chns[2], self.ft_chns[2], self.dropout[2], self.bilinear) 
+        self.up3 = UpBlock(self.ft_chns[2], self.ft_chns[1], self.ft_chns[1], self.dropout[1], self.bilinear) 
+        self.up4 = UpBlock(self.ft_chns[1], self.ft_chns[0], self.ft_chns[0], self.dropout[0], self.bilinear) 
+        self.out_conv = nn.Conv3d(self.ft_chns[0], self.n_class, kernel_size = 1)
+
+    def forward(self, x):
+        if(len(self.ft_chns) == 5):
+            assert(len(x) == 5)
+            x0, x1, x2, x3, x4 = x 
+            x_d3 = self.up1(x4, x3)
+        else:
+            assert(len(x) == 4)
+            x0, x1, x2, x3 = x 
+            x_d3 = x3
+        x_d2 = self.up2(x_d3, x2)
+        x_d1 = self.up3(x_d2, x1)
+        x_d0 = self.up4(x_d1, x0)
+        output = self.out_conv(x_d0)
+        return output
 
 class UNet3D(nn.Module):
     """

pymic/net/net3d/unet3d_dual_branch.py

@@ -0,0 +1,46 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+
+import torch
+import torch.nn as nn
+from pymic.net.net3d.unet3d import *
+
+class UNet3D_DualBranch(nn.Module):
+    """
+    A dual branch network using UNet3D as backbone.
+
+    * Reference: Xiangde Luo, Minhao Hu, Wenjun Liao, Shuwei Zhai, Tao Song, Guotai Wang,
+      Shaoting Zhang. ScribblScribble-Supervised Medical Image Segmentation via 
+      Dual-Branch Network and Dynamically Mixed Pseudo Labels Supervision.
+      `MICCAI 2022. <https://arxiv.org/abs/2203.02106>`_ 
+
+    The parameters for the backbone should be given in the `params` dictionary. 
+    See :mod:`pymic.net.net3d.unet3d.UNet3D` for details. 
+    In addition, the following field should be included:
+
+    :param output_mode: (str) How to obtain the result during the inference. 
+      `average`: taking average of the two branches. 
+      `first`: takeing the result in the first branch. 
+      `second`: taking the result in the second branch.
+    """
+    def __init__(self, params):
+        super(UNet3D_DualBranch, self).__init__()
+        self.output_mode = params.get("output_mode", "average")
+        self.encoder  = Encoder(params)
+        self.decoder1 = Decoder(params)    
+        self.decoder2 = Decoder(params)        
+
+    def forward(self, x):
+        f = self.encoder(x)
+        output1 = self.decoder1(f)
+        output2 = self.decoder2(f)
+
+        if(self.training):
+          return output1, output2
+        else:
+          if(self.output_mode == "average"):
+            return (output1 + output2)/2
+          elif(self.output_mode == "first"):
+            return output1
+          else:
+            return output2

pymic/net/net_dict_seg.py

@@ -24,6 +24,7 @@
 from pymic.net.net3d.unet2d5 import UNet2D5
 from pymic.net.net3d.unet3d import UNet3D
 from pymic.net.net3d.unet3d_scse import UNet3D_ScSE
+from pymic.net.net3d.unet3d_dual_branch import UNet3D_DualBranch
 
 SegNetDict = {
 	'UNet2D': UNet2D,
@@ -35,5 +36,7 @@
 	'UNet2D_ScSE': UNet2D_ScSE,
 	'UNet2D5': UNet2D5,
 	'UNet3D': UNet3D,
-	'UNet3D_ScSE': UNet3D_ScSE
+	'UNet3D_ScSE': UNet3D_ScSE,
+	'UNet3D_DualBranch': UNet3D_DualBranch
+
 	}

pymic/net_run/agent_abstract.py

@@ -259,7 +259,7 @@ def worker_init_fn(worker_id):
 
             bn_train = self.config['dataset']['train_batch_size']
             bn_valid = self.config['dataset'].get('valid_batch_size', 1)
-            num_worker = self.config['dataset'].get('num_workder', 16)
+            num_worker = self.config['dataset'].get('num_worker', 16)
             g_train, g_valid = torch.Generator(), torch.Generator()
             g_train.manual_seed(self.random_seed)
             g_valid.manual_seed(self.random_seed)

pymic/net_run/agent_seg.py

@@ -162,10 +162,6 @@ def training(self):
             loss = self.get_loss_value(data, outputs, labels_prob)
             loss.backward()
             self.optimizer.step()
-            if(self.scheduler is not None and \
-                not isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
-                self.scheduler.step()
-
             train_loss = train_loss + loss.item()
             # get dice evaluation for each class
             if(isinstance(outputs, tuple) or isinstance(outputs, list)):
@@ -177,9 +173,9 @@ def training(self):
             train_dice_list.append(dice_list.cpu().numpy())
         train_avg_loss = train_loss / iter_valid
         train_cls_dice = np.asarray(train_dice_list).mean(axis = 0)
-        train_avg_dice = train_cls_dice.mean()
+        train_avg_dice = train_cls_dice[1:].mean()
 
-        train_scalers = {'loss': train_avg_loss, 'avg_dice':train_avg_dice,\
+        train_scalers = {'loss': train_avg_loss, 'avg_fg_dice':train_avg_dice,\
             'class_dice': train_cls_dice}
         return train_scalers
 
@@ -218,18 +214,14 @@ def validation(self):
 
         valid_avg_loss = np.asarray(valid_loss_list).mean()
         valid_cls_dice = np.asarray(valid_dice_list).mean(axis = 0)
-        valid_avg_dice = valid_cls_dice.mean()
-        
-        if(isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
-            self.scheduler.step(valid_avg_dice)
-
-        valid_scalers = {'loss': valid_avg_loss, 'avg_dice': valid_avg_dice,\
+        valid_avg_dice = valid_cls_dice[1:].mean()
+        valid_scalers = {'loss': valid_avg_loss, 'avg_fg_dice': valid_avg_dice,\
             'class_dice': valid_cls_dice}
         return valid_scalers
 
     def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
         loss_scalar ={'train':train_scalars['loss'], 'valid':valid_scalars['loss']}
-        dice_scalar ={'train':train_scalars['avg_dice'], 'valid':valid_scalars['avg_dice']}
+        dice_scalar ={'train':train_scalars['avg_fg_dice'], 'valid':valid_scalars['avg_fg_dice']}
         self.summ_writer.add_scalars('loss', loss_scalar, glob_it)
         self.summ_writer.add_scalars('dice', dice_scalar, glob_it)
         self.summ_writer.add_scalars('lr', {"lr": lr_value}, glob_it)
@@ -239,11 +231,11 @@ def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
                 'valid':valid_scalars['class_dice'][c]}
             self.summ_writer.add_scalars('class_{0:}_dice'.format(c), cls_dice_scalar, glob_it)
 
-        logging.info('train loss {0:.4f}, avg dice {1:.4f} '.format(
-            train_scalars['loss'], train_scalars['avg_dice']) + "[" + \
+        logging.info('train loss {0:.4f}, avg foreground dice {1:.4f} '.format(
+            train_scalars['loss'], train_scalars['avg_fg_dice']) + "[" + \
             ' '.join("{0:.4f}".format(x) for x in train_scalars['class_dice']) + "]")        
-        logging.info('valid loss {0:.4f}, avg dice {1:.4f} '.format(
-            valid_scalars['loss'], valid_scalars['avg_dice']) + "[" + \
+        logging.info('valid loss {0:.4f}, avg foreground dice {1:.4f} '.format(
+            valid_scalars['loss'], valid_scalars['avg_fg_dice']) + "[" + \
             ' '.join("{0:.4f}".format(x) for x in valid_scalars['class_dice']) + "]")        
 
     def train_valid(self):
@@ -300,16 +292,20 @@ def train_valid(self):
             t0 = time.time()
             train_scalars = self.training()
             t1 = time.time()
-            
             valid_scalars = self.validation()
             t2 = time.time()
+            if(isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
+                self.scheduler.step(valid_scalars['avg_fg_dice'])
+            else:
+                self.scheduler.step()
+
             self.glob_it = it + iter_valid
             logging.info("\n{0:} it {1:}".format(str(datetime.now())[:-7], self.glob_it))
             logging.info('learning rate {0:}'.format(lr_value))
             logging.info("training/validation time: {0:.2f}s/{1:.2f}s".format(t1-t0, t2-t1))
             self.write_scalars(train_scalars, valid_scalars, lr_value, self.glob_it)
-            if(valid_scalars['avg_dice'] > self.max_val_dice):
-                self.max_val_dice = valid_scalars['avg_dice']
+            if(valid_scalars['avg_fg_dice'] > self.max_val_dice):
+                self.max_val_dice = valid_scalars['avg_fg_dice']
                 self.max_val_it   = self.glob_it
                 if(len(device_ids) > 1):
                     self.best_model_wts = copy.deepcopy(self.net.module.state_dict())
@@ -320,7 +316,7 @@ def train_valid(self):
                 self.glob_it - self.max_val_it > early_stop_it) else False
             if ((self.glob_it in iter_save_list) or stop_now):
                 save_dict = {'iteration': self.glob_it,
-                             'valid_pred': valid_scalars['avg_dice'],
+                             'valid_pred': valid_scalars['avg_fg_dice'],
                              'model_state_dict': self.net.module.state_dict() \
                                  if len(device_ids) > 1 else self.net.state_dict(),
                              'optimizer_state_dict': self.optimizer.state_dict()}