Anpassungen für HF-Upload

segformer_plusplus/configuration_segformer_plusplus.py

@@ -0,0 +1,10 @@
+from transformers import PretrainedConfig
+
+class SegformerPlusPlusConfig(PretrainedConfig):
+    model_type = "segformer_plus_plus"
+
+    def __init__(self, backbone_cfg=None, decode_head_cfg=None, out_channels=19, **kwargs):
+        super().__init__(**kwargs)
+        self.backbone_cfg = backbone_cfg
+        self.decode_head_cfg = decode_head_cfg
+        self.out_channels = out_channels

segformer_plusplus/modeling_segformer_plusplus.py

@@ -1,44 +1,13 @@
-from typing import Optional, Tuple
-import torch.nn as nn
-from transformers import PreTrainedModel, PretrainedConfig
-from segformer_plusplus.model import create_model
+from transformers import PreTrainedModel
+from .configuration_segformer_plusplus import SegformerPlusPlusConfig
+from .build_model import SegFormer
 
-
-class SegformerPlusPlusConfig(PretrainedConfig):
-    model_type = "segformer_plusplus"
-
-    def __init__(
-        self,
-        backbone: str = "b5",
-        tome_strategy: Optional[str] = "bsm_hq",
-        num_labels: int = 19,
-        id2label: Optional[dict] = None,
-        label2id: Optional[dict] = None,
-        **kwargs,
-    ):
-        self.backbone = backbone
-        self.tome_strategy = tome_strategy
-        self.num_labels = num_labels
-
-        if id2label is None:
-            id2label = {i: f"class_{i}" for i in range(num_labels)}
-        if label2id is None:
-            label2id = {v: k for k, v in id2label.items()}
-
-        self.id2label = id2label
-        self.label2id = label2id
-
-        super().__init__(**kwargs)
-
-
-class SegformerPlusPlusForSemanticSegmentation(PreTrainedModel):
+class SegformerPlusPlusModel(PreTrainedModel):
     config_class = SegformerPlusPlusConfig
 
     def __init__(self, config: SegformerPlusPlusConfig):
         super().__init__(config)
-        self.segformer = create_model(
-            backbone=config.backbone,
-            tome_strategy=config.tome_strategy,
-            out_channels=config.num_labels,
-            pretrained=False,
-        )
+        self.model = SegFormer(config)
+
+    def forward(self, x, **kwargs):
+        return self.model(x)