Upload GPReconResNet

GPModelConfigs.py

@@ -0,0 +1,86 @@
+from transformers import PretrainedConfig
+
+class GPUNetConfig(PretrainedConfig):
+    model_type = "GPUNet"
+    def __init__(
+            self,
+            in_channels=1, 
+            n_classes=3,
+            depth=3, 
+            wf=6, 
+            padding=True,
+            batch_norm=False, 
+            up_mode="sinc", 
+            dropout=True, 
+            Relu="Relu", 
+            out_act="None",
+            **kwargs):
+        self.in_channels = in_channels
+        self.n_classes = n_classes
+        self.depth = depth
+        self.wf = wf
+        self.padding = padding
+        self.batch_norm = batch_norm
+        self.up_mode = up_mode
+        self.dropout = dropout
+        self.Relu = Relu
+        self.out_act = out_act
+        super().__init__(**kwargs)
+        
+class GPReconResNetConfig(PretrainedConfig):
+    model_type = "GPReconResNet"
+    def __init__(
+            self,
+            in_channels=1, 
+            n_classes=3, 
+            res_blocks=14, 
+            starting_nfeatures=64, 
+            updown_blocks=2, 
+            is_relu_leaky=True, 
+            do_batchnorm=False, 
+            res_drop_prob=0.5,
+            out_act="None", 
+            forwardV=0, 
+            upinterp_algo='sinc', 
+            post_interp_convtrans=False, 
+            is3D=False,
+            **kwargs):
+        self.in_channels = in_channels
+        self.n_classes = n_classes
+        self.res_blocks = res_blocks
+        self.starting_nfeatures = starting_nfeatures
+        self.updown_blocks = updown_blocks
+        self.is_relu_leaky = is_relu_leaky
+        self.do_batchnorm = do_batchnorm
+        self.res_drop_prob = res_drop_prob
+        self.out_act = out_act
+        self.forwardV = forwardV
+        self.upinterp_algo = upinterp_algo
+        self.post_interp_convtrans = post_interp_convtrans
+        self.is3D = is3D
+        super().__init__(**kwargs)
+
+class GPShuffleUNetConfig(PretrainedConfig):
+    model_type = "GPShuffleUNet"
+    def __init__(
+            self,
+            d=2, 
+            in_ch=1, 
+            num_features=64, 
+            n_levels=3, 
+            out_ch=3, 
+            kernel_size=3, 
+            stride=1, 
+            dropout=True, 
+            out_act="None",
+            **kwargs):
+        self.d = d
+        self.in_ch = in_ch
+        self.num_features = num_features
+        self.n_levels = n_levels
+        self.out_ch = out_ch
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.dropout = dropout
+        self.out_act = out_act
+        super().__init__(**kwargs)

GPModels.py

@@ -0,0 +1,64 @@
+from transformers import PreTrainedModel
+
+from .GP_UNet import GP_UNet
+from .GP_ReconResNet import GP_ReconResNet
+from .GP_ShuffleUNet import GP_ShuffleUNet
+
+from .GPModelConfigs import GPUNetConfig, GPReconResNetConfig, GPShuffleUNetConfig
+
+class GPUNet(PreTrainedModel):
+    config_class = GPUNetConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = GP_UNet(
+            in_channels=config.in_channels, 
+            n_classes=config.n_classes,
+            depth=config.depth, 
+            wf=config.wf, 
+            padding=config.padding,
+            batch_norm=config.batch_norm, 
+            up_mode=config.up_mode, 
+            dropout=config.dropout, 
+            Relu=config.Relu, 
+            out_act=config.out_act)
+    def forward(self, x):
+        return self.model(x)
+    
+    
+class GPReconResNet(PreTrainedModel):    
+    config_class = GPReconResNetConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = GP_ReconResNet(
+            in_channels=config.in_channels, 
+            n_classes=config.n_classes, 
+            res_blocks=config.res_blocks, 
+            starting_nfeatures=config.starting_nfeatures, 
+            updown_blocks=config.updown_blocks, 
+            is_relu_leaky=config.is_relu_leaky, 
+            do_batchnorm=config.do_batchnorm, 
+            res_drop_prob=config.res_drop_prob,
+            out_act=config.out_act, 
+            forwardV=config.forwardV, 
+            upinterp_algo=config.upinterp_algo, 
+            post_interp_convtrans=config.post_interp_convtrans, 
+            is3D=config.is3D)
+    def forward(self, x):
+        return self.model(x)
+    
+class GPShuffleUNet(PreTrainedModel):
+    config_class = GPShuffleUNetConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = GP_ShuffleUNet(
+            d=config.d, 
+            in_ch=config.in_ch, 
+            num_features=config.num_features, 
+            n_levels=config.n_levels, 
+            out_ch=config.out_ch, 
+            kernel_size=config.kernel_size, 
+            stride=config.stride, 
+            dropout=config.dropout, 
+            out_act=config.out_act)
+    def forward(self, x):
+        return self.model(x)

GP_ReconResNet.py

@@ -0,0 +1,270 @@
+# Adapted from https://raw.githubusercontent.com/soumickmj/NCC1701/main/Bridge/models/ResNet/MickResNet.py
+
+import torch
+import torch.nn as nn
+import numpy as np
+import sys
+import torch.nn.functional as F
+from tricorder.torch.transforms import Interpolator
+
+__author__ = "Soumick Chatterjee"
+__copyright__ = "Copyright 2019, Soumick Chatterjee & OvGU:ESF:MEMoRIAL"
+__credits__ = ["Soumick Chatterjee"]
+
+__license__ = "GPL"
+__version__ = "1.0.0"
+__email__ = "[email protected]"
+__status__ = "Published"
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_features, drop_prob=0.2): #drop_prob=0.2
+        super(ResidualBlock, self).__init__()
+
+        conv_block = [  layer_pad(1),
+                        layer_conv(in_features, in_features, 3),
+                        layer_norm(in_features),
+                        act_relu(),
+                        layer_drop(p=drop_prob, inplace=True),
+                        layer_pad(1),
+                        layer_conv(in_features, in_features, 3) ,
+                        layer_norm(in_features) ]
+
+        self.conv_block = nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        return x + self.conv_block(x)
+
+class DownsamplingBlock(nn.Module):
+    def __init__(self, in_features, out_features):
+        super(DownsamplingBlock, self).__init__()
+
+        conv_block = [  layer_conv(in_features, out_features, 3, stride=2, padding=1),
+                        layer_norm(out_features),
+                        act_relu()  ]
+        self.conv_block = nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        return self.conv_block(x)
+
+class UpsamplingBlock(nn.Module):
+    def __init__(self, in_features, out_features, mode="upconv", interpolator=None, post_interp_convtrans=False):
+        super(UpsamplingBlock, self).__init__()
+
+        self.interpolator = interpolator
+        self.mode = mode
+        self.post_interp_convtrans = post_interp_convtrans
+        if self.post_interp_convtrans:
+            self.post_conv = layer_conv(out_features, out_features, 1)
+
+        if mode == "upconv":
+            conv_block = [  layer_convtrans(in_features, out_features, 3, stride=2, padding=1, output_padding=1),   ]
+        else:
+            conv_block = [  layer_pad(1),
+                            layer_conv(in_features, out_features, 3),   ]
+        conv_block += [ layer_norm(out_features),
+                        act_relu()  ]
+        self.conv_block = nn.Sequential(*conv_block)
+
+    def forward(self, x, out_shape=None):
+        if self.mode == "upconv":
+            if self.post_interp_convtrans:
+                x = self.conv_block(x)
+                if x.shape[2:] != out_shape:
+                    return self.post_conv(self.interpolator(x, out_shape))
+                else:
+                    return x
+            else:
+                return self.conv_block(x)
+        else:
+            return self.conv_block(self.interpolator(x, out_shape))
+
+class GP_ReconResNet(nn.Module):
+    def __init__(self, in_channels=1, n_classes=1, res_blocks=14, starting_nfeatures=64, updown_blocks=2, is_relu_leaky=True, do_batchnorm=False, res_drop_prob=0.2,    #res_drop_prob=0.2
+                        out_act="softmax", forwardV=0, upinterp_algo='upconv', post_interp_convtrans=False, is3D=False): #should use 14 as that gives number of trainable parameters close to number of possible pixel values in a image 256x256 
+        super(GP_ReconResNet, self).__init__()
+
+        layers = {}
+        if is3D:
+            sys.exit("ResNet: for implemented for 3D, ReflectionPad3d code is required")
+            layers["layer_conv"] = nn.Conv3d
+            layers["layer_convtrans"] = nn.ConvTranspose3d
+            if do_batchnorm:
+                layers["layer_norm"] = nn.BatchNorm3d
+            else:
+                layers["layer_norm"] = nn.InstanceNorm3d
+            layers["layer_drop"] = nn.Dropout3d
+            layers["layer_pad"] = ReflectionPad3d
+            layers["interp_mode"] = 'trilinear'
+        else:
+            layers["layer_conv"] = nn.Conv2d
+            layers["layer_convtrans"] = nn.ConvTranspose2d
+            if do_batchnorm:
+                layers["layer_norm"] = nn.BatchNorm2d
+            else:
+                layers["layer_norm"] = nn.InstanceNorm2d
+            layers["layer_drop"] = nn.Dropout2d
+            layers["layer_pad"] = nn.ReflectionPad2d
+            layers["interp_mode"] = 'bilinear'            
+        if is_relu_leaky:
+            layers["act_relu"] = nn.PReLU
+        else:
+            layers["act_relu"] = nn.ReLU
+        globals().update(layers)
+
+        self.forwardV = forwardV
+        self.upinterp_algo = upinterp_algo
+
+        interpolator = Interpolator(mode=layers["interp_mode"] if self.upinterp_algo == "upconv" else self.upinterp_algo)
+
+        in_channels = in_channels
+        out_channels = n_classes
+        # Initial convolution block
+        intialConv = [  layer_pad(3),
+                        layer_conv(in_channels, starting_nfeatures, 7),
+                        layer_norm(starting_nfeatures),
+                        act_relu() ]
+
+        # Downsampling [need to save the shape for upsample]
+        downsam = []
+        in_features = starting_nfeatures
+        out_features = in_features*2
+        for _ in range(updown_blocks):
+            downsam.append(DownsamplingBlock(in_features, out_features))
+            in_features = out_features
+            out_features = in_features*2
+
+        # Residual blocks
+        resblocks = []
+        for _ in range(res_blocks):
+            resblocks += [ResidualBlock(in_features, res_drop_prob)]
+
+        # Upsampling
+        upsam = []
+        out_features = in_features//2
+        for _ in range(updown_blocks):
+            upsam.append(UpsamplingBlock(in_features, out_features, self.upinterp_algo, interpolator, post_interp_convtrans))
+            in_features = out_features
+            out_features = in_features//2
+
+        # Output layer
+        finalconv = [   layer_conv(starting_nfeatures, out_channels, 1),    ] #kernel size changed from 7 to 1 to make GMP work
+
+        if out_act == "sigmoid":
+            finalconv += [   nn.Sigmoid(), ]
+        elif out_act == "relu":
+            finalconv += [   act_relu(), ]
+        elif out_act == "tanh":
+            finalconv += [   nn.Tanh(), ]
+        elif out_act == "softmax":
+            finalconv += [   nn.Softmax2d(), ]
+
+
+        self.intialConv = nn.Sequential(*intialConv)
+        self.downsam = nn.ModuleList(downsam)
+        self.resblocks = nn.Sequential(*resblocks)
+        self.upsam = nn.ModuleList(upsam)
+        self.finalconv = nn.Sequential(*finalconv)
+
+        ### For Classification, following Florian's GP-UNet
+        self.GMP = nn.AdaptiveMaxPool2d((1, 1))
+
+        if self.forwardV == 0:
+            self.forward = self.forwardV0
+        elif self.forwardV == 1:
+            sys.exit("ResNet: its identical to V0 in case of GP_ResNet")
+        elif self.forwardV == 2:
+            self.forward = self.forwardV2
+        elif self.forwardV == 3:
+            self.forward = self.forwardV3
+        elif self.forwardV == 4:
+            self.forward = self.forwardV4
+        elif self.forwardV == 5:
+            self.forward = self.forwardV5
+
+    def final_step(self, x):    
+        if self.training:
+            x = self.GMP(x)
+            return self.finalconv(x).view(x.shape[0],-1)
+        else:
+            mask = self.finalconv(x)
+            x = self.GMP(x)
+            pred = self.finalconv(x).view(x.shape[0],-1)
+            return pred, mask
+            
+    def forwardV0(self, x):
+        #v0: Original Version
+        x = self.intialConv(x)
+        shapes = []
+        for downblock in self.downsam:
+            shapes.append(x.shape[2:])
+            x = downblock(x)
+        x = self.resblocks(x)
+        for i, upblock in enumerate(self.upsam):
+            x = upblock(x, shapes[-1-i])
+        return self.final_step(x)
+
+    def forwardV2(self, x):
+        #v2: residual of v1 + input to the residual blocks added back with the output
+        out = self.intialConv(x)
+        shapes = []
+        for downblock in self.downsam:
+            shapes.append(out.shape[2:])
+            out = downblock(out)
+        out = out + self.resblocks(out)
+        for i, upblock in enumerate(self.upsam):
+            out = upblock(out, shapes[-1-i])
+        return self.final_step(out)
+
+    def forwardV3(self, x):
+        #v3: residual of v2 + input of the initial conv added back with the output
+        out = x + self.intialConv(x)
+        shapes = []
+        for downblock in self.downsam:
+            shapes.append(out.shape[2:])
+            out = downblock(out)
+        out = out + self.resblocks(out)
+        for i, upblock in enumerate(self.upsam):
+            out = upblock(out, shapes[-1-i])
+        return self.final_step(out)
+
+    def forwardV4(self, x):
+        #v4: residual of v3 + output of the initial conv added back with the input of final conv
+        iniconv = x + self.intialConv(x)
+        shapes = []
+        if len(self.downsam) > 0:
+            for i, downblock in enumerate(self.downsam):
+                if i == 0:
+                    shapes.append(iniconv.shape[2:])
+                    out = downblock(iniconv)
+                else:
+                    shapes.append(out.shape[2:])
+                    out = downblock(out)
+        else:
+            out = iniconv
+        out = out + self.resblocks(out)
+        for i, upblock in enumerate(self.upsam):
+            out = upblock(out, shapes[-1-i])
+        out = iniconv + out
+        return self.final_step(out)
+
+    def forwardV5(self, x):
+        #v5: residual of v4 + individual down blocks with individual up blocks
+        outs = [x + self.intialConv(x)]
+        shapes = []
+        for i, downblock in enumerate(self.downsam):
+            shapes.append(outs[-1].shape[2:])
+            outs.append(downblock(outs[-1]))
+        outs[-1] = outs[-1] + self.resblocks(outs[-1])
+        for i, upblock in enumerate(self.upsam):
+            outs[-1] = upblock(outs[-1], shapes[-1-i])
+            outs[-1] = outs[-2] + outs.pop()
+        return self.final_step(outs.pop())
+
+#to run it here from this script, uncomment the following
+
+if __name__ == "__main__":                #to run it
+    image = torch.rand(2, 1, 240, 240)    #specify your image: batch size, Channel, height, width
+    model = GP_ReconResNet(in_channels=1, n_classes=3, upinterp_algo='sinc')                        #Initialize the model
+    # model.eval()
+    out = model(image)
+    print(model(image))

GP_ShuffleUNet.py

@@ -0,0 +1,187 @@
+##Dropout and out_act was added by hadya
+
+import sys
+import torch
+import torch.nn as nn
+
+from . import GP_ShuffleUNet_pixel_shuffle, GP_ShuffleUNet_pixel_unshuffle
+# import pixel_shuffle, pixel_unshuffle
+
+# -------------------------------------------------------------------------------------------------------------------------------------------------##
+
+class _double_conv(nn.Module):
+    """
+    Double Convolution Block
+    """
+
+    def __init__(self, in_channels, out_channels, k_size, stride, bias=True, conv_layer=nn.Conv3d):
+        super(_double_conv, self).__init__()
+        self.conv_1 = conv_layer(in_channels=in_channels, out_channels=out_channels, kernel_size=k_size,
+                      stride=stride, padding=k_size // 2, bias=bias)
+        self.conv_2 = conv_layer(in_channels=out_channels, out_channels=out_channels, kernel_size=k_size,
+                      stride=stride, padding=k_size // 2, bias=bias)
+
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.conv_1(x)
+        x = self.relu((x))
+        x = self.conv_2(x)
+        x = self.relu((x))
+
+        return x
+
+
+class _conv_decomp(nn.Module):
+    """
+    Convolutional Decomposition Block
+    """
+
+    def __init__(self, in_channels, out_channels, k_size, stride, bias=True, conv_layer=nn.Conv3d):
+        super(_conv_decomp, self).__init__()
+        self.conv1 = conv_layer(in_channels=in_channels, out_channels=out_channels, kernel_size=k_size,
+                      stride=stride, padding=k_size // 2, bias=bias)
+        self.conv2 = conv_layer(in_channels=in_channels, out_channels=out_channels, kernel_size=k_size,
+                      stride=stride, padding=k_size // 2, bias=bias)
+        self.conv3 = conv_layer(in_channels=in_channels, out_channels=out_channels, kernel_size=k_size,
+                      stride=stride, padding=k_size // 2, bias=bias)
+        self.conv4 = conv_layer(in_channels=in_channels, out_channels=out_channels, kernel_size=k_size,
+                      stride=stride, padding=k_size // 2, bias=bias)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x1 = self.conv1(x)
+        x1 = self.relu((x1))
+        x2 = self.conv2(x)
+        x2 = self.relu((x2))
+        x3 = self.conv3(x)
+        x3 = self.relu((x3))
+        x4 = self.conv4(x)
+        x4 = self.relu((x4))
+        return x1, x2, x3, x4
+
+
+class _concat(nn.Module):
+    """
+    Skip-Addition block
+    """
+
+    def __init__(self):
+        super(_concat, self).__init__()
+
+    def forward(self, e1, e2, e3, e4, d1, d2, d3, d4):
+        self.X1 = e1 + d1
+        self.X2 = e2 + d2
+        self.X3 = e3 + d3
+        self.X4 = e4 + d4
+        x = torch.cat([self.X1, self.X2, self.X3, self.X4], dim=1)
+
+        return x
+
+# -------------------------------------------------------------------------------------------------------------------------------------------------##
+
+class GP_ShuffleUNet(nn.Module):
+
+    def __init__(self, d=3, in_ch=1, num_features=64, n_levels=3, out_ch=1, kernel_size=3, stride=1, dropout=False, out_act="softmax"):
+        super(GP_ShuffleUNet, self).__init__()
+
+        self.n_levels = n_levels
+        self.dropout = nn.Dropout2d() if dropout else nn.Sequential()       #added by Hadya
+
+        num_features = num_features
+        filters = [num_features]
+        for _ in range(n_levels):
+            filters.append(filters[-1]*2)
+
+        if d==3:
+            conv_layer = nn.Conv3d
+            ps_fact = (2 ** 2)
+        elif d==2:
+            conv_layer = nn.Conv2d
+            ps_fact = 2
+        else:
+            sys.exit("Invalid d")
+
+        # Input
+        self.conv_inp = _double_conv(in_ch, filters[0], kernel_size, stride, conv_layer=conv_layer)     
+
+        #Contraction path
+        self.wave_down = nn.ModuleList()
+        self.pix_unshuff = nn.ModuleList()
+        self.conv_enc = nn.ModuleList()
+        for i in range(0, n_levels):
+            self.wave_down.append(_conv_decomp(filters[i], filters[i], kernel_size, stride, conv_layer=conv_layer))
+            self.pix_unshuff.append(pixel_unshuffle.PixelUnshuffle(num_features * (2**i), num_features * (2**i), kernel_size, stride, d=d))       
+            self.conv_enc.append(_double_conv(filters[i], filters[i+1], kernel_size, stride, conv_layer=conv_layer))     
+
+        #Expansion path
+        self.cat = _concat()
+        self.pix_shuff = nn.ModuleList()
+        self.wave_up = nn.ModuleList()
+        self.convup = nn.ModuleList()
+        for i in range(n_levels-1,-1,-1):
+            self.pix_shuff.append(pixel_shuffle.PixelShuffle(num_features * (2**(i+1)), num_features * (2**(i+1)) * ps_fact, kernel_size, stride, d=d))
+            self.wave_up.append(_conv_decomp(filters[i], filters[i], kernel_size, stride, conv_layer=conv_layer))
+            self.convup.append(_double_conv(filters[i] * 5, filters[i], kernel_size, stride, conv_layer=conv_layer))
+
+        ###added For Classification, following Florian's GP-UNet
+        self.GMP = nn.AdaptiveMaxPool2d((1, 1))
+
+        #FC
+        if out_act == "softmax":            #added by Hadya
+            self.last = nn.Sequential(
+                      conv_layer(filters[0], out_ch, kernel_size=1, stride=1, padding=0, bias=True),
+                      nn.Softmax2d()
+                    )
+        else:
+            self.out = conv_layer(filters[0], out_ch, kernel_size=1, stride=1, padding=0, bias=True)  #original
+
+        #Weight init
+        for m in self.modules():
+            if isinstance(m, conv_layer):
+                weight = nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
+                m.weight.data.copy_(weight)
+                if m.bias is not None:
+                    m.bias.data.zero_()
+
+    def forward(self, x):
+        encs = [self.conv_inp(x)]
+
+        waves = []
+        for i in range(self.n_levels):
+            waves.append(self.wave_down[i](encs[-1]))
+            _tmp = self.pix_unshuff[i](waves[-1][-1])
+            encs.append(self.conv_enc[i](_tmp))
+
+        dec = encs.pop()
+
+        dec = self.dropout(dec) #added by hadya
+
+        for i in range(self.n_levels):
+            _tmp = self.pix_shuff[i](dec)
+            _tmp_waves = self.wave_up[i](_tmp) + waves.pop()
+            _tmp_cat = self.cat(*_tmp_waves)
+            dec = self.convup[i](torch.cat([encs.pop(), _tmp_cat], dim=1))
+
+        ###added section to make it GP-UNet
+        if self.training:
+            x = self.GMP(dec)
+            return self.out(x).view(x.shape[0],-1)
+        else:
+            mask = self.out(dec)
+            x = self.GMP(dec)
+            pred = self.out(x).view(x.shape[0],-1)
+            return pred, mask
+
+
+        # return self.out(dec) #####replace by line 154-161 to make it GP_ShuffleUNet
+
+
+#to run it here from this script, uncomment the following
+
+if __name__ == "__main__":                #to run it
+    image = torch.rand(2, 1, 512, 512)    #specify your image: batch size, Channel, height, width
+    model = GP_ShuffleUNet(d=2, in_ch=1, num_features=64, n_levels=3, out_ch=3, kernel_size=3, stride=1)    #Initialize the model, d=3 default is for dimensionality conv2d or 3d, default out channel = 1 but in gp we need 3
+    model.eval()
+    out = model(image)
+    print(model(image))

GP_UNet.py

@@ -0,0 +1,189 @@
+# Adapted from https://github.com/soumickmj/FTSuperResDynMRI/blob/main/models/unet2d.py
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torchcomplex.nn.functional as cF
+
+
+
+__author__ = "Soumick Chatterjee"
+__copyright__ = "Copyright 2020, Faculty of Computer Science, Otto von Guericke University Magdeburg, Germany"
+__credits__ = ["Soumick Chatterjee"]
+__license__ = "GPL"
+__version__ = "1.0.0"
+__maintainer__ = "Soumick Chatterjee"
+__email__ = "[email protected]"
+__status__ = "Production"
+
+
+class GP_UNet(nn.Module):
+    """
+    Implementation of
+    U-Net: Convolutional Networks for Biomedical Image Segmentation
+    (Ronneberger et al., 2015)
+    https://arxiv.org/abs/1505.04597
+
+    Using the default arguments will yield the exact version used
+    in the original paper
+
+    Args:
+        in_channels (int): number of input channels
+        n_classes (int): number of output channels
+        depth (int): depth of the network
+        wf (int): number of filters in the first layer is 2**wf
+        padding (bool): if True, apply padding such that the input shape
+                        is the same as the output.
+                        This may introduce artifacts
+        batch_norm (bool): Use BatchNorm after layers with an
+                            activation function
+        up_mode (str): one of 'upconv' or 'upsample'.
+                        'upconv' will use transposed convolutions for
+                        learned upsampling.
+                        'upsample_Bi' will use bilinear upsampling.
+                        'upsample_Sinc' will use sinc upsampling.
+    """
+    def __init__(self, in_channels=1, n_classes=1, depth=3, wf=6, padding=True,
+                 batch_norm=False, up_mode='upconv', dropout=False, Relu = "Relu", out_act="None"): #dropout=False
+        super(GP_UNet, self).__init__()
+        assert up_mode in ('upconv', 'bilinear', 'sinc', "upsample_Sinc")
+        assert out_act in ("softmax", "None", "sigmoid", "relu")
+        self.padding = padding
+        self.depth = depth
+        self.Relu = Relu
+        self.dropout = nn.Dropout2d() if dropout else nn.Sequential()
+        prev_channels = in_channels
+        self.down_path = nn.ModuleList()
+        for i in range(depth):
+            self.down_path.append(UNetConvBlock(prev_channels, 2**(wf+i),
+                                                padding, batch_norm, Relu))
+            prev_channels = 2**(wf+i)
+
+        self.up_path = nn.ModuleList()
+        for i in reversed(range(depth - 1)):
+            self.up_path.append(UNetUpBlock(prev_channels, 2**(wf+i), up_mode,
+                                            padding, batch_norm, Relu))
+            prev_channels = 2**(wf+i)
+
+        if out_act == "softmax":
+            self.last = nn.Sequential(
+                      nn.Conv2d(prev_channels, n_classes, kernel_size=1),
+                      nn.Softmax2d()
+                    )
+        
+        elif out_act == "sigmoid":
+            self.last = nn.Sequential(
+                      nn.Conv2d(prev_channels, n_classes, kernel_size=1),
+                      nn.Sigmoid()
+                    )
+        
+        elif out_act == "relu":
+            self.last = nn.Sequential(
+                      nn.Conv2d(prev_channels, n_classes, kernel_size=1),
+                      nn.ReLU()
+                    )
+
+        else:
+            self.last = nn.Conv2d(prev_channels, n_classes, kernel_size=1)
+        
+
+        ### For Classification, following Florian's GP-UNet
+        self.GMP = nn.AdaptiveMaxPool2d((1, 1))
+
+    def forward(self, x):
+        blocks = []
+        for i, down in enumerate(self.down_path):
+            x = down(x)
+            if i != len(self.down_path)-1:
+                blocks.append(x)
+                #x = nn.AvgPool2d(x, 2)
+                x = F.avg_pool2d(x, 2)
+        x = self.dropout(x)
+
+        for i, up in enumerate(self.up_path):
+            x = up(x, blocks[-i-1])
+
+        if self.training:
+            x = self.GMP(x)
+            return self.last(x).view(x.shape[0],-1)
+        else:
+            mask = self.last(x)
+            x = self.GMP(x)
+            pred = self.last(x).view(x.shape[0],-1)
+            return pred, mask
+
+class UNetConvBlock(nn.Module):
+    def __init__(self, in_size, out_size, padding, batch_norm, Relu):
+        super(UNetConvBlock, self).__init__()
+        block = []
+
+        block.append(nn.Conv2d(in_size, out_size, kernel_size=3,
+                               padding=int(padding)))
+        if Relu == "Relu":
+            block.append(nn.ReLU())
+        else:
+            block.append(nn.PReLU())
+
+        if batch_norm:
+            block.append(nn.BatchNorm2d(out_size))
+
+        block.append(nn.Conv2d(out_size, out_size, kernel_size=3,
+                               padding=int(padding)))
+        
+        if Relu == "Relu":
+            block.append(nn.ReLU())
+        else:
+            block.append(nn.PReLU())
+
+        if batch_norm:
+            block.append(nn.BatchNorm2d(out_size))
+
+        self.block = nn.Sequential(*block)
+
+    def forward(self, x):
+        out = self.block(x)
+        return out
+
+
+class UNetUpBlock(nn.Module):
+    def __init__(self, in_size, out_size, up_mode, padding, batch_norm, Relu):
+        super(UNetUpBlock, self).__init__()
+
+        self.up_mode = up_mode
+
+        if up_mode == 'upconv':
+            self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2,
+                                         stride=2)
+        elif up_mode == 'bilinear':
+            self.up = nn.Sequential(nn.Upsample(mode='bilinear', scale_factor=2),      #'trilinear'
+                                    nn.Conv2d(in_size, out_size, kernel_size=1))
+        elif 'inc' in up_mode:   
+            self.up = nn.Conv2d(in_size, out_size, kernel_size=1)
+
+
+        self.conv_block = UNetConvBlock(in_size, out_size, padding, batch_norm, Relu)
+
+    def forward(self, x, bridge):
+        if self.up_mode == 'upconv': #  'upconv'
+            up = self.up(x)
+        elif self.up_mode == 'bilinear':
+            up = self.up(x)
+        elif 'inc' in self.up_mode: 
+            x = cF._sinc_interpolate(x, size=[int(x.shape[2]*2), int(x.shape[3]*2)]) #'sinc' ###sth wrong
+            up = self.up(x)
+
+        # bridge = self.center_crop(bridge, up.shape[2:]) #sending shape ignoring 2 digit, so target size start with 0,1,2
+        up = F.interpolate(up, size=bridge.shape[2:], mode='bilinear')
+        out = torch.cat([up, bridge], 1)
+        out = self.conv_block(out)
+
+        return out
+
+#to run it here from this script, uncomment the following
+
+if __name__ == "__main__":                #to run it
+    image = torch.rand(2, 4, 240, 240)    #specify your image: batch size, Channel, height, width
+    model = GP_UNet(in_channels=4, n_classes=3, depth=4, wf=6, up_mode="upsample_Sinc", Relu = "Relu")          #Initialize the model, up_mode = "upconv" or "upsample1" == interpolate mode Bilinear or "upsample" == interpolate mode sinc
+    model.eval()
+    out = model(image)
+    print(model(image))

config.json

@@ -0,0 +1,25 @@
+{
+  "architectures": [
+    "GPReconResNet"
+  ],
+  "auto_map": {
+    "AutoConfig": "GPModelConfigs.GPReconResNetConfig",
+    "AutoModel": "GPModels.GPReconResNet"
+  },
+  "do_batchnorm": false,
+  "forwardV": 0,
+  "in_channels": 1,
+  "is3D": false,
+  "is_relu_leaky": true,
+  "model_type": "GPReconResNet",
+  "n_classes": 3,
+  "out_act": "None",
+  "post_interp_convtrans": false,
+  "res_blocks": 14,
+  "res_drop_prob": 0.5,
+  "starting_nfeatures": 64,
+  "torch_dtype": "float32",
+  "transformers_version": "4.44.2",
+  "updown_blocks": 2,
+  "upinterp_algo": "sinc"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f86cb2396ae9c392ad99f6de61459f8764a368a3073a1244c780dfecb05ced54
+size 69063232