melage 0.0.65__py3-none-any.whl

melage/widgets/DeepLModels/InfantSegment/Unet.py

@@ -0,0 +1,464 @@
+from functools import partial
+
+import numpy as np
+import torch.nn as nn
+import torch
+import math
+# from .model_utils import *
+import math
+import torch
+from functools import partial
+import torch.nn as nn
+from einops import repeat, rearrange
+
+
+
+class BlockLayer(nn.Module):
+    def __init__(self, num_blcks, block_layer, planes_in, planes_out, kernel_size=3, first_layer=False,
+                 input_size=None, norm_type='layer'):
+        super(BlockLayer, self).__init__()
+
+        self.blocks = nn.ModuleList()
+        for i in range(num_blcks):
+            if i == 0:
+                self.blocks.append(block_layer(planes_in, planes_out, kernel_size=kernel_size, first_layer=first_layer,
+                                               input_size=input_size, norm_type=norm_type))
+            else:
+                self.blocks.append(block_layer(planes_in, planes_out, kernel_size=kernel_size, first_layer=False,
+                                               input_size=input_size, norm_type=norm_type))
+            planes_in = planes_out
+
+    def forward(self, x):
+        for i, block in enumerate(self.blocks):
+            x = block(x)
+        return x
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, planes_in, planes_out, kernel_size=3, first_layer=False, input_size=128, norm_type='layer'):
+        super(ResidualBlock, self).__init__()
+
+        self.conv1 = ConvolutionalBlock(planes_in=planes_in, planes_out=planes_out, first_layer=first_layer,
+                                        kernel_size=kernel_size, dilation=1,
+                                        activation=nn.ELU, input_size=input_size, norm_type=norm_type)
+        self.conv2 = ConvolutionalBlock(planes_in=planes_out, planes_out=planes_out, first_layer=False,
+                                        kernel_size=1,
+                                        dilation=1, activation=nn.ELU, input_size=input_size, norm_type=norm_type)
+        if planes_in != planes_out:
+            self.sample = nn.Conv3d(planes_in, planes_out, (1, 1, 1), stride=(1, 1, 1), dilation=(1, 1, 1),
+                                    bias=True)  #
+        else:
+            self.sample = None
+
+    def forward(self, x):
+        identity = x.clone()
+
+        x = self.conv1(x)
+        x = self.conv2(x)
+
+        if self.sample is not None:
+            identity = self.sample(identity)
+
+        x += identity
+
+        return x
+
+
+class UnetEncoder(nn.Module):
+    def __init__(self, in_channel, base_inc_channel=8, layer=BlockLayer, block=None, layer_blocks=None,
+                 downsampling_stride=None, feature_dilation=1.5, layer_widths=None, kernel_size=3,
+                 norm_type='layer'):
+        super(UnetEncoder, self).__init__()
+
+        self.layers = nn.ModuleList()
+        self.downsampling_convolutions = nn.ModuleList()
+        self.attention_modules = nn.ModuleList()
+        self.downsampling_zarib = []
+        in_channel_layer = in_channel
+        input_size = 192
+        self._layers_with = []
+        # self._layers_with.append(base_inc_channel)
+        for i, num_blcks in enumerate(layer_blocks):
+            if layer_widths is not None:
+                out_channel_layer = layer_widths[i]
+            else:
+                out_channel_layer = base_inc_channel * int(feature_dilation ** (i + 1)) // 2
+
+            if i == 0:
+                first_layer = True
+            else:
+                first_layer = False
+            self.layers.append(layer(num_blcks=num_blcks, block_layer=block,
+                                     planes_in=in_channel_layer, planes_out=out_channel_layer,
+                                     kernel_size=kernel_size,
+                                     first_layer=first_layer, input_size=input_size,
+                                     norm_type=norm_type))
+            # self.attention_modules.append(Attention(out_channel_layer))
+            if i != len(layer_blocks) - 1:
+                # padding = kernel_size // 2  # constant size
+                downsampling_conv = nn.Conv3d(out_channel_layer, out_channel_layer, (3, 3, 3), padding=3 // 2,
+                                              stride=(downsampling_stride, downsampling_stride, downsampling_stride),
+                                              bias=True)
+                # downsampling_conv = nn.MaxPool3d(kernel_size=2, stride=2)
+
+                self.downsampling_convolutions.append(downsampling_conv)
+
+                input_size = input_size // 2
+            print("Encoder {}:".format(i), in_channel_layer, out_channel_layer)
+            self._layers_with.append(out_channel_layer)
+            in_channel_layer = out_channel_layer
+        self.out_channel_layer = in_channel_layer // 2
+        self.last_downsampling_conv = nn.Conv3d(out_channel_layer, out_channel_layer, (3, 3, 3),
+                                                padding=3 // 2,
+                                                stride=(downsampling_stride, downsampling_stride, downsampling_stride),
+                                                bias=True)
+        self.output_size = input_size
+
+    def forward(self, x):
+        outputs = list()
+        # outputs.insert(0, x)
+        for layer, downsampling in zip(self.layers[:-1], self.downsampling_convolutions):
+            x = layer(x)
+
+            outputs.insert(0, x)
+
+            x = downsampling(x)
+        outputs.insert(0, x)
+        x = self.layers[-1](x)
+        x = self.last_downsampling_conv(x)
+        #outputs.insert(0, x)  # bottle neck layer
+        return x, outputs
+
+
+class ConvolutionalBlock(nn.Module):
+    def __init__(self, planes_in, planes_out, first_layer=False, kernel_size=3, dilation=1, activation=None,
+                 input_size=None, norm_type='layer'):
+        super(ConvolutionalBlock, self).__init__()
+        if dilation == 1:
+            padding = kernel_size // 2  # constant size
+        else:
+            # (In + 2*padding - dilation * (kernel_size - 1) - 1)/stride + 1
+            if kernel_size == 3:
+                if dilation == 2:
+                    padding = 2
+                elif dilation == 4:
+                    padding = 4
+                elif dilation == 3:
+                    padding = 3
+                else:
+                    padding = None
+            elif kernel_size == 1:
+                padding = 0
+        self.activation = None
+        self.norm = None
+        if first_layer:
+            self.norm = nn.InstanceNorm3d(planes_in)
+            self.activation = activation()
+            self.conv = nn.Conv3d(planes_in, planes_out, (kernel_size, kernel_size, kernel_size),
+                                  padding=padding, bias=True,
+                                  dilation=(dilation, dilation, dilation))
+        else:
+            if activation is not None:
+                if norm_type.lower() == 'layer':
+                    self.norm = nn.LayerNorm([input_size, input_size, input_size])
+                elif norm_type.lower() == 'group':
+                    valid_num_groups = np.array([16, 8, 4, 2])
+                    valid_num_groups = valid_num_groups[valid_num_groups < planes_in]
+                    num_groups = None
+                    for num_groups in valid_num_groups:
+                        if planes_in % num_groups != 0:
+                            break
+                    if num_groups is None:
+                        raise exit('Num groups can not be determined')
+                    self.norm = nn.GroupNorm(num_groups=num_groups, num_channels=planes_in)
+                elif norm_type.lower() == 'batch':
+                    self.norm = nn.BatchNorm3d(planes_in)
+                elif norm_type.lower() == 'instance':
+                    self.norm = nn.InstanceNorm3d(planes_in)
+                else:
+                    self.norm = None
+
+                self.activation = activation()
+                self.conv = nn.Conv3d(planes_in, planes_out, (kernel_size, kernel_size, kernel_size),
+                                      padding=padding, bias=True,
+                                      dilation=(dilation, dilation, dilation))
+
+            else:
+                if norm_type.lower() == 'layer':
+                    if input_size < 120:
+                        self.norm = nn.LayerNorm([input_size, input_size, input_size])
+                    else:
+                        self.norm = nn.InstanceNorm3d(planes_in)
+                elif norm_type.lower() == 'group':
+                    valid_num_groups = [16, 8, 4, 2]
+                    valid_num_groups = valid_num_groups[valid_num_groups < planes_in]
+                    num_groups = None
+                    for num_groups in valid_num_groups:
+                        if planes_in % num_groups != 0:
+                            break
+                    if num_groups is None:
+                        raise exit('Num groups can not be determined')
+                    self.norm = nn.GroupNorm(num_groups=planes_in, num_channels=planes_in)
+                elif norm_type.lower() == 'batch':
+                    self.norm = nn.BatchNorm3d(planes_in)
+                elif norm_type.lower() == 'instance':
+                    self.norm = nn.InstanceNorm3d(planes_in)
+                else:
+                    self.norm = None
+
+                self.conv = nn.Conv3d(planes_in, planes_out, (kernel_size, kernel_size, kernel_size),
+                                      padding=padding, bias=True,
+                                      dilation=(dilation, dilation, dilation))
+
+    def forward(self, x):
+        if self.norm is not None:
+            x = self.norm(x)
+
+        if self.activation is not None:
+            x = self.activation(x)
+
+        x = self.conv(x)
+
+        return x
+
+
+class UnetDecoder(nn.Module):
+    def __init__(self, in_channel, base_inc_channel=64, layer=BlockLayer, block=None, layer_blocks=[1, 1, 1, 1],
+                 feature_dilation=2, upsampling_stride=2, layer_widths=None, kernel_size=3,
+                 upsampling_mode="trilinear", align_corners=False, use_transposed_convolutions=False,
+                 last_cov_channels=256,
+                 norm_type='layer'
+                 ):
+        super(UnetDecoder, self).__init__()
+        self.layers = nn.ModuleList()
+
+        self.upsampling_blocks = nn.ModuleList()
+
+        self.attention_modules = nn.ModuleList()
+        in_channel_layer = in_channel
+        # input_size = 24
+        input_size = 16
+
+        for i, num_blcks in enumerate(layer_blocks):
+            if layer_widths is not None:
+                out_channel_layer = layer_widths[i]
+            else:
+                out_channel_layer = base_inc_channel // (feature_dilation ** (i))
+
+            if i == 0:
+                first_layer = True
+                self.layers.append(layer(num_blcks=num_blcks, block_layer=block,
+                                         planes_in=last_cov_channels, planes_out=out_channel_layer,
+                                         kernel_size=kernel_size,
+                                         first_layer=first_layer, input_size=input_size, norm_type=norm_type))
+            else:
+                first_layer = False
+
+                self.layers.append(layer(num_blcks=num_blcks, block_layer=block,
+                                         planes_in=in_channel_layer + layer_widths[i - 1], planes_out=out_channel_layer,
+                                         kernel_size=kernel_size,
+                                         first_layer=first_layer, input_size=input_size, norm_type=norm_type))
+
+
+            # self.upsampling_blocks.append(nn.ConvTranspose3d(out_channel_layer, out_channel_layer, kernel_size=2,
+            #                                                         stride=upsampling_stride, padding=0))
+            self.upsampling_blocks.append(nn.Upsample(scale_factor=2, mode='trilinear'))
+
+            input_size = input_size * 2
+            last_cov_channels = in_channel_layer  # last_cov_channels//2
+            print("Decoder {}:".format(i), in_channel_layer, out_channel_layer)
+            in_channel_layer = out_channel_layer
+        self.out_channel_layer = in_channel_layer
+
+    def forward(self, y, x):
+        i = 0
+        outputs = list()
+        #y = x[0]
+        for up, lay in zip(self.upsampling_blocks, self.layers[:-1]):
+            if i == 0:
+                y = lay(y)
+            else:
+                y = lay(y)
+            outputs.insert(0, y)
+            y = up(y)
+            y = torch.cat([y, x[i]], 1)
+
+            # y = att(y)
+            # y = torch.cat([y, x[i + 1]],1)
+            i += 1
+        outputs.insert(0, y)
+        y = self.layers[-1](y)
+        y = up(y)
+        outputs.insert(0, y)
+        return y, outputs
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads=4, dim_head=16):
+        super().__init__()
+        self.scale = dim_head ** -0.5
+        self.heads = heads
+        hidden_dim = dim_head * heads
+
+        self.to_qkv = nn.Conv3d(dim, hidden_dim * 3, 1, bias=False)
+        self.to_out = nn.Conv3d(hidden_dim, dim // 2, 1)
+
+    def forward(self, x, mask=None):
+        b, c, h, w, z = x.shape
+        qkv = self.to_qkv(x).chunk(3, dim=1)
+        q, k, v = map(lambda t: rearrange(t, 'b (h c) x y z -> b h c (x y z)', h=self.heads), qkv)
+
+        scaled_dot_prod = torch.einsum('... i d , ... j d -> ... i j', q, k) * self.scale
+        attention = torch.softmax(scaled_dot_prod, dim=-1)
+        v = v / (h * w * z)
+        atv = torch.einsum('... i j , ... j d -> ... i d', attention, v)
+        out = rearrange(atv, "b h c (x y z) -> b (h c) x y z", h=self.heads, x=h, y=w, z=z)
+        return self.to_out(out)
+
+
+class UnetGen(nn.Module):
+    def __init__(self, base_inc_channel=8,
+                 feature_dilation=2, downsampling_stride=2,
+                 encoder_class=UnetEncoder, layer_widths=None, block=None,
+                 kernel_size=3, interpolation_mode="trilinear", decoder_class=None,
+                 use_transposed_convolutions=True, norm_type='layer', outChannels=3):
+        super(UnetGen, self).__init__()
+
+        use_transposed_convolutions = self.use_tr_conv
+        inblock = 16
+        base_inc_channel = inblock
+        self.base_inc_channel = base_inc_channel
+
+        # encoder_blocks = [1, 1, 1, 1, 1, 1]
+
+        # decoder_blocks = [1,1,1,1, 1, 1]
+        #encoder_blocks = [1, 1, 1]
+
+        #decoder_blocks = [1, 1, 1]
+
+        encoder_blocks = [1, 1, 1, 1, 1]
+
+        decoder_blocks = [1, 1, 1, 1, 1]
+
+        padding = kernel_size // 2  # constant size
+        self.before_encoder = nn.Conv3d(1, inblock, kernel_size=(3, 3, 3),
+                                        stride=(1, 1, 1), padding=3 // 2,
+                                        bias=True)
+
+        self.encoder = encoder_class(in_channel=inblock, base_inc_channel=base_inc_channel, layer_blocks=encoder_blocks,
+                                     block=block,
+                                     feature_dilation=feature_dilation, downsampling_stride=downsampling_stride,
+                                     layer_widths=layer_widths, kernel_size=kernel_size,
+                                     norm_type=norm_type)
+
+        layer_widths = self.encoder._layers_with
+        in_channel = layer_widths[-1]
+        self.BottleNeck = BlockLayer(num_blcks=1, block_layer=block,
+                                     planes_in=in_channel, planes_out=in_channel // 2,
+                                     kernel_size=kernel_size,
+                                     first_layer=False, input_size=self.encoder.output_size, norm_type=norm_type)
+
+        # self.BottleNeck_att = Attention(in_channel)
+
+        layer_widths = layer_widths[::-1]  # [1:]
+        layer_widths[0] = layer_widths[0] // 2
+        layer_widths[-1] = layer_widths[-2]
+
+        in_channel = in_channel // 2
+        self.decoder = decoder_class(in_channel=in_channel, base_inc_channel=base_inc_channel * 8,
+                                     layer_blocks=decoder_blocks,
+                                     block=block, last_cov_channels=self.encoder.out_channel_layer,
+                                     upsampling_mode=interpolation_mode, layer_widths=layer_widths,
+                                     use_transposed_convolutions=use_transposed_convolutions,
+                                     kernel_size=kernel_size, norm_type=norm_type,
+                                     )
+
+        kernel_size = 3
+
+        #self.last_convolution_rec = BlockLayer(num_blcks=1, block_layer=block,
+        #                                       planes_in=inblock * 2, planes_out=inblock * 4,
+        #                                       kernel_size=kernel_size,
+        #                                       first_layer=False, input_size=192, norm_type=norm_type)
+
+        """
+
+        self.decoder2 = decoder_class(in_channel=in_channel, base_inc_channel=base_inc_channel*8, layer_blocks=decoder_blocks,
+                                     block=block, last_cov_channels = self.encoder.out_channel_layer,
+                                     upsampling_mode=interpolation_mode, layer_widths=layer_widths,
+                                     use_transposed_convolutions=use_transposed_convolutions,
+                                     kernel_size=kernel_size, norm_type=norm_type,
+                                     )
+        self.last_convolution_rec2 = BlockLayer(num_blcks=1, block_layer=block,
+                                         planes_in=inblock*2, planes_out=inblock//2,
+                                         kernel_size=kernel_size,
+                                         first_layer=False, input_size=192, norm_type=norm_type)
+        self.final_convolution_rec2 = nn.Conv3d(inblock//2, 1, kernel_size=(kernel_size, kernel_size, kernel_size),
+                                           stride=(1, 1, 1), bias=True, padding=kernel_size // 2)
+        """
+        kernel_size = 1
+        self.final_convolution_rec = nn.Conv3d(inblock * 3, outChannels, kernel_size=(kernel_size, kernel_size, kernel_size),
+                                               stride=(1, 1, 1), bias=True, padding=kernel_size // 2)
+        #self.seg_layer = nn.Conv3d(inblock * 3, outChannels, kernel_size=(kernel_size, kernel_size, kernel_size),
+        #                                       stride=(1, 1, 1), bias=True, padding=kernel_size // 2)
+        self.activation = nn.Softmax(dim=1)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, y):
+        y = self.before_encoder(y)
+
+        y1, x = self.encoder(y)
+        y1 = self.BottleNeck(y1)
+        # x[0] = self.BottleNeck_att(x[0])
+
+        x1, _ = self.decoder(y1, x)
+
+        x1 = torch.cat([x1, y], 1)
+        # x = (x * mask)
+        #z1 = self.last_convolution_rec(x1)
+        z1 = self.final_convolution_rec(x1)
+        z2 = 0# self.seg_layer(x1)
+
+        """
+
+        x2, _ = self.decoder2(x)
+
+        x2 = torch.cat([x2, y], 1)
+        #x = (x * mask)
+        z2 = self.last_convolution_rec2(x2)
+        z2 = self.final_convolution_rec2(z2)
+        """
+        return z1, z2
+
+
+class MGA_NET(UnetGen):
+    def __init__(self, channels=3,  *args, encoder_class=UnetEncoder, **kwargs):
+        self.use_tr_conv = False
+
+        norm_type = "instance"
+        super().__init__(encoder_class=encoder_class, decoder_class=UnetDecoder,
+                         block=ResidualBlock, outChannels=channels, norm_type=norm_type, **kwargs)
+
+        self.channels = channels
+        self.netName = 'MGA_NET'
+
+    def name(self):
+        return 'MGA_NET'
+
+
+class Segmentor(nn.Module):
+    """
+    Segmentor
+    """
+
+    def __init__(self,
+                 channels=30):
+        """
+
+        """
+        super().__init__()
+        self.seg_model = MGA_NET(channels=channels)
+
+    def forward(self, source, shape=None):
+        pred_logits, pred_logits_seg = self.seg_model(source)
+
+        return pred_logits, pred_logits_seg

melage/widgets/DeepLModels/NPP/dataset/mri_dataset_affine.py

@@ -0,0 +1,149 @@
+import os.path
+from torch.utils.data import Dataset
+import nibabel as nib
+from multiprocessing import Manager
+import random
+import torchio as tio
+import pickle
+import torch
+
+error_list = [5298, 5894,  393,   66,   68, 6576,   76, 6653,  347, 6218,  804,  844,
+        5445,  751, 6075, 5501,  368,  236,  269, 5472,  494, 6711, 5571, 6543,
+        5837, 5586, 6663,  791, 6113,  318,  299,  688,  323,   26, 6496, 5746]
+def Generate_dataset():
+
+    # This if else sentence is used to decide whether enables dataset cache funciton
+    if False:
+        cache = DatasetCache(None,use_cache=False)
+        manager2 = Manager()
+        cache2 = DatasetCache(manager2,use_cache=True)
+    else:
+        manager = None
+        cache = DatasetCache(manager,use_cache=False)
+        manager2 = None
+        cache2 = DatasetCache(manager2,use_cache=False)
+
+    # Loading all availible files
+    if os.path.exists('dataset/pretrain_files_list.pkl'):
+        with open('dataset/pretrain_files_list.pkl', 'rb') as file:
+            files_list = pickle.load(file)
+
+        # splitting the training and testing datasets. The Oasis will be treated as testing while the rest of them will be trated as training dataset.
+        testing_set = files_list[0:1]
+
+        training_sets = files_list
+        del training_sets[0]
+
+        training_sets[2] = [i.replace('GSP','GSP/FS_4p5') for i in training_sets[2]]
+        training_sets_path = sum(training_sets, [])
+        for index in sorted(error_list, reverse=True):
+            del training_sets_path[index]
+        training_sets_path= ['/scratch/datasets/xh278/orig'+i for i in training_sets_path]
+        testing_set_path = sum(testing_set, [])
+
+        train_dataset = VoxelDataset(training_sets_path, cache=cache, train=True)
+        val_dataset = VoxelDataset(testing_set_path, cache=cache2, train=False)
+
+    return train_dataset,val_dataset
+
+class DatasetCache(object):
+    def __init__(self, manager, use_cache=True):
+        self.use_cache = use_cache
+        self.manager = manager
+        if self.manager is not None:
+            self._dict = manager.dict()
+
+    def is_cached(self, key):
+        if not self.use_cache:
+            return False
+        return str(key) in self._dict
+
+    def reset(self):
+        self._dict.clear()
+
+    def get(self, key):
+        if not self.use_cache:
+            raise AttributeError('Data caching is disabled and get funciton is unavailable! Check your config.')
+        return self._dict[str(key)]
+
+    def cache(self, key, img, lbl):
+        # only store if full data in memory is enabled
+        if not self.use_cache:
+            return
+        # only store if not already cached
+        if str(key) in self._dict:
+            return
+        self._dict[str(key)] = (img, lbl)
+
+
+class VoxelDataset(Dataset):
+    def __init__(self, norm_file_path,  rescale_sdf=True, cache = None,train=True):
+        # note that input image paths are for processed images rather than unprocessed
+        self.norm_file_path = [i.replace('proc/ForAdrian_Talairach','orig').replace('talairach/','') for i in norm_file_path]
+        self.orig_file_path = [i.replace('norm.mgz','orig.mgz').replace('talairach/','') for i in self.norm_file_path]
+
+        self.intensity_spatial_norm_file_path = [i.replace('orig.mgz','talairach/norm.mgz') for i in self.orig_file_path]
+        self.intensity_norm_file_path =[i.replace('orig.mgz','norm.mgz') for i in self.orig_file_path]
+        self.seg_path =[i.replace('orig.mgz','talairach/aseg.mgz') for i in self.orig_file_path]
+
+        removal = []
+        for i in range(len(self.seg_path)):
+            if not os.path.exists(self.seg_path[i]):
+                removal.append(i)
+        removal = []
+
+        for i in reversed(removal):
+            self.seg_path.pop(i)
+            self.intensity_spatial_norm_file_path.pop(i)
+            self.orig_file_path.pop(i)
+        self.rescale_sdf = rescale_sdf
+        self.cache = cache
+        self.train = train
+        self.transform = tio.Compose([
+            tio.RandomGamma(log_gamma = (-0.3,0.3),p=0.3),
+            ]
+        )
+    def __len__(self):
+        return len(self.orig_file_path)
+
+    def __getitem__(self, index):
+        image_resolution= 256
+        normalization=255
+        output_index = [index]
+        for cur_index in output_index:
+            array = nib.load(self.orig_file_path[cur_index])
+            input = array.get_fdata()
+            input = torch.FloatTensor(input)
+            input = input.unsqueeze(0).unsqueeze(0)
+            if image_resolution != 256:
+                input = torch.nn.functional.interpolate(input, size=[image_resolution, image_resolution, image_resolution], mode='trilinear',align_corners=False)[0]
+            else:
+                input = torch.Tensor(input)[0]
+
+            array = nib.load(self.intensity_spatial_norm_file_path[cur_index])
+            intensity_spatial_norm = array.get_fdata()
+            intensity_spatial_norm = torch.FloatTensor(intensity_spatial_norm)
+            intensity_spatial_norm = intensity_spatial_norm.unsqueeze(0).unsqueeze(0)
+            if image_resolution != 256:
+                intensity_spatial_norm = torch.nn.functional.interpolate(intensity_spatial_norm, size=[image_resolution, image_resolution, image_resolution], mode='trilinear',align_corners=False)[0]
+            else:
+                intensity_spatial_norm = torch.Tensor(intensity_spatial_norm)[0]
+
+            seg = intensity_spatial_norm.unsqueeze(0)
+            if image_resolution != 256:
+                seg = torch.nn.functional.interpolate(seg, size=[image_resolution, image_resolution, image_resolution], mode='nearest')[0]
+            else:
+                seg = torch.Tensor(seg)[0]
+        if self.train:
+            subject = tio.Subject(image = tio.ScalarImage(tensor=input))
+            transformed = self.transform(subject)
+            input = transformed['image'].data
+            input = input.clip(0, normalization) / normalization
+            intensity_spatial_norm = intensity_spatial_norm.clip(0, normalization) / normalization
+        else:
+            input = input.clip(0, normalization) / normalization
+            intensity_spatial_norm = intensity_spatial_norm.clip(0, normalization) / normalization
+        return input,intensity_spatial_norm,seg,cur_index
+
+
+