Upload with huggingface_hub

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+AllinonSAM/eval/lits/output_demo.nii filter=lfs diff=lfs merge=lfs -text
+AllinonSAM/wandb/run-20241018_210810-zrrx3qz9/run-zrrx3qz9.wandb filter=lfs diff=lfs merge=lfs -text
+AllinonSAM/wandb/run-20241018_162125-i4stmvih/run-i4stmvih.wandb filter=lfs diff=lfs merge=lfs -text
+AllinonSAM/wandb/run-20240915_215641-1usjns7w/run-1usjns7w.wandb filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,14 @@
+*.pyc
+*.cpython-38.pyc
+*.pth
+*.gz
+*.zip
+*.png
+*.jpg
+*.JPG
+*.tif
+*.bmp
+*.out
+*.txt
+AllinonSAM/wandb/
+__pycache__

AllinonSAM/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Ahmed Heakl
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

AllinonSAM/axialnet.py

@@ -0,0 +1,730 @@
+import pdb
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from utils import *
+import pdb
+import matplotlib.pyplot as plt
+ 
+import random
+
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class AxialAttention(nn.Module):
+    def __init__(self, in_planes, out_planes, groups=8, kernel_size=56,
+                 stride=1, bias=False, width=False):
+        assert (in_planes % groups == 0) and (out_planes % groups == 0)
+        super(AxialAttention, self).__init__()
+        self.in_planes = in_planes
+        self.out_planes = out_planes
+        self.groups = groups
+        self.group_planes = out_planes // groups
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.bias = bias
+        self.width = width
+
+        # Multi-head self attention
+        self.qkv_transform = qkv_transform(in_planes, out_planes * 2, kernel_size=1, stride=1,
+                                           padding=0, bias=False)
+        self.bn_qkv = nn.BatchNorm1d(out_planes * 2)
+        self.bn_similarity = nn.BatchNorm2d(groups * 3)
+
+        self.bn_output = nn.BatchNorm1d(out_planes * 2)
+
+        # Position embedding
+        self.relative = nn.Parameter(torch.randn(self.group_planes * 2, kernel_size * 2 - 1), requires_grad=True)
+        query_index = torch.arange(kernel_size).unsqueeze(0)
+        key_index = torch.arange(kernel_size).unsqueeze(1)
+        relative_index = key_index - query_index + kernel_size - 1
+        self.register_buffer('flatten_index', relative_index.view(-1))
+        if stride > 1:
+            self.pooling = nn.AvgPool2d(stride, stride=stride)
+
+        self.reset_parameters()
+
+    def forward(self, x):
+        # pdb.set_trace()
+        if self.width:
+            x = x.permute(0, 2, 1, 3)
+        else:
+            x = x.permute(0, 3, 1, 2)  # N, W, C, H
+        N, W, C, H = x.shape
+        x = x.contiguous().view(N * W, C, H)
+
+        # Transformations
+        qkv = self.bn_qkv(self.qkv_transform(x))
+        q, k, v = torch.split(qkv.reshape(N * W, self.groups, self.group_planes * 2, H), [self.group_planes // 2, self.group_planes // 2, self.group_planes], dim=2)
+
+        # Calculate position embedding
+        all_embeddings = torch.index_select(self.relative, 1, self.flatten_index).view(self.group_planes * 2, self.kernel_size, self.kernel_size)
+        q_embedding, k_embedding, v_embedding = torch.split(all_embeddings, [self.group_planes // 2, self.group_planes // 2, self.group_planes], dim=0)
+        
+        qr = torch.einsum('bgci,cij->bgij', q, q_embedding)
+        kr = torch.einsum('bgci,cij->bgij', k, k_embedding).transpose(2, 3)
+        
+        qk = torch.einsum('bgci, bgcj->bgij', q, k)
+        
+        stacked_similarity = torch.cat([qk, qr, kr], dim=1)
+        stacked_similarity = self.bn_similarity(stacked_similarity).view(N * W, 3, self.groups, H, H).sum(dim=1)
+        #stacked_similarity = self.bn_qr(qr) + self.bn_kr(kr) + self.bn_qk(qk)
+        # (N, groups, H, H, W)
+        similarity = F.softmax(stacked_similarity, dim=3)
+        sv = torch.einsum('bgij,bgcj->bgci', similarity, v)
+        sve = torch.einsum('bgij,cij->bgci', similarity, v_embedding)
+        stacked_output = torch.cat([sv, sve], dim=-1).view(N * W, self.out_planes * 2, H)
+        output = self.bn_output(stacked_output).view(N, W, self.out_planes, 2, H).sum(dim=-2)
+
+        if self.width:
+            output = output.permute(0, 2, 1, 3)
+        else:
+            output = output.permute(0, 2, 3, 1)
+
+        if self.stride > 1:
+            output = self.pooling(output)
+
+        return output
+
+    def reset_parameters(self):
+        self.qkv_transform.weight.data.normal_(0, math.sqrt(1. / self.in_planes))
+        #nn.init.uniform_(self.relative, -0.1, 0.1)
+        nn.init.normal_(self.relative, 0., math.sqrt(1. / self.group_planes))
+
+class AxialAttention_dynamic(nn.Module):
+    def __init__(self, in_planes, out_planes, groups=8, kernel_size=56,
+                 stride=1, bias=False, width=False):
+        assert (in_planes % groups == 0) and (out_planes % groups == 0)
+        super(AxialAttention_dynamic, self).__init__()
+        self.in_planes = in_planes
+        self.out_planes = out_planes
+        self.groups = groups
+        self.group_planes = out_planes // groups
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.bias = bias
+        self.width = width
+
+        # Multi-head self attention
+        self.qkv_transform = qkv_transform(in_planes, out_planes * 2, kernel_size=1, stride=1,
+                                           padding=0, bias=False)
+        self.bn_qkv = nn.BatchNorm1d(out_planes * 2)
+        self.bn_similarity = nn.BatchNorm2d(groups * 3)
+        self.bn_output = nn.BatchNorm1d(out_planes * 2)
+
+        # Priority on encoding
+
+        ## Initial values 
+
+        self.f_qr = nn.Parameter(torch.tensor(0.1),  requires_grad=False) 
+        self.f_kr = nn.Parameter(torch.tensor(0.1),  requires_grad=False)
+        self.f_sve = nn.Parameter(torch.tensor(0.1),  requires_grad=False)
+        self.f_sv = nn.Parameter(torch.tensor(1.0),  requires_grad=False)
+
+
+        # Position embedding
+        self.relative = nn.Parameter(torch.randn(self.group_planes * 2, kernel_size * 2 - 1), requires_grad=True)
+        query_index = torch.arange(kernel_size).unsqueeze(0)
+        key_index = torch.arange(kernel_size).unsqueeze(1)
+        relative_index = key_index - query_index + kernel_size - 1
+        self.register_buffer('flatten_index', relative_index.view(-1))
+        if stride > 1:
+            self.pooling = nn.AvgPool2d(stride, stride=stride)
+
+        self.reset_parameters()
+        # self.print_para()
+
+    def forward(self, x):
+        if self.width:
+            x = x.permute(0, 2, 1, 3)
+        else:
+            x = x.permute(0, 3, 1, 2)  # N, W, C, H
+        N, W, C, H = x.shape
+        x = x.contiguous().view(N * W, C, H)
+
+        # Transformations
+        qkv = self.bn_qkv(self.qkv_transform(x))
+        q, k, v = torch.split(qkv.reshape(N * W, self.groups, self.group_planes * 2, H), [self.group_planes // 2, self.group_planes // 2, self.group_planes], dim=2)
+
+        # Calculate position embedding
+        all_embeddings = torch.index_select(self.relative, 1, self.flatten_index).view(self.group_planes * 2, self.kernel_size, self.kernel_size)
+        q_embedding, k_embedding, v_embedding = torch.split(all_embeddings, [self.group_planes // 2, self.group_planes // 2, self.group_planes], dim=0)
+        qr = torch.einsum('bgci,cij->bgij', q, q_embedding)
+        kr = torch.einsum('bgci,cij->bgij', k, k_embedding).transpose(2, 3)
+        qk = torch.einsum('bgci, bgcj->bgij', q, k)
+
+
+        # multiply by factors
+        qr = torch.mul(qr, self.f_qr)
+        kr = torch.mul(kr, self.f_kr)
+
+        stacked_similarity = torch.cat([qk, qr, kr], dim=1)
+        stacked_similarity = self.bn_similarity(stacked_similarity).view(N * W, 3, self.groups, H, H).sum(dim=1)
+        #stacked_similarity = self.bn_qr(qr) + self.bn_kr(kr) + self.bn_qk(qk)
+        # (N, groups, H, H, W)
+        similarity = F.softmax(stacked_similarity, dim=3)
+        sv = torch.einsum('bgij,bgcj->bgci', similarity, v)
+        sve = torch.einsum('bgij,cij->bgci', similarity, v_embedding)
+
+        # multiply by factors
+        sv = torch.mul(sv, self.f_sv)
+        sve = torch.mul(sve, self.f_sve)
+
+        stacked_output = torch.cat([sv, sve], dim=-1).view(N * W, self.out_planes * 2, H)
+        output = self.bn_output(stacked_output).view(N, W, self.out_planes, 2, H).sum(dim=-2)
+
+        if self.width:
+            output = output.permute(0, 2, 1, 3)
+        else:
+            output = output.permute(0, 2, 3, 1)
+
+        if self.stride > 1:
+            output = self.pooling(output)
+
+        return output
+    def reset_parameters(self):
+        self.qkv_transform.weight.data.normal_(0, math.sqrt(1. / self.in_planes))
+        #nn.init.uniform_(self.relative, -0.1, 0.1)
+        nn.init.normal_(self.relative, 0., math.sqrt(1. / self.group_planes))
+
+class AxialAttention_wopos(nn.Module):
+    def __init__(self, in_planes, out_planes, groups=8, kernel_size=56,
+                 stride=1, bias=False, width=False):
+        assert (in_planes % groups == 0) and (out_planes % groups == 0)
+        super(AxialAttention_wopos, self).__init__()
+        self.in_planes = in_planes
+        self.out_planes = out_planes
+        self.groups = groups
+        self.group_planes = out_planes // groups
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.bias = bias
+        self.width = width
+
+        # Multi-head self attention
+        self.qkv_transform = qkv_transform(in_planes, out_planes * 2, kernel_size=1, stride=1,
+                                           padding=0, bias=False)
+        self.bn_qkv = nn.BatchNorm1d(out_planes * 2)
+        self.bn_similarity = nn.BatchNorm2d(groups )
+
+        self.bn_output = nn.BatchNorm1d(out_planes * 1)
+
+        if stride > 1:
+            self.pooling = nn.AvgPool2d(stride, stride=stride)
+
+        self.reset_parameters()
+
+    def forward(self, x):
+        if self.width:
+            x = x.permute(0, 2, 1, 3)
+        else:
+            x = x.permute(0, 3, 1, 2)  # N, W, C, H
+        N, W, C, H = x.shape
+        x = x.contiguous().view(N * W, C, H)
+
+        # Transformations
+        qkv = self.bn_qkv(self.qkv_transform(x))
+        q, k, v = torch.split(qkv.reshape(N * W, self.groups, self.group_planes * 2, H), [self.group_planes // 2, self.group_planes // 2, self.group_planes], dim=2)
+
+        qk = torch.einsum('bgci, bgcj->bgij', q, k)
+
+        stacked_similarity = self.bn_similarity(qk).reshape(N * W, 1, self.groups, H, H).sum(dim=1).contiguous()
+
+        similarity = F.softmax(stacked_similarity, dim=3)
+        sv = torch.einsum('bgij,bgcj->bgci', similarity, v)
+
+        sv = sv.reshape(N*W,self.out_planes * 1, H).contiguous()
+        output = self.bn_output(sv).reshape(N, W, self.out_planes, 1, H).sum(dim=-2).contiguous()
+
+
+        if self.width:
+            output = output.permute(0, 2, 1, 3)
+        else:
+            output = output.permute(0, 2, 3, 1)
+
+        if self.stride > 1:
+            output = self.pooling(output)
+
+        return output
+
+    def reset_parameters(self):
+        self.qkv_transform.weight.data.normal_(0, math.sqrt(1. / self.in_planes))
+        #nn.init.uniform_(self.relative, -0.1, 0.1)
+        # nn.init.normal_(self.relative, 0., math.sqrt(1. / self.group_planes))
+
+#end of attn definition
+
+class AxialBlock(nn.Module):
+    expansion = 2
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None, kernel_size=56):
+        super(AxialBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.))
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv_down = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.hight_block = AxialAttention(width, width, groups=groups, kernel_size=kernel_size)
+        self.width_block = AxialAttention(width, width, groups=groups, kernel_size=kernel_size, stride=stride, width=True)
+        self.conv_up = conv1x1(width, planes * self.expansion)
+        self.bn2 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv_down(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+        # print(out.shape)
+        out = self.hight_block(out)
+        out = self.width_block(out)
+        out = self.relu(out)
+
+        out = self.conv_up(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+class AxialBlock_dynamic(nn.Module):
+    expansion = 2
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None, kernel_size=56):
+        super(AxialBlock_dynamic, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.))
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv_down = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.hight_block = AxialAttention_dynamic(width, width, groups=groups, kernel_size=kernel_size)
+        self.width_block = AxialAttention_dynamic(width, width, groups=groups, kernel_size=kernel_size, stride=stride, width=True)
+        self.conv_up = conv1x1(width, planes * self.expansion)
+        self.bn2 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv_down(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.hight_block(out)
+        out = self.width_block(out)
+        out = self.relu(out)
+
+        out = self.conv_up(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+class AxialBlock_wopos(nn.Module):
+    expansion = 2
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None, kernel_size=56):
+        super(AxialBlock_wopos, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        # print(kernel_size)
+        width = int(planes * (base_width / 64.))
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv_down = conv1x1(inplanes, width)
+        self.conv1 = nn.Conv2d(width, width, kernel_size = 1)
+        self.bn1 = norm_layer(width)
+        self.hight_block = AxialAttention_wopos(width, width, groups=groups, kernel_size=kernel_size)
+        self.width_block = AxialAttention_wopos(width, width, groups=groups, kernel_size=kernel_size, stride=stride, width=True)
+        self.conv_up = conv1x1(width, planes * self.expansion)
+        self.bn2 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        # pdb.set_trace()
+
+        out = self.conv_down(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+        # print(out.shape)
+        out = self.hight_block(out)
+        out = self.width_block(out)
+
+        out = self.relu(out)
+
+        out = self.conv_up(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+#end of block definition
+
+
+class ResAxialAttentionUNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=2, zero_init_residual=True,
+                 groups=8, width_per_group=64, replace_stride_with_dilation=None,
+                 norm_layer=None, s=0.125, img_size = 128,imgchan = 3):
+        super(ResAxialAttentionUNet, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.inplanes = int(64 * s)
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(imgchan, self.inplanes, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.conv2 = nn.Conv2d(self.inplanes, 128, kernel_size=3, stride=1, padding=1, bias=False)
+        self.conv3 = nn.Conv2d(128, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = norm_layer(self.inplanes)
+        self.bn2 = norm_layer(128)
+        self.bn3 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        # self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, int(128 * s), layers[0], kernel_size= (img_size//2))
+        self.layer2 = self._make_layer(block, int(256 * s), layers[1], stride=2, kernel_size=(img_size//2),
+                                       dilate=replace_stride_with_dilation[0])
+        self.layer3 = self._make_layer(block, int(512 * s), layers[2], stride=2, kernel_size=(img_size//4),
+                                       dilate=replace_stride_with_dilation[1])
+        self.layer4 = self._make_layer(block, int(1024 * s), layers[3], stride=2, kernel_size=(img_size//8),
+                                       dilate=replace_stride_with_dilation[2])
+        
+        # Decoder
+        self.decoder1 = nn.Conv2d(int(1024 *2*s)      ,        int(1024*2*s), kernel_size=3, stride=2, padding=1)
+        self.decoder2 = nn.Conv2d(int(1024  *2*s)     , int(1024*s), kernel_size=3, stride=1, padding=1)
+        self.decoder3 = nn.Conv2d(int(1024*s),  int(512*s), kernel_size=3, stride=1, padding=1)
+        self.decoder4 = nn.Conv2d(int(512*s) ,  int(256*s), kernel_size=3, stride=1, padding=1)
+        self.decoder5 = nn.Conv2d(int(256*s) , int(128*s) , kernel_size=3, stride=1, padding=1)
+        self.adjust   = nn.Conv2d(int(128*s) , num_classes, kernel_size=1, stride=1, padding=0)
+        self.soft     = nn.Softmax(dim=1)
+
+
+    def _make_layer(self, block, planes, blocks, kernel_size=56, stride=1, dilate=False):
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, groups=self.groups,
+                            base_width=self.base_width, dilation=previous_dilation, 
+                            norm_layer=norm_layer, kernel_size=kernel_size))
+        self.inplanes = planes * block.expansion
+        if stride != 1:
+            kernel_size = kernel_size // 2
+
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, groups=self.groups,
+                                base_width=self.base_width, dilation=self.dilation,
+                                norm_layer=norm_layer, kernel_size=kernel_size))
+
+        return nn.Sequential(*layers)
+
+    def _forward_impl(self, x):
+        
+        # AxialAttention Encoder
+        # pdb.set_trace()
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+        x = self.conv3(x)
+        x = self.bn3(x)
+        x = self.relu(x)
+
+        x1 = self.layer1(x)
+
+        x2 = self.layer2(x1)
+        # print(x2.shape)
+        x3 = self.layer3(x2)
+        # print(x3.shape)
+        x4 = self.layer4(x3)
+
+        x = F.relu(F.interpolate(self.decoder1(x4), scale_factor=(2,2), mode ='bilinear'))
+        x = torch.add(x, x4)
+        x = F.relu(F.interpolate(self.decoder2(x) , scale_factor=(2,2), mode ='bilinear'))
+        x = torch.add(x, x3)
+        x = F.relu(F.interpolate(self.decoder3(x) , scale_factor=(2,2), mode ='bilinear'))
+        x = torch.add(x, x2)
+        x = F.relu(F.interpolate(self.decoder4(x) , scale_factor=(2,2), mode ='bilinear'))
+        x = torch.add(x, x1)
+        x = F.relu(F.interpolate(self.decoder5(x) , scale_factor=(2,2), mode ='bilinear'))
+        x = self.adjust(F.relu(x))
+        # pdb.set_trace()
+        return x
+
+    def forward(self, x):
+        return self._forward_impl(x)
+
+class medt_net(nn.Module):
+
+    def __init__(self, block, block_2, layers, num_classes=2, zero_init_residual=True,
+                 groups=8, width_per_group=64, replace_stride_with_dilation=None,
+                 norm_layer=None, s=0.125, img_size = 128,imgchan = 3):
+        super(medt_net, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.inplanes = int(64 * s)
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(imgchan, self.inplanes, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.conv2 = nn.Conv2d(self.inplanes, 128, kernel_size=3, stride=1, padding=1, bias=False)
+        self.conv3 = nn.Conv2d(128, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = norm_layer(self.inplanes)
+        self.bn2 = norm_layer(128)
+        self.bn3 = norm_layer(self.inplanes)
+        # self.conv1 = nn.Conv2d(1, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        # self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, int(128 * s), layers[0], kernel_size= (img_size//2))
+        self.layer2 = self._make_layer(block, int(256 * s), layers[1], stride=2, kernel_size=(img_size//2),
+                                       dilate=replace_stride_with_dilation[0])
+        # self.layer3 = self._make_layer(block, int(512 * s), layers[2], stride=2, kernel_size=(img_size//4),
+        #                                dilate=replace_stride_with_dilation[1])
+        # self.layer4 = self._make_layer(block, int(1024 * s), layers[3], stride=2, kernel_size=(img_size//8),
+        #                                dilate=replace_stride_with_dilation[2])
+        
+        # Decoder
+        # self.decoder1 = nn.Conv2d(int(1024 *2*s)      ,        int(1024*2*s), kernel_size=3, stride=2, padding=1)
+        # self.decoder2 = nn.Conv2d(int(1024  *2*s)     , int(1024*s), kernel_size=3, stride=1, padding=1)
+        # self.decoder3 = nn.Conv2d(int(1024*s),  int(512*s), kernel_size=3, stride=1, padding=1)
+        self.decoder4 = nn.Conv2d(int(512*s) ,  int(256*s), kernel_size=3, stride=1, padding=1)
+        self.decoder5 = nn.Conv2d(int(256*s) , int(128*s) , kernel_size=3, stride=1, padding=1)
+        self.adjust   = nn.Conv2d(int(128*s) , num_classes, kernel_size=1, stride=1, padding=0)
+        self.soft     = nn.Softmax(dim=1)
+
+
+        self.conv1_p = nn.Conv2d(imgchan, self.inplanes, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.conv2_p = nn.Conv2d(self.inplanes,128, kernel_size=3, stride=1, padding=1,
+                               bias=False)
+        self.conv3_p = nn.Conv2d(128, self.inplanes, kernel_size=3, stride=1, padding=1,
+                               bias=False)
+        # self.conv1 = nn.Conv2d(1, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1_p = norm_layer(self.inplanes)
+        self.bn2_p = norm_layer(128)
+        self.bn3_p = norm_layer(self.inplanes)
+
+        self.relu_p = nn.ReLU(inplace=True)
+
+        img_size_p = img_size // 4
+
+        self.layer1_p = self._make_layer(block_2, int(128 * s), layers[0], kernel_size= (img_size_p//2))
+        self.layer2_p = self._make_layer(block_2, int(256 * s), layers[1], stride=2, kernel_size=(img_size_p//2),
+                                       dilate=replace_stride_with_dilation[0])
+        self.layer3_p = self._make_layer(block_2, int(512 * s), layers[2], stride=2, kernel_size=(img_size_p//4),
+                                       dilate=replace_stride_with_dilation[1])
+        self.layer4_p = self._make_layer(block_2, int(1024 * s), layers[3], stride=2, kernel_size=(img_size_p//8),
+                                       dilate=replace_stride_with_dilation[2])
+        
+        # Decoder
+        self.decoder1_p = nn.Conv2d(int(1024 *2*s)      ,        int(1024*2*s), kernel_size=3, stride=2, padding=1)
+        self.decoder2_p = nn.Conv2d(int(1024  *2*s)     , int(1024*s), kernel_size=3, stride=1, padding=1)
+        self.decoder3_p = nn.Conv2d(int(1024*s),  int(512*s), kernel_size=3, stride=1, padding=1)
+        self.decoder4_p = nn.Conv2d(int(512*s) ,  int(256*s), kernel_size=3, stride=1, padding=1)
+        self.decoder5_p = nn.Conv2d(int(256*s) , int(128*s) , kernel_size=3, stride=1, padding=1)
+
+        self.decoderf = nn.Conv2d(int(128*s) , int(128*s) , kernel_size=3, stride=1, padding=1)
+        self.adjust_p   = nn.Conv2d(int(128*s) , num_classes, kernel_size=1, stride=1, padding=0)
+        self.soft_p     = nn.Softmax(dim=1)
+
+
+    def _make_layer(self, block, planes, blocks, kernel_size=56, stride=1, dilate=False):
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, groups=self.groups,
+                            base_width=self.base_width, dilation=previous_dilation, 
+                            norm_layer=norm_layer, kernel_size=kernel_size))
+        self.inplanes = planes * block.expansion
+        if stride != 1:
+            kernel_size = kernel_size // 2
+
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, groups=self.groups,
+                                base_width=self.base_width, dilation=self.dilation,
+                                norm_layer=norm_layer, kernel_size=kernel_size))
+
+        return nn.Sequential(*layers)
+
+    def _forward_impl(self, x):
+
+        xin = x.clone()
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+        x = self.conv3(x)
+        x = self.bn3(x)
+        # x = F.max_pool2d(x,2,2)
+        x = self.relu(x)
+        
+        # x = self.maxpool(x)
+        # pdb.set_trace()
+        x1 = self.layer1(x)
+        # print(x1.shape)
+        x2 = self.layer2(x1)
+        # print(x2.shape)
+        # x3 = self.layer3(x2)
+        # # print(x3.shape)
+        # x4 = self.layer4(x3)
+        # # print(x4.shape)
+        # x = F.relu(F.interpolate(self.decoder1(x4), scale_factor=(2,2), mode ='bilinear'))
+        # x = torch.add(x, x4)
+        # x = F.relu(F.interpolate(self.decoder2(x4) , scale_factor=(2,2), mode ='bilinear'))
+        # x = torch.add(x, x3)
+        # x = F.relu(F.interpolate(self.decoder3(x3) , scale_factor=(2,2), mode ='bilinear'))
+        # x = torch.add(x, x2)
+        x = F.relu(F.interpolate(self.decoder4(x2) , scale_factor=(2,2), mode ='bilinear'))
+        x = torch.add(x, x1)
+        x = F.relu(F.interpolate(self.decoder5(x) , scale_factor=(2,2), mode ='bilinear'))
+        # print(x.shape)
+        
+        # end of full image training 
+
+        # y_out = torch.ones((1,2,128,128))
+        x_loc = x.clone()
+        # x = F.relu(F.interpolate(self.decoder5(x) , scale_factor=(2,2), mode ='bilinear'))
+        #start 
+        for i in range(0,4):
+            for j in range(0,4):
+
+                x_p = xin[:,:,32*i:32*(i+1),32*j:32*(j+1)]
+                # begin patch wise
+                x_p = self.conv1_p(x_p)
+                x_p = self.bn1_p(x_p)
+                # x = F.max_pool2d(x,2,2)
+                x_p = self.relu(x_p)
+
+                x_p = self.conv2_p(x_p)
+                x_p = self.bn2_p(x_p)
+                # x = F.max_pool2d(x,2,2)
+                x_p = self.relu(x_p)
+                x_p = self.conv3_p(x_p)
+                x_p = self.bn3_p(x_p)
+                # x = F.max_pool2d(x,2,2)
+                x_p = self.relu(x_p)
+                
+                # x = self.maxpool(x)
+                # pdb.set_trace()
+                x1_p = self.layer1_p(x_p)
+                # print(x1.shape)
+                x2_p = self.layer2_p(x1_p)
+                # print(x2.shape)
+                x3_p = self.layer3_p(x2_p)
+                # # print(x3.shape)
+                x4_p = self.layer4_p(x3_p)
+                
+                x_p = F.relu(F.interpolate(self.decoder1_p(x4_p), scale_factor=(2,2), mode ='bilinear'))
+                x_p = torch.add(x_p, x4_p)
+                x_p = F.relu(F.interpolate(self.decoder2_p(x_p) , scale_factor=(2,2), mode ='bilinear'))
+                x_p = torch.add(x_p, x3_p)
+                x_p = F.relu(F.interpolate(self.decoder3_p(x_p) , scale_factor=(2,2), mode ='bilinear'))
+                x_p = torch.add(x_p, x2_p)
+                x_p = F.relu(F.interpolate(self.decoder4_p(x_p) , scale_factor=(2,2), mode ='bilinear'))
+                x_p = torch.add(x_p, x1_p)
+                x_p = F.relu(F.interpolate(self.decoder5_p(x_p) , scale_factor=(2,2), mode ='bilinear'))
+                
+                x_loc[:,:,32*i:32*(i+1),32*j:32*(j+1)] = x_p
+
+        x = torch.add(x,x_loc)
+        x = F.relu(self.decoderf(x))
+        
+        x = self.adjust(F.relu(x))
+
+        # pdb.set_trace()
+        return x
+
+    def forward(self, x, text_dummy):
+        return self.soft(self._forward_impl(x)),0
+
+
+def axialunet(pretrained=False, **kwargs):
+    model = ResAxialAttentionUNet(AxialBlock, [1, 2, 4, 1], s= 0.125, **kwargs)
+    return model
+
+def gated(pretrained=False, **kwargs):
+    model = ResAxialAttentionUNet(AxialBlock_dynamic, [1, 2, 4, 1], s= 0.125, **kwargs)
+    return model
+
+def MedT(pretrained=False, **kwargs):
+    model = medt_net(AxialBlock_dynamic,AxialBlock_wopos, [1, 2, 4, 1], s= 0.125,  **kwargs)
+    return model
+
+def logo(pretrained=False, **kwargs):
+    model = medt_net(AxialBlock,AxialBlock, [1, 2, 4, 1], s= 0.125, **kwargs)
+    return model
+
+# EOF

AllinonSAM/baselines.py

@@ -0,0 +1,630 @@
+import torch
+import torch.nn as nn
+from backbones_unet.model.unet import Unet
+import torch.nn.functional as F
+from utils import *
+__all__ = ['UNext']
+
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+import math
+
+class UNet(nn.Module):
+    def __init__(self, in_channels = 3, out_channels = 1, init_features = 32, pretrained=True , back_bone=None):
+        super().__init__()
+        if back_bone is None:
+            self.model = torch.hub.load(
+                'mateuszbuda/brain-segmentation-pytorch', 'unet', in_channels=in_channels, out_channels=out_channels, 
+                init_features=init_features, pretrained=pretrained
+            )
+        else:
+            self.model = UNet(
+                in_channels= in_channels,
+                out_channels= out_channels,
+                backbone=back_bone
+            )
+
+        self.soft = nn.Softmax(dim =1)
+    def forward(self, x, text_dummy):
+        return self.soft(self.model(x)),0
+
+
+def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=1, bias=False)
+
+class shiftmlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0., shift_size=5):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.dim = in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.dwconv = DWConv(hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+        self.shift_size = shift_size
+        self.pad = shift_size // 2
+
+        
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+    
+
+    def forward(self, x, H, W):
+        # pdb.set_trace()
+        B, N, C = x.shape
+
+        xn = x.transpose(1, 2).view(B, C, H, W).contiguous()
+        xn = F.pad(xn, (self.pad, self.pad, self.pad, self.pad) , "constant", 0)
+        xs = torch.chunk(xn, self.shift_size, 1)
+        x_shift = [torch.roll(x_c, shift, 2) for x_c, shift in zip(xs, range(-self.pad, self.pad+1))]
+        x_cat = torch.cat(x_shift, 1)
+        x_cat = torch.narrow(x_cat, 2, self.pad, H)
+        x_s = torch.narrow(x_cat, 3, self.pad, W)
+
+
+        x_s = x_s.reshape(B,C,H*W).contiguous()
+        x_shift_r = x_s.transpose(1,2)
+
+
+        x = self.fc1(x_shift_r)
+
+        x = self.dwconv(x, H, W)
+        x = self.act(x) 
+        x = self.drop(x)
+
+        xn = x.transpose(1, 2).view(B, C, H, W).contiguous()
+        xn = F.pad(xn, (self.pad, self.pad, self.pad, self.pad) , "constant", 0)
+        xs = torch.chunk(xn, self.shift_size, 1)
+        x_shift = [torch.roll(x_c, shift, 3) for x_c, shift in zip(xs, range(-self.pad, self.pad+1))]
+        x_cat = torch.cat(x_shift, 1)
+        x_cat = torch.narrow(x_cat, 2, self.pad, H)
+        x_s = torch.narrow(x_cat, 3, self.pad, W)
+        x_s = x_s.reshape(B,C,H*W).contiguous()
+        x_shift_c = x_s.transpose(1,2)
+
+        x = self.fc2(x_shift_c)
+        x = self.drop(x)
+        return x
+
+
+
+class shiftedBlock(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1):
+        super().__init__()
+
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = shiftmlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+
+    def forward(self, x, H, W):
+
+        x = x + self.drop_path(self.mlp(self.norm2(x), H, W))
+        return x
+
+
+class DWConv(nn.Module):
+    def __init__(self, dim=768):
+        super(DWConv, self).__init__()
+        self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)
+
+    def forward(self, x, H, W):
+        B, N, C = x.shape
+        x = x.transpose(1, 2).view(B, C, H, W)
+        x = self.dwconv(x)
+        x = x.flatten(2).transpose(1, 2)
+
+        return x
+
+class OverlapPatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+
+    def __init__(self, img_size=224, patch_size=7, stride=4, in_chans=3, embed_dim=768):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]
+        self.num_patches = self.H * self.W
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride,
+                              padding=(patch_size[0] // 2, patch_size[1] // 2))
+        self.norm = nn.LayerNorm(embed_dim)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+
+    def forward(self, x):
+        x = self.proj(x)
+        _, _, H, W = x.shape
+        x = x.flatten(2).transpose(1, 2)
+        x = self.norm(x)
+
+        return x, H, W
+
+
+class UNext(nn.Module):
+
+    ## Conv 3 + MLP 2 + shifted MLP
+    
+    def __init__(self,  num_classes, input_channels=3, deep_supervision=False,img_size=256, patch_size=16, in_chans=3,  embed_dims=[ 128, 160, 256],
+                 num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0.,
+                 attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm,
+                 depths=[1, 1, 1], sr_ratios=[8, 4, 2, 1], **kwargs):
+        super().__init__()
+        
+        self.encoder1 = nn.Conv2d(3, 16, 3, stride=1, padding=1)  
+        self.encoder2 = nn.Conv2d(16, 32, 3, stride=1, padding=1)  
+        self.encoder3 = nn.Conv2d(32, 128, 3, stride=1, padding=1)
+
+        self.ebn1 = nn.BatchNorm2d(16)
+        self.ebn2 = nn.BatchNorm2d(32)
+        self.ebn3 = nn.BatchNorm2d(128)
+        
+        self.norm3 = norm_layer(embed_dims[1])
+        self.norm4 = norm_layer(embed_dims[2])
+
+        self.dnorm3 = norm_layer(160)
+        self.dnorm4 = norm_layer(128)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+
+        self.block1 = nn.ModuleList([shiftedBlock(
+            dim=embed_dims[1], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[0], norm_layer=norm_layer,
+            sr_ratio=sr_ratios[0])])
+
+        self.block2 = nn.ModuleList([shiftedBlock(
+            dim=embed_dims[2], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[1], norm_layer=norm_layer,
+            sr_ratio=sr_ratios[0])])
+
+        self.dblock1 = nn.ModuleList([shiftedBlock(
+            dim=embed_dims[1], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[0], norm_layer=norm_layer,
+            sr_ratio=sr_ratios[0])])
+
+        self.dblock2 = nn.ModuleList([shiftedBlock(
+            dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[1], norm_layer=norm_layer,
+            sr_ratio=sr_ratios[0])])
+
+        self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_chans=embed_dims[0],
+                                              embed_dim=embed_dims[1])
+        self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_chans=embed_dims[1],
+                                              embed_dim=embed_dims[2])
+
+        self.decoder1 = nn.Conv2d(256, 160, 3, stride=1,padding=1)  
+        self.decoder2 =   nn.Conv2d(160, 128, 3, stride=1, padding=1)  
+        self.decoder3 =   nn.Conv2d(128, 32, 3, stride=1, padding=1) 
+        self.decoder4 =   nn.Conv2d(32, 16, 3, stride=1, padding=1)
+        self.decoder5 =   nn.Conv2d(16, 16, 3, stride=1, padding=1)
+
+        self.dbn1 = nn.BatchNorm2d(160)
+        self.dbn2 = nn.BatchNorm2d(128)
+        self.dbn3 = nn.BatchNorm2d(32)
+        self.dbn4 = nn.BatchNorm2d(16)
+        
+        self.final = nn.Conv2d(16, num_classes, kernel_size=1)
+
+        self.soft = nn.Softmax(dim =1)
+
+    def forward(self, x, text_dummy):
+        
+        B = x.shape[0]
+        ### Encoder
+        ### Conv Stage
+
+        ### Stage 1
+        out = F.relu(F.max_pool2d(self.ebn1(self.encoder1(x)),2,2))
+        t1 = out
+        ### Stage 2
+        out = F.relu(F.max_pool2d(self.ebn2(self.encoder2(out)),2,2))
+        t2 = out
+        ### Stage 3
+        out = F.relu(F.max_pool2d(self.ebn3(self.encoder3(out)),2,2))
+        t3 = out
+
+        ### Tokenized MLP Stage
+        ### Stage 4
+
+        out,H,W = self.patch_embed3(out)
+        for i, blk in enumerate(self.block1):
+            out = blk(out, H, W)
+        out = self.norm3(out)
+        out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+        t4 = out
+
+        ### Bottleneck
+
+        out ,H,W= self.patch_embed4(out)
+        for i, blk in enumerate(self.block2):
+            out = blk(out, H, W)
+        out = self.norm4(out)
+        out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+
+        ### Stage 4
+
+        out = F.relu(F.interpolate(self.dbn1(self.decoder1(out)),scale_factor=(2,2),mode ='bilinear'))
+        
+        out = torch.add(out,t4)
+        _,_,H,W = out.shape
+        out = out.flatten(2).transpose(1,2)
+        for i, blk in enumerate(self.dblock1):
+            out = blk(out, H, W)
+
+        ### Stage 3
+        
+        out = self.dnorm3(out)
+        out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+        out = F.relu(F.interpolate(self.dbn2(self.decoder2(out)),scale_factor=(2,2),mode ='bilinear'))
+        out = torch.add(out,t3)
+        _,_,H,W = out.shape
+        out = out.flatten(2).transpose(1,2)
+        
+        for i, blk in enumerate(self.dblock2):
+            out = blk(out, H, W)
+
+        out = self.dnorm4(out)
+        out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+
+        out = F.relu(F.interpolate(self.dbn3(self.decoder3(out)),scale_factor=(2,2),mode ='bilinear'))
+        out = torch.add(out,t2)
+        out = F.relu(F.interpolate(self.dbn4(self.decoder4(out)),scale_factor=(2,2),mode ='bilinear'))
+        out = torch.add(out,t1)
+        out = F.relu(F.interpolate(self.decoder5(out),scale_factor=(2,2),mode ='bilinear'))
+
+        return self.soft(self.final(out)),0
+
+
+class UNext_S(nn.Module):
+
+    ## Conv 3 + MLP 2 + shifted MLP w less parameters
+    
+    def __init__(self,  num_classes, input_channels=3, deep_supervision=False,img_size=256, patch_size=16, in_chans=3,  embed_dims=[32, 64, 128, 512],
+                 num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0.,
+                 attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm,
+                 depths=[1, 1, 1], sr_ratios=[8, 4, 2, 1], **kwargs):
+        super().__init__()
+        
+        self.encoder1 = nn.Conv2d(3, 8, 3, stride=1, padding=1)  
+        self.encoder2 = nn.Conv2d(8, 16, 3, stride=1, padding=1)  
+        self.encoder3 = nn.Conv2d(16, 32, 3, stride=1, padding=1)
+
+        self.ebn1 = nn.BatchNorm2d(8)
+        self.ebn2 = nn.BatchNorm2d(16)
+        self.ebn3 = nn.BatchNorm2d(32)
+        
+        self.norm3 = norm_layer(embed_dims[1])
+        self.norm4 = norm_layer(embed_dims[2])
+
+        self.dnorm3 = norm_layer(64)
+        self.dnorm4 = norm_layer(32)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+
+        self.block1 = nn.ModuleList([shiftedBlock(
+            dim=embed_dims[1], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[0], norm_layer=norm_layer,
+            sr_ratio=sr_ratios[0])])
+
+        self.block2 = nn.ModuleList([shiftedBlock(
+            dim=embed_dims[2], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[1], norm_layer=norm_layer,
+            sr_ratio=sr_ratios[0])])
+
+        self.dblock1 = nn.ModuleList([shiftedBlock(
+            dim=embed_dims[1], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[0], norm_layer=norm_layer,
+            sr_ratio=sr_ratios[0])])
+
+        self.dblock2 = nn.ModuleList([shiftedBlock(
+            dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[1], norm_layer=norm_layer,
+            sr_ratio=sr_ratios[0])])
+
+        self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_chans=embed_dims[0],
+                                              embed_dim=embed_dims[1])
+        self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_chans=embed_dims[1],
+                                              embed_dim=embed_dims[2])
+
+        self.decoder1 = nn.Conv2d(128, 64, 3, stride=1,padding=1)  
+        self.decoder2 =   nn.Conv2d(64, 32, 3, stride=1, padding=1)  
+        self.decoder3 =   nn.Conv2d(32, 16, 3, stride=1, padding=1) 
+        self.decoder4 =   nn.Conv2d(16, 8, 3, stride=1, padding=1)
+        self.decoder5 =   nn.Conv2d(8, 8, 3, stride=1, padding=1)
+
+        self.dbn1 = nn.BatchNorm2d(64)
+        self.dbn2 = nn.BatchNorm2d(32)
+        self.dbn3 = nn.BatchNorm2d(16)
+        self.dbn4 = nn.BatchNorm2d(8)
+        
+        self.final = nn.Conv2d(8, num_classes, kernel_size=1)
+
+        self.soft = nn.Softmax(dim =1)
+
+    def forward(self, x, text_dummy):
+        
+        B = x.shape[0]
+        ### Encoder
+        ### Conv Stage
+
+        ### Stage 1
+        out = F.relu(F.max_pool2d(self.ebn1(self.encoder1(x)),2,2))
+        t1 = out
+        ### Stage 2
+        out = F.relu(F.max_pool2d(self.ebn2(self.encoder2(out)),2,2))
+        t2 = out
+        ### Stage 3
+        out = F.relu(F.max_pool2d(self.ebn3(self.encoder3(out)),2,2))
+        t3 = out
+
+        ### Tokenized MLP Stage
+        ### Stage 4
+
+        out,H,W = self.patch_embed3(out)
+        for i, blk in enumerate(self.block1):
+            out = blk(out, H, W)
+        out = self.norm3(out)
+        out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+        t4 = out
+
+        ### Bottleneck
+
+        out ,H,W= self.patch_embed4(out)
+        for i, blk in enumerate(self.block2):
+            out = blk(out, H, W)
+        out = self.norm4(out)
+        out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+
+        ### Stage 4
+
+        out = F.relu(F.interpolate(self.dbn1(self.decoder1(out)),scale_factor=(2,2),mode ='bilinear'))
+        
+        out = torch.add(out,t4)
+        _,_,H,W = out.shape
+        out = out.flatten(2).transpose(1,2)
+        for i, blk in enumerate(self.dblock1):
+            out = blk(out, H, W)
+
+        ### Stage 3
+        
+        out = self.dnorm3(out)
+        out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+        out = F.relu(F.interpolate(self.dbn2(self.decoder2(out)),scale_factor=(2,2),mode ='bilinear'))
+        out = torch.add(out,t3)
+        _,_,H,W = out.shape
+        out = out.flatten(2).transpose(1,2)
+        
+        for i, blk in enumerate(self.dblock2):
+            out = blk(out, H, W)
+
+        out = self.dnorm4(out)
+        out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+
+        out = F.relu(F.interpolate(self.dbn3(self.decoder3(out)),scale_factor=(2,2),mode ='bilinear'))
+        out = torch.add(out,t2)
+        out = F.relu(F.interpolate(self.dbn4(self.decoder4(out)),scale_factor=(2,2),mode ='bilinear'))
+        out = torch.add(out,t1)
+        out = F.relu(F.interpolate(self.decoder5(out),scale_factor=(2,2),mode ='bilinear'))
+
+        return self.final(out)
+
+
+class medt_net(nn.Module):
+
+    def __init__(self, block, block_2, layers, num_classes=2, zero_init_residual=True,
+                 groups=8, width_per_group=64, replace_stride_with_dilation=None,
+                 norm_layer=None, s=0.125, img_size = 128,imgchan = 3):
+        super(medt_net, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.inplanes = int(64 * s)
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(imgchan, self.inplanes, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.conv2 = nn.Conv2d(self.inplanes, 128, kernel_size=3, stride=1, padding=1, bias=False)
+        self.conv3 = nn.Conv2d(128, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = norm_layer(self.inplanes)
+        self.bn2 = norm_layer(128)
+        self.bn3 = norm_layer(self.inplanes)
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.layer1 = self._make_layer(block, int(128 * s), layers[0], kernel_size= (img_size//2))
+        self.layer2 = self._make_layer(block, int(256 * s), layers[1], stride=2, kernel_size=(img_size//2),
+                                       dilate=replace_stride_with_dilation[0])
+        
+        self.decoder4 = nn.Conv2d(int(512*s) ,  int(256*s), kernel_size=3, stride=1, padding=1)
+        self.decoder5 = nn.Conv2d(int(256*s) , int(128*s) , kernel_size=3, stride=1, padding=1)
+        self.adjust   = nn.Conv2d(int(128*s) , num_classes, kernel_size=1, stride=1, padding=0)
+        self.soft     = nn.Softmax(dim=1)
+
+
+        self.conv1_p = nn.Conv2d(imgchan, self.inplanes, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.conv2_p = nn.Conv2d(self.inplanes,128, kernel_size=3, stride=1, padding=1,
+                               bias=False)
+        self.conv3_p = nn.Conv2d(128, self.inplanes, kernel_size=3, stride=1, padding=1,
+                               bias=False)
+        self.bn1_p = norm_layer(self.inplanes)
+        self.bn2_p = norm_layer(128)
+        self.bn3_p = norm_layer(self.inplanes)
+
+        self.relu_p = nn.ReLU(inplace=True)
+
+        img_size_p = img_size // 4
+
+        self.layer1_p = self._make_layer(block_2, int(128 * s), layers[0], kernel_size= (img_size_p//2))
+        self.layer2_p = self._make_layer(block_2, int(256 * s), layers[1], stride=2, kernel_size=(img_size_p//2),
+                                       dilate=replace_stride_with_dilation[0])
+        self.layer3_p = self._make_layer(block_2, int(512 * s), layers[2], stride=2, kernel_size=(img_size_p//4),
+                                       dilate=replace_stride_with_dilation[1])
+        self.layer4_p = self._make_layer(block_2, int(1024 * s), layers[3], stride=2, kernel_size=(img_size_p//8),
+                                       dilate=replace_stride_with_dilation[2])
+        
+        # Decoder
+        self.decoder1_p = nn.Conv2d(int(1024 *2*s)      ,        int(1024*2*s), kernel_size=3, stride=2, padding=1)
+        self.decoder2_p = nn.Conv2d(int(1024  *2*s)     , int(1024*s), kernel_size=3, stride=1, padding=1)
+        self.decoder3_p = nn.Conv2d(int(1024*s),  int(512*s), kernel_size=3, stride=1, padding=1)
+        self.decoder4_p = nn.Conv2d(int(512*s) ,  int(256*s), kernel_size=3, stride=1, padding=1)
+        self.decoder5_p = nn.Conv2d(int(256*s) , int(128*s) , kernel_size=3, stride=1, padding=1)
+
+        self.decoderf = nn.Conv2d(int(128*s) , int(128*s) , kernel_size=3, stride=1, padding=1)
+        self.adjust_p   = nn.Conv2d(int(128*s) , num_classes, kernel_size=1, stride=1, padding=0)
+        self.soft_p     = nn.Softmax(dim=1)
+
+
+    def _make_layer(self, block, planes, blocks, kernel_size=56, stride=1, dilate=False):
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, groups=self.groups,
+                            base_width=self.base_width, dilation=previous_dilation, 
+                            norm_layer=norm_layer, kernel_size=kernel_size))
+        self.inplanes = planes * block.expansion
+        if stride != 1:
+            kernel_size = kernel_size // 2
+
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, groups=self.groups,
+                                base_width=self.base_width, dilation=self.dilation,
+                                norm_layer=norm_layer, kernel_size=kernel_size))
+
+        return nn.Sequential(*layers)
+
+    def _forward_impl(self, x):
+
+        xin = x.clone()
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+        x = self.conv3(x)
+        x = self.bn3(x)
+        x = self.relu(x)
+        
+        x1 = self.layer1(x)
+        x2 = self.layer2(x1)
+
+        x = F.relu(F.interpolate(self.decoder4(x2) , scale_factor=(2,2), mode ='bilinear'))
+        x = torch.add(x, x1)
+        x = F.relu(F.interpolate(self.decoder5(x) , scale_factor=(2,2), mode ='bilinear'))
+        
+        # end of full image training 
+
+        x_loc = x.clone()
+        #start 
+        for i in range(0,4):
+            for j in range(0,4):
+
+                x_p = xin[:,:,32*i:32*(i+1),32*j:32*(j+1)]
+                # begin patch wise
+                x_p = self.conv1_p(x_p)
+                x_p = self.bn1_p(x_p)
+                x_p = self.relu(x_p)
+
+                x_p = self.conv2_p(x_p)
+                x_p = self.bn2_p(x_p)
+                x_p = self.relu(x_p)
+                x_p = self.conv3_p(x_p)
+                x_p = self.bn3_p(x_p)
+                x_p = self.relu(x_p)
+                
+                x1_p = self.layer1_p(x_p)
+                x2_p = self.layer2_p(x1_p)
+                x3_p = self.layer3_p(x2_p)
+                x4_p = self.layer4_p(x3_p)
+                
+                x_p = F.relu(F.interpolate(self.decoder1_p(x4_p), scale_factor=(2,2), mode ='bilinear'))
+                x_p = torch.add(x_p, x4_p)
+                x_p = F.relu(F.interpolate(self.decoder2_p(x_p) , scale_factor=(2,2), mode ='bilinear'))
+                x_p = torch.add(x_p, x3_p)
+                x_p = F.relu(F.interpolate(self.decoder3_p(x_p) , scale_factor=(2,2), mode ='bilinear'))
+                x_p = torch.add(x_p, x2_p)
+                x_p = F.relu(F.interpolate(self.decoder4_p(x_p) , scale_factor=(2,2), mode ='bilinear'))
+                x_p = torch.add(x_p, x1_p)
+                x_p = F.relu(F.interpolate(self.decoder5_p(x_p) , scale_factor=(2,2), mode ='bilinear'))
+                
+                x_loc[:,:,32*i:32*(i+1),32*j:32*(j+1)] = x_p
+
+        x = torch.add(x,x_loc)
+        x = F.relu(self.decoderf(x))
+        
+        x = self.adjust(F.relu(x))
+
+        return x
+
+    def forward(self, x, text_dummy):
+        return self._forward_impl(x)