update weakly supervised leraning

add Mumford Shah loss

Author: taigw

pymic/loss/seg/mumford_shah.py

@@ -0,0 +1,95 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+
+import torch
+import torch.nn as nn
+
+class DiceLoss(nn.Module):
+    def __init__(self, params = None):
+        super(DiceLoss, self).__init__()
+        if(params is None):
+            self.softmax = True
+        else:
+            self.softmax = params.get('loss_softmax', True)
+
+    def forward(self, loss_input_dict):
+        predict = loss_input_dict['prediction']
+        soft_y  = loss_input_dict['ground_truth']
+        
+        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) 
+        dice_score = get_classwise_dice(predict, soft_y)
+        dice_loss  = 1.0 - dice_score.mean()
+        return dice_loss
+
+class MumfordShahLoss(nn.Module):
+    """
+    Implementation of Mumford Shah Loss in this paper:
+        Boah Kim and Jong Chul Ye, Mumford–Shah Loss Functional 
+        for Image Segmentation With Deep Learning. IEEE TIP, 2019.
+    The oringial implementation is availabel at:
+    https://github.com/jongcye/CNN_MumfordShah_Loss 
+    
+    currently only 2D version is supported.
+    """
+    def __init__(self, params = None):
+        super(MumfordShahLoss, self).__init__()
+        if(params is None):
+            params = {}
+        self.softmax = params.get('loss_softmax', True)
+        self.penalty = params.get('MumfordShahLoss_penalty', "l1")
+        self.grad_w  = params.get('MumfordShahLoss_lambda', 1.0)
+
+    def get_levelset_loss(self, output, target):
+        """
+        output: softmax output of a network
+        target: the input image
+        """
+        outshape = output.shape
+        tarshape = target.shape
+        loss = 0.0
+        for ich in range(tarshape[1]):
+            target_ = torch.unsqueeze(target[:,ich], 1)
+            target_ = target_.expand(tarshape[0], outshape[1], tarshape[2], tarshape[3])
+            pcentroid = torch.sum(target_ * output, (2,3))/torch.sum(output, (2,3))
+            pcentroid = pcentroid.view(tarshape[0], outshape[1], 1, 1)
+            plevel = target_ - pcentroid.expand(tarshape[0], outshape[1], tarshape[2], tarshape[3])
+            pLoss = plevel * plevel * output
+            loss += torch.sum(pLoss)
+        return loss
+
+    def get_gradient_loss(self, pred, penalty = "l2"):
+        dH = torch.abs(pred[:, :, 1:, :] - pred[:, :, :-1, :])
+        dW = torch.abs(pred[:, :, :, 1:] - pred[:, :, :, :-1])
+        if penalty == "l2":
+            dH = dH * dH
+            dW = dW * dW
+        loss = torch.sum(dH) + torch.sum(dW)
+        return loss
+
+    def forward(self, loss_input_dict):
+        predict = loss_input_dict['prediction']
+        image   = loss_input_dict['image']
+        if(isinstance(predict, (list, tuple))):
+            predict = predict[0]
+        if(self.softmax):
+            predict = nn.Softmax(dim = 1)(predict) 
+
+        pred_shape  = list(predict.shape)
+        if(len(pred_shape) == 5):
+                [N, C, D, H, W] = pred_shape
+                new_shape  = [N*D, C, H, W]
+                predict = torch.transpose(predict, 1, 2)
+                predict = torch.reshape(predict, new_shape)
+                [N, C, D, H, W] = list(image.shape)
+                new_shape    = [N*D, C, H, W]
+                image = torch.transpose(image, 1, 2)
+                image = torch.reshape(image, new_shape)
+        loss0 = self.get_levelset_loss(predict, image)
+        loss1 = self.get_gradient_loss(predict, self.penalty)
+        loss = loss0 + self.grad_w * loss1
+        return loss/torch.numel(predict)

pymic/net_run_wsl/wsl_main.py

@@ -7,12 +7,14 @@
 from pymic.util.parse_config import *
 from pymic.net_run_wsl.wsl_em import WSL_EntropyMinimization
 from pymic.net_run_wsl.wsl_gatedcrf import WSL_GatedCRF
+from pymic.net_run_wsl.wsl_mumford_shah import WSL_MumfordShah
 from pymic.net_run_wsl.wsl_tv import WSL_TotalVariation
 from pymic.net_run_wsl.wsl_ustm import WSL_USTM
 from pymic.net_run_wsl.wsl_dmpls import WSL_DMPLS
 
 WSLMethodDict = {'EntropyMinimization': WSL_EntropyMinimization,
     'GatedCRF': WSL_GatedCRF,
+    'MumfordShah': WSL_MumfordShah,
     'TotalVariation': WSL_TotalVariation,
     'USTM': WSL_USTM,
     'DMPLS': WSL_DMPLS}

pymic/net_run_wsl/wsl_mumford_shah.py

@@ -0,0 +1,91 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+import logging
+import numpy as np
+import random
+import torch
+import torchvision.transforms as transforms
+from pymic.io.nifty_dataset import NiftyDataset
+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
+from pymic.loss.seg.mumford_shah import MumfordShahLoss
+from pymic.net_run.agent_seg import SegmentationAgent
+from pymic.net_run_wsl.wsl_em import WSL_EntropyMinimization
+from pymic.util.ramps import sigmoid_rampup
+
+class WSL_MumfordShah(WSL_EntropyMinimization):
+    """
+    Training and testing agent for semi-supervised segmentation
+    """
+    def __init__(self, config, stage = 'train'):
+        super(WSL_MumfordShah, self).__init__(config, stage)
+
+    def training(self):
+        class_num   = self.config['network']['class_num']
+        iter_valid  = self.config['training']['iter_valid']
+        wsl_cfg     = self.config['weakly_supervised_learning']
+        train_loss  = 0
+        train_loss_sup = 0
+        train_loss_reg = 0
+        train_dice_list = []
+
+        reg_loss_calculator = MumfordShahLoss(wsl_cfg)
+        self.net.train()
+        for it in range(iter_valid):
+            try:
+                data = next(self.trainIter)
+            except StopIteration:
+                self.trainIter = iter(self.train_loader)
+                data = next(self.trainIter)
+            
+            # get the inputs
+            inputs = self.convert_tensor_type(data['image'])
+            y      = self.convert_tensor_type(data['label_prob'])  
+                         
+            inputs, y = inputs.to(self.device), y.to(self.device)
+            
+            # zero the parameter gradients
+            self.optimizer.zero_grad()
+                
+            # forward + backward + optimize
+            outputs = self.net(inputs)
+            loss_sup = self.get_loss_value(data, outputs, y)
+            loss_dict = {"prediction":outputs, 'image':inputs}
+            loss_reg = reg_loss_calculator(loss_dict)
+            
+            iter_max = self.config['training']['iter_max']
+            ramp_up_length = wsl_cfg.get('ramp_up_length', iter_max)
+            regular_w = 0.0
+            if(self.glob_it > wsl_cfg.get('iter_sup', 0)):
+                regular_w = wsl_cfg.get('regularize_w', 0.1)
+                if(ramp_up_length is not None and self.glob_it < ramp_up_length):
+                    regular_w = regular_w * sigmoid_rampup(self.glob_it, ramp_up_length)
+            loss = loss_sup + regular_w*loss_reg
+            # if (self.config['training']['use'])
+            loss.backward()
+            self.optimizer.step()
+            self.scheduler.step()
+
+            train_loss = train_loss + loss.item()
+            train_loss_sup = train_loss_sup + loss_sup.item()
+            train_loss_reg = train_loss_reg + loss_reg.item() 
+            # get dice evaluation for each class in annotated images
+            if(isinstance(outputs, tuple) or isinstance(outputs, list)):
+                outputs = outputs[0] 
+            p_argmax = torch.argmax(outputs, dim = 1, keepdim = True)
+            p_soft   = get_soft_label(p_argmax, class_num, self.tensor_type)
+            p_soft, y = reshape_prediction_and_ground_truth(p_soft, y) 
+            dice_list   = get_classwise_dice(p_soft, y)
+            train_dice_list.append(dice_list.cpu().numpy())
+        train_avg_loss = train_loss / iter_valid
+        train_avg_loss_sup = train_loss_sup / iter_valid
+        train_avg_loss_reg = train_loss_reg / iter_valid
+        train_cls_dice = np.asarray(train_dice_list).mean(axis = 0)
+        train_avg_dice = train_cls_dice.mean()
+
+        train_scalers = {'loss': train_avg_loss, 'loss_sup':train_avg_loss_sup,
+            'loss_reg':train_avg_loss_reg, 'regular_w':regular_w,
+            'avg_dice':train_avg_dice,     'class_dice': train_cls_dice}
+        return train_scalers
+        

pymic/util/evaluation_seg.py

@@ -222,6 +222,7 @@ def evaluation(config_file):
         for i in range(item_num):
             gt_name  = image_items.iloc[i, 0]
             seg_name = image_items.iloc[i, 1]
+            # seg_name = seg_name.replace(".nii.gz", "_pred.nii.gz")
             gt_full_name  = gt_root  + '/' + gt_name
             seg_full_name = seg_root_n + '/' + seg_name