Upload model

RRDB.py

@@ -0,0 +1,118 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+
+class make_dense(nn.Module):
+    def __init__(self,nChannels,GrowthRate,kernel_size=3):
+        super(make_dense,self).__init__()
+        self.conv = nn.Conv3d(nChannels,GrowthRate,kernel_size=kernel_size,padding=(kernel_size-1)//2,bias=True)
+        # self.norm = nn.BatchNorm3d(nChannels)
+    def forward(self,x):
+        # out = self.norm(x)
+        out = F.relu(self.conv(x))
+        out = torch.cat([x,out],dim=1)
+        return out
+
+# class RDB(nn.Module):
+#     def __init__(self,inChannels,outChannels,nDenseLayer,GrowthRate,KernelSize = 3):
+#         super(RDB,self).__init__()
+#         nChannels_ = inChannels
+#         modules = []
+#         for i in range (nDenseLayer):
+#             modules.append(make_dense(nChannels_,GrowthRate,kernel_size=KernelSize))
+#             nChannels_ += GrowthRate
+#         self.dense_layers = nn.Sequential(*modules)
+#         self.conv_1x1 = nn.Conv3d(nChannels_,outChannels,kernel_size=1,padding=0,bias = False)
+#     def forward(self,x):
+#         out = self.dense_layers(x)
+#         out = self.conv_1x1(out)
+#         # out = out + x
+#         return out
+
+class RDB(nn.Module):
+    def __init__(self,inChannels,outChannels,nDenseLayer,GrowthRate,KernelSize = 3,
+                 block_dropout = True, block_dropout_rate = 0.2):
+        super(RDB,self).__init__()
+        nChannels_ = inChannels
+        modules = []
+        for i in range (nDenseLayer):
+            modules.append(make_dense(nChannels_,GrowthRate,kernel_size=KernelSize))
+            nChannels_ += GrowthRate
+        if block_dropout:
+            modules.append(nn.Dropout3d(block_dropout_rate))
+        self.dense_layers = nn.Sequential(*modules)
+        self.conv_1x1 = nn.Conv3d(nChannels_,outChannels,kernel_size=1,padding=0,bias = False)
+    def forward(self,x):
+        out = self.dense_layers(x)
+        out = self.conv_1x1(out)
+        # out = out + x
+        return out
+
+
+class RRDB(nn.Module):
+    def __init__(self,nChannels,nDenseLayers,nInitFeat,GrowthRate,featureFusion=True,kernel_config = [3,3,3,3]):
+        super(RRDB,self).__init__()
+        nChannels_ = nChannels
+        nDenseLayers_ = nDenseLayers
+        nInitFeat_ = nInitFeat
+        GrowthRate_ = GrowthRate
+        self.featureFusion = featureFusion
+
+        #First Convolution
+        self.C1 = nn.Conv3d(nChannels_,nInitFeat_,kernel_size=kernel_config[0],padding=(kernel_config[0]-1)//2,bias=True)
+        # Initialize RDB
+        if self.featureFusion:
+            self.RDB1 = RDB(nInitFeat_,nInitFeat_,nDenseLayers_,GrowthRate_,kernel_config[1])
+            # print(f"RDB1 =========================================== \n {self.RDB1}")
+            self.RDB2 = RDB(nInitFeat_*2,nInitFeat_, nDenseLayers_, GrowthRate_,kernel_config[2])
+            # print(f"RDB2 =========================================== \n {self.RDB2}")
+            self.RDB3 = RDB(nInitFeat_*3,nInitFeat_, nDenseLayers_, GrowthRate_,kernel_config[3])
+            # print(f"RDB3 =========================================== \n {self.RDB3}")
+            self.FF_1X1 = nn.Conv3d(nInitFeat_*4, 1, kernel_size=1, padding=0, bias=True)
+            # print(f"FF1x1 =========================================== \n {self.FF_1X1}")
+        else:
+            self.RDB1 = RDB(nInitFeat_, nDenseLayers_, GrowthRate_, kernel_config[1])
+            self.RDB2 = RDB(nInitFeat_, nDenseLayers_, GrowthRate_, kernel_config[2])
+            self.RDB3 = RDB(nInitFeat_, nDenseLayers_, GrowthRate_, kernel_config[3])
+            self.FF_1X1 = nn.Conv3d(nInitFeat_, 1, kernel_size=1, padding=0, bias=True)
+
+
+        # Feature Fusion
+
+
+        # self.FF_3X3 = nn.Conv3d(nInitFeat_,nInitFeat_,kernel_size=3,padding=1,bias=True)
+
+        # self.final_layer = nn.Conv3d(nInitFeat_,nChannels_,kernel_size=1,padding=0,bias=False)
+
+    def forward(self,x):
+        First = F.relu(self.C1(x))
+        R_1 = self.RDB1(First)
+
+        if self.featureFusion:
+            FF0 = torch.cat([First,R_1],dim = 1)
+            R_2 = self.RDB2(FF0)
+            FF1 = torch.cat([First,R_1,R_2],dim=1)
+            R_3 = self.RDB3(FF1)
+            FF2 = torch.cat([First,R_1, R_2, R_3], dim=1)
+            FF1X1 = self.FF_1X1(FF2)
+        else:
+            R_2 = self.RDB2(R_1)
+            R_3 = self.RDB3(R_2)
+            FF1X1 = self.FF_1X1(R_3)
+
+        # FF2 = torch.cat([R_1,R_2,R_3],dim=1)
+        # FF1X1 = self.FF_1X1(FF2)
+        # FF3X3 = self.FF_3X3(FF1X1)
+        # output = self.final_layer(FF3X3)
+
+        return FF1X1
+
+if __name__ == '__main__':
+    model = RRDB(nChannels=1,nDenseLayers=6,nInitFeat=6,GrowthRate=12,featureFusion=True,kernel_config = [3,3,3,3]).cuda()
+    dimensions = 1, 1, 64, 64, 64
+    x = torch.rand(dimensions)
+    x = x.cuda()
+    out = model(x)
+    print(out.shape)

SRMRIModels.py

@@ -0,0 +1,41 @@
+from transformers import PreTrainedModel
+from .RRDB import RRDB
+from .unet3DMSS import UNetMSS
+from .SRMRIModelsConfigs import RRDBConfiguration, UNetMSSConfiguration
+
+class SRMRIModelUNetMSS(PreTrainedModel):
+    config_class = UNetMSSConfiguration
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = UNetMSS(
+            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, 
+            mss_level=config.mss_level,
+            mss_fromlatent=config.mss_fromlatent, 
+            mss_up=config.mss_up, 
+            mss_interpb4=config.mss_interpb4)
+    def forward(self, x):
+        return self.model.forward(x)
+    
+
+class SRMRIModelRRDB(PreTrainedModel):
+    config_class = RRDBConfiguration
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = RRDB(
+            nChannels=config.nChannels, 
+            nDenseLayers=config.nDenseLayers, 
+            nInitFeat=config.nInitFeat, 
+            GrowthRate=config.GrowthRate, 
+            featureFusion=config.featureFusion,
+            kernel_config=config.kernel_config,
+        )
+
+    def forward(self, x):
+        return self.model.forward(x)

SRMRIModelsConfigs.py

@@ -0,0 +1,52 @@
+from transformers import PretrainedConfig
+from typing import List
+
+class RRDBConfiguration(PretrainedConfig):
+    model_type = "SRMRIModelRRDB"
+    def __init__(
+            self,
+            nChannels=1, 
+            nDenseLayers=6, 
+            nInitFeat=6, 
+            GrowthRate=12, 
+            featureFusion=True,
+            kernel_config=[3, 3, 3, 3],
+            **kwargs):
+        self.nChannels = nChannels
+        self.nDenseLayers = nDenseLayers
+        self.nInitFeat = nInitFeat
+        self.GrowthRate = GrowthRate
+        self.featureFusion = featureFusion
+        self.kernel_config = kernel_config
+        super().__init__(**kwargs)
+
+class UNetMSSConfiguration(PretrainedConfig):
+    model_type = "SRMRIModelUNetMSS"
+    def __init__(
+            self,
+            in_channels=1, 
+            n_classes=1, 
+            depth=3, 
+            wf=6, 
+            padding=True,
+            batch_norm=False, 
+            up_mode='upconv', 
+            dropout=False, 
+            mss_level=2,
+            mss_fromlatent=True, 
+            mss_up="trilinear", 
+            mss_interpb4=True,
+            **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.mss_level = mss_level
+        self.mss_fromlatent = mss_fromlatent
+        self.mss_up = mss_up
+        self.mss_interpb4 = mss_interpb4
+        super().__init__(**kwargs)

config.json

@@ -0,0 +1,23 @@
+{
+  "GrowthRate": 12,
+  "architectures": [
+    "SRMRIModelRRDB"
+  ],
+  "auto_map": {
+    "AutoConfig": "SRMRIModelsConfigs.RRDBConfiguration",
+    "AutoModel": "SRMRIModels.SRMRIModelRRDB"
+  },
+  "featureFusion": true,
+  "kernel_config": [
+    3,
+    3,
+    3,
+    3
+  ],
+  "model_type": "SRMRIModelRRDB",
+  "nChannels": 1,
+  "nDenseLayers": 6,
+  "nInitFeat": 6,
+  "torch_dtype": "float32",
+  "transformers_version": "4.44.0"
+}

model.safetensors

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

unet3DMSS.py

@@ -0,0 +1,189 @@
+# Adapted from https://discuss.pytorch.org/t/unet-implementation/426
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+# import torchcomplex.nn.functional as cF
+
+__author__ = "Soumick Chatterjee, Chompunuch Sarasaen"
+__copyright__ = "Copyright 2020, Faculty of Computer Science, Otto von Guericke University Magdeburg, Germany"
+__credits__ = ["Soumick Chatterjee", "Chompunuch Sarasaen"]
+__license__ = "GPL"
+__version__ = "1.0.0"
+__maintainer__ = "Soumick Chatterjee"
+__email__ = "[email protected]"
+__status__ = "Production"
+
+
+class UNetMSS(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' will use bilinear upsampling.
+    """
+    def __init__(self, in_channels=1, n_classes=1, depth=3, wf=6, padding=True,
+                 batch_norm=False, up_mode='upconv', dropout=False, mss_level=2, mss_fromlatent=True, 
+                 mss_up="trilinear", mss_interpb4=False):
+        super(UNetMSS, self).__init__()
+        assert up_mode in ('upconv', 'upsample')
+        self.padding = padding
+        self.depth = depth
+        self.dropout = nn.Dropout3d() if dropout else nn.Sequential()
+        prev_channels = in_channels
+        self.down_path = nn.ModuleList()
+        up_out_features = []
+        for i in range(depth):
+            self.down_path.append(UNetConvBlock(prev_channels, 2**(wf+i),
+                                                padding, batch_norm))
+            prev_channels = 2**(wf+i)
+
+        if mss_fromlatent:
+            mss_features = [prev_channels]
+        else:
+            mss_features = []
+
+        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))
+            prev_channels = 2**(wf+i)
+            up_out_features.append(prev_channels)
+
+        self.last = nn.Conv3d(prev_channels, n_classes, kernel_size=1)
+
+        mss_features += up_out_features[len(up_out_features)-1-mss_level if not mss_fromlatent 
+                                                                        else len(up_out_features)-1-mss_level+1:-1]
+
+        self.mss_level = mss_level
+        self.mss_up = mss_up
+        self.mss_fromlatent = mss_fromlatent
+        self.mss_interpb4 = mss_interpb4
+        self.mss_convs = nn.ModuleList()
+        for i in range(self.mss_level):
+            self.mss_convs.append(nn.Conv3d(mss_features[i], n_classes, kernel_size=1))
+        if self.mss_level == 1:
+            self.mss_coeff = [0.5]
+        else:
+            lmbda = []
+            for i in range(self.mss_level-1, -1, -1):
+                lmbda.append(2**i)
+            self.mss_coeff = []
+            fact = 1.0 / sum(lmbda)
+            for i in range(self.mss_level-1):
+                self.mss_coeff.append(fact*lmbda[i])
+            self.mss_coeff.append(1.0 - sum(self.mss_coeff))          
+            self.mss_coeff.reverse()
+
+
+    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 = F.avg_pool3d(x, 2)
+        x = self.dropout(x)
+            
+        if self.mss_fromlatent:
+            mss = [x]
+        else:
+            mss = []
+
+        for i, up in enumerate(self.up_path):
+            x = up(x, blocks[-i-1])
+            if self.training and ((len(self.up_path)-1-i <= self.mss_level) and not(i+1 == len(self.up_path))):
+                mss.append(x)
+
+        if self.training:
+            for i in range(len(mss)):
+                if not self.mss_interpb4:
+                    mss[i] = F.interpolate(self.mss_convs[i](mss[i]), size=x.shape[2:], mode=self.mss_up) 
+                else:
+                    mss[i] = self.mss_convs[i](F.interpolate(mss[i], size=x.shape[2:], mode=self.mss_up)) 
+            
+            return self.last(x), mss
+        else:
+            return self.last(x)
+
+class UNetConvBlock(nn.Module):
+    def __init__(self, in_size, out_size, padding, batch_norm):
+        super(UNetConvBlock, self).__init__()
+        block = []
+
+        block.append(nn.Conv3d(in_size, out_size, kernel_size=3,
+                               padding=int(padding)))
+        block.append(nn.ReLU())
+        if batch_norm:
+            block.append(nn.BatchNorm3d(out_size))
+
+        block.append(nn.Conv3d(out_size, out_size, kernel_size=3,
+                               padding=int(padding)))
+        block.append(nn.ReLU())
+        if batch_norm:
+            block.append(nn.BatchNorm3d(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):
+        super(UNetUpBlock, self).__init__()
+        if up_mode == 'upconv':
+            self.up = nn.ConvTranspose3d(in_size, out_size, kernel_size=2,
+                                         stride=2)
+        elif up_mode == 'upsample':
+            self.up = nn.Sequential(nn.Upsample(mode='trilinear', scale_factor=2),
+                                    nn.Conv3d(in_size, out_size, kernel_size=1))
+
+        self.conv_block = UNetConvBlock(in_size, out_size, padding, batch_norm)
+
+    def center_crop(self, layer, target_size):
+        _, _, layer_depth, layer_height, layer_width = layer.size()
+        diff_z = (layer_depth - target_size[0]) // 2
+        diff_y = (layer_height - target_size[1]) // 2
+        diff_x = (layer_width - target_size[2]) // 2
+        return layer[:, :, diff_z:(diff_z + target_size[0]), diff_y:(diff_y + target_size[1]), diff_x:(diff_x + target_size[2])]
+        #  _, _, layer_height, layer_width = layer.size() #for 2D data
+        # diff_y = (layer_height - target_size[0]) // 2
+        # diff_x = (layer_width - target_size[1]) // 2
+        # return layer[:, :, diff_y:(diff_y + target_size[0]), diff_x:(diff_x + target_size[1])]
+
+    def forward(self, x, bridge):
+        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='trilinear')
+        out = torch.cat([up, bridge], 1)
+        out = self.conv_block(out)
+
+        return out
+
+
+if __name__ == "__main__":
+    model = UNetMSS(in_channels=1, n_classes=1, depth=4, wf=6, padding=True,
+                    batch_norm=False, up_mode='upconv', dropout=True, mss_level=3,
+                    mss_fromlatent=True, mss_up="trilinear", mss_interpb4=True).cuda()
+
+    print(model)