Upload model.py with huggingface_hub

model.py

@@ -5,6 +5,19 @@ import torchvision
 import torch.nn.functional as F
 from torchvision.transforms import functional as FF
 
+######################################################################
+# Configuration File
+######################################################################
+class VesselSegmentConfig:
+  def __init__(self, num_classes, input_channels, image_size, features, attention_dims, patch_size, batch_size):
+    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
+
 ######################################################################
 # IMAGE DOWN SAMPLING
 ######################################################################
@@ -208,33 +221,33 @@ class Decoder(nn.Module):
     return x
 
 class VesselSegmentModel(nn.Module):
-  def __init__(self, input_channel: int, feature: list, attention_dim: list, output_channel: int=1):
+  def __init__(self, config: VesselSegmentConfig):
     super().__init__()
-
+    self.config = config
     # image patch
-    self.img_patch = ImagePatching(patch_size=PATCH_SIZE)
+    self.img_patch = ImagePatching(patch_size=config.patch_size)
 
     # image downsampling
-    self.img_down_sampling_1 = ImageDownSampling(height=PATCH_SIZE, width=PATCH_SIZE, scale=2)
-    self.img_down_sampling_2 = ImageDownSampling(height=PATCH_SIZE, width=PATCH_SIZE, scale=4)
+    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(input_channel, feature[0], enc_fet_ch=feature[0], max_pool_size=2, is_concate=False)
-    self.encoder_layer_2 = Encoder(input_channel, feature[1], enc_fet_ch=feature[0]*2, max_pool_size=2, is_concate=True)
-    self.encoder_layer_3 = Encoder(input_channel, feature[2], enc_fet_ch=feature[0]*4, max_pool_size=2, is_concate=True)
+    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=feature[2]*2, out_ch=feature[2]*4)
+    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=feature[-1]*4, tensor_dim_encoder=feature[-1]*2, tensor_dim_mid=attention_dim[0], up_conv_in_ch=feature[-1]*4, up_conv_out_ch=feature[-1]*2, up_conv_scale=2, dconv_in_feature=feature[-1]*4, dconv_out_feature=feature[-1]*2, is_concat=True)
-    self.decoder_layer_2 = Decoder(tensor_dim_decoder=feature[-1]*2, tensor_dim_encoder=feature[-1], tensor_dim_mid=attention_dim[1], up_conv_in_ch=feature[-1]*2, up_conv_out_ch=feature[-1], up_conv_scale=2, dconv_in_feature=feature[-1]*2, dconv_out_feature=feature[-1], is_concat=True)
-    self.decoder_layer_3 = Decoder(tensor_dim_decoder=feature[-1], tensor_dim_encoder=feature[-2], tensor_dim_mid=attention_dim[2], up_conv_in_ch=feature[-1], up_conv_out_ch=feature[-2], up_conv_scale=2, dconv_in_feature=feature[-1], dconv_out_feature=feature[-2], is_concat=True)
+    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=feature[-3], out_channels=output_channel, kernel_size=1, padding=0, stride=1),
-        ImageFolding(image_size=IMAGE_SIZE[0], patch_size=PATCH_SIZE, batch_size=BATCH_SIZE)
+        nn.Conv2d(in_channels=config.features[-3], out_channels=config.num_classes, kernel_size=1, padding=0, stride=1),
+        ImageFolding(image_size=IMAGE_SIZE[0], patch_size=config.patch_size, batch_size=config.batch_size)
     )
 
   def forward(self, x):