Merge pull request #146 from vzinche/super-dev

Added some functionality for 2D input with channels

inferno/extensions/layers/convolutional.py

@@ -16,7 +16,9 @@
            'BNReLUConv2D', 'BNReLUConv3D',
            'BNReLUDepthwiseConv2D',
            'ConvSELU2D', 'ConvSELU3D',
-           'ConvReLU2D', 'ConvReLU3D']
+           'ConvReLU2D', 'ConvReLU3D',
+           'BNReLUDilatedConv2D', 'DilatedConv2D',
+           'GlobalConv2D']
 _all = __all__
 
 
@@ -236,6 +238,17 @@ def __init__(self, in_channels, out_channels, kernel_size, dilation=2):
                                                activation='ELU',
                                                initialization=OrthogonalWeightsZeroBias())
 
+class DilatedConv2D(ConvActivation):
+    """2D dilated convolutional layer with 'SAME' padding, no activation and orthogonal weight initialization."""
+    def __init__(self, in_channels, out_channels, kernel_size, dilation=2):
+        super(DilatedConv2D, self).__init__(in_channels=in_channels,
+                                               out_channels=out_channels,
+                                               kernel_size=kernel_size,
+                                               dilation=dilation,
+                                               dim=2,
+                                               activation=None,
+                                               initialization=OrthogonalWeightsZeroBias())
+
 
 class ConvReLU2D(ConvActivation):
     """2D Convolutional layer with 'SAME' padding, ReLU and Kaiming normal weight initialization."""
@@ -342,6 +355,21 @@ def __init__(self, in_channels, out_channels, kernel_size, stride=1):
                                            initialization=KaimingNormalWeightsZeroBias(0))
         self.batchnorm = nn.BatchNorm2d(in_channels)
 
+class BNReLUDilatedConv2D(_BNReLUSomeConv,ConvActivation):
+    """
+    2D dilated convolutional layer with 'SAME' padding, Batch norm,  Relu and He
+    weight initialization.
+    """
+    def __init__(self, in_channels, out_channels, kernel_size, dilation=2):
+        super(BNReLUDilatedConv2D, self).__init__(in_channels=in_channels,
+                                               out_channels=out_channels,
+                                               kernel_size=kernel_size,
+                                               dilation=dilation,
+                                               dim=2,
+                                               activation=nn.ReLU(inplace=True),
+                                               initialization=KaimingNormalWeightsZeroBias(0))
+        self.batchnorm = nn.BatchNorm2d(in_channels)  
+
 
 class BNReLUConv3D(_BNReLUSomeConv, ConvActivation):
     """
@@ -441,3 +469,41 @@ def __init__(self, in_channels, out_channels, kernel_size):
                                          dim=3,
                                          activation=activation,
                                          initialization=SELUWeightsZeroBias())
+
+class GlobalConv2D(nn.Module):
+    """From https://arxiv.org/pdf/1703.02719.pdf
+    Main idea: we can have a bigger kernel size computationally acceptable
+    if we separate 2D-conv in 2 1D-convs """
+    def __init__(self, in_channels, out_channels, kernel_size, local_conv_type, activation=None, use_BN=False, **kwargs):
+        super(GlobalConv2D, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        assert isinstance(kernel_size, (int, list, tuple))
+        if isinstance(kernel_size, int):
+	        kernel_size = (kernel_size,)*2
+        self.kwargs=kwargs
+        self.conv1a = local_conv_type(in_channels=self.in_channels, out_channels=self.out_channels,
+                               kernel_size=(kernel_size[0], 1), **kwargs)
+        self.conv1b = local_conv_type(in_channels=self.out_channels, out_channels=self.out_channels,
+                               kernel_size=(1, kernel_size[1]), **kwargs)
+        self.conv2a = local_conv_type(in_channels=self.in_channels, out_channels=self.out_channels,
+                               kernel_size=(1, kernel_size[1]), **kwargs)
+        self.conv2b = local_conv_type(in_channels=self.out_channels, out_channels=self.out_channels,
+                               kernel_size=(kernel_size[0], 1), **kwargs)
+        if use_BN: 
+            self.batchnorm = nn.BatchNorm2d(self.out_channels)
+        else:
+            self.batchnorm = None
+        self.activation = activation
+    def forward(self, input_):
+        out1 = self.conv1a(input_)
+        out1 = self.conv1b(out1)
+        out2 = self.conv2a(input_)
+        out2 = self.conv2b(out2)
+        out = out1.add(1,out2)        
+        if self.activation is not None:
+            out = self.activation(out)
+        if self.batchnorm is not None:
+            out = self.batchnorm(out) 
+        return out

inferno/extensions/layers/sampling.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 
-__all__ = ['AnisotropicUpsample', 'AnisotropicPool', 'Upsample']
+__all__ = ['AnisotropicUpsample', 'AnisotropicPool', 'Upsample', 'AnisotropicUpsample2D', 'AnisotropicPool2D']
 
 
 # torch is deprecating nn.Upsample in favor of nn.functional.interpolate
@@ -56,4 +56,29 @@ def __init__(self, downscale_factor):
                                               stride=(1, ds, ds),
                                               padding=(0, 1, 1))
 
+class AnisotropicUpsample2D(nn.Module):
+    def __init__(self, scale_factor):
+        super(AnisotropicUpsample2D, self).__init__()
+        self.upsampler = nn.Upsample(scale_factor=scale_factor)
+
+    def forward(self, input):
+        # input is 2D of shape NCDW (or NCDH, egal)
+        N, C, D, W = input.size()
+        # Fold C and D axes in one
+        folded = input.view(N, C * D, W)
+        # Upsample
+        upsampled = self.upsampler(folded)
+        # Unfold out the C and D axes
+        unfolded = upsampled.view(N, C, D,
+                                  self.upsampler.scale_factor * W)
+        # Done
+        return unfolded
+
+
+class AnisotropicPool2D(nn.MaxPool2d):
+    def __init__(self, downscale_factor):
+        ds = downscale_factor
+        super(AnisotropicPool2D, self).__init__(kernel_size=(1, ds + 1),
+                                              stride=(1, ds),
+                                              padding=(0, 1))
 

inferno/io/transform/base.py

@@ -174,4 +174,4 @@ class DTypeMapping(object):
                      'byte': 'uint8',
                      'uint8': 'uint8',
                      'int': 'int32',
-                     'int32': 'int32'}
+                     'int32': 'int32'}

inferno/io/transform/image.py

@@ -566,14 +566,15 @@ def build_random_variables(self):
 
     def batch_function(self, image):
         scale = self.get_random_variable('seg_scale')
-        image_shape = np.array(image[0].shape[1:])
-
+        input_image, segmentation = image
+        image_shape = np.array(input_image.shape[1:])
+        if input_image.ndim == segmentation.ndim + 1:
+            segmentation = segmentation[None]
         with catch_warnings():
             simplefilter('ignore')
-            img = np.stack([zoom(x, scale, order=3) for x in image[0]])
-            seg = np.stack([zoom(x, scale, order=0) for x in image[1]])
+            img = np.stack([zoom(x, scale, order=3) for x in input_image])
+            seg = np.stack([zoom(x, scale, order=0) for x in segmentation])
         new_shape = np.array(img.shape[1:])
-
         if self.resize:
             if scale > 1.:
                 # pad image to original size
@@ -589,7 +590,5 @@ def batch_function(self, image):
                 pad_r = image_shape - new_shape - pad_l
                 padding = [(0,0)] + list(zip(pad_l, pad_r))
                 img = np.pad(img, padding, 'constant', constant_values=self.pad_const)
-
-                seg = np.pad(seg, padding, 'constant', constant_values=self.pad_const)
-
-        return img, seg
+                seg = np.pad(seg, padding, 'constant', constant_values=self.pad_const)     
+        return img, seg

inferno/io/transform/volume.py

@@ -27,6 +27,8 @@ def volume_function(self, volume):
             volume = volume[::-1, :, :]
         return volume
 
+
+
 class RandomRot3D(Transform):
     def __init__(self, rot_range, p=0.125,  only_one=True, **super_kwargs):
         super(RandomRot3D, self).__init__(**super_kwargs)
@@ -129,3 +131,30 @@ def volume_function(self, volume):
         x1, y1, z1 = self.crop_left
         x2, y2, z2 = (np.array(volume.shape) - np.array(self.crop_right)).astype('uint32')
         return volume[x1:x2, y1:y2, z1:z2]
+
+
+
+class Slices2Channels(Transform):
+    """ Needed for training 2D network with slices above/below as additional channels
+        For the input data transforms one dimension (x, y or z) into channels 
+        For the target data just takes the central slice and discards all the rest"""
+    def __init__(self, num_channels, downsampling = 1, **super_kwargs):
+        super(Slices2Channels, self).__init__(**super_kwargs)
+        self.channels = num_channels
+        self.downsampling = downsampling
+    def batch_function(self, batch):
+        try:
+            self.axis = batch[0].shape.index(self.channels)
+        except ValueError:
+            print ("The axis has the shape of the desired channels number!")
+        half = int(self.channels/2)
+        new_input = np.moveaxis(batch[0], self.axis, 0)
+        #take every nth slice to the both directions of the central slice
+        indices = []
+        for i in range (self.channels):
+            if i%self.downsampling == half%self.downsampling:
+                indices.append(i)
+        new_input = new_input[indices]   #num_chan after - int (num_chan/(2*downsample)) * 2 + 1
+        new_target = np.moveaxis(batch[1], self.axis, 0)
+        new_target = new_target[half]
+        return (new_input, new_target)

inferno/io/volumetric/volume.py

@@ -69,7 +69,10 @@ def pad_volume(self, padding=None):
         padding = self.padding if padding is None else padding
         if padding is None:
             return self.volume
-        else:
+        else:            
+            #for symmertic padding only one int can be passed for each axis
+            assert_(all(isinstance(pad, (int, tuple, list)) for pad in self.padding), "Expect int or iterable", TypeError)
+            self.padding = [[pad, pad] if isinstance(pad, int) else pad for pad in self.padding]
             self.volume = np.pad(self.volume,
                                  pad_width=self.padding,
                                  mode=self.padding_mode)