Upload Model.py with huggingface_hub

Model.py

@@ -208,6 +208,21 @@ class Decoder(nn.Module):
     x = self.double_conv(x)
     return x
 
+######################################################################
+# SEGMENTATION HEAD
+######################################################################
+class SegmentationHead(nn.Module):
+  def __init__(self, feature_dim, num_classes, config:VesselSegmentConfig = VesselSegmentConfig()):
+    super().__init__()
+    self.config = config
+    self.conv = nn.Conv2d(in_channels=feature_dim, out_channels=num_classes, kernel_size=1, stride=1, padding=0)
+
+  def forward(self, x, batch_size):
+    x1 = self.conv(x)
+    x1 = ImageFolding(image_size=self.config.image_size[0], patch_size=self.config.patch_size, batch_size=batch_size)(x1)
+    return x1
+
+
 class VesselSegmentModel(PreTrainedModel):
   config_class = VesselSegmentConfig
   def __init__(self, config: VesselSegmentConfig=VesselSegmentConfig()):
@@ -233,12 +248,14 @@ class VesselSegmentModel(PreTrainedModel):
     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)
-    )
-
+    self.segmenation_head = SegmentationHead(feature_dim=config.features[-3], num_classes=config.num_classes)
+    # 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):
+    B,C,H,W = x.shape
     IMG_1 = self.img_patch(x)
     IMG_2 = self.img_down_sampling_1(IMG_1)
     IMG_3 = self.img_down_sampling_2(IMG_2)
@@ -257,6 +274,6 @@ class VesselSegmentModel(PreTrainedModel):
     d3 = self.decoder_layer_3(sk1, d2)
 
     # head
-    head = self.segmenation_head(d3)
+    head = self.segmenation_head(d3, B)
 
     return head