Merge branch 'dev'

Author: taigw

.gitignore

@@ -1 +1,6 @@
 *.pyc
+build/*
+dist/*
+*egg*/*
+*stop*
+files.txt

pymic/io/h5_dataset.py

@@ -0,0 +1,100 @@
+# -*- coding: utf-8 -*-
+
+from logging import root
+import os
+from re import S 
+import torch
+import random
+import h5py 
+import pandas as pd
+from scipy import ndimage
+from torch.utils.data import Dataset
+from torch.utils.data.sampler import Sampler
+
+class H5DataSets(Dataset):
+    """
+    Dataset for loading images stored in h5 format. It generates 
+    4D tensors with dimention order [C, D, H, W] for 3D images, and 
+    3D tensors with dimention order [C, H, W] for 2D images
+    """
+    def __init__(self, root_dir, sample_list_name, transform = None):
+        self.root_dir = root_dir 
+        self.transform  = transform
+        with open(sample_list_name, 'r') as f:
+            lines = f.readlines()
+            self.sample_list = [item.replace('\n', '') for item in lines]
+
+    def __len__(self):
+        return len(self.sample_list)
+
+    def __getitem__(self, idx):
+        sample_name = self.sample_list[idx]
+        h5f = h5py.File(self.root_dir + '/' +  sample_name, 'r')
+        image = h5f['image'][:]
+        label = h5f['label'][:]
+        sample = {'image': image, 'label': label}
+        if self.transform:
+            sample = self.transform(sample)
+        # sample["idx"] = idx
+        return sample
+    
+class TwoStreamBatchSampler(Sampler):
+    """Iterate two sets of indices
+
+    An 'epoch' is one iteration through the primary indices.
+    During the epoch, the secondary indices are iterated through
+    as many times as needed.
+    """
+
+    def __init__(self, primary_indices, secondary_indices, batch_size, secondary_batch_size):
+        self.primary_indices = primary_indices
+        self.secondary_indices = secondary_indices
+        self.secondary_batch_size = secondary_batch_size
+        self.primary_batch_size = batch_size - secondary_batch_size
+
+        assert len(self.primary_indices) >= self.primary_batch_size > 0
+        assert len(self.secondary_indices) >= self.secondary_batch_size > 0
+
+    def __iter__(self):
+        primary_iter = iterate_once(self.primary_indices)
+        secondary_iter = iterate_eternally(self.secondary_indices)
+        return (
+            primary_batch + secondary_batch
+            for (primary_batch, secondary_batch)
+            in zip(grouper(primary_iter, self.primary_batch_size),
+                   grouper(secondary_iter, self.secondary_batch_size))
+        )
+
+    def __len__(self):
+        return len(self.primary_indices) // self.primary_batch_size
+
+
+def iterate_once(iterable):
+    return np.random.permutation(iterable)
+
+
+def iterate_eternally(indices):
+    def infinite_shuffles():
+        while True:
+            yield np.random.permutation(indices)
+    return itertools.chain.from_iterable(infinite_shuffles())
+
+
+def grouper(iterable, n):
+    "Collect data into fixed-length chunks or blocks"
+    # grouper('ABCDEFG', 3) --> ABC DEF"
+    args = [iter(iterable)] * n
+    return zip(*args)
+
+
+if __name__ == "__main__":
+    root_dir = "/home/guotai/disk2t/projects/semi_supervise/SSL4MIS/data/ACDC/data/slices"
+    file_name = "/home/guotai/disk2t/projects/semi_supervise/slices.txt"
+    dataset = H5DataSets(root_dir, file_name)
+    train_loader = torch.utils.data.DataLoader(dataset, 
+                batch_size = 4, shuffle=True, num_workers= 1)
+    for sample in train_loader:
+        image = sample['image']
+        label = sample['label']
+        print(image.shape, label.shape)
+        print(image.min(), image.max(), label.max())

pymic/io/image_read_write.py

@@ -4,10 +4,8 @@
 import os
 import numpy as np
 import SimpleITK as sitk
-
 from PIL import Image
 
-
 def load_nifty_volume_as_4d_array(filename):
     """Read a nifty image and return a dictionay storing data array, spacing and direction
     output['data_array'] 4d array with shape [C, D, H, W]

pymic/io/nifty_dataset.py

@@ -5,7 +5,6 @@
 import torch
 import pandas as pd
 import numpy as np
-
 from torch.utils.data import Dataset, DataLoader
 from torchvision import transforms, utils
 from pymic.io.image_read_write import load_image_as_nd_array
@@ -95,7 +94,6 @@ def __init__(self, root_dir, csv_file, modal_num = 1, class_num = 2,
         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):
         csv_keys = list(self.csv_items.keys())
@@ -129,6 +127,4 @@ def __getitem__(self, idx):
             sample['image_weight'] = self.__getweight__(idx) 
         if self.transform:
             sample = self.transform(sample)
-
         return sample
-

pymic/loss/cls/ce.py

@@ -29,15 +29,13 @@ class SigmoidCELoss(nn.Module):
     Args:
     predict has a shape of [N, C] where C is the class number
     labels  has a shape of [N, C] with binary values
-    
     note that predict is the digit output of a network, before using sigmoid."""
     def __init__(self, params):
         super(SigmoidCELoss, self).__init__()
 
     def forward(self, loss_input_dict):
         predict = loss_input_dict['prediction']
         labels  = loss_input_dict['ground_truth']
-        # for numeric stability
         predict = nn.Sigmoid()(predict) * 0.999 + 5e-4
         loss = - labels * torch.log(predict) - (1 - labels) * torch.log( 1 - predict)
         loss = loss.mean()

pymic/loss/cls/l1.py

@@ -6,6 +6,9 @@
 from pymic.loss.cls.util import get_soft_label
 
 class L1Loss(nn.Module):
+    """
+    L1 (MAE) loss for classification
+    """
     def __init__(self, params):
         super(L1Loss, self).__init__()
         self.l1_loss = nn.L1Loss()
@@ -20,3 +23,30 @@ def forward(self, loss_input_dict):
         soft_y = get_soft_label(labels, num_class, data_type)
         loss = self.l1_loss(predict, soft_y)
         return loss
+
+class RectifiedLoss(nn.Module):
+    def __init__(self, params):
+        super(RectifiedLoss, self).__init__()
+        # self.l1_loss = nn.L1Loss()
+    
+    def forward(self, loss_input_dict):
+        predict = loss_input_dict['prediction']
+        labels  = loss_input_dict['ground_truth'][:, None] # reshape to N, 1
+        
+        # softmax = nn.Softmax(dim = 1)
+        # predict = softmax(predict)
+        num_class  = list(predict.size())[1]
+        data_type = 'float' if(predict.dtype is torch.float32) else 'double'
+        soft_y = get_soft_label(labels, num_class, data_type)
+        g = 2* soft_y - 1
+        loss = torch.exp((g*1.5- predict) * g)
+        mask = predict < g 
+        if (data_type  == 'float'):
+            mask = mask.float() 
+        else:
+            mask = mask.double() 
+        w = (mask - 0.5) * g + 0.5    
+        loss = w * loss + 0.1*(g - predict) * (g - predict)
+        loss = loss.mean()
+        # loss = self.l1_loss(predict, soft_y)
+        return loss

pymic/loss/loss_dict_seg.py

@@ -2,15 +2,15 @@
 from __future__ import print_function, division
 import torch.nn as nn 
 from pymic.loss.seg.ce import CrossEntropyLoss, GeneralizedCrossEntropyLoss
-from pymic.loss.seg.dice import DiceLoss, MultiScaleDiceLoss
-from pymic.loss.seg.dice import FocalDiceLoss, NoiseRobustDiceLoss
+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,
-    'MultiScaleDiceLoss': MultiScaleDiceLoss,
     'FocalDiceLoss': FocalDiceLoss,
     'NoiseRobustDiceLoss': NoiseRobustDiceLoss,
     'ExpLogLoss': ExpLogLoss,

pymic/loss/seg/ce.py

@@ -8,40 +8,55 @@
 class CrossEntropyLoss(nn.Module):
     def __init__(self, params):
         super(CrossEntropyLoss, self).__init__()
-        self.enable_pix_weight = params.get('CrossEntropyLoss_Enable_Pixel_Weight'.lower(), False)
-        self.enable_cls_weight = params.get('CrossEntropyLoss_Enable_Class_Weight'.lower(), False)
+        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']
-        pix_w   = loss_input_dict['pixel_weight']
-        cls_w   = loss_input_dict['class_weight']
-        softmax = loss_input_dict['softmax']
+        pix_w   = loss_input_dict.get('pixel_weight', None)
 
         if(isinstance(predict, (list, tuple))):
             predict = predict[0]
-        if(softmax):
+        if(self.softmax):
             predict = nn.Softmax(dim = 1)(predict)
         predict = reshape_tensor_to_2D(predict)
         soft_y  = reshape_tensor_to_2D(soft_y)
 
         # for numeric stability
         predict = predict * 0.999 + 5e-4
         ce = - soft_y* torch.log(predict)
-        if(self.enable_cls_weight):
-            if(cls_w is None):
-                raise ValueError("Class weight is enabled but not defined")
-            ce = torch.sum(ce * cls_w, dim = 1)
-        else:
-            ce = torch.sum(ce, dim = 1) # shape is [N]
-        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) # shape is [N, 1]
-            pix_w = torch.squeeze(pix_w)        # squeeze to [N]
-            ce    = torch.sum(ce * pix_w) / torch.sum(pix_w)
-        else:
+        ce = torch.sum(ce, dim = 1) # shape is [N]
+        if(pix_w is None):
             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) 
+        return ce
+
+class PartialCrossEntropyLoss(nn.Module):
+    def __init__(self, params):
+        super(CrossEntropyLoss, self).__init__()
+        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)
+
+        # for numeric stability
+        predict = predict * 0.999 + 5e-4
+        ce = - soft_y* torch.log(predict)
+        ce = torch.sum(ce, dim = 1) # shape is [N]
+        ce = torch.mean(ce)  
         return ce
 
 class GeneralizedCrossEntropyLoss(nn.Module):
@@ -85,4 +100,4 @@ def forward(self, loss_input_dict):
             gce    = torch.sum(gce * pix_w) / torch.sum(pix_w)
         else:
             gce = torch.mean(gce)
-        return gce 
+        return gce 

pymic/loss/seg/deep_sup.py

@@ -0,0 +1,29 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+
+import torch.nn as nn
+
+class DeepSuperviseLoss(nn.Module):
+    def __init__(self, params):
+        super(DeepSuperviseLoss, self).__init__()
+        self.deep_sup_weight = params.get('deep_suervise_weight', None)
+        self.base_loss = params['base_loss']
+
+    def forward(self, loss_input_dict):
+        predict = loss_input_dict['prediction']
+        if(not isinstance(predict, (list,tuple))):
+            raise ValueError("""For deep supervision, the prediction should
+                be a list or a tuple""")
+        predict_num = len(predict)
+        if(self.deep_sup_weight is None):
+            self.deep_sup_weight = [1.0] * predict_num
+        else:
+            assert(predict_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]
+            temp_loss   = self.base_loss(loss_input_dict)
+            loss_sum   += temp_loss * self.deep_sup_weight[i]
+            weight_sum += self.deep_sup_weight[i]
+        loss = loss_sum/weight_sum
+        return loss