Upload HFUnetPlusPlus

config.json

@@ -0,0 +1,21 @@
+{
+  "architectures": [
+    "HFUnetPlusPlus"
+  ],
+  "auto_map": {
+    "AutoConfig": "hf_config.UnetConfig",
+    "AutoModelForImageSegmentation": "hf_model.HFUnetPlusPlus"
+  },
+  "decoder_channels": [
+    1024,
+    512,
+    256,
+    128,
+    64
+  ],
+  "encoder_name": "resnet18",
+  "input_channels": 1,
+  "num_classes": 16,
+  "torch_dtype": "float32",
+  "transformers_version": "4.24.0"
+}

hf_config.py

@@ -0,0 +1,17 @@
+from transformers import PretrainedConfig
+
+class UnetConfig(PretrainedConfig):
+
+    def __init__(
+        self,
+        encoder_name: str = "resnet18",
+        num_classes: int = 16,
+        input_channels: int = 1,
+        decoder_channels: tuple = (1024, 512, 256, 128, 64),
+        **kwargs
+    ):
+        self.encoder_name = encoder_name
+        self.num_classes = num_classes
+        self.input_channels = input_channels
+        self.decoder_channels = decoder_channels
+        super().__init__(**kwargs)

hf_model.py

@@ -0,0 +1,21 @@
+import segmentation_models_pytorch as smp
+from hf_config import UnetConfig
+from transformers import PreTrainedModel
+
+
+class HFUnetPlusPlus(PreTrainedModel):
+    config_class = UnetConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.model = smp.UnetPlusPlus(
+            encoder_name=config.encoder_name,
+            encoder_weights="imagenet",
+            decoder_channels=config.decoder_channels,
+            in_channels=config.input_channels,
+            classes=config.num_classes,
+            decoder_attention_type="scse")
+
+    def forward(self, tensor):
+        return self.model(tensor)

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:11241ff685300ecf0d1314eabc7bf74223d2071bae43b7b3ce7859a01f599efd
+size 164204733