Delete railnet_model.py

railnet_model.py

@@ -1,975 +0,0 @@
-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