Upload ReconResNet

ReconResNet.py

@@ -0,0 +1,25 @@
+from transformers import PreTrainedModel
+from .ReconResNetBase import ReconResNetBase
+from .ReconResNetConfig import ReconResNetConfig
+
+class ReconResNet(PreTrainedModel):
+    config_class = ReconResNetConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = ReconResNetBase(
+            in_channels=config.in_channels,
+            out_channels=config.out_channels,
+            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,
+            is_replicatepad=config.is_replicatepad,
+            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)

ReconResNetBase.py

@@ -0,0 +1,267 @@
+#!/usr/bin/env python
+
+# This model is part of the paper "ReconResNet: Regularised Residual Learning for MR Image Reconstruction of Undersampled Cartesian and Radial Data" (https://doi.org/10.1016/j.compbiomed.2022.105321)
+# and has been published on GitHub: https://github.com/soumickmj/NCC1701/blob/main/Bridge/WarpDrives/ReconResNet/ReconResNet.py
+
+import torch.nn as nn
+from tricorder.torch.transforms import Interpolator
+
+__author__ = "Soumick Chatterjee"
+__copyright__ = "Copyright 2019, Soumick Chatterjee & OvGU:ESF:MEMoRIAL"
+__credits__ = ["Soumick Chatterjee"]
+
+__license__ = "apache-2.0"
+__version__ = "1.0.0"
+__email__ = "soumick@live.com"
+__status__ = "Published"
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_features, 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="convtrans", 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 == "convtrans":
+            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 == "convtrans":
+            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 ReconResNetBase(nn.Module):
+    def __init__(self, in_channels=1, out_channels=1, res_blocks=14, starting_nfeatures=64, updown_blocks=2, is_relu_leaky=True, do_batchnorm=False, res_drop_prob=0.2,
+                 is_replicatepad=0, out_act="sigmoid", forwardV=0, upinterp_algo='convtrans', 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(ReconResNetBase, self).__init__()
+
+        layers = {}
+        if is3D:
+            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
+            if is_replicatepad == 0:
+                layers["layer_pad"] = nn.ReflectionPad3d
+            elif is_replicatepad == 1:
+                layers["layer_pad"] = nn.ReplicationPad3d
+            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
+            if is_replicatepad == 0:
+                layers["layer_pad"] = nn.ReflectionPad2d
+            elif is_replicatepad == 1:
+                layers["layer_pad"] = nn.ReplicationPad2d
+            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 == "convtrans" else self.upinterp_algo)
+
+        # 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_pad(3),
+                     layer_conv(starting_nfeatures, out_channels, 7), ]
+
+        if out_act == "sigmoid":
+            finalconv += [nn.Sigmoid(), ]
+        elif out_act == "relu":
+            finalconv += [act_relu(), ]
+        elif out_act == "tanh":
+            finalconv += [nn.Tanh(), ]
+
+        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)
+
+        if self.forwardV == 0:
+            self.forward = self.forwardV0
+        elif self.forwardV == 1:
+            self.forward = self.forwardV1
+        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 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.finalconv(x)
+
+    def forwardV1(self, x):
+        # v1: input is added to the final output
+        out = self.intialConv(x)
+        shapes = []
+        for downblock in self.downsam:
+            shapes.append(out.shape[2:])
+            out = downblock(out)
+        out = self.resblocks(out)
+        for i, upblock in enumerate(self.upsam):
+            out = upblock(out, shapes[-1-i])
+        return x + self.finalconv(out)
+
+    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 x + self.finalconv(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 x + self.finalconv(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 x + self.finalconv(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 x + self.finalconv(outs.pop())

ReconResNetConfig.py

@@ -0,0 +1,37 @@
+from transformers import PretrainedConfig
+from typing import List
+
+class ReconResNetConfig(PretrainedConfig):
+    model_type = "ReconResNet"
+    def __init__(
+            self,
+            in_channels=1, 
+            out_channels=1, 
+            res_blocks=14, 
+            starting_nfeatures=64, 
+            updown_blocks=2, 
+            is_relu_leaky=True, 
+            do_batchnorm=False, 
+            res_drop_prob=0.2,
+            is_replicatepad=0, 
+            out_act="sigmoid", 
+            forwardV=0, 
+            upinterp_algo='convtrans', 
+            post_interp_convtrans=False, 
+            is3D=False,
+            **kwargs):
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        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.is_replicatepad = is_replicatepad
+        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)

config.json

@@ -0,0 +1,26 @@
+{
+  "architectures": [
+    "ReconResNet"
+  ],
+  "auto_map": {
+    "AutoConfig": "ReconResNetConfig.ReconResNetConfig",
+    "AutoModel": "ReconResNet.ReconResNet"
+  },
+  "do_batchnorm": false,
+  "forwardV": 0,
+  "in_channels": 1,
+  "is3D": false,
+  "is_relu_leaky": true,
+  "is_replicatepad": 0,
+  "model_type": "ReconResNet",
+  "out_act": "sigmoid",
+  "out_channels": 1,
+  "post_interp_convtrans": false,
+  "res_blocks": 14,
+  "res_drop_prob": 0.2,
+  "starting_nfeatures": 64,
+  "torch_dtype": "float32",
+  "transformers_version": "4.44.2",
+  "updown_blocks": 2,
+  "upinterp_algo": "convtrans"
+}

model.safetensors

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