noise-learning

update pymic for noise-learning

Author: taigw

pymic/io/nifty_dataset.py

@@ -95,6 +95,7 @@ def __init__(self, root_dir, csv_file, modal_num = 1, class_num = 2,
             csv_file, modal_num, with_label, transform)
         self.class_num = class_num
         print("class number for ClassificationDataset", self.class_num)
+        print("self.transform", self.transform)
 
     def __getlabel__(self, idx):
         csv_keys = list(self.csv_items.keys())
@@ -126,7 +127,7 @@ def __getitem__(self, idx):
             sample['label'] = self.__getlabel__(idx) 
         if (self.image_weight_idx is not None):
             sample['image_weight'] = self.__getweight__(idx) 
+        print("***transform", self.transform)
         if self.transform:
             sample = self.transform(sample)
         return sample
-

pymic/loss/loss_dict_seg.py

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

pymic/loss/seg/gatedcrf_backup.py

@@ -0,0 +1,48 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+
+import torch
+import torch.nn as nn
+import numpy as np 
+from pymic.loss.seg.ce import CrossEntropyLoss
+from pymic.loss.seg.gatedcrf_util import ModelLossSemsegGatedCRF
+
+class GatedCRFLoss(nn.Module):
+    def __init__(self, params):
+        super(GatedCRFLoss, self).__init__()
+        self.gcrf_loss = ModelLossSemsegGatedCRF()
+        self.softmax = params.get('loss_softmax', True)
+        w0 = params['GatedCRFLoss_W0'.lower()]
+        xy0= params['GatedCRFLoss_XY0'.lower()]
+        rgb= params['GatedCRFLoss_rgb'.lower()]
+        w1 = params['GatedCRFLoss_W1'.lower()]
+        xy1= params['GatedCRFLoss_XY1'.lower()]
+        kernel0 = {'weight': w0, 'xy': xy0, 'rgb': rgb}
+        kernel1 = {'weight': w1, 'xy': xy1}
+        self.kernels = [kernel0, kernel1]
+        self.radius  = params['GatedCRFLoss_Radius'.lower()]
+    
+    def forward(self, loss_input_dict):
+        predict = loss_input_dict['prediction']
+        image   = loss_input_dict['image']  # should be normalized by mean, std
+        scribble= loss_input_dict['scribbles']
+        validity_mask = loss_input_dict['validity_mask']
+
+        if(self.softmax):
+            predict = nn.Softmax(dim = 1)(predict)
+
+        batch_dict = {'rgb': image, 
+            'semseg_scribbles': scribble}
+        x_shape = list(predict.shape)
+        l_crf = {'loss': 0}
+        if(self.gcrf_w > 0):
+            l_crf = self.gcrf_loss(predict,
+                    self.kernels,
+                    self.radius,
+                    batch_dict,
+                    x_shape[-2],
+                    x_shape[-1],
+                    mask_src=validity_mask,
+                    out_kernels_vis=False,
+                )
+        return l_crf['loss']

pymic/loss/seg/slsr.py

@@ -0,0 +1,50 @@
+# -*- coding: utf-8 -*-
+"""
+Spatial Label Smoothing Regularization (SLSR) loss for learning from
+noisy annotatins according to the following paper:
+    Minqing Zhang, Jiantao Gao et al., Characterizing Label Errors: 
+    Confident Learning for Noisy-Labeled Image Segmentation, MICCAI 2020.
+"""
+from __future__ import print_function, division
+
+import torch
+import torch.nn as nn
+from pymic.loss.seg.util import reshape_tensor_to_2D
+
+class SLSRLoss(nn.Module):
+    def __init__(self, params):
+        super(SLSRLoss, self).__init__()
+        if(params is None):
+            params = {}
+        self.softmax = params.get('loss_softmax', True)
+        self.epsilon = params.get('slsrloss_softmax', 0.25)
+    
+    def forward(self, loss_input_dict):
+        predict = loss_input_dict['prediction']
+        soft_y  = loss_input_dict['ground_truth']
+        pix_w   = loss_input_dict.get('pixel_weight', None)
+        # the pixel wight here is actually the confidence mask
+        # i.e., if the value is one, it means the label of corresponding 
+        # pixel is noisy and should be replaced with smoothed label.
+
+        if(isinstance(predict, (list, tuple))):
+            predict = predict[0]
+        if(self.softmax):
+            predict = nn.Softmax(dim = 1)(predict)
+        predict = reshape_tensor_to_2D(predict)
+        soft_y  = reshape_tensor_to_2D(soft_y)
+        if(pix_w is not None):
+            pix_w   = reshape_tensor_to_2D(pix_w > 0).float()
+
+            # smooth labels for pixels in the unconfident mask 
+            smooth_y = (soft_y - 0.5) * (0.5 - self.epsilon) / 0.5 + 0.5
+            smooth_y = pix_w * smooth_y + (1 - pix_w) * soft_y
+        else:
+            smooth_y = soft_y
+
+        # for numeric stability
+        predict = predict * 0.999 + 5e-4
+        ce = - smooth_y* torch.log(predict)
+        ce = torch.sum(ce, dim = 1) # shape is [N]
+        ce = torch.mean(ce)  
+        return ce

pymic/net/cls/torch_pretrained_net.py

@@ -24,10 +24,9 @@ class ResNet18(nn.Module):
     def __init__(self, params):
         super(ResNet18, self).__init__()
         self.params    = params
-        net_name = params['net_type']
         cls_num  = params['class_num']
-        in_chns  = params['input_chns']
-        self.pretrain = params['pretrain']
+        in_chns  = params.get('input_chns', 3)
+        self.pretrain = params.get('pretrain', True)
         self.update_layers = params.get('update_layers', 0)
         self.net = models.resnet18(pretrained = self.pretrain)
 
@@ -51,10 +50,9 @@ class VGG16(nn.Module):
     def __init__(self, params):
         super(VGG16, self).__init__()
         self.params    = params
-        net_name = params['net_type']
         cls_num  = params['class_num']
-        in_chns  = params['input_chns']
-        self.pretrain = params['pretrain']
+        in_chns  = params.get('input_chns', 3)
+        self.pretrain = params.get('pretrain', True)
         self.update_layers = params.get('update_layers', 0)
         self.net = models.vgg16(pretrained = self.pretrain)
 
@@ -78,10 +76,9 @@ 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']
+        in_chns  = params.get('input_chns', 3)
+        self.pretrain = params.get('pretrain', True)
         self.update_layers = params.get('update_layers', 0)
         self.net = models.mobilenet_v2(pretrained = self.pretrain)
 

pymic/net/net3d/unet3d.py

@@ -96,6 +96,7 @@ def __init__(self, params):
         self.n_class   = self.params['class_num']
         self.trilinear = self.params['trilinear']
         self.deep_sup  = self.params['deep_supervise']
+        self.stage     = self.params['stage']
         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])
@@ -133,7 +134,7 @@ def forward(self, x):
         x_d1 = self.up3(x_d2, x1)
         x_d0 = self.up4(x_d1, x0)
         output = self.out_conv(x_d0)
-        if(self.deep_sup):
+        if(self.deep_sup and self.stage == "train"):
             out_shape = list(output.shape)[2:]
             output1 = self.out_conv1(x_d1)
             output1 = interpolate(output1, out_shape, mode = 'trilinear')

pymic/net_run/agent_abstract.py

@@ -144,7 +144,7 @@ def worker_init_fn(worker_id):
             bn_test = self.config['dataset'].get('test_batch_size', 1)
             if(self.test_set  is None):
                 self.test_set  = self.get_stage_dataset_from_config('test')
-            self.test_loder = torch.utils.data.DataLoader(self.test_set, 
+            self.test_loader = torch.utils.data.DataLoader(self.test_set, 
                 batch_size = bn_test, shuffle=False, num_workers= bn_test)
 
     def create_optimizer(self, params):

pymic/net_run/agent_cls.py

@@ -207,7 +207,7 @@ def train_valid(self):
         self.net.to(self.device)
 
         ckpt_dir    = self.config['training']['ckpt_save_dir']
-        ckpt_prefx  = self.config['training']['ckpt_save_prefix']
+        ckpt_prefx  = ckpt_dir.split('/')[-1]
         iter_start  = self.config['training']['iter_start']
         iter_max    = self.config['training']['iter_max']
         iter_valid  = self.config['training']['iter_valid']
@@ -290,9 +290,12 @@ def infer(self):
         out_prob_list   = []
         out_lab_list    = []
         with torch.no_grad():
-            for data in self.test_loder:
+            for data in self.test_loader:
                 names  = data['names']
+                if(names[0] != "data3_process/20190711_1005487059.png"):
+                    continue
                 inputs = self.convert_tensor_type(data['image'])
+                print("intensity mean and std", inputs.mean(), inputs.std())
                 inputs = inputs.to(device) 
 
                 start_time = time.time()

pymic/net_run/agent_seg.py

@@ -40,7 +40,7 @@ def __init__(self, config, stage = 'train'):
     def get_stage_dataset_from_config(self, stage):
         assert(stage in ['train', 'valid', 'test'])
         root_dir  = self.config['dataset']['root_dir']
-        modal_num = self.config['dataset']['modal_num']
+        modal_num = self.config['dataset'].get('modal_num', 1)
 
         transform_key = stage +  '_transform'
         if(stage == "valid" and transform_key not in self.config['dataset']):
@@ -61,7 +61,7 @@ def get_stage_dataset_from_config(self, stage):
             data_transform = transforms.Compose(self.transform_list)
 
         csv_file = self.config['dataset'].get(stage + '_csv', None)
-        dataset  = NiftyDataset(root_dir=root_dir,
+        dataset  = NiftyDataset(root_dir  = root_dir,
                                 csv_file  = csv_file,
                                 modal_num = modal_num,
                                 with_label= not (stage == 'test'),
@@ -286,7 +286,7 @@ def train_valid(self):
             self.device = torch.device("cuda:{0:}".format(device_ids[0]))
         self.net.to(self.device)
         ckpt_dir    = self.config['training']['ckpt_save_dir']
-        ckpt_prefx = ckpt_dir.split('/')[-1]
+        ckpt_prefx  = ckpt_dir.split('/')[-1]
         iter_start  = self.config['training']['iter_start']
         iter_max    = self.config['training']['iter_max']
         iter_valid  = self.config['training']['iter_valid']
@@ -397,7 +397,7 @@ def test_time_dropout(m):
         infer_obj = Inferer(self.net, infer_cfg)
         infer_time_list = []
         with torch.no_grad():
-            for data in self.test_loder:
+            for data in self.test_loader:
                 images = self.convert_tensor_type(data['image'])
                 images = images.to(device)
 
@@ -444,7 +444,7 @@ def infer_with_multiple_checkpoints(self):
         infer_obj = Inferer(self.net, infer_cfg)
         infer_time_list = []
         with torch.no_grad():
-            for data in self.test_loder:
+            for data in self.test_loader:
                 images = self.convert_tensor_type(data['image'])
                 images = images.to(device)
 

pymic/net_run/infer_func.py

@@ -137,22 +137,22 @@ def run(self, image):
             outputs = self.__infer(image)
         elif(tta_mode == 1): # test time augmentation with flip in 2D
             outputs1 = self.__infer(image)
-            outputs2 = self.__infer(torch.flip(image, [-1]))
-            outputs3 = self.__infer(torch.flip(image, [-2]))
-            outputs4 = self.__infer(torch.flip(image, [-2, -1]))
+            outputs2 = self.__infer(torch.flip(image, [-2]))
+            outputs3 = self.__infer(torch.flip(image, [-3]))
+            outputs4 = self.__infer(torch.flip(image, [-2, -3]))
             if(isinstance(outputs1, (tuple, list))):
                 outputs = []
                 for i in range(len(outputs)):
                     temp_out1 = outputs1[i]
-                    temp_out2 = torch.flip(outputs2[i], [-1])
-                    temp_out3 = torch.flip(outputs3[i], [-2])
-                    temp_out4 = torch.flip(outputs4[i], [-2, -1])
+                    temp_out2 = torch.flip(outputs2[i], [-2])
+                    temp_out3 = torch.flip(outputs3[i], [-3])
+                    temp_out4 = torch.flip(outputs4[i], [-2, -3])
                     temp_mean = (temp_out1 + temp_out2 + temp_out3 + temp_out4) / 4
                     outputs.append(temp_mean)
             else:
-                outputs2 = torch.flip(outputs2, [-1])
-                outputs3 = torch.flip(outputs3, [-2])
-                outputs4 = torch.flip(outputs4, [-2, -1])
+                outputs2 = torch.flip(outputs2, [-2])
+                outputs3 = torch.flip(outputs3, [-3])
+                outputs4 = torch.flip(outputs4, [-2, -3])
                 outputs = (outputs1 + outputs2 + outputs3 + outputs4) / 4
         else:
             raise ValueError("Undefined tta_mode {0:}".format(tta_mode))