Upload LayoutFIDNetV3

config.json

@@ -8,16 +8,16 @@
   },
   "d_model": 256,
   "id2label": {
-    "0": "LABEL_0",
-    "1": "LABEL_1",
-    "2": "LABEL_2",
-    "3": "LABEL_3"
+    "0": "logo",
+    "1": "text",
+    "2": "underlay",
+    "3": "embellishment"
   },
   "label2id": {
-    "LABEL_0": 0,
-    "LABEL_1": 1,
-    "LABEL_2": 2,
-    "LABEL_3": 3
+    "embellishment": 3,
+    "logo": 0,
+    "text": 1,
+    "underlay": 2
   },
   "max_bbox": 10,
   "model_type": "layoutdm_fidnet_v3",

modeling_layout_fidnet_v3.py

@@ -1,6 +1,6 @@
 import logging
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -95,7 +95,9 @@ class LayoutFIDNetV3(PreTrainedModel):
         self.fc_out_cls = nn.Linear(config.d_model, config.num_labels)
         self.fc_out_bbox = nn.Linear(config.d_model, 4)
 
-    def extract_features(self, bbox, label, padding_mask):
+    def extract_features(
+        self, bbox: torch.Tensor, label: torch.Tensor, padding_mask: torch.Tensor
+    ) -> torch.Tensor:
         b = self.fc_bbox(bbox)
         l = self.emb_label(label)
         x = self.enc_fc_in(torch.cat([b, l], dim=-1))
@@ -103,7 +105,13 @@ class LayoutFIDNetV3(PreTrainedModel):
         x = self.enc_transformer(x, padding_mask)
         return x[0]
 
-    def forward(self, bbox, label, padding_mask):
+    def forward(
+        self,
+        bbox: torch.Tensor,
+        label: torch.Tensor,
+        padding_mask: torch.Tensor,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, LayoutFIDNetV3Output]:
         B, N, _ = bbox.size()
         x = self.extract_features(bbox, label, padding_mask)
 
@@ -122,6 +130,9 @@ class LayoutFIDNetV3(PreTrainedModel):
         logit_cls = self.fc_out_cls(x)
         bbox_pred = torch.sigmoid(self.fc_out_bbox(x))
 
+        if not return_dict:
+            return logit_disc, logit_cls, bbox_pred
+
         return LayoutFIDNetV3Output(
             logit_disc=logit_disc, logit_cls=logit_cls, bbox_pred=bbox_pred
         )