update transform and lr scheduler

add PolynomialLR and Transpose

Author: taigw

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_run/agent_seg.py

@@ -173,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
 
@@ -214,14 +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()
-        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)
@@ -231,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):
@@ -295,7 +295,7 @@ def train_valid(self):
             valid_scalars = self.validation()
             t2 = time.time()
             if(isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
-                self.scheduler.step(valid_scalars['avg_dice'])
+                self.scheduler.step(valid_scalars['avg_fg_dice'])
             else:
                 self.scheduler.step()
 
@@ -304,8 +304,8 @@ def train_valid(self):
             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())
@@ -316,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()}

pymic/net_run/get_optimizer.py

@@ -39,8 +39,9 @@ def get_optimizer(name, net_params, optim_params):
 
 
 def get_lr_scheduler(optimizer, sched_params):
-    name = sched_params["lr_scheduler"]
-    val_it = sched_params["iter_valid"]
+    name       = sched_params["lr_scheduler"]
+    val_it     = sched_params["iter_valid"]
+    epoch_last = sched_params["last_iter"] / val_it
     if(name is None):
         return None
     if(keyword_match(name, "ReduceLROnPlateau")):
@@ -52,16 +53,19 @@ def get_lr_scheduler(optimizer, sched_params):
     elif(keyword_match(name, "MultiStepLR")):
         lr_milestones = sched_params["lr_milestones"]
         lr_milestones = [int(item / val_it) for item in lr_milestones]
-        epoch_last    = sched_params["last_iter"] / val_it
         lr_gamma  = sched_params["lr_gamma"]
         scheduler = lr_scheduler.MultiStepLR(optimizer,
                     lr_milestones, lr_gamma, epoch_last)
     elif(keyword_match(name, "CosineAnnealingLR")):
         epoch_max  = sched_params["iter_max"] / val_it
-        epoch_last = sched_params["last_iter"] / val_it
         lr_min     = sched_params.get("lr_min", 0)
         scheduler = lr_scheduler.CosineAnnealingLR(optimizer,
                     epoch_max, lr_min, epoch_last)
+    elif(keyword_match(name, "PolynomialLR")):
+        epoch_max  = sched_params["iter_max"] / val_it
+        power      = sched_params["lr_power"]
+        scheduler = lr_scheduler.PolynomialLR(optimizer,
+                    epoch_max, power, epoch_last)
     else:
         raise ValueError("unsupported lr scheduler {0:}".format(name))
     return scheduler

pymic/net_run/infer_func.py

@@ -15,7 +15,8 @@ class Inferer(object):
     :param `sliding_window_stride`: (optional, list) The sliding window stride. 
     :param `tta_mode`: (optional, int) The test time augmentation mode. Default
         is 0 (no test time augmentation). The other option is 1 (augmentation 
-        with horinzontal and vertical flipping).
+        with horinzontal and vertical flipping) and 2 (ensemble of inference
+        in axial, sagittal and coronal views for 2D networks applied to 3D volumes)
     """
     def __init__(self, config):
         self.config = config
@@ -170,6 +171,22 @@ def run(self, model, image):
                 outputs3 = torch.flip(outputs3, [-1])
                 outputs4 = torch.flip(outputs4, [-2, -1])
                 outputs = (outputs1 + outputs2 + outputs3 + outputs4) / 4
+        elif(tta_mode == 2):
+            outputs1 = self.__infer(image)
+            outputs2 = self.__infer(torch.transpose(image, -1, -3))
+            outputs3 = self.__infer(torch.transpose(image, -2, -3))
+            if(isinstance(outputs1, (tuple, list))):
+                outputs = []
+                for i in range(len(outputs1)):
+                    temp_out1 = outputs1[i]
+                    temp_out2 = torch.transpose(outputs2[i], -1, -3)
+                    temp_out3 = torch.transpose(outputs3[i], -2, -3)
+                    temp_mean = (temp_out1 + temp_out2 + temp_out3) / 3
+                    outputs.append(temp_mean)
+            else:
+                outputs2 = torch.transpose(outputs2, -1, -3)
+                outputs3 = torch.transpose(outputs3, -2, -3)
+                outputs = (outputs1 + outputs2 + outputs3) / 3
         else:
             raise ValueError("Undefined tta_mode {0:}".format(tta_mode))
         return outputs

pymic/net_run/net_run.py

@@ -3,6 +3,7 @@
 import logging
 import os
 import sys
+import shutil
 from pymic.util.parse_config import *
 from pymic.net_run.agent_cls import ClassificationAgent
 from pymic.net_run.agent_seg import SegmentationAgent
@@ -22,8 +23,15 @@ def main():
     log_dir  = config['training']['ckpt_save_dir']
     if(not os.path.exists(log_dir)):
         os.makedirs(log_dir, exist_ok=True)
-    logging.basicConfig(filename=log_dir+"/log_{0:}.txt".format(stage), level=logging.INFO,
-                        format='%(message)s')
+    if(stage == "train"):
+        dst_cfg = cfg_file if "/" not in cfg_file else cfg_file.split("/")[-1]
+        shutil.copy(cfg_file, log_dir + "/" + dst_cfg)
+    if sys.version.startswith("3.9"):
+        logging.basicConfig(filename=log_dir+"/log_{0:}.txt".format(stage), level=logging.INFO,
+                            format='%(message)s', force=True) # for python 3.9
+    else:
+        logging.basicConfig(filename=log_dir+"/log_{0:}.txt".format(stage), level=logging.INFO,
+                            format='%(message)s') # for python 3.6
     logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
     logging_config(config)
     task     = config['dataset']['task_type']

pymic/test/test_nifty_dataset.py

@@ -1,50 +1,68 @@
 # -*- coding: utf-8 -*-
 from __future__ import print_function, division
+import sys
+import numpy as np 
+from pymic.util.parse_config import *
+from pymic.net_run.agent_cls import ClassificationAgent
+from pymic.net_run.agent_seg import SegmentationAgent
+import SimpleITK as sitk
 
-import os
-import torch
-import pandas as pd
-import numpy as np
-from skimage import io, transform
-from torch.utils.data import Dataset, DataLoader
-from torchvision import transforms, utils
-from pymic.io.image_read_write import *
-from pymic.io.nifty_dataset import NiftyDataset
-from pymic.io.transform3d import *
+def save_array_as_nifty_volume(data, image_name, reference_name = None):
+    """
+    Save a numpy array as nifty image
 
-if __name__ == "__main__":
-    root_dir = '/home/guotai/data/brats/BraTS2018_Training'
-    csv_file = '/home/guotai/projects/torch_brats/brats/config/brats18_train_train.csv'
-    
-    crop1 = CropWithBoundingBox(start = None, output_size = [4, 144, 180, 144])
-    norm  = ChannelWiseNormalize(mean = None, std = None, zero_to_random = True)
-    labconv = LabelConvert([0, 1, 2, 4], [0, 1, 2, 3])
-    crop2 = RandomCrop([128, 128, 128])
-    rescale =Rescale([64, 64, 64])
-    transform_list = [crop1, norm, labconv, crop2,rescale, ToTensor()]
-    transformed_dataset = NiftyDataset(root_dir=root_dir,
-                                    csv_file  = csv_file,
-                                    modal_num = 4,
-                                    transform = transforms.Compose(transform_list))
-    dataloader = DataLoader(transformed_dataset, batch_size=4,
-                            shuffle=True, num_workers=4)
-    # Helper function to show a batch
+    :param data:  (numpy.ndarray) A numpy array with shape [Depth, Height, Width].
+    :param image_name:  (str) The ouput file name.
+    :param reference_name:  (str) File name of the reference image of which 
+        meta information is used.
+    """
+    img = sitk.GetImageFromArray(data)
+    if(reference_name is not None):
+        img_ref = sitk.ReadImage(reference_name)
+        #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 main():
+    """
+    The main function for running a network for training or inference.
+    """
+    if(len(sys.argv) < 3):
+        print('Number of arguments should be 3. e.g.')
+        print('python test_nifty_dataset.py train config.cfg')
+        exit()
+    stage    = str(sys.argv[1])
+    cfg_file = str(sys.argv[2])
+    config   = parse_config(cfg_file)
+    config   = synchronize_config(config)
+    # task     = config['dataset']['task_type']
+    # assert task in ['cls', 'cls_nexcl', 'seg']
+    # if(task == 'cls' or task == 'cls_nexcl'):
+    #     agent = ClassificationAgent(config, stage)
+    # else:
+    #     agent = SegmentationAgent(config, stage)
+    agent = SegmentationAgent(config, stage)
+    agent.create_dataset()
+    data_loader = agent.train_loader if stage == "train" else agent.test_loader
+    it = 0
+    for data in data_loader:
+        inputs      = agent.convert_tensor_type(data['image'])
+        labels_prob = agent.convert_tensor_type(data['label_prob']) 
+        for i in range(inputs.shape[0]):
+            image_i = inputs[i][0]
+            label_i = np.argmax(labels_prob[i], axis = 0)
+            print(image_i.shape, label_i.shape)
+            image_name = "temp/image_{0:}_{1:}.nii.gz".format(it, i)
+            label_name = "temp/label_{0:}_{1:}.nii.gz".format(it, i)
+            save_array_as_nifty_volume(image_i, image_name, reference_name = None)
+            save_array_as_nifty_volume(label_i, label_name, reference_name = None)
+        it = it + 1
+        if(it == 10):
+            break
 
-    for i_batch, sample_batched in enumerate(dataloader):
-        print(i_batch, sample_batched['image'].size(),
-            sample_batched['label'].size())
+if __name__ == "__main__":
+    main()
+    
 
-        # # observe 4th batch and stop.
-        modals = ['flair', 't1ce', 't1', 't2']
-        if i_batch == 0:
-            image =  sample_batched['image'].numpy()
-            label =  sample_batched['label'].numpy()
-            for i in range(image.shape[0]):
-                for mod in range(4):
-                    image_i = image[i][mod]
-                    label_i = label[i][0]
-                    image_name = "temp/image_{0:}_{1:}.nii.gz".format(i, modals[mod])
-                    label_name = "temp/label_{0:}.nii.gz".format(i)
-                    save_array_as_nifty_volume(image_i, image_name, reference_name = None)
-                    save_array_as_nifty_volume(label_i, label_name, reference_name = None)