Upload railnet_model.py

railnet_model.py

@@ -0,0 +1,975 @@
+import os
+os.environ['KMP_DUPLICATE_LIB_OK']='True'
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub import PyTorchModelHubMixin
+
+import numpy as np
+import nibabel as nib
+from skimage import morphology
+
+import math
+from scipy import ndimage
+from medpy import metric
+
+from huggingface_hub import hf_hub_download
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, n_stages, n_filters_in, n_filters_out, normalization='none'):
+        super(ConvBlock, self).__init__()
+
+        ops = []
+        for i in range(n_stages):
+            if i == 0:
+                input_channel = n_filters_in
+            else:
+                input_channel = n_filters_out
+
+            ops.append(nn.Conv3d(input_channel, n_filters_out, 3, padding=1))
+            if normalization == 'batchnorm':
+                ops.append(nn.BatchNorm3d(n_filters_out))
+            elif normalization == 'groupnorm':
+                ops.append(nn.GroupNorm(num_groups=16, num_channels=n_filters_out))
+            elif normalization == 'instancenorm':
+                ops.append(nn.InstanceNorm3d(n_filters_out))
+            elif normalization != 'none':
+                assert False
+            ops.append(nn.ReLU(inplace=True))
+
+        self.conv = nn.Sequential(*ops)
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+class DownsamplingConvBlock(nn.Module):
+    def __init__(self, n_filters_in, n_filters_out, stride=2, normalization='none'):
+        super(DownsamplingConvBlock, self).__init__()
+
+        ops = []
+        if normalization != 'none':
+            ops.append(nn.Conv3d(n_filters_in, n_filters_out, stride, padding=0, stride=stride))
+            if normalization == 'batchnorm':
+                ops.append(nn.BatchNorm3d(n_filters_out))
+            elif normalization == 'groupnorm':
+                ops.append(nn.GroupNorm(num_groups=16, num_channels=n_filters_out))
+            elif normalization == 'instancenorm':
+                ops.append(nn.InstanceNorm3d(n_filters_out))
+            else:
+                assert False
+        else:
+            ops.append(nn.Conv3d(n_filters_in, n_filters_out, stride, padding=0, stride=stride))
+
+        ops.append(nn.ReLU(inplace=True))
+
+        self.conv = nn.Sequential(*ops)
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+class UpsamplingDeconvBlock(nn.Module):
+    def __init__(self, n_filters_in, n_filters_out, stride=2, normalization='none'):
+        super(UpsamplingDeconvBlock, self).__init__()
+
+        ops = []
+        if normalization != 'none':
+            ops.append(nn.ConvTranspose3d(n_filters_in, n_filters_out, stride, padding=0, stride=stride))
+            if normalization == 'batchnorm':
+                ops.append(nn.BatchNorm3d(n_filters_out))
+            elif normalization == 'groupnorm':
+                ops.append(nn.GroupNorm(num_groups=16, num_channels=n_filters_out))
+            elif normalization == 'instancenorm':
+                ops.append(nn.InstanceNorm3d(n_filters_out))
+            else:
+                assert False
+        else:
+            ops.append(nn.ConvTranspose3d(n_filters_in, n_filters_out, stride, padding=0, stride=stride))
+
+        ops.append(nn.ReLU(inplace=True))
+
+        self.conv = nn.Sequential(*ops)
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+class Upsampling(nn.Module):
+    def __init__(self, n_filters_in, n_filters_out, stride=2, normalization='none'):
+        super(Upsampling, self).__init__()
+
+        ops = []
+        ops.append(nn.Upsample(scale_factor=stride, mode='trilinear', align_corners=False))
+        ops.append(nn.Conv3d(n_filters_in, n_filters_out, kernel_size=3, padding=1))
+        if normalization == 'batchnorm':
+            ops.append(nn.BatchNorm3d(n_filters_out))
+        elif normalization == 'groupnorm':
+            ops.append(nn.GroupNorm(num_groups=16, num_channels=n_filters_out))
+        elif normalization == 'instancenorm':
+            ops.append(nn.InstanceNorm3d(n_filters_out))
+        elif normalization != 'none':
+            assert False
+        ops.append(nn.ReLU(inplace=True))
+
+        self.conv = nn.Sequential(*ops)
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+class ConnectNet(nn.Module):
+    def __init__(self, in_channels, out_channels, input_size):
+        super(ConnectNet, self).__init__()
+        self.encoder = nn.Sequential(
+            nn.Conv3d(in_channels, 128, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.MaxPool3d(kernel_size=2, stride=2),
+            nn.Conv3d(128, 64, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.MaxPool3d(kernel_size=2, stride=2)
+        )
+
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose3d(64, 128, kernel_size=2, stride=2),
+            nn.ReLU(),
+            nn.ConvTranspose3d(128, out_channels, kernel_size=2, stride=2),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        encoded = self.encoder(x)
+        decoded = self.decoder(encoded)
+        return decoded
+
+
+class VNet(nn.Module):
+    def __init__(self, n_channels=3, n_classes=2, n_filters=16, normalization='none', has_dropout=False):
+        super(VNet, self).__init__()
+        self.has_dropout = has_dropout
+
+        self.block_one = ConvBlock(1, n_channels, n_filters, normalization=normalization)
+        self.block_one_dw = DownsamplingConvBlock(n_filters, 2 * n_filters, normalization=normalization)
+
+        self.block_two = ConvBlock(2, n_filters * 2, n_filters * 2, normalization=normalization)
+        self.block_two_dw = DownsamplingConvBlock(n_filters * 2, n_filters * 4, normalization=normalization)
+
+        self.block_three = ConvBlock(3, n_filters * 4, n_filters * 4, normalization=normalization)
+        self.block_three_dw = DownsamplingConvBlock(n_filters * 4, n_filters * 8, normalization=normalization)
+
+        self.block_four = ConvBlock(3, n_filters * 8, n_filters * 8, normalization=normalization)
+        self.block_four_dw = DownsamplingConvBlock(n_filters * 8, n_filters * 16, normalization=normalization)
+
+        self.block_five = ConvBlock(3, n_filters * 16, n_filters * 16, normalization=normalization)
+        self.block_five_up = UpsamplingDeconvBlock(n_filters * 16, n_filters * 8, normalization=normalization)
+
+        self.block_six = ConvBlock(3, n_filters * 8, n_filters * 8, normalization=normalization)
+        self.block_six_up = UpsamplingDeconvBlock(n_filters * 8, n_filters * 4, normalization=normalization)
+
+        self.block_seven = ConvBlock(3, n_filters * 4, n_filters * 4, normalization=normalization)
+        self.block_seven_up = UpsamplingDeconvBlock(n_filters * 4, n_filters * 2, normalization=normalization)
+
+        self.block_eight = ConvBlock(2, n_filters * 2, n_filters * 2, normalization=normalization)
+        self.block_eight_up = UpsamplingDeconvBlock(n_filters * 2, n_filters, normalization=normalization)
+
+        self.block_nine = ConvBlock(1, n_filters, n_filters, normalization=normalization)
+        self.out_conv = nn.Conv3d(n_filters, n_classes, 1, padding=0)
+
+        self.dropout = nn.Dropout3d(p=0.5, inplace=False)
+
+        self.__init_weight()
+
+    def encoder(self, input):
+        x1 = self.block_one(input)
+        x1_dw = self.block_one_dw(x1)
+
+        x2 = self.block_two(x1_dw)
+        x2_dw = self.block_two_dw(x2)
+
+        x3 = self.block_three(x2_dw)
+        x3_dw = self.block_three_dw(x3)
+
+        x4 = self.block_four(x3_dw)
+        x4_dw = self.block_four_dw(x4)
+
+        x5 = self.block_five(x4_dw)
+        if self.has_dropout:
+            x5 = self.dropout(x5)
+
+        res = [x1, x2, x3, x4, x5]
+
+        return res
+
+    def decoder(self, features):
+        x1 = features[0]
+        x2 = features[1]
+        x3 = features[2]
+        x4 = features[3]
+        x5 = features[4]
+
+        x5_up = self.block_five_up(x5)
+        x5_up = x5_up + x4
+
+        x6 = self.block_six(x5_up)
+        x6_up = self.block_six_up(x6)
+        x6_up = x6_up + x3
+
+        x7 = self.block_seven(x6_up)
+        x7_up = self.block_seven_up(x7)
+        x7_up = x7_up + x2
+
+        x8 = self.block_eight(x7_up)
+        x8_up = self.block_eight_up(x8)
+        x8_up = x8_up + x1
+        x9 = self.block_nine(x8_up)
+        if self.has_dropout:
+            x9 = self.dropout(x9)
+        out = self.out_conv(x9)
+        return out
+
+    def forward(self, input, turnoff_drop=False):
+        if turnoff_drop:
+            has_dropout = self.has_dropout
+            self.has_dropout = False
+        features = self.encoder(input)
+        out = self.decoder(features)
+        if turnoff_drop:
+            self.has_dropout = has_dropout
+        return out
+
+    def __init_weight(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv3d) or isinstance(m, nn.ConvTranspose3d):
+                torch.nn.init.kaiming_normal_(m.weight)
+            elif isinstance(m, nn.BatchNorm3d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+
+class VNet_roi(nn.Module):
+    def __init__(self, n_channels=3, n_classes=2, n_filters=16, normalization='none', has_dropout=False):
+        super(VNet_roi, self).__init__()
+        self.has_dropout = has_dropout
+
+        self.block_one = ConvBlock(1, n_channels, n_filters, normalization=normalization)
+        self.block_one_dw = DownsamplingConvBlock(n_filters, 2 * n_filters, normalization=normalization)
+
+        self.block_two = ConvBlock(2, n_filters * 2, n_filters * 2, normalization=normalization)
+        self.block_two_dw = DownsamplingConvBlock(n_filters * 2, n_filters * 4, normalization=normalization)
+
+        self.block_three = ConvBlock(3, n_filters * 4, n_filters * 4, normalization=normalization)
+        self.block_three_dw = DownsamplingConvBlock(n_filters * 4, n_filters * 8, normalization=normalization)
+
+        self.block_four = ConvBlock(3, n_filters * 8, n_filters * 8, normalization=normalization)
+        self.block_four_dw = DownsamplingConvBlock(n_filters * 8, n_filters * 16, normalization=normalization)
+
+        self.block_five = ConvBlock(3, n_filters * 16, n_filters * 16, normalization=normalization)
+        self.block_five_up = UpsamplingDeconvBlock(n_filters * 16, n_filters * 8, normalization=normalization)
+
+        self.block_six = ConvBlock(3, n_filters * 8, n_filters * 8, normalization=normalization)
+        self.block_six_up = UpsamplingDeconvBlock(n_filters * 8, n_filters * 4, normalization=normalization)
+
+        self.block_seven = ConvBlock(3, n_filters * 4, n_filters * 4, normalization=normalization)
+        self.block_seven_up = UpsamplingDeconvBlock(n_filters * 4, n_filters * 2, normalization=normalization)
+
+        self.block_eight = ConvBlock(2, n_filters * 2, n_filters * 2, normalization=normalization)
+        self.block_eight_up = UpsamplingDeconvBlock(n_filters * 2, n_filters, normalization=normalization)
+
+        self.block_nine = ConvBlock(1, n_filters, n_filters, normalization=normalization)
+        self.out_conv = nn.Conv3d(n_filters, n_classes, 1, padding=0)
+
+        self.dropout = nn.Dropout3d(p=0.5, inplace=False)
+        # self.__init_weight()
+
+    def encoder(self, input):
+        x1 = self.block_one(input)
+        x1_dw = self.block_one_dw(x1)
+
+        x2 = self.block_two(x1_dw)
+        x2_dw = self.block_two_dw(x2)
+
+        x3 = self.block_three(x2_dw)
+        x3_dw = self.block_three_dw(x3)
+
+        x4 = self.block_four(x3_dw)
+        x4_dw = self.block_four_dw(x4)
+
+        x5 = self.block_five(x4_dw)
+        # x5 = F.dropout3d(x5, p=0.5, training=True)
+        if self.has_dropout:
+            x5 = self.dropout(x5)
+
+        res = [x1, x2, x3, x4, x5]
+
+        return res
+
+    def decoder(self, features):
+        x1 = features[0]
+        x2 = features[1]
+        x3 = features[2]
+        x4 = features[3]
+        x5 = features[4]
+
+        x5_up = self.block_five_up(x5)
+        x5_up = x5_up + x4
+
+        x6 = self.block_six(x5_up)
+        x6_up = self.block_six_up(x6)
+        x6_up = x6_up + x3
+
+        x7 = self.block_seven(x6_up)
+        x7_up = self.block_seven_up(x7)
+        x7_up = x7_up + x2
+
+        x8 = self.block_eight(x7_up)
+        x8_up = self.block_eight_up(x8)
+        x8_up = x8_up + x1
+        x9 = self.block_nine(x8_up)
+        # x9 = F.dropout3d(x9, p=0.5, training=True)
+        if self.has_dropout:
+            x9 = self.dropout(x9)
+        out = self.out_conv(x9)
+        return out
+
+
+    def forward(self, input, turnoff_drop=False):
+        if turnoff_drop:
+            has_dropout = self.has_dropout
+            self.has_dropout = False
+        features = self.encoder(input)
+        out = self.decoder(features)
+        if turnoff_drop:
+            self.has_dropout = has_dropout
+        return out
+
+
+class ResVNet(nn.Module):
+    def __init__(self, n_channels=1, n_classes=2, n_filters=16, normalization='instancenorm', has_dropout=False):
+        super(ResVNet, self).__init__()
+        self.resencoder = resnet34()
+        self.has_dropout = has_dropout
+
+        self.block_one = ConvBlock(1, n_channels, n_filters, normalization=normalization)
+        self.block_one_dw = DownsamplingConvBlock(n_filters, 2 * n_filters, normalization=normalization)
+
+        self.block_two = ConvBlock(2, n_filters * 2, n_filters * 2, normalization=normalization)
+        self.block_two_dw = DownsamplingConvBlock(n_filters * 2, n_filters * 4, normalization=normalization)
+
+        self.block_three = ConvBlock(3, n_filters * 4, n_filters * 4, normalization=normalization)
+        self.block_three_dw = DownsamplingConvBlock(n_filters * 4, n_filters * 8, normalization=normalization)
+
+        self.block_four = ConvBlock(3, n_filters * 8, n_filters * 8, normalization=normalization)
+        self.block_four_dw = DownsamplingConvBlock(n_filters * 8, n_filters * 16, normalization=normalization)
+
+        self.block_five = ConvBlock(3, n_filters * 16, n_filters * 16, normalization=normalization)
+        self.block_five_up = UpsamplingDeconvBlock(n_filters * 16, n_filters * 8, normalization=normalization)
+
+        self.block_six = ConvBlock(3, n_filters * 8, n_filters * 8, normalization=normalization)
+        self.block_six_up = UpsamplingDeconvBlock(n_filters * 8, n_filters * 4, normalization=normalization)
+
+        self.block_seven = ConvBlock(3, n_filters * 4, n_filters * 4, normalization=normalization)
+        self.block_seven_up = UpsamplingDeconvBlock(n_filters * 4, n_filters * 2, normalization=normalization)
+
+        self.block_eight = ConvBlock(2, n_filters * 2, n_filters * 2, normalization=normalization)
+        self.block_eight_up = UpsamplingDeconvBlock(n_filters * 2, n_filters, normalization=normalization)
+
+
+        self.block_nine = ConvBlock(1, n_filters, n_filters, normalization=normalization)
+        self.out_conv = nn.Conv3d(n_filters, n_classes, 1, padding=0)
+
+
+        if has_dropout:
+            self.dropout = nn.Dropout3d(p=0.5)
+        self.branchs = nn.ModuleList()
+        for i in range(1):
+            if has_dropout:
+                seq = nn.Sequential(
+                    ConvBlock(1, n_filters, n_filters, normalization=normalization),
+                    nn.Dropout3d(p=0.5),
+                    nn.Conv3d(n_filters, n_classes, 1, padding=0)
+                )
+            else:
+                seq = nn.Sequential(
+                    ConvBlock(1, n_filters, n_filters, normalization=normalization),
+                    nn.Conv3d(n_filters, n_classes, 1, padding=0)
+                )
+            self.branchs.append(seq)
+
+    def encoder(self, input):
+        x1 = self.block_one(input)
+        x1_dw = self.block_one_dw(x1)
+
+        x2 = self.block_two(x1_dw)
+        x2_dw = self.block_two_dw(x2)
+
+        x3 = self.block_three(x2_dw)
+        x3_dw = self.block_three_dw(x3)
+
+        x4 = self.block_four(x3_dw)
+        x4_dw = self.block_four_dw(x4)
+
+        x5 = self.block_five(x4_dw)
+
+        if self.has_dropout:
+            x5 = self.dropout(x5)
+
+        res = [x1, x2, x3, x4, x5]
+
+        return res
+
+    def decoder(self, features):
+        x1 = features[0]
+        x2 = features[1]
+        x3 = features[2]
+        x4 = features[3]
+        x5 = features[4]
+
+        x5_up = self.block_five_up(x5)
+        x5_up = x5_up + x4
+
+        x6 = self.block_six(x5_up)
+        x6_up = self.block_six_up(x6)
+        x6_up = x6_up + x3
+
+        x7 = self.block_seven(x6_up)
+        x7_up = self.block_seven_up(x7)
+        x7_up = x7_up + x2
+
+        x8 = self.block_eight(x7_up)
+        x8_up = self.block_eight_up(x8)
+        x8_up = x8_up + x1
+
+
+        x9 = self.block_nine(x8_up)
+
+        out = self.out_conv(x9)
+
+
+        return out
+
+    def forward(self, input, turnoff_drop=False):
+        if turnoff_drop:
+            has_dropout = self.has_dropout
+            self.has_dropout = False
+        features = self.resencoder(input)
+        out = self.decoder(features)
+        if turnoff_drop:
+            self.has_dropout = has_dropout
+        return out
+
+
+__all__ = ['ResNet',  'resnet34']
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv3d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
+
+
+def conv3x3_bn_relu(in_planes, out_planes, stride=1):
+    return nn.Sequential(
+        conv3x3(in_planes, out_planes, stride),
+        nn.InstanceNorm3d(out_planes),
+        nn.ReLU()
+    )
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None,
+                 groups=1, base_width=64, dilation=-1):
+        super(BasicBlock, self).__init__()
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.InstanceNorm3d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.InstanceNorm3d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None,
+                 groups=1, base_width=64, dilation=1):
+        super(Bottleneck, self).__init__()
+        width = int(planes * (base_width / 64.)) * groups
+        self.conv1 = nn.Conv3d(inplanes, width, kernel_size=1, bias=False)
+        self.bn1 = nn.InstanceNorm3d(width)
+        self.conv2 = nn.Conv3d(width, width, kernel_size=3, stride=stride, dilation=dilation,
+                               padding=dilation, groups=groups, bias=False)
+        self.bn2 = nn.InstanceNorm3d(width)
+        self.conv3 = nn.Conv3d(width, planes * self.expansion, kernel_size=1, bias=False)
+        self.bn3 = nn.InstanceNorm3d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, in_channel=1, width=1,
+                 groups=1, width_per_group=64,
+                 mid_dim=1024, low_dim=128,
+                 avg_down=False, deep_stem=False,
+                 head_type='mlp_head', layer4_dilation=1):
+        super(ResNet, self).__init__()
+        self.avg_down = avg_down
+        self.inplanes = 16 * width
+        self.base = int(16 * width)
+        self.groups = groups
+        self.base_width = width_per_group
+
+        mid_dim = self.base * 8 * block.expansion
+
+        if deep_stem:
+            self.conv1 = nn.Sequential(
+                conv3x3_bn_relu(in_channel, 32, stride=2),
+                conv3x3_bn_relu(32, 32, stride=1),
+                conv3x3(32, 64, stride=1)
+            )
+        else:
+            self.conv1 = nn.Conv3d(in_channel, self.inplanes, kernel_size=7, stride=1, padding=3, bias=False)
+
+        self.bn1 = nn.InstanceNorm3d(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.maxpool = nn.MaxPool3d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, self.base*2, layers[0],stride=2)
+        self.layer2 = self._make_layer(block, self.base * 4, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, self.base * 8, layers[2], stride=2)
+        if layer4_dilation == 1:
+            self.layer4 = self._make_layer(block, self.base * 16, layers[3], stride=2)
+        elif layer4_dilation == 2:
+            self.layer4 = self._make_layer(block, self.base * 16, layers[3], stride=1, dilation=2)
+        else:
+            raise NotImplementedError
+        self.avgpool = nn.AvgPool3d(7, stride=1)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilation=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            if self.avg_down:
+                downsample = nn.Sequential(
+                    nn.AvgPool3d(kernel_size=stride, stride=stride),
+                    nn.Conv3d(self.inplanes, planes * block.expansion,
+                              kernel_size=1, stride=1, bias=False),
+                    nn.InstanceNorm3d(planes * block.expansion),
+                )
+            else:
+                downsample = nn.Sequential(
+                    nn.Conv3d(self.inplanes, planes * block.expansion,
+                              kernel_size=1, stride=stride, bias=False),
+                    nn.InstanceNorm3d(planes * block.expansion),
+                )
+
+        layers = [block(self.inplanes, planes, stride, downsample, self.groups, self.base_width, dilation)]
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, groups=self.groups, base_width=self.base_width, dilation=dilation))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        #c2 = self.maxpool(x)
+        c2 = self.layer1(x)
+        c3 = self.layer2(c2)
+        c4 = self.layer3(c3)
+        c5 = self.layer4(c4)
+
+
+        return [x,c2,c3,c4,c5]
+
+
+def resnet34(**kwargs):
+    return ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
+
+
+def label_rescale(image_label, w_ori, h_ori, z_ori, flag):
+    w_ori, h_ori, z_ori = int(w_ori), int(h_ori), int(z_ori)
+    # resize label map (int)
+    if flag == 'trilinear':
+        teeth_ids = np.unique(image_label)
+        image_label_ori = np.zeros((w_ori, h_ori, z_ori))
+        
+        
+        image_label = torch.from_numpy(image_label).cuda(0)
+        
+        
+        # image_label = torch.from_numpy(image_label).to("cpu")
+        for label_id in range(len(teeth_ids)):
+            image_label_bn = (image_label == teeth_ids[label_id]).float()
+            image_label_bn = image_label_bn[None, None, :, :, :]
+            image_label_bn = torch.nn.functional.interpolate(image_label_bn, size=(w_ori, h_ori, z_ori),
+                                                             mode='trilinear', align_corners=False)
+            image_label_bn = image_label_bn[0, 0, :, :, :]
+            image_label_bn = image_label_bn.cpu().data.numpy()
+            image_label_ori[image_label_bn > 0.5] = teeth_ids[label_id]
+        image_label = image_label_ori
+
+    if flag == 'nearest':
+
+        
+        image_label = torch.from_numpy(image_label).cuda(0)
+        
+        
+        # image_label = torch.from_numpy(image_label).to("cpu")
+        image_label = image_label[None, None, :, :, :].float()
+        image_label = torch.nn.functional.interpolate(image_label, size=(w_ori, h_ori, z_ori), mode='nearest')
+        image_label = image_label[0, 0, :, :, :].cpu().data.numpy()
+    return image_label
+
+
+def img_crop(image_bbox):
+    if image_bbox.sum() > 0:
+
+        x_min = np.nonzero(image_bbox)[0].min() - 8
+        x_max = np.nonzero(image_bbox)[0].max() + 8
+
+        y_min = np.nonzero(image_bbox)[1].min() - 16
+        y_max = np.nonzero(image_bbox)[1].max() + 16
+
+        z_min = np.nonzero(image_bbox)[2].min() - 16
+        z_max = np.nonzero(image_bbox)[2].max() + 16
+
+        if x_min < 0:
+            x_min = 0
+        if y_min < 0:
+            y_min = 0
+        if z_min < 0:
+            z_min = 0
+        if x_max > image_bbox.shape[0]:
+            x_max = image_bbox.shape[0]
+        if y_max > image_bbox.shape[1]:
+            y_max = image_bbox.shape[1]
+        if z_max > image_bbox.shape[2]:
+            z_max = image_bbox.shape[2]
+
+        if (x_max - x_min) % 16 != 0: 
+            x_max -= (x_max - x_min) % 16
+        if (y_max - y_min) % 16 != 0:
+            y_max -= (y_max - y_min) % 16
+        if (z_max - z_min) % 16 != 0:
+            z_max -= (z_max - z_min) % 16
+
+    if image_bbox.sum() == 0:
+        x_min, x_max, y_min, y_max, z_min, z_max = -1, image_bbox.shape[0], 0, image_bbox.shape[1], 0, image_bbox.shape[
+            2]
+    return x_min, x_max, y_min, y_max, z_min, z_max
+
+
+def roi_extraction(image, net_roi, ids):
+    w, h, d = image.shape
+    # roi binary segmentation parameters, the input spacing is 0.4 mm
+    print('---run the roi binary segmentation.')
+
+    stride_xy = 32
+    stride_z = 16
+    patch_size_roi_stage = (112, 112, 80)
+
+    label_roi = roi_detection(net_roi, image[0:w:2, 0:h:2, 0:d:2], stride_xy, stride_z,
+                              patch_size_roi_stage)  # (400,400,200)
+    print(label_roi.shape, np.max(label_roi))
+    label_roi = label_rescale(label_roi, w, h, d, 'trilinear')  # (800,800,400)
+
+    label_roi = morphology.remove_small_objects(label_roi.astype(bool), 5000, connectivity=3).astype(float)
+
+    label_roi = ndimage.grey_dilation(label_roi, size=(5, 5, 5))
+
+    label_roi = morphology.remove_small_objects(label_roi.astype(bool), 400000, connectivity=3).astype(
+        float)  
+
+    label_roi = ndimage.grey_erosion(label_roi, size=(5, 5, 5))
+
+    # crop image
+    x_min, x_max, y_min, y_max, z_min, z_max = img_crop(label_roi)
+    if x_min == -1:  # non-foreground label
+        whole_label = np.zeros((w, h, d))
+        return whole_label
+    image = image[x_min:x_max, y_min:y_max, z_min:z_max]  
+    print("image shape(after roi): ", image.shape)
+
+    return image, x_min, x_max, y_min, y_max, z_min, z_max
+
+
+def roi_detection(net, image, stride_xy, stride_z, patch_size):
+    w, h, d = image.shape  # (400,400,200)
+
+    # if the size of image is less than patch_size, then padding it
+    add_pad = False
+    if w < patch_size[0]:
+        w_pad = patch_size[0] - w
+        add_pad = True
+    else:
+        w_pad = 0
+    if h < patch_size[1]:
+        h_pad = patch_size[1] - h
+        add_pad = True
+    else:
+        h_pad = 0
+    if d < patch_size[2]:
+        d_pad = patch_size[2] - d
+        add_pad = True
+    else:
+        d_pad = 0
+    wl_pad, wr_pad = w_pad // 2, w_pad - w_pad // 2
+    hl_pad, hr_pad = h_pad // 2, h_pad - h_pad // 2
+    dl_pad, dr_pad = d_pad // 2, d_pad - d_pad // 2
+    if add_pad:
+        image = np.pad(image, [(wl_pad, wr_pad), (hl_pad, hr_pad), (dl_pad, dr_pad)], mode='constant',
+                       constant_values=0)
+    ww, hh, dd = image.shape
+
+    sx = math.ceil((ww - patch_size[0]) / stride_xy) + 1  # 2
+    sy = math.ceil((hh - patch_size[1]) / stride_xy) + 1  # 2
+    sz = math.ceil((dd - patch_size[2]) / stride_z) + 1  # 2
+    score_map = np.zeros((2,) + image.shape).astype(np.float32) 
+    cnt = np.zeros(image.shape).astype(np.float32) 
+    count = 0
+    for x in range(0, sx):
+        xs = min(stride_xy * x, ww - patch_size[0])
+        for y in range(0, sy):
+            ys = min(stride_xy * y, hh - patch_size[1])
+            for z in range(0, sz):
+                zs = min(stride_z * z, dd - patch_size[2])
+                test_patch = image[xs:xs + patch_size[0], ys:ys + patch_size[1],
+                             zs:zs + patch_size[2]]  
+                test_patch = np.expand_dims(np.expand_dims(test_patch, axis=0), axis=0).astype(
+                    np.float32) 
+                
+                
+                test_patch = torch.from_numpy(test_patch).cuda(0)
+                
+                
+                # test_patch = torch.from_numpy(test_patch).to("cpu")
+                with torch.no_grad():
+                    y1 = net(test_patch)  # （1,2,256,256,160）
+                y = F.softmax(y1, dim=1)  # （1,2,256,256,160）
+                y = y.cpu().data.numpy()
+                y = y[0, :, :, :, :]  # （2,256,256,160）
+                score_map[:, xs:xs + patch_size[0], ys:ys + patch_size[1], zs:zs + patch_size[2]] \
+                    = score_map[:, xs:xs + patch_size[0], ys:ys + patch_size[1],
+                      zs:zs + patch_size[2]] + y  # (2,400,400,200)
+                cnt[xs:xs + patch_size[0], ys:ys + patch_size[1], zs:zs + patch_size[2]] \
+                    = cnt[xs:xs + patch_size[0], ys:ys + patch_size[1], zs:zs + patch_size[2]] + 1  # (400,400,200)
+                count = count + 1
+    score_map = score_map / np.expand_dims(cnt, axis=0)
+
+    label_map = np.argmax(score_map, axis=0)  # (400,400,200),0/1
+    if add_pad:
+        label_map = label_map[wl_pad:wl_pad + w, hl_pad:hl_pad + h, dl_pad:dl_pad + d]
+        score_map = score_map[:, wl_pad:wl_pad + w, hl_pad:hl_pad + h, dl_pad:dl_pad + d]
+    return label_map
+
+
+def test_single_case_array(model_array, image=None, stride_xy=None, stride_z=None, patch_size=None, num_classes=1):
+    w, h, d = image.shape
+
+    # if the size of image is less than patch_size, then padding it
+    add_pad = False
+    if w < patch_size[0]:
+        w_pad = patch_size[0]-w
+        add_pad = True
+    else:
+        w_pad = 0
+    if h < patch_size[1]:
+        h_pad = patch_size[1]-h
+        add_pad = True
+    else:
+        h_pad = 0
+    if d < patch_size[2]:
+        d_pad = patch_size[2]-d
+        add_pad = True
+    else:
+        d_pad = 0
+    wl_pad, wr_pad = w_pad//2,w_pad-w_pad//2
+    hl_pad, hr_pad = h_pad//2,h_pad-h_pad//2
+    dl_pad, dr_pad = d_pad//2,d_pad-d_pad//2
+    if add_pad:
+        image = np.pad(image, [(wl_pad,wr_pad),(hl_pad,hr_pad), (dl_pad, dr_pad)], mode='constant', constant_values=0)
+
+    ww,hh,dd = image.shape
+
+    sx = math.ceil((ww - patch_size[0]) / stride_xy) + 1
+    sy = math.ceil((hh - patch_size[1]) / stride_xy) + 1
+    sz = math.ceil((dd - patch_size[2]) / stride_z) + 1
+    score_map = np.zeros((num_classes, ) + image.shape).astype(np.float32)
+    cnt = np.zeros(image.shape).astype(np.float32)
+
+    for x in range(0, sx):
+        xs = min(stride_xy*x, ww-patch_size[0])
+        for y in range(0, sy):
+            ys = min(stride_xy * y,hh-patch_size[1])
+            for z in range(0, sz):
+                zs = min(stride_z * z, dd-patch_size[2])
+                test_patch = image[xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]]
+                test_patch = np.expand_dims(np.expand_dims(test_patch,axis=0),axis=0).astype(np.float32)
+
+
+                test_patch = torch.from_numpy(test_patch).cuda()
+
+
+                # test_patch = torch.from_numpy(test_patch).to("cpu")
+                for model in model_array:
+                    output = model(test_patch)
+                    y_temp = F.softmax(output, dim=1)
+                    y_temp = y_temp.cpu().data.numpy()
+                    y += y_temp[0,:,:,:,:]
+                y /= len(model_array)
+                score_map[:, xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] \
+                  = score_map[:, xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] + y
+                cnt[xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] \
+                  = cnt[xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] + 1
+    score_map = score_map/np.expand_dims(cnt,axis=0)
+
+    label_map = np.argmax(score_map, axis = 0)
+    if add_pad:
+        label_map = label_map[wl_pad:wl_pad+w,hl_pad:hl_pad+h,dl_pad:dl_pad+d]
+        score_map = score_map[:,wl_pad:wl_pad+w,hl_pad:hl_pad+h,dl_pad:dl_pad+d]
+    return label_map, score_map
+
+def calculate_metric_percase(pred, gt):
+    dice = metric.binary.dc(pred, gt)
+    jc = metric.binary.jc(pred, gt)
+    hd = metric.binary.hd95(pred, gt)
+    asd = metric.binary.asd(pred, gt)
+
+    return dice, jc, hd, asd
+
+
+class RailNetSystem(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, n_channels: int, n_classes: int, normalization: str):
+        super().__init__()
+
+        self.num_classes = 2
+
+
+        self.net_roi = VNet_roi(n_channels = n_channels, n_classes = n_classes, normalization = normalization, has_dropout=False).cuda()
+        
+        
+        # self.net_roi = VNet_roi(n_channels = n_channels, n_classes = n_classes, normalization = normalization, has_dropout=False).to("cpu")
+
+        self.model_array = []
+        for i in range(4):
+            if i < 2:
+
+
+                model = VNet(n_channels = n_channels, n_classes = n_classes, normalization = normalization, has_dropout=True).cuda()
+                
+                
+                # model = VNet(n_channels = n_channels, n_classes = n_classes, normalization = normalization, has_dropout=True).to("cpu")
+            else:
+
+
+                model = ResVNet(n_channels = n_channels, n_classes = n_classes, normalization = normalization, has_dropout=True).cuda()
+                
+                
+                # model = ResVNet(n_channels = n_channels, n_classes = n_classes, normalization = normalization, has_dropout=True).to("cpu")
+            self.model_array.append(model)
+
+    def load_weights(self, weight_dir=".", from_hub=False, repo_id=None):
+        def load(file_name):
+            if from_hub:
+                return hf_hub_download(repo_id=repo_id, filename=f"model weights/{file_name}")
+            else:
+                return os.path.join(weight_dir, "model weights", file_name)
+
+        
+        # self.net_roi.load_state_dict(torch.load(os.path.join(weight_dir, "model weights", "roi_best_model.pth"), map_location="cuda", weights_only=True))
+        
+        
+        # self.net_roi.load_state_dict(torch.load(os.path.join(weight_dir, "model weights", "roi_best_model.pth"), map_location="cpu", weights_only=True))
+        self.net_roi.load_state_dict(torch.load(load("roi_best_model.pth"), map_location="cuda", weights_only=True))
+        self.net_roi.eval()
+
+        model_files = [
+            "rail_0_iter_7995_best.pth",
+            "rail_1_iter_7995_best.pth",
+            "rail_2_iter_7995_best.pth",
+            "rail_3_iter_7995_best.pth",
+        ]
+        for i, file in enumerate(model_files):
+            
+            
+            # self.model_array[i].load_state_dict(torch.load(os.path.join(weight_dir, "model weights", file), map_location="cuda", weights_only=True))
+           
+           
+            # self.model_array[i].load_state_dict(torch.load(os.path.join(weight_dir, "model weights", file), map_location="cpu", weights_only=True))
+            self.model_array[i].load_state_dict(torch.load(load(file), map_location="cuda", weights_only=True))
+            self.model_array[i].eval()
+
+    def forward(self, image, label, save_path="./output", name="case"):
+        if not os.path.exists(save_path):
+            os.makedirs(save_path)
+        nib.save(nib.Nifti1Image(image.astype(np.float32), np.eye(4)), os.path.join(save_path, f"{name}_img.nii.gz"))
+
+        w, h, d = image.shape
+
+        image, x_min, x_max, y_min, y_max, z_min, z_max = roi_extraction(image, self.net_roi, name)
+
+        prediction, _ = test_single_case_array(self.model_array, image, stride_xy=64, stride_z=32, patch_size=(112, 112, 80), num_classes=self.num_classes)
+
+        prediction = morphology.remove_small_objects(prediction.astype(bool), 3000, connectivity=3).astype(float)
+
+        new_prediction = np.zeros((w, h, d))
+        new_prediction[x_min:x_max, y_min:y_max, z_min:z_max] = prediction
+
+        dice, jc, hd, asd = calculate_metric_percase(new_prediction, label[:])
+
+        nib.save(nib.Nifti1Image(new_prediction.astype(np.float32), np.eye(4)), os.path.join(save_path, f"{name}_pred.nii.gz"))
+
+        return new_prediction, dice, jc, hd, asd