Delete model.py

model.py

@@ -1,278 +0,0 @@
-
-import torch
-import torch.nn as nn
-import torchvision
-import torch.nn.functional as F
-from torchvision.transforms import functional as FF
-from torchvision import transforms
-from transformers import PretrainedConfig, PreTrainedModel
-
-######################################################################
-# Configuration File
-######################################################################
-class VesselSegmentConfig(PretrainedConfig):
-  model_type = "SegformerForSemanticSegmentation"
-  def __init__(self, num_classes=1, input_channels=1, image_size=[512,512], features=[64,64,128], attention_dims=[64,32,16], patch_size=256, batch_size=4, **kwargs):
-    self.num_classes = num_classes
-    self.input_channels = input_channels
-    self.image_size = image_size
-    self.features = features
-    self.attention_dims = attention_dims
-    self.patch_size = patch_size
-    self.batch_size = batch_size
-    super().__init__(**kwargs)
-
-######################################################################
-# IMAGE DOWN SAMPLING
-######################################################################
-class ImageDownSampling(nn.Module):
-  def __init__(self, height, width, scale):
-    super().__init__()
-    self.resize = transforms.Resize(size=(height//scale, width//scale))
-
-  def forward(self, x):
-    return self.resize(x)
-
-######################################################################
-# IMAGE SHARPENING
-######################################################################
-class ImageSharp(nn.Module):
-  def __init__(self):
-    super(ImageSharp, self).__init__()
-
-  def forward(self, x):
-    B, C, H, W = x.shape
-    device = x.device
-    # Sharpening kernel: basic 3x3
-    kernel = torch.tensor([[[[0, -1,  0],
-                             [-1, 5, -1],
-                             [0, -1,  0]]]], dtype=torch.float32, device=device)  # (1, 1, 3, 3)
-    # Apply the kernel using group convolution (one group per channel)
-    kernel = kernel.repeat(C, 1, 1, 1)  # (C, 1, 3, 3) --> here C=1, so it's still (1, 1, 3, 3)
-
-    # Apply convolution
-    sharpened = F.conv2d(x, kernel, padding=1, groups=C)  # padding=1 keeps same spatial size
-
-    # Clamp to stay within valid image range
-    sharpened = torch.clamp(sharpened, 0, 1)
-
-    return sharpened
-
-######################################################################
-# IMAGE PATCHING
-######################################################################
-class ImagePatching(nn.Module):
-  def __init__(self, patch_size: int):
-    super(ImagePatching, self).__init__()
-    self.patch_size = patch_size
-    self.image_patch = nn.Unfold(kernel_size=patch_size, stride=patch_size)
-    self.image_sharp = ImageSharp()
-
-  def forward(self, x):
-    batch_size, channels, height, width = x.shape
-    x = self.image_sharp(x)
-    x = self.image_patch(x)
-    x = x.transpose(1, 2).contiguous()
-    x = x.view(-1, height // self.patch_size, width // self.patch_size, channels, self.patch_size, self.patch_size)
-    x = x.view(-1, channels, self.patch_size, self.patch_size)
-    return x
-
-######################################################################
-# DOUBLE CONVOLUTION LAYER
-######################################################################
-class DoubleConvLayer(nn.Module):
-  def __init__(self, in_feature: int, out_feature: int):
-    super(DoubleConvLayer, self).__init__()
-    self.double_conv_layer = nn.Sequential(
-        nn.Conv2d(in_channels=in_feature, out_channels=out_feature, kernel_size=3, padding=1),
-        nn.InstanceNorm2d(num_features=out_feature),
-        nn.LeakyReLU(inplace=True),
-        nn.Conv2d(in_channels=out_feature, out_channels=out_feature, kernel_size=3, padding=1),
-        nn.InstanceNorm2d(num_features=out_feature),
-        nn.LeakyReLU(inplace=True)
-    )
-
-  def forward(self, x):
-    return self.double_conv_layer(x)
-
-######################################################################
-# FEATURE EXTRACTION FROM ENCODER PART
-######################################################################
-class EncoderFetureExtraction(nn.Module):
-  def __init__(self, feature: int):
-    super(EncoderFetureExtraction, self).__init__()
-
-    self.feature_extraction = nn.Sequential(
-        nn.Conv2d(in_channels=feature, out_channels=1, kernel_size=1, stride=1),
-        nn.InstanceNorm2d(num_features=1),
-        nn.LeakyReLU(inplace=True),
-        nn.Sigmoid()
-    )
-
-    self.relu = nn.LeakyReLU()
-
-  def forward(self, x):
-    x1 = self.feature_extraction(x)
-    return x * x1
-
-
-######################################################################
-# BOTTLENECK LAYER OF THE MODEL
-######################################################################
-class BottleNeck(nn.Module):
-  def __init__(self, in_ch, out_ch):
-    super(BottleNeck, self).__init__()
-    self.bottleneck = nn.Sequential(
-        nn.Conv2d(in_channels=in_ch, out_channels=out_ch, kernel_size=3, padding=1),
-        nn.InstanceNorm2d(num_features=out_ch),
-        nn.LeakyReLU(inplace=True)
-    )
-
-  def forward(self, x):
-    return self.bottleneck(x)
-
-
-######################################################################
-# SOFT-ATTENTION IN DECODER LAYER
-######################################################################
-class AttentionGate(nn.Module):
-  def __init__(self, dim_g, dim_x, dim_l):
-    super(AttentionGate, self).__init__()
-    self.Wg = nn.Sequential(
-        nn.Conv2d(in_channels=dim_g, out_channels=dim_l, kernel_size=1, stride=1),
-        nn.BatchNorm2d(num_features=dim_l))
-
-    self.Wx = nn.Sequential(
-        nn.Conv2d(in_channels=dim_x, out_channels=dim_l, kernel_size=1, stride=1),
-        nn.BatchNorm2d(num_features=dim_l))
-
-    self.alpha_conv = nn.Sequential(
-        nn.Conv2d(in_channels=dim_l, out_channels=1, kernel_size=1, stride=1),
-        nn.BatchNorm2d(num_features=1),
-        nn.Sigmoid())
-
-    self.up_conv = nn.ConvTranspose2d(in_channels=dim_g, out_channels=dim_g,
-                                      kernel_size=2, stride=2)
-
-    self.relu = nn.ReLU()
-
-  def forward(self, encoder_tensor, decoder_tensor):
-    # g > x, g is decoder, x is encoder
-    g = self.up_conv(decoder_tensor) # [b, 512, 32, 32]
-    w_x = self.Wx(encoder_tensor)    # [b, 128, 32 ,32]
-    w_g = self.Wg(g)           # [b, 128, 32, 32]
-
-    alpha = self.alpha_conv(self.relu(w_x + w_g))
-
-    return encoder_tensor * alpha
-
-
-######################################################################
-# IMAGE RECONSTRUCTION FROM PATCH
-######################################################################
-class ImageFolding(nn.Module):
-  def __init__(self, image_size: int, patch_size: int, batch_size: int):
-    super(ImageFolding, self).__init__()
-    self.num_patches = image_size // patch_size
-    self.batch_size = batch_size
-    self.folding = nn.Fold(output_size=(image_size, image_size),
-                           kernel_size=(patch_size, patch_size),
-                           stride=(patch_size, patch_size))
-
-  def forward(self, x):
-    x1 = x.view(self.batch_size, self.num_patches * self.num_patches, -1)
-    x1 = x1.transpose(1, 2).contiguous()
-    x1 = self.folding(x1)
-    return x1
-
-######################################################################
-# ENCODER LAYERS
-######################################################################
-class Encoder(nn.Module):
-  def __init__(self, in_channel, out_channel, enc_fet_ch, max_pool_size, is_concate=False):
-    super().__init__()
-    self.double_conv = DoubleConvLayer(in_feature=in_channel, out_feature=out_channel)
-    self.enc_feature_extraction = EncoderFetureExtraction(feature=enc_fet_ch)
-    self.pooling_layer = nn.MaxPool2d(kernel_size=max_pool_size, stride=max_pool_size)
-    self.concat = is_concate
-
-  def forward(self, x, concat_tensor=None):
-    x = self.double_conv(x)
-    if self.concat:
-      x = torch.cat([concat_tensor, x], dim=1)
-    skip_connection = self.enc_feature_extraction(x)
-    x = self.pooling_layer(x)
-    return x, skip_connection
-
-
-######################################################################
-# Decoder LAYERS
-######################################################################
-class Decoder(nn.Module):
-  def __init__(self, tensor_dim_encoder, tensor_dim_decoder, tensor_dim_mid, up_conv_in_ch, up_conv_out_ch, up_conv_scale, dconv_in_feature, dconv_out_feature, is_concat=False):
-    super().__init__()
-    self.soft_attention = AttentionGate(dim_g=tensor_dim_decoder, dim_x=tensor_dim_encoder, dim_l=tensor_dim_mid)
-    self.up_conv = nn.ConvTranspose2d(in_channels=up_conv_in_ch, out_channels=up_conv_out_ch, kernel_size=up_conv_scale, stride=up_conv_scale)
-    self.double_conv = DoubleConvLayer(in_feature=dconv_in_feature, out_feature=dconv_out_feature)
-    self.concat = is_concat
-
-  def forward(self, encoder_tensor, decoder_tensor):
-    x = self.soft_attention(encoder_tensor, decoder_tensor)
-    y = self.up_conv(decoder_tensor)
-    if self.concat:
-      x = torch.cat([x, y], dim=1)
-    x = self.double_conv(x)
-    return x
-
-class VesselSegmentModel(PreTrainedModel):
-  config_class = VesselSegmentConfig
-  def __init__(self, config: VesselSegmentConfig=VesselSegmentConfig()):
-    super().__init__(config)
-    # image patch
-    self.img_patch = ImagePatching(patch_size=config.patch_size)
-
-    # image downsampling
-    self.img_down_sampling_1 = ImageDownSampling(height=config.patch_size, width=config.patch_size, scale=2)
-    self.img_down_sampling_2 = ImageDownSampling(height=config.patch_size, width=config.patch_size, scale=4)
-
-    # encoder layers
-    self.encoder_layer_1 = Encoder(config.input_channels, config.features[0], enc_fet_ch=config.features[0], max_pool_size=2, is_concate=False)
-    self.encoder_layer_2 = Encoder(config.input_channels, config.features[1], enc_fet_ch=config.features[0]*2, max_pool_size=2, is_concate=True)
-    self.encoder_layer_3 = Encoder(config.input_channels, config.features[2], enc_fet_ch=config.features[0]*4, max_pool_size=2, is_concate=True)
-
-    # bottle-neck layer
-    self.bottleneck = BottleNeck(in_ch=config.features[2]*2, out_ch=config.features[2]*4)
-
-    # decoder layers
-    self.decoder_layer_1 = Decoder(tensor_dim_decoder=config.features[-1]*4, tensor_dim_encoder=config.features[-1]*2, tensor_dim_mid=config.features[0], up_conv_in_ch=config.features[-1]*4, up_conv_out_ch=config.features[-1]*2, up_conv_scale=2, dconv_in_feature=config.features[-1]*4, dconv_out_feature=config.features[-1]*2, is_concat=True)
-    self.decoder_layer_2 = Decoder(tensor_dim_decoder=config.features[-1]*2, tensor_dim_encoder=config.features[-1], tensor_dim_mid=config.features[1], up_conv_in_ch=config.features[-1]*2, up_conv_out_ch=config.features[-1], up_conv_scale=2, dconv_in_feature=config.features[-1]*2, dconv_out_feature=config.features[-1], is_concat=True)
-    self.decoder_layer_3 = Decoder(tensor_dim_decoder=config.features[-1], tensor_dim_encoder=config.features[-2], tensor_dim_mid=config.features[2], up_conv_in_ch=config.features[-1], up_conv_out_ch=config.features[-2], up_conv_scale=2, dconv_in_feature=config.features[-1], dconv_out_feature=config.features[-2], is_concat=True)
-
-    # Segmentation Head
-    self.segmenation_head = nn.Sequential(
-        nn.Conv2d(in_channels=config.features[-3], out_channels=config.num_classes, kernel_size=1, padding=0, stride=1),
-        ImageFolding(image_size=config.image_size[0], patch_size=config.patch_size, batch_size=config.batch_size)
-    )
-
-  def forward(self, x):
-    IMG_1 = self.img_patch(x)
-    IMG_2 = self.img_down_sampling_1(IMG_1)
-    IMG_3 = self.img_down_sampling_2(IMG_2)
-
-    # encoder
-    e1, sk1 = self.encoder_layer_1(IMG_1, None)
-    e2, sk2 = self.encoder_layer_2(IMG_2, e1)
-    e3, sk3 = self.encoder_layer_3(IMG_3, e2)
-
-    # bottleneck
-    b = self.bottleneck(e3)
-
-    # decoder
-    d1 = self.decoder_layer_1(sk3, b)
-    d2 = self.decoder_layer_2(sk2, d1)
-    d3 = self.decoder_layer_3(sk1, d2)
-
-    # head
-    head = self.segmenation_head(d3)
-
-    return head