update net run agent and mse loss

Author: taigw

pymic/io/nifty_dataset.py

@@ -16,25 +16,46 @@ class NiftyDataset(Dataset):
     with dimention order [C, H, W] for 2D images"""
 
     def __init__(self, root_dir, csv_file, modal_num = 1, 
-            with_label = False, transform=None):
+            with_label = False, with_weight = None, transform=None):
         """
         Args:
             root_dir (string): Directory with all the images.
             csv_file (string): Path to the csv file with image names.
             modal_num (int): Number of modalities. 
             with_label (bool): Load the data with segmentation ground truth.
+            with_weight(bool): Load pixel-wise weight map.
             transform (callable, optional): Optional transform to be applied
                 on a sample.
         """
         self.root_dir   = root_dir
         self.csv_items  = pd.read_csv(csv_file)
         self.modal_num  = modal_num
         self.with_label = with_label
+        self.with_weight= with_weight
         self.transform  = transform
 
+        if(self.with_label):
+            self.label_idx = list(self.csv_items.keys()).index('label')
+        if(self.with_weight):
+            self.weight_idx = list(self.csv_items.keys()).index('weight')
+
     def __len__(self):
         return len(self.csv_items)
 
+    def __getlabel__(self, idx):
+        label_name = "{0:}/{1:}".format(self.root_dir, 
+            self.csv_items.iloc[idx, self.label_idx])
+        label = load_image_as_nd_array(label_name)['data_array']
+        label = np.asarray(label, np.int32)
+        return label
+
+    def __getweight__(self, idx):
+        weight_name = "{0:}/{1:}".format(self.root_dir, 
+            self.csv_items.iloc[idx, self.weight_idx])
+        weight = load_image_as_nd_array(weight_name)['data_array']
+        weight = np.asarray(weight, np.float32)
+        return weight        
+
     def __getitem__(self, idx):
         names_list, image_list = [], []
         for i in range (self.modal_num):
@@ -46,17 +67,32 @@ def __getitem__(self, idx):
             image_list.append(image_data)
         image = np.concatenate(image_list, axis = 0)
         image = np.asarray(image, np.float32)    
-        
         sample = {'image': image, 'names' : names_list[0], 
                  'origin':image_dict['origin'],
                  'spacing': image_dict['spacing'],
                  'direction':image_dict['direction']}
-        if (self.with_label):    
-            label_name = "{0:}/{1:}".format(self.root_dir, self.csv_items.iloc[idx, -1])
-            label = load_image_as_nd_array(label_name)['data_array']
-            label = np.asarray(label, np.int32)
-            sample['label'] = label
+        if (self.with_label):   
+            sample['label'] = self.__getlabel__(idx) 
+            assert(image.shape[1:] == sample['label'].shape[1:])
+        if (self.with_weight):
+            sample['weight'] = self.__getweight__(idx) 
+            assert(image.shape[1:] == sample['weight'].shape[1:])
         if self.transform:
             sample = self.transform(sample)
 
         return sample
+
+
+class ClassificationDataset(NiftyDataset):
+    def __init__(self, root_dir, csv_file, modal_num = 1, class_num = 2, 
+            with_label = False, transform=None):
+        super(ClassificationDataset, self).__init__(root_dir, 
+            csv_file, modal_num, with_label, transform)
+        self.class_num = class_num
+        print("class number for ClassificationDataset", self.class_num)
+
+    def __getlabel__(self, idx):
+        label_idx = self.csv_items.iloc[idx, -1]
+        label = np.zeros((self.class_num, ))
+        label[label_idx] = 1
+        return label

pymic/loss/loss_dict.py

@@ -4,11 +4,14 @@
 from pymic.loss.dice import DiceLoss, MultiScaleDiceLoss
 from pymic.loss.dice import DiceWithCrossEntropyLoss, NoiseRobustDiceLoss
 from pymic.loss.exp_log import ExpLogLoss
+from pymic.loss.mse import MSELoss, MAELoss
 
 LossDict = {'CrossEntropyLoss': CrossEntropyLoss,
     'GeneralizedCrossEntropyLoss': GeneralizedCrossEntropyLoss,
     'DiceLoss': DiceLoss,
     'MultiScaleDiceLoss': MultiScaleDiceLoss,
     'DiceWithCrossEntropyLoss': DiceWithCrossEntropyLoss,
     'NoiseRobustDiceLoss': NoiseRobustDiceLoss,
-    'ExpLogLoss': ExpLogLoss}
+    'ExpLogLoss': ExpLogLoss,
+    'MSELoss': MSELoss,
+    'MAELoss': MAELoss}

pymic/loss/mse.py

@@ -0,0 +1,50 @@
+import torch
+import torch.nn as nn
+from pymic.loss.util import reshape_tensor_to_2D
+
+class MSELoss(nn.Module):
+    def __init__(self, params):
+        super(MSELoss, self).__init__()
+        self.enable_pix_weight = params['MSELoss_Enable_Pixel_Weight'.lower()]
+        self.enable_cls_weight = params['MSELoss_Enable_Class_Weight'.lower()]
+    
+    def forward(self, loss_input_dict):
+        predict = loss_input_dict['prediction']
+        soft_y  = loss_input_dict['ground_truth']
+        pix_w   = loss_input_dict['pixel_weight']
+        cls_w   = loss_input_dict['class_weight']
+        softmax = loss_input_dict['softmax']
+
+        if(softmax):
+            predict = nn.Softmax(dim = 1)(predict)
+        predict = reshape_tensor_to_2D(predict)
+        soft_y  = reshape_tensor_to_2D(soft_y)
+        se   = self.get_prediction_error(predict, soft_y)
+        if(self.enable_cls_weight):
+            if(cls_w is None):
+                raise ValueError("Class weight is enabled but not defined")
+            mse = torch.sum(se * cls_w, dim = 1) / torch.sum(cls_w)
+        else:
+            mse = torch.mean(se, dim = 1)
+        if(self.enable_pix_weight):
+            if(pix_w is None):
+                raise ValueError("Pixel weight is enabled but not defined")
+            pix_w = reshape_tensor_to_2D(pix_w)
+            mse   = torch.sum(mse * pix_w) / torch.sum(pix_w)
+        else:
+            mse = torch.mean(mse)  
+        return mse 
+    
+    def get_prediction_error(self, predict, soft_y):
+        diff = predict - soft_y
+        error = diff * diff 
+        return error 
+    
+class MAELoss(nn.MSELoss):
+    def __init__(self, params):
+        super(MAELoss, self).__init__(params)
+    
+    def get_prediction_error(self, predict, soft_y):
+        diff  = predict - soft_y
+        error = torch.abs(diff)
+        return error 

pymic/net_run/net_run_agent.py

@@ -66,26 +66,30 @@ def get_stage_dataset_from_config(self, stage):
 
         if(stage == "train" or stage == "valid"):
             transform_names = self.config['dataset']['train_transform']
-            with_weight = self.config['dataset']['load_pixelwise_weight']
+            with_weight = self.config['dataset'].get('load_pixelwise_weight', False)
         elif(stage == "test"):
             transform_names = self.config['dataset']['test_transform']
             with_weight = False 
         else:
             raise ValueError("Incorrect value for stage: {0:}".format(stage))
         self.transform_list  = []
-        for name in transform_names:
-            if(name not in self.transform_dict):
-                raise(ValueError("Undefined transform {0:}".format(name))) 
-            one_transform = self.transform_dict[name](self.config['dataset'])
-            self.transform_list.append(one_transform)
+        if(transform_names is None or len(transform_names) == 0):
+            data_transform = None 
+        else:
+            for name in transform_names:
+                if(name not in self.transform_dict):
+                    raise(ValueError("Undefined transform {0:}".format(name))) 
+                one_transform = self.transform_dict[name](self.config['dataset'])
+                self.transform_list.append(one_transform)
+            data_transform = transforms.Compose(self.transform_list)
 
         csv_file = self.config['dataset'].get(stage + '_csv', None)
         dataset  = NiftyDataset(root_dir=root_dir,
                                 csv_file  = csv_file,
                                 modal_num = modal_num,
                                 with_label= not (stage == 'test'),
                                 with_weight = with_weight,
-                                transform = transforms.Compose(self.transform_list))
+                                transform = data_transform )
         return dataset
 
     def create_dataset(self):
@@ -141,7 +145,6 @@ def convert_tensor_type(self, input_tensor):
     def train(self):
         device = torch.device(self.config['training']['device_name'])
         self.net.to(device)
-        
         class_num   = self.config['network']['class_num']
         summ_writer = SummaryWriter(self.config['training']['summary_dir'])
         chpt_prefx  = self.config['training']['checkpoint_prefix']
@@ -171,7 +174,7 @@ def train(self):
         if(loss_name not in self.loss_dict):
             raise ValueError("Undefined loss function {0:}".format(loss_name))
         self.loss_calculater = self.loss_dict[loss_name](self.config['training'])
-
+        
         trainIter  = iter(self.train_loader)
         train_loss = 0
         train_dice_list = []

pymic/util/evaluation.py

@@ -4,6 +4,7 @@
 import os
 import sys
 import math
+import pandas as pd
 import random
 import GeodisTK
 import configparser
@@ -180,7 +181,7 @@ def get_evaluation_score(s_volume, g_volume, spacing, metric):
 
     return score
 
-def evaluation(config_file):
+def evaluation_backup(config_file):
     config = parse_config(config_file)['evaluation']
     metric = config['metric']
     labels = config['label_list']
@@ -218,7 +219,9 @@ def evaluation(config_file):
 
             for g_folder in g_folder_list:
                 g_name = os.path.join(g_folder, patient_names[i] + g_postfix_long)
-                if(os.path.isfile(g_name)):
+                if(not os.path.isfile(g_name)):
+                    g_name = g_name.replace(patient_names[i], patient_names[i] + '/' + patient_names[i])
+                if(not os.path.isfile(g_name)):  
                     break
             s_dict = load_image_as_nd_array(s_name)
             g_dict = load_image_as_nd_array(g_name)
@@ -257,6 +260,64 @@ def evaluation(config_file):
         print("{0:} mean ".format(metric), score_mean)
         print("{0:} std  ".format(metric), score_std) 
 
+def evaluation(config_file):
+    config = parse_config(config_file)['evaluation']
+    metric = config['metric']
+    labels = config['label_list']
+    organ_name = config['organ_name']
+    gt_root    = config['ground_truth_folder_root']
+    seg_root   = config['segmentation_folder_root']
+    image_pair_csv = config['evaluation_image_pair']
+    ground_truth_label_convert_source = config.get('ground_truth_label_convert_source', None)
+    ground_truth_label_convert_target = config.get('ground_truth_label_convert_target', None)
+    segmentation_label_convert_source = config.get('segmentation_label_convert_source', None)
+    segmentation_label_convert_target = config.get('segmentation_label_convert_target', None)
+
+    image_items = pd.read_csv(image_pair_csv)
+    item_num = len(image_items)
+    score_all_data = []
+    for i in range(item_num):
+        gt_name  = image_items.iloc[i, 0]
+        seg_name = image_items.iloc[i, 1]
+        gt_full_name  = gt_root  + '/' + gt_name
+        seg_full_name = seg_root + '/' + seg_name
+        
+        s_dict = load_image_as_nd_array(seg_full_name)
+        g_dict = load_image_as_nd_array(gt_full_name)
+        s_volume = s_dict["data_array"]; s_spacing = s_dict["spacing"]
+        g_volume = g_dict["data_array"]; g_spacing = g_dict["spacing"]
+        # for dim in range(len(s_spacing)):
+        #     assert(s_spacing[dim] == g_spacing[dim])
+        if((ground_truth_label_convert_source is not None) and \
+            ground_truth_label_convert_target is not None):
+            g_volume = convert_label(g_volume, ground_truth_label_convert_source, \
+                ground_truth_label_convert_target)
+
+        if((segmentation_label_convert_source is not None) and \
+            segmentation_label_convert_target is not None):
+            s_volume = convert_label(s_volume, segmentation_label_convert_source, \
+                segmentation_label_convert_target)
+
+        # fuse multiple labels
+        s_volume_sub = np.zeros_like(s_volume)
+        g_volume_sub = np.zeros_like(g_volume)
+        for lab in labels:
+            s_volume_sub = s_volume_sub + np.asarray(s_volume == lab, np.uint8)
+            g_volume_sub = g_volume_sub + np.asarray(g_volume == lab, np.uint8)
+        
+        # get evaluation score
+        temp_score = get_evaluation_score(s_volume_sub > 0, g_volume_sub > 0,
+                    s_spacing, metric)
+        score_all_data.append(temp_score)
+        print(seg_name, temp_score)
+    score_all_data = np.asarray(score_all_data)
+    score_mean = [score_all_data.mean(axis = 0)]
+    score_std  = [score_all_data.std(axis = 0)]
+    np.savetxt("{0:}/{1:}_{2:}_all.txt".format(seg_root, organ_name, metric), score_all_data)
+    np.savetxt("{0:}/{1:}_{2:}_mean.txt".format(seg_root, organ_name, metric), score_mean)
+    np.savetxt("{0:}/{1:}_{2:}_std.txt".format(seg_root, organ_name, metric), score_std)
+    print("{0:} mean ".format(metric), score_mean)
+    print("{0:} std  ".format(metric), score_std) 
 
 def main():
     if(len(sys.argv) < 2):

pymic/util/image_process.py

@@ -2,6 +2,7 @@
 from __future__ import absolute_import, print_function
 
 import numpy as np
+import SimpleITK as sitk
 from scipy import ndimage
 
 def get_ND_bounding_box(volume, margin = None):
@@ -16,8 +17,8 @@ def get_ND_bounding_box(volume, margin = None):
     idx_min = []
     idx_max = []
     for i in range(len(input_shape)):
-        idx_min.append(indxes[i].min())
-        idx_max.append(indxes[i].max() + 1)
+        idx_min.append(int(indxes[i].min()))
+        idx_max.append(int(indxes[i].max()) + 1)
 
     for i in range(len(input_shape)):
         idx_min[i] = max(idx_min[i] - margin[i], 0)
@@ -162,3 +163,19 @@ def convert_label(label, source_list, target_list):
         label_temp = label_temp * target_list[i]
         label_converted = label_converted + label_temp
     return label_converted
+
+def resample_sitk_image_to_given_spacing(image, spacing, order):
+    """
+    image: an sitk image object
+    spacing: 3D tuple / list for spacing along x, y, z direction
+    order: order for interpolation
+    """
+    spacing0 = image.GetSpacing()
+    data = sitk.GetArrayFromImage(image)
+    zoom = [spacing0[i] / spacing[i] for i in range(3)]
+    zoom = [zoom[2], zoom[0], zoom[1]]
+    data = ndimage.interpolation.zoom(data, zoom, order = order)
+    out_img = sitk.GetImageFromArray(data)
+    out_img.SetSpacing(spacing)
+    out_img.SetDirection(image.GetDirection())
+    return out_img