Upload SegformerForSemanticSegmentation

config.json

@@ -0,0 +1,106 @@
+{
+  "architectures": [
+    "SegformerForSemanticSegmentation"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "auto_map": {
+    "AutoModelForImageSegmentation": "model.SegformerForSemanticSegmentation"
+  },
+  "classifier_dropout_prob": 0.1,
+  "decoder_hidden_size": 256,
+  "depths": [
+    2,
+    2,
+    2,
+    2
+  ],
+  "drop_path_rate": 0.1,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.0,
+  "hidden_sizes": [
+    32,
+    64,
+    160,
+    256
+  ],
+  "id2label": {
+    "0": "skin",
+    "1": "l_brow",
+    "2": "r_brow",
+    "3": "l_eye",
+    "4": "r_eye",
+    "5": "eye_g",
+    "6": "l_ear",
+    "7": "r_ear",
+    "8": "ear_r",
+    "9": "nose",
+    "10": "mouth",
+    "11": "u_lip",
+    "12": "l_lip",
+    "13": "neck",
+    "14": "neck_l",
+    "15": "cloth",
+    "16": "hair",
+    "17": "hat"
+  },
+  "initializer_range": 0.02,
+  "label2id": {
+    "cloth": 15,
+    "ear_r": 8,
+    "eye_g": 5,
+    "hair": 16,
+    "hat": 17,
+    "l_brow": 1,
+    "l_ear": 6,
+    "l_eye": 3,
+    "l_lip": 12,
+    "mouth": 10,
+    "neck": 13,
+    "neck_l": 14,
+    "nose": 9,
+    "r_brow": 2,
+    "r_ear": 7,
+    "r_eye": 4,
+    "skin": 0,
+    "u_lip": 11
+  },
+  "layer_norm_eps": 1e-06,
+  "mlp_ratios": [
+    4,
+    4,
+    4,
+    4
+  ],
+  "model_type": "segformer",
+  "num_attention_heads": [
+    1,
+    2,
+    5,
+    8
+  ],
+  "num_channels": 3,
+  "num_classes": 18,
+  "num_encoder_blocks": 4,
+  "patch_sizes": [
+    7,
+    3,
+    3,
+    3
+  ],
+  "reshape_last_stage": true,
+  "semantic_loss_ignore_index": 255,
+  "sr_ratios": [
+    8,
+    4,
+    2,
+    1
+  ],
+  "strides": [
+    4,
+    2,
+    2,
+    2
+  ],
+  "torch_dtype": "float32",
+  "transformers_version": "4.36.2"
+}

model.py

@@ -0,0 +1,111 @@
+import torch
+import transformers
+from torch import nn
+from transformers.modeling_outputs import SemanticSegmenterOutput
+
+
+def encode_down(c_in: int, c_out: int):
+    return nn.Sequential(
+        nn.Conv2d(in_channels=c_in, out_channels=c_out, kernel_size=3, padding=1),
+        nn.BatchNorm2d(num_features=c_out),
+        nn.ReLU(inplace=True),
+        nn.Conv2d(in_channels=c_out, out_channels=c_out, kernel_size=3, padding=1),
+        nn.BatchNorm2d(num_features=c_out),
+        nn.ReLU(inplace=True),
+    )
+
+
+def decode_up(c: int):
+    return nn.ConvTranspose2d(
+        in_channels=c,
+        out_channels=int(c / 2),
+        kernel_size=2,
+        stride=2,
+    )
+
+
+class FaceUNet(nn.Module):
+    def __init__(self, num_classes: int):
+        super().__init__()
+        self.num_classes = num_classes
+        
+        self.down_1 = nn.Conv2d(
+            in_channels=3,
+            out_channels=64,
+            kernel_size=3,
+            padding=1,
+        )
+        self.down_2 = encode_down(64, 128)
+        self.down_3 = encode_down(128, 256)
+        self.down_4 = encode_down(256, 512)
+        self.down_5 = encode_down(512, 1024)
+
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+
+        self.up_1 = decode_up(1024)
+        self.up_c1 = encode_down(1024, 512)
+        self.up_2 = decode_up(512)
+        self.up_c2 = encode_down(512, 256)
+        self.up_3 = decode_up(256)
+        self.up_c3 = encode_down(256, 128)
+        self.up_4 = decode_up(128)
+        self.up_c4 = encode_down(128, 64)
+
+        self.segment = nn.Conv2d(
+            in_channels=64,
+            out_channels=self.num_classes,
+            kernel_size=3,
+            padding=1,
+        )
+
+    def forward(self, x):
+        d1 = self.down_1(x)
+        d2 = self.pool(d1)
+        d3 = self.down_2(d2)
+        d4 = self.pool(d3)
+        d5 = self.down_3(d4)
+        d6 = self.pool(d5)
+        d7 = self.down_4(d6)
+        d8 = self.pool(d7)
+        d9 = self.down_5(d8)
+
+        u1 = self.up_1(d9)
+        x = self.up_c1(torch.cat([d7, u1], 1))
+        u2 = self.up_2(x)
+        x = self.up_c2(torch.cat([d5, u2], 1))
+        u3 = self.up_3(x)
+        x = self.up_c3(torch.cat([d3, u3], 1))
+        u4 = self.up_4(x)
+        x = self.up_c4(torch.cat([d1, u4], 1))
+
+        x = self.segment(x)
+        return x
+
+
+class Segformer(transformers.PreTrainedModel):
+    config_class = transformers.SegformerConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.model = FaceUNet(num_classes=config.num_classes)
+
+    def forward(self, tensor):
+        return self.model.forward_features(tensor)
+
+
+class SegformerForSemanticSegmentation(transformers.PreTrainedModel):
+    config_class = transformers.SegformerConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.model = FaceUNet(num_classes=config.num_classes)
+
+    def forward(self, pixel_values, labels=None):
+        logits = self.model(pixel_values)
+        values = {"logits": logits}
+        if labels is not None:
+            loss = torch.nn.cross_entropy(logits, labels)
+            values["loss"] = loss
+        return SemanticSegmenterOutput(**values)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:92930e2231ef4b99841c68ab826b59621934f91a27c7ed7e62c849be7a7b6d64
+size 124124040