update transform

Author: taigw

pymic/transform/abstract_transform.py

@@ -0,0 +1,12 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+
+class AbstractTransform(object):
+    def __init__(self, params):
+        pass
+
+    def __call__(self, sample):
+        return sample
+
+    def inverse_transform_for_prediction(self, sample):
+        raise(ValueError("not implemented"))

pymic/transform/crop.py

@@ -0,0 +1,200 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+
+import torch
+import json
+import math
+import random
+import numpy as np
+from scipy import ndimage
+from pymic.transform.abstract_transform import AbstractTransform
+from pymic.util.image_process import *
+
+
+class CropWithBoundingBox(AbstractTransform):
+    """Crop the image (shape [C, D, H, W] or [C, H, W]) based on bounding box
+    """
+    def __init__(self, params):
+        """
+        start (None or tuple/list): The start index along each spatial axis.
+            if None, calculate the start index automatically so that 
+            the cropped region is centered at the non-zero region.
+        output_size (None or tuple/list): Desired spatial output size.
+            if None, set it as the size of bounding box of non-zero region 
+        """
+        self.start       = params['CropWithBoundingBox_start'.lower()]
+        self.output_size = params['CropWithBoundingBox_output_size'.lower()]
+        self.inverse = params['CropWithBoundingBox_inverse'.lower()]
+        
+    def __call__(self, sample):
+        image = sample['image']
+        input_shape = image.shape
+        input_dim   = len(input_shape) - 1
+        bb_min, bb_max = get_ND_bounding_box(image)
+        bb_min, bb_max = bb_min[1:], bb_max[1:]
+        if(self.start is None):
+            if(self.output_size is None):
+                crop_min, crop_max = bb_min, bb_max
+            else:
+                assert(len(self.output_size) == input_dim)
+                crop_min = [int((bb_min[i] + bb_max[i] + 1)/2) - int(self.output_size[i]/2) \
+                    for i in range(input_dim)]
+                crop_min = [max(0, crop_min[i]) for i in range(input_dim)]
+                crop_max = [crop_min[i] + self.output_size[i] for i in range(input_dim)]
+        else:
+            assert(len(self.start) == input_dim)
+            crop_min = self.start
+            if(self.output_size is None):
+                assert(len(self.output_size) == input_dim)
+                crop_max = [crop_min[i] + bb_max[i] - bb_min[i] \
+                    for i in range(input_dim)]
+            else:
+                crop_max =  [crop_min[i] + self.output_size[i] for i in range(input_dim)]
+        crop_min = [0] + crop_min
+        crop_max = list(input_shape[0:1]) + crop_max
+        sample['CropWithBoundingBox_Param'] = json.dumps((input_shape, crop_min, crop_max))
+
+        image_t = crop_ND_volume_with_bounding_box(image, crop_min, crop_max)
+        sample['image'] = image_t
+        
+        if('label' in sample and sample['label'].shape[1:] == image.shape[1:]):
+            label = sample['label']
+            crop_max[0] = label.shape[0]
+            label = crop_ND_volume_with_bounding_box(label, crop_min, crop_max)
+            sample['label'] = label
+        if('weight' in sample and sample['weight'].shape[1:] == image.shape[1:]):
+            weight = sample['weight']
+            crop_max[0] = weight.shape[0]
+            weight = crop_ND_volume_with_bounding_box(weight, crop_min, crop_max)
+            sample['weight'] = weight
+        return sample
+
+    def inverse_transform_for_prediction(self, sample):
+        ''' rescale sample['predict'] (5D or 4D) to the original spatial shape.
+         assume batch size is 1, otherwise scale may be different for 
+         different elemenets in the batch.
+
+        origin_shape is a 4D or 3D vector as saved in __call__().'''
+        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]
+        predict = sample['predict']
+        if(isinstance(predict, tuple) or isinstance(predict, list)):
+            output_predict = []
+            for predict_i in predict:
+                origin_shape     = list(predict_i.shape[:2]) + origin_shape[1:]
+                output_predict_i = np.zeros(origin_shape, predict_i.dtype)
+                crop_min = [0, 0] + crop_min[1:]
+                crop_max = list(predict_i.shape[:2]) + crop_max[1:]
+                output_predict_i = set_ND_volume_roi_with_bounding_box_range(output_predict_i,
+                    crop_min, crop_max, predict_i)
+                output_predict.append(output_predict_i)
+        else:
+            origin_shape   = list(predict.shape[:2]) + origin_shape[1:]
+            output_predict = np.zeros(origin_shape, predict.dtype)
+            crop_min = [0, 0] + crop_min[1:]
+            crop_max = list(predict.shape[:2]) + crop_max[1:]
+            output_predict = set_ND_volume_roi_with_bounding_box_range(output_predict,
+                crop_min, crop_max, predict)
+
+        sample['predict'] = output_predict
+        return sample
+
+class RandomCrop(object):
+    """Randomly crop the input image (shape [C, D, H, W] or [C, H, W]) 
+    """
+    def __init__(self, params):
+        """
+        output_size (tuple or list): Desired output size [D, H, W] or [H, W].
+            the output channel is the same as the input channel.
+        foreground_focus (bool): If true, allow crop around the foreground.
+        foreground_ratio (float): Specifying the probability of foreground 
+            focus cropping when foreground_focus is true.
+        mask_label (None, or tuple / list): Specifying the foreground labels for foreground 
+            focus cropping
+        """
+        self.output_size = params['RandomCrop_output_size'.lower()]
+        self.fg_focus    = params['RandomCrop_foreground_focus'.lower()]
+        self.fg_ratio    = params['RandomCrop_foreground_ratio'.lower()]
+        self.mask_label  = params['RandomCrop_mask_label'.lower()]
+        self.inverse     = params['RandomCrop_inverse'.lower()]
+        assert isinstance(self.output_size, (list, tuple))
+        if(self.mask_label is not None):
+            assert isinstance(self.mask_label, (list, tuple))
+
+    def __call__(self, sample):
+        image = sample['image']
+        input_shape = image.shape
+        input_dim   = len(input_shape) - 1
+
+        assert(input_dim == len(self.output_size))
+        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)
+            if(mask.sum() == 0):
+                bb_min = [0] * (input_dim + 1)
+                bb_max = mask.shape
+            else:
+                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
+        crop_max = list(input_shape[0:1]) + crop_max
+        sample['RandomCrop_Param'] = json.dumps((input_shape, crop_min, crop_max))
+
+        image_t = crop_ND_volume_with_bounding_box(image, crop_min, crop_max)
+        sample['image'] = image_t
+        
+        if('label' in sample and sample['label'].shape[1:] == image.shape[1:]):
+            label = sample['label']
+            crop_max[0] = label.shape[0]
+            label = crop_ND_volume_with_bounding_box(label, crop_min, crop_max)
+            sample['label'] = label
+        if('weight' in sample and sample['weight'].shape[1:] == image.shape[1:]):
+            weight = sample['weight']
+            crop_max[0] = weight.shape[0]
+            weight = crop_ND_volume_with_bounding_box(weight, crop_min, crop_max)
+            sample['weight'] = weight
+        return sample
+
+    def inverse_transform_for_prediction(self, sample):
+        ''' rescale sample['predict'] (5D or 4D) to the original spatial shape.
+         assume batch size is 1, otherwise scale may be different for 
+         different elemenets in the batch.
+
+        origin_shape is a 4D or 3D vector as saved in __call__().'''
+        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]
+        predict = sample['predict']
+        origin_shape   = list(predict.shape[:2]) + origin_shape[1:]
+        output_predict = np.zeros(origin_shape, predict.dtype)
+        crop_min = [0, 0] + crop_min[1:]
+        crop_max = list(predict.shape[:2]) + crop_max[1:]
+        output_predict = set_ND_volume_roi_with_bounding_box_range(output_predict,
+            crop_min, crop_max, predict)
+
+        sample['predict'] = output_predict
+        return sample

pymic/transform/flip.py

@@ -0,0 +1,68 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+
+import torch
+import json
+import math
+import random
+import numpy as np
+from scipy import ndimage
+from pymic.transform.abstract_transform import AbstractTransform
+from pymic.util.image_process import *
+
+
+class RandomFlip(AbstractTransform):
+    """ random flip the image (shape [C, D, H, W] or [C, H, W]) """
+    def __init__(self, params):
+        """
+            flip_depth (bool) : random flip along depth axis or not, only used for 3D images
+            flip_height (bool): random flip along height axis or not
+            flip_width (bool) : random flip along width axis or not
+        """
+        self.flip_depth  = params['RandomFlip_flip_depth'.lower()]
+        self.flip_height = params['RandomFlip_flip_height'.lower()]
+        self.flip_width  = params['RandomFlip_flip_width'.lower()]
+        self.inverse = params['RandomFlip_inverse'.lower()]
+
+    def __call__(self, sample):
+        image = sample['image']
+        input_shape = image.shape
+        input_dim = len(input_shape) - 1
+        flip_axis = []
+        if(self.flip_width):
+            if(random.random() > 0.5):
+                flip_axis.append(-1)
+        if(self.flip_height):
+            if(random.random() > 0.5):
+                flip_axis.append(-2)
+        if(input_dim == 3 and self.flip_depth):
+            if(random.random() > 0.5):
+                flip_axis.append(-3)
+
+        sample['RandomFlip_Param'] = json.dumps(flip_axis)
+        if(len(flip_axis) > 0):
+            # use .copy() to avoid negative strides of numpy array
+            # current pytorch does not support negative strides
+            image_t = np.flip(image, flip_axis).copy()
+            sample['image'] = image_t
+            if('label' in sample and sample['label'].shape[1:] == image.shape[1:]):
+                sample['label'] = np.flip(sample['label'] , flip_axis).copy()
+            if('weight' in sample and sample['weight'].shape[1:] == image.shape[1:]):
+                sample['weight'] = np.flip(sample['weight'] , flip_axis).copy()
+            
+        return sample
+
+    def  inverse_transform_for_prediction(self, sample):
+        ''' flip sample['predict'] (5D or 4D) to the original direction.
+         assume batch size is 1, otherwise flip parameter may be different for 
+         different elemenets in the batch.
+
+        flip_axis is a list as saved in __call__().'''
+        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).copy()
+        return sample

pymic/transform/gamma_correction.py

@@ -0,0 +1,39 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+
+import torch
+import json
+import math
+import random
+import numpy as np
+from scipy import ndimage
+from pymic.transform.abstract_transform import AbstractTransform
+from pymic.util.image_process import *
+
+
+class ChannelWiseGammaCorrection(AbstractTransform):
+    """
+    apply random gamma correction to each channel
+    """
+    def __init__(self, params):
+        """
+        (gamma_min, gamma_max) specify the range of gamma
+        """
+        self.gamma_min = params['ChannelWiseGammaCorrection_gamma_min'.lower()]
+        self.gamma_max = params['ChannelWiseGammaCorrection_gamma_max'.lower()]
+        self.inverse = params['ChannelWiseGammaCorrection_inverse'.lower()]
+    
+    def __call__(self, sample):
+        image= sample['image']
+        for chn in range(image.shape[0]):
+            gamma_c = random.random() * (self.gamma_max - self.gamma_min) + self.gamma_min
+            img_c = image[chn]
+            v_min = img_c.min()
+            v_max = img_c.max()
+            img_c = (img_c - v_min)/(v_max - v_min)
+            img_c = np.power(img_c, gamma_c)*(v_max - v_min) + v_min
+            image[chn] = img_c
+
+        sample['image'] = image
+        return sample
+