update files for DMSPS

Author: taigw

pymic/io/image_read_write.py

@@ -99,7 +99,7 @@ def save_array_as_nifty_volume(data, image_name, reference_name = None, spacing
     :param spacing: (list or tuple) the spacing of a volume data when `reference_name` is not provided.  
     """
     img = sitk.GetImageFromArray(data)
-    if(reference_name is not None):
+    if((reference_name is not None) and (not reference_name.endswith(".h5"))):
         img_ref = sitk.ReadImage(reference_name)
         #img.CopyInformation(img_ref)
         img.SetSpacing(img_ref.GetSpacing())
@@ -141,11 +141,15 @@ def save_nd_array_as_image(data, image_name, reference_name = None, spacing = [1
     """
     data_dim = len(data.shape)
     assert(data_dim == 2 or data_dim == 3)
+    if(image_name.endswith(".h5")):
+        if(data_dim == 3):
+            image_name = image_name.replace(".h5", ".nii.gz")
+        else:
+            image_name = image_name.replace(".h5", ".png")
     if (image_name.endswith(".nii.gz") or image_name.endswith(".nii") or
         image_name.endswith(".mha")):
         assert(data_dim == 3)
         save_array_as_nifty_volume(data, image_name, reference_name, spacing)
-
     elif(image_name.endswith(".jpg") or image_name.endswith(".jpeg") or
          image_name.endswith(".tif") or image_name.endswith(".png")):
          assert(data_dim == 2)

pymic/loss/loss_dict_seg.py

@@ -26,6 +26,7 @@
 from pymic.loss.seg.dice import DiceLoss, FocalDiceLoss, \
      NoiseRobustDiceLoss, BinaryDiceLoss, GroupDiceLoss
 from pymic.loss.seg.exp_log import ExpLogLoss
+from pymic.loss.seg.ars_tversky import ARSTverskyLoss
 from pymic.loss.seg.mse import MSELoss, MAELoss
 from pymic.loss.seg.slsr import SLSRLoss
 
@@ -35,6 +36,7 @@
     'DiceLoss': DiceLoss,
     'BinaryDiceLoss': BinaryDiceLoss,
     'FocalDiceLoss': FocalDiceLoss,
+    'ARSTverskyLoss': ARSTverskyLoss,
     'NoiseRobustDiceLoss': NoiseRobustDiceLoss,
     'GroupDiceLoss': GroupDiceLoss,
     'ExpLogLoss': ExpLogLoss,

pymic/loss/seg/ce.py

@@ -36,7 +36,7 @@ def forward(self, loss_input_dict):
         soft_y  = reshape_tensor_to_2D(soft_y)
 
         # for numeric stability
-        predict = predict * 0.999 + 5e-4
+        # predict = predict  * (1-1e-10) + 0.5e-10
         ce = - soft_y* torch.log(predict)
         if(cls_w is not None):
             ce = torch.sum(ce*cls_w, dim = 1)
@@ -46,7 +46,7 @@ def forward(self, loss_input_dict):
             ce = torch.mean(ce)  
         else:
             pix_w = torch.squeeze(reshape_tensor_to_2D(pix_w))
-            ce = torch.sum(pix_w * ce) / (pix_w.sum() + 1e-5) 
+            ce = torch.sum(pix_w * ce) / (pix_w.sum() + 1e-10) 
         return ce
 
 class GeneralizedCELoss(AbstractSegLoss):

pymic/net/net2d/unet2d_multi_decoder.py

@@ -63,15 +63,16 @@ def forward(self, x):
             output2 = torch.reshape(output2, new_shape)
             output2 = torch.transpose(output2, 1, 2)
 
-        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
+        return output1, output2
+        # 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
 
 class UNet2D_TriBranch(nn.Module):
     """

pymic/net_run/weak_sup/wsl_dmsps.py

@@ -1,11 +1,13 @@
 # -*- coding: utf-8 -*-
 from __future__ import print_function, division
 import logging
+import os
 import numpy as np
 import random
 import time
 import torch
-from PIL import Image
+import scipy
+from pymic.io.image_read_write import save_nd_array_as_image
 from pymic.loss.seg.util import get_soft_label
 from pymic.loss.seg.util import reshape_prediction_and_ground_truth
 from pymic.loss.seg.util import get_classwise_dice
@@ -133,4 +135,68 @@ def training(self):
             'avg_fg_dice':train_avg_dice,     'class_dice': train_cls_dice,
             'data_time': data_time, 'forward_time':gpu_time, 
             'loss_time':loss_time, 'backward_time':back_time}
-        return train_scalers
+        return train_scalers
+    
+    def save_outputs(self, data):
+        """
+        Save prediction output. 
+
+        :param data: (dictionary) A data dictionary with prediciton result and other 
+            information such as input image name. 
+        """
+        output_dir = self.config['testing']['output_dir']
+        test_mode  = self.config['testing'].get('dmsps_test_mode', 0)
+        uct_threshold = self.config['testing'].get('dmsps_uncertainty_threshold', 0.1)
+        # DMSPS_test_mode == 0: only save the segmentation label for the main decoder
+        # DMSPS_test_mode == 1: save all the results, including the the probability map of each decoder,
+        # the uncertainty map, and the confident predictions
+        if(not os.path.exists(output_dir)):
+            os.makedirs(output_dir, exist_ok=True)
+
+        names, pred  = data['names'], data['predict']
+        pred0, pred1 = pred
+        prob0   = scipy.special.softmax(pred0, axis = 1) 
+        prob1   = scipy.special.softmax(pred1, axis = 1) 
+        prob_mean = (prob0 + prob1) / 2
+        lab0     = np.asarray(np.argmax(prob0,  axis = 1), np.uint8)
+        lab1     = np.asarray(np.argmax(prob1,  axis = 1), np.uint8)
+        lab_mean = np.asarray(np.argmax(prob_mean,  axis = 1), np.uint8)
+
+        # save the output and (optionally) probability predictions
+        test_dir = self.config['dataset'].get('test_dir', None)
+        if(test_dir is None):
+            test_dir = self.config['dataset']['train_dir']
+        img_name  = names[0][0].split('/')[-1]
+        print(img_name)
+        lab0_name = img_name
+        if(".h5" in lab0_name):
+            lab0_name = lab0_name.replace(".h5", ".nii.gz")
+        save_nd_array_as_image(lab0[0], output_dir + "/" + lab0_name, test_dir + '/' + names[0][0])
+        if(test_mode == 1):
+            lab1_name = lab0_name.replace(".nii.gz", "_predaux.nii.gz")
+            save_nd_array_as_image(lab1[0], output_dir + "/" + lab1_name, test_dir + '/' + names[0][0])
+            C = pred0.shape[1]
+            uct = -1.0 * np.sum(prob_mean * np.log(prob_mean), axis=1, keepdims=False)/ np.log(C)
+            uct_name = lab0_name.replace(".nii.gz", "_uncertainty.nii.gz")
+            save_nd_array_as_image(uct[0], output_dir + "/" + uct_name, test_dir + '/' + names[0][0])
+            conf_mask = uct < uct_threshold
+            conf_lab  = conf_mask * lab_mean + (1 - conf_mask)*4
+            conf_lab_name = lab0_name.replace(".nii.gz", "_seeds_expand.nii.gz")
+            
+            # get the largest connected component in each slice for each class
+            D, H, W = conf_lab[0].shape
+            from pymic.util.image_process import get_largest_k_components
+            for d in range(D):
+                lab2d = conf_lab[0][d]
+                for c in range(C):
+                    lab2d_c = lab2d == c 
+                    mask_c  = get_largest_k_components(lab2d_c, k = 1)
+                    diff = lab2d_c != mask_c
+                    if(np.sum(diff) > 0):
+                        lab2d[diff] = C 
+                conf_lab[0][d] = lab2d
+            save_nd_array_as_image(conf_lab[0], output_dir + "/" + conf_lab_name, test_dir + '/' + img_name)
+            
+            
+            
+            

pymic/transform/rescale.py

@@ -72,12 +72,22 @@ def inverse_transform_for_prediction(self, sample):
             origin_shape = json.loads(sample['Rescale_origin_shape'])
         origin_dim   = len(origin_shape) - 1
         predict = sample['predict']
-        input_shape = predict.shape
-        scale = [(origin_shape[1:][i] + 0.0)/input_shape[2:][i] for \
-                i in range(origin_dim)]
-        scale = [1.0, 1.0] + scale
 
-        output_predict = ndimage.interpolation.zoom(predict, scale, order = 1)
+        if(isinstance(predict, tuple) or isinstance(predict, list)):
+            output_predict = []
+            for predict_i in predict:
+                input_shape = predict_i.shape
+                scale = [(origin_shape[1:][i] + 0.0)/input_shape[2:][i] for \
+                        i in range(origin_dim)]
+                scale = [1.0, 1.0] + scale
+                output_predict_i = ndimage.interpolation.zoom(predict_i, scale, order = 1)
+                output_predict.append(output_predict_i)                
+        else:
+            input_shape = predict.shape
+            scale = [(origin_shape[1:][i] + 0.0)/input_shape[2:][i] for \
+                    i in range(origin_dim)]
+            scale = [1.0, 1.0] + scale
+            output_predict = ndimage.interpolation.zoom(predict, scale, order = 1)
         sample['predict'] = output_predict
         return sample
 

pymic/util/evaluation_seg.py

@@ -212,8 +212,10 @@ def get_binary_evaluation_score(s_volume, g_volume, spacing, metric):
         score = binary_iou(s_volume,g_volume)
     elif(metric_lower == 'assd'):
         score = binary_assd(s_volume, g_volume, spacing)
+        score = min(score, 20) # to reject outliers
     elif(metric_lower == "hd95"):
         score = binary_hd95(s_volume, g_volume, spacing)
+        score = min(score, 50) # to reject outliers
     elif(metric_lower == "rve"):
         score = binary_relative_volume_error(s_volume, g_volume)
     elif(metric_lower == "volume"):
@@ -269,8 +271,8 @@ def evaluation(config):
     :param label_fuse: (option, bool) If true, fuse the labels in the `label_list`
         as the foreground, and other labels as the background. Default is False. 
     :param organ_name: (str) The name of the organ for segmentation.
-    :param ground_truth_folder_root: (str) The root dir of ground truth images. 
-    :param segmentation_folder_root: (str or list) The root dir of segmentation images. 
+    :param ground_truth_folder: (str) The root dir of ground truth images. 
+    :param segmentation_folder: (str or list) The root dir of segmentation images. 
         When a list is given, each list element should be the root dir of the results of one method. 
     :param evaluation_image_pair: (str) The csv file that provide the segmentation 
         images and the corresponding ground truth images. 
@@ -366,23 +368,23 @@ def main():
 
     """
     parser = argparse.ArgumentParser()
-    parser.add_argument("-cfg", help="configuration file for evaluation", 
+    parser.add_argument("--cfg", help="configuration file for evaluation", 
                         required=False, default=None)
-    parser.add_argument("-metric", help="evaluation metrics, e.g., dice, or [dice, assd]", 
+    parser.add_argument("--metric", help="evaluation metrics, e.g., dice, or [dice, assd]", 
                         required=False, default=None)
-    parser.add_argument("-cls_num", help="number of classes", 
+    parser.add_argument("--cls_num", help="number of classes", 
                         required=False, default=None)
-    parser.add_argument("-cls_index", help="The class index for evaluation, e.g., 255, [1, 2]", 
+    parser.add_argument("--cls_index", help="The class index for evaluation, e.g., 255, [1, 2]", 
                         required=False, default=None)
-    parser.add_argument("-gt_dir", help="path of folder for ground truth", 
+    parser.add_argument("--gt_dir", help="path of folder for ground truth", 
                         required=False, default=None)
-    parser.add_argument("-seg_dir", help="path of folder for segmentation", 
+    parser.add_argument("--seg_dir", help="path of folder for segmentation", 
                         required=False, default=None)
-    parser.add_argument("-name_pair", help="the .csv file for name mapping in case"
+    parser.add_argument("--name_pair", help="the .csv file for name mapping in case"
                         " the names of one case are different in the gt_dir "
                         " and seg_dir", 
                         required=False, default=None)
-    parser.add_argument("-out", help="the output .csv file name", 
+    parser.add_argument("--out", help="the output .csv file name", 
                         required=False, default=None)
     args = parser.parse_args()
     print(args)
@@ -402,5 +404,3 @@ def main():
 
 if __name__ == '__main__':
     main()
-
-    main()

pymic/util/parse_config.py

@@ -96,12 +96,22 @@ def parse_config(args):
                 args_key = getattr(args, key)
                 if(args_key is not None):
                     val_str = args_key
+            print(section, key, val_str)
             if(len(val_str)>0):
                 val = parse_value_from_string(val_str)
                 output[section][key] = val
             else:
                 val = None
-            print(section, key, val)
+            
+    for key in ["train_dir", "train_csv", "valid_csv", "test_dir", "test_csv"]:
+        if key in args and getattr(args, key) is not None:
+            output["dataset"][key] = parse_value_from_string(getattr(args, key))
+    for key in ["ckpt_dir", "iter_max", "gpus"]:
+        if key in args and getattr(args, key) is not None:
+            output["training"][key] = parse_value_from_string(getattr(args, key))
+    for key in ["output_dir", "ckpt_mode", "ckpt_name"]:
+        if key in args and getattr(args, key) is not None:
+            output["testing"][key] = parse_value_from_string(getattr(args, key))
     return output
 
 def synchronize_config(config):
@@ -133,6 +143,17 @@ def synchronize_config(config):
         if('RandomResizedCrop' in transform and \
             'RandomResizedCrop_output_size'.lower() not in data_cfg):
             data_cfg['RandomResizedCrop_output_size'.lower()] = patch_size
+        if('testing' in config):
+            test_cfg  = config['testing']
+            sliding_window_enable = test_cfg.get("sliding_window_enable", False)
+            if(sliding_window_enable):
+                sliding_window_size = test_cfg.get("sliding_window_size", None)
+                if(sliding_window_size is None):
+                    test_cfg["sliding_window_size"] = patch_size
+                sliding_window_stride = test_cfg.get("sliding_window_stride", None)
+                if(sliding_window_stride is None):
+                    test_cfg["sliding_window_stride"] = [item // 2 for item in patch_size]
+            config['testing'] = test_cfg
     config['dataset'] = data_cfg
     # config['network'] = net_cfg
     return config