update transform and image process functions

Author: taigw

pymic/io/transform3d.py

@@ -3,6 +3,7 @@
 
 import torch
 import json
+import math
 import random
 import numpy as np
 
@@ -59,7 +60,11 @@ def inverse_transform_for_prediction(self, sample):
          different elemenets in the batch.
 
         origin_shape is a 4D or 3D vector as saved in __call__().'''
-        origin_shape = json.loads(sample['Rescale_origin_shape'][0])
+        if(isinstance(sample['Rescale_origin_shape'], list) or \
+            isinstance(sample['Rescale_origin_shape'], tuple)):
+            origin_shape = json.loads(sample['Rescale_origin_shape'][0])
+        else:
+            origin_shape = json.loads(sample['Rescale_origin_shape'])
         origin_dim   = len(origin_shape) - 1
         predict = sample['predict']
         input_shape = predict.shape
@@ -116,7 +121,11 @@ def  inverse_transform_for_prediction(self, sample):
          different elemenets in the batch.
 
         flip_axis is a list as saved in __call__().'''
-        flip_axis = json.loads(sample['RandomFlip_Param'][0]) 
+        if(isinstance(sample['RandomFlip_Param'], list) or \
+            isinstance(sample['RandomFlip_Param'], tuple)):
+            flip_axis = json.loads(sample['RandomFlip_Param'][0]) 
+        else:
+            flip_axis = json.loads(sample['RandomFlip_Param']) 
         if(len(flip_axis) > 0):
             sample['predict']  = np.flip(sample['predict'] , flip_axis)
         return sample
@@ -180,7 +189,11 @@ def  inverse_transform_for_prediction(self, sample):
 
         transform_param_list is a list as saved in __call__().'''
         # get the paramters for invers transformation
-        transform_param_list = json.loads(sample['RandomRotate_Param'][0]) 
+        if(isinstance(sample['RandomRotate_Param'], list) or \
+            isinstance(sample['RandomRotate_Param'], tuple)):
+            transform_param_list = json.loads(sample['RandomRotate_Param'][0]) 
+        else:
+            transform_param_list = json.loads(sample['RandomRotate_Param']) 
         transform_param_list.reverse()
         for i in range(len(transform_param_list)):
             transform_param_list[i][0] = - transform_param_list[i][0]
@@ -196,8 +209,9 @@ class Pad(object):
        output_size (tuple/list): the size along each spatial axis. 
        
     """
-    def __init__(self, output_size, inverse = True):
+    def __init__(self, output_size, ceil_mode = False, inverse = True):
         self.output_size = output_size
+        self.ceil_mode   = ceil_mode
         self.inverse = inverse
 
 
@@ -206,7 +220,14 @@ def __call__(self, sample):
         input_shape = image.shape
         input_dim = len(input_shape) - 1
         assert(len(self.output_size) == input_dim)
-        margin = [max(0, self.output_size[i] - input_shape[1+i]) \
+        if(self.ceil_mode):
+            multiple = [int(math.ceil(float(input_shape[1+i])/self.output_size[i]))\
+                for i in range(input_dim)]
+            output_size = [multiple[i] * self.output_size[i] \
+                for i in range(input_dim)]
+        else:
+            output_size = self.output_size
+        margin = [max(0, output_size[i] - input_shape[1+i]) \
             for i in range(input_dim)]
 
         margin_lower = [int(margin[i] / 2) for i in range(input_dim)]
@@ -233,7 +254,10 @@ def inverse_transform_for_prediction(self, sample):
 
         origin_shape is a 4D or 3D vector as saved in __call__().'''
         # raise ValueError("not implemented")
-        params = json.loads(sample['Pad_Param'][0]) 
+        if(isinstance(sample['Pad_Param'], list) or isinstance(sample['Pad_Param'], tuple)):
+            params = json.loads(sample['Pad_Param'][0]) 
+        else:
+            params = json.loads(sample['Pad_Param']) 
         margin_lower = params[0]
         margin_upper = params[1]
         predict = sample['predict']
@@ -310,7 +334,11 @@ def inverse_transform_for_prediction(self, sample):
          different elemenets in the batch.
 
         origin_shape is a 4D or 3D vector as saved in __call__().'''
-        params = json.loads(sample['CropWithBoundingBox_Param'][0]) 
+        if(isinstance(sample['CropWithBoundingBox_Param'], list) or \
+            isinstance(sample['CropWithBoundingBox_Param'], tuple)):
+            params = json.loads(sample['CropWithBoundingBox_Param'][0]) 
+        else:
+            params = json.loads(sample['CropWithBoundingBox_Param']) 
         origin_shape = params[0]
         crop_min     = params[1]
         crop_max     = params[2]
@@ -333,10 +361,13 @@ class RandomCrop(object):
             the output channel is the same as the input channel.
     """
 
-    def __init__(self, output_size, inverse = True):
+    def __init__(self, output_size, fg_focus = False, fg_ratio = 0.0, mask_label = None,  inverse = True):
         assert isinstance(output_size, (list, tuple))
         self.output_size = output_size
-        self.inverse = inverse
+        self.inverse  = inverse
+        self.fg_focus = fg_focus
+        self.fg_ratio = fg_ratio
+        self.mask_label = mask_label
 
     def __call__(self, sample):
         image = sample['image']
@@ -347,6 +378,19 @@ def __call__(self, sample):
         crop_margin = [input_shape[i + 1] - self.output_size[i]\
             for i in range(input_dim)]
         crop_min = [random.randint(0, item) for item in crop_margin]
+        if(self.fg_focus and random.random() < self.fg_ratio):
+            label = sample['label']
+            mask  = np.zeros_like(label)
+            for temp_lab in self.mask_label:
+                mask = np.maximum(mask, label == temp_lab)
+            bb_min, bb_max = get_ND_bounding_box(mask)
+            bb_min, bb_max = bb_min[1:], bb_max[1:]
+            crop_min = [random.randint(bb_min[i], bb_max[i]) - int(self.output_size[i]/2) \
+                for i in range(input_dim)]
+            crop_min = [max(0, item) for item in crop_min]
+            crop_min = [min(crop_min[i], input_shape[i+1] - self.output_size[i]) \
+                for i in range(input_dim)]
+
         crop_max = [crop_min[i] + self.output_size[i] \
             for i in range(input_dim)]
         crop_min = [0] + crop_min
@@ -368,7 +412,11 @@ def inverse_transform_for_prediction(self, sample):
          different elemenets in the batch.
 
         origin_shape is a 4D or 3D vector as saved in __call__().'''
-        params = json.loads(sample['RandomCrop_Param'][0]) 
+        if(isinstance(sample['RandomCrop_Param'], list) or \
+            isinstance(sample['RandomCrop_Param'], tuple)):
+            params = json.loads(sample['RandomCrop_Param'][0]) 
+        else:
+            params = json.loads(sample['RandomCrop_Param']) 
         origin_shape = params[0]
         crop_min     = params[1]
         crop_max     = params[2]
@@ -416,27 +464,40 @@ class ChannelWiseNormalize(object):
         mean (None or tuple/list): The mean values along each channel.
         std  (None or tuple/list): The std values along each channel.
             if mean and std are None, calculate them from non-zero region
-        zero_to_random (bool): If true, replace zero values with random values.
+        chns (None, or tuple/list): The list of channel indices
+        zero_to_random (bool, or tuple/list or bool): indicate whether zero values
+             in each channel is replaced  with random values.
     """
-    def __init__(self, mean, std, zero_to_random = False, inverse = False):
+    def __init__(self, mean, std, chns = None, zero_to_random = False, inverse = False):
         self.mean = mean
         self.std  = std
+        self.chns = chns
         self.zero_to_random = zero_to_random
         self.inverse = inverse
 
     def __call__(self, sample):
         image= sample['image']
         mask = image[0] > 0
-        for chn in range(image.shape[0]):
+        chns = self.chns
+        if(chns is None):
+            chns = range(image.shape[0])
+        zero_to_random = self.zero_to_random
+        if(isinstance(zero_to_random, bool)):
+            zero_to_random = [zero_to_random]*len(chns)
+        if(not(self.mean is None and self.std is None)):
+            assert(len(self.mean) == len(self.std))
+            assert(len(self.mean) == len(chns))
+        for i in range(len(chns)):
+            chn = chns[i]
             if(self.mean is None and self.std is None):
                 pixels = image[chn][mask > 0]
                 chn_mean = pixels.mean()
                 chn_std  = pixels.std()
             else:
-                chn_mean = self.mean[chn]
-                chn_std  = self.std[chn]
+                chn_mean = self.mean[i]
+                chn_std  = self.std[i]
             chn_norm = (image[chn] - chn_mean)/chn_std
-            if(self.zero_to_random):
+            if(zero_to_random[i]):
                 chn_random = np.random.normal(0, 1, size = chn_norm.shape)
                 chn_norm[mask == 0] = chn_random[mask == 0]
             image[chn] = chn_norm
@@ -490,6 +551,49 @@ def __call__(self, sample):
     def inverse_transform_for_prediction(self, sample):
         raise(ValueError("not implemented"))
 
+class LabelToProbability(object):
+    """
+        Convert one-channel label map to multi-channel probability map
+    Args:
+        class_num (int): the class number in the label map
+    """
+    def __init__(self, class_num, inverse = False):
+        self.class_num = class_num
+        self.inverse   = inverse
+    
+    def __call__(self, sample):
+        label = sample['label'][0]
+        label_prob = []
+        for i in range(self.class_num):
+            temp_prob = label == i*np.ones_like(label)
+            label_prob.append(temp_prob)
+        label_prob = np.asarray(label_prob, np.float32)
+   
+        sample['label_prob'] = label_prob
+        return sample
+    
+    def inverse_transform_for_prediction(self, sample):
+        raise(ValueError("not implemented"))
+
+class ProbabilityToDistance(object):
+    """
+     get distance transform for each label
+    """
+    def __init__(self, inverse = False):
+        self.inverse = inverse
+
+    
+    def __call__(self, sample):
+        label_prob = sample['label_prob']
+        label_distance = []
+        for i in range(label_prob.shape[0]):
+            temp_lab = label_prob[i]
+            temp_dis = get_euclidean_distance(temp_lab, dim = 3, spacing = [1.0, 1.0, 1.0])
+            label_distance.append(temp_dis)
+        label_distance = np.asarray(label_distance)
+        sample['label_distance'] = label_distance
+        return sample
+
 class RegionSwop(object):
     """
     Swop a subregion randomly between two images and their corresponding label
@@ -567,8 +671,9 @@ def get_transform(name, params):
 
     elif(name == "Pad"):
         output_size = params["Pad_output_size".lower()]
+        ceil_mode   = params["Pad_ceil_mode".lower()]
         inverse = params["Pad_inverse".lower()]
-        return Pad(output_size, inverse)
+        return Pad(output_size, ceil_mode, inverse)
 
     elif(name == "ChannelWiseGammaCorrection"):
         gamma_min = params['ChannelWiseGammaCorrection_gamma_min'.lower()]
@@ -577,11 +682,12 @@ def get_transform(name, params):
         return ChannelWiseGammaCorrection(gamma_min, gamma_max, inverse)
 
     elif (name == 'ChannelWiseNormalize'):
+        chns = params['ChannelWiseNormalize_channels'.lower()]
         mean = params['ChannelWiseNormalize_mean'.lower()]
         std  = params['ChannelWiseNormalize_std'.lower()]
         zero_to_random = params['ChannelWiseNormalize_zero_to_random'.lower()]
         inverse = params['ChannelWiseNormalize_inverse'.lower()]
-        return ChannelWiseNormalize(mean, std, zero_to_random, inverse)
+        return ChannelWiseNormalize(mean, std, chns, zero_to_random, inverse)
 
     elif(name == 'ChannelWiseThreshold'):
         threshold = params['ChannelWiseThreshold_threshold'.lower()]
@@ -594,10 +700,22 @@ def get_transform(name, params):
         inverse = params['LabelConvert_inverse'.lower()]
         return LabelConvert(source_list, target_list, inverse)
 
+    elif(name == 'LabelToProbability'):
+        class_num = params['LabelToProbability_class_num'.lower()]
+        inverse   = params['LabelToProbability_inverse'.lower()]
+        return LabelToProbability(class_num, inverse)
+
+    elif(name == 'ProbabilityToDistance'):
+        inverse   = params['ProbabilityToDistance_inverse'.lower()]
+        return ProbabilityToDistance(inverse)
+
     elif(name == 'RandomCrop'):
         output_size = params['RandomCrop_output_size'.lower()]
+        fg_focus    = params['RandomCrop_foreground_focus'.lower()]
+        fg_ratio    = params['RandomCrop_foreground_ratio'.lower()]
+        mask_label  = params['RandomCrop_mask_label'.lower()]
         inverse     = params['RandomCrop_inverse'.lower()]
-        return RandomCrop(output_size, inverse)
+        return RandomCrop(output_size, fg_focus, fg_ratio, mask_label,  inverse)
 
     elif(name == 'RegionSwop'):
         spatial_axes = params['RegionSwop_spatial_axes'.lower()]

pymic/layer/convolution.py

@@ -38,3 +38,69 @@ def forward(self, x):
         if(self.acti_func is not None):
             f = self.acti_func(f)
         return f
+
+class ConvolutionSepAll3DLayer(nn.Module):
+    """
+    A compose layer with the following components:
+    convolution -> (batch_norm) -> activation -> (dropout)
+    batch norm and dropout are optional
+    """
+    def __init__(self, in_channels, out_channels, kernel_size, dim = 3,
+            stride = 1, padding = 0, dilation =1, groups = 1, bias = True, 
+            batch_norm = True, acti_func = None):
+        super(ConvolutionSepAll3DLayer, self).__init__()
+        self.n_in_chns  = in_channels
+        self.n_out_chns = out_channels
+        self.batch_norm = batch_norm
+        self.acti_func  = acti_func
+
+        assert(dim == 3)
+        chn = min(in_channels, out_channels)
+
+        self.conv_intra_plane1 = nn.Conv2d(chn, chn,
+            kernel_size, stride, padding, dilation, chn, bias)
+
+        self.conv_intra_plane2 = nn.Conv2d(chn, chn,
+            kernel_size, stride, padding, dilation, chn, bias)
+
+        self.conv_intra_plane3 = nn.Conv2d(chn, chn,
+            kernel_size, stride, padding, dilation, chn, bias)
+
+        self.conv_space_wise  = nn.Conv2d(in_channels, out_channels,
+            1, stride, 0, dilation, 1, bias)
+        
+        if(self.batch_norm):
+                self.bn = nn.modules.BatchNorm3d(out_channels)
+
+    def forward(self, x):
+        in_shape = list(x.shape)
+        assert(len(in_shape) == 5)
+        [B, C, D, H, W] = in_shape
+        f0 = x.permute(0, 2, 1, 3, 4)  #[B, D, C, H, W] 
+        f0 = f0.contiguous().view([B*D, C, H, W])
+
+        Cc = min(self.n_in_chns, self.n_out_chns)
+        Co = self.n_out_chns
+        if(self.n_in_chns > self.n_out_chns):
+            f0 = self.conv_space_wise(f0)  #[B*D, Cc, H, W]  
+   
+        f1 = self.conv_intra_plane1(f0)
+        f2 = f1.contiguous().view([B, D, Cc, H, W])
+        f2 = f2.permute(0, 3, 2, 1, 4) #[B, H, Cc, D, W]
+        f2 = f2.contiguous().view([B*H, Cc, D, W])
+        f2 = self.conv_intra_plane2(f2)
+        f3 = f2.contiguous().view([B, H, Cc, D, W])
+        f3 = f3.permute(0, 4, 2, 3, 1) #[B, W, Cc, D, H]
+        f3 = f3.contiguous().view([B*W, Cc, D, H])
+        f3 = self.conv_intra_plane3(f3)
+        if(self.n_in_chns <= self.n_out_chns):
+            f3 = self.conv_space_wise(f3)  #[B*W, Co, D, H] 
+
+        f3 = f3.contiguous().view([B, W, Co, D, H])
+        f3 = f3.permute([0, 2, 3, 4, 1]) #[B, Co, D, H, W]
+
+        if(self.batch_norm):
+            f3 = self.bn(f3)
+        if(self.acti_func is not None):
+            f3 = self.acti_func(f3)
+        return f3

pymic/net2d/unet2d.py

@@ -81,6 +81,12 @@ def __init__(self, params):
             kernel_size = 3, padding = 1)
 
     def forward(self, x):
+        x_shape = list(x.shape)
+        if(len(x_shape) == 5):
+          [N, C, D, H, W] = x_shape
+          new_shape = [N*D, C, H, W]
+          x = torch.transpose(x, 1, 2)
+          x = torch.reshape(x, new_shape)
         f1 = self.block1(x);  d1 = self.down1(f1)
         f2 = self.block2(d1); d2 = self.down2(f2)
         f3 = self.block3(d2); d3 = self.down3(f3)
@@ -104,6 +110,10 @@ def forward(self, x):
         f9    = self.block9(f1cat)
 
         output = self.conv(f9)
+        if(len(x_shape) == 5):
+            new_shape = [N, D] + list(output.shape)[1:]
+            output = torch.reshape(output, new_shape)
+            output = torch.transpose(output, 1, 2)
         return output
 
 if __name__ == "__main__":
@@ -114,7 +124,7 @@ def forward(self, x):
     Net = UNet2D(params)
     Net = Net.double()
 
-    x  = np.random.rand(4, 4, 96, 96)
+    x  = np.random.rand(4, 4, 10, 96, 96)
     xt = torch.from_numpy(x)
     xt = torch.tensor(xt)