Upload Model.py with huggingface_hub

Model.py

@@ -0,0 +1,261 @@
+
+import torch
+import torch.nn as nn
+import torchvision
+from torchvision import transforms
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+######################################################################
+# 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