Update app.py

app.py

@@ -11,34 +11,77 @@ from transformers import LayoutLMv3ForTokenClassification
 from transformers import AutoProcessor
 from transformers import AutoModelForTokenClassification
 
-reading_order_model = AutoModelForTokenClassification.from_pretrained("omarelsayeed/LayoutReader85Large")
+finetuned_fully = LayoutLMv3ForTokenClassification.from_pretrained("omarelsayeed/finetuned_pretrained_model")
+
 processor = AutoProcessor.from_pretrained("microsoft/layoutlmv3-base", 
                                           apply_ocr=False)
 
-def predict_reading_order(boxes,image_path):
-    words = ["<unk>"]*len(boxes)
-    print(boxes)
-    encoding = processor(image_path , text = words
-                            , boxes=boxes
-                            ,return_tensors="pt" ,
-                            return_offsets_mapping=True)
-    offset_mapping = encoding.pop('offset_mapping')
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    for k,v in encoding.items():
-      encoding[k] = v.to(device)
-    outputs = reading_order_model(**encoding)
-    predictions = outputs.logits.argmax(-1).squeeze().tolist()
-    token_boxes = encoding.bbox.squeeze().tolist()
-    is_subword = np.array(offset_mapping.squeeze().tolist())[:,0] != 0
-    # true_predictions = [id2label[pred] for idx, pred in enumerate(predictions) if not is_subword[idx]]
-    predictions = predictions[1:-1]
-    return predictions
-    
-def get_orders(image_path, boxes):
-    b = scale_and_normalize_boxes(boxes)
-    orders = predict_reading_order(b,  image_path)
-    return orders
+MAX_LEN = 70
+CLS_TOKEN_ID = 0
+UNK_TOKEN_ID = 3
+EOS_TOKEN_ID = 2
+import torch
+def boxes2inputs(boxes):
+    bbox = [[0, 0, 0, 0]] + boxes + [[0, 0, 0, 0]]
+    input_ids = [CLS_TOKEN_ID] + [UNK_TOKEN_ID] * len(boxes) + [EOS_TOKEN_ID]
+    attention_mask = [1] + [1] * len(boxes) + [1]
+    return {
+        "bbox": torch.tensor([bbox]),
+        "attention_mask": torch.tensor([attention_mask]),
+        "input_ids": torch.tensor([input_ids]),
+    }
+def parse_logits(logits: torch.Tensor, length):
+    """
+    parse logits to orders
 
+    :param logits: logits from model
+    :param length: input length
+    :return: orders
+    """
+    logits = logits[1 : length + 1, :length]
+    orders = logits.argsort(descending=False).tolist()
+    ret = [o.pop() for o in orders]
+    while True:
+        order_to_idxes = defaultdict(list)
+        for idx, order in enumerate(ret):
+            order_to_idxes[order].append(idx)
+        # filter idxes len > 1
+        order_to_idxes = {k: v for k, v in order_to_idxes.items() if len(v) > 1}
+        if not order_to_idxes:
+            break
+        # filter
+        for order, idxes in order_to_idxes.items():
+            # find original logits of idxes
+            idxes_to_logit = {}
+            for idx in idxes:
+                idxes_to_logit[idx] = logits[idx, order]
+            idxes_to_logit = sorted(
+                idxes_to_logit.items(), key=lambda x: x[1], reverse=True
+            )
+            # keep the highest logit as order, set others to next candidate
+            for idx, _ in idxes_to_logit[1:]:
+                ret[idx] = orders[idx].pop()
+    return ret
+
+    
+def prepare_inputs(
+    inputs, model
+):
+    ret = {}
+    for k, v in inputs.items():
+        v = v.to(model.device)
+        if torch.is_floating_point(v):
+            v = v.to(model.dtype)
+        ret[k] = v
+    return ret
+    
+    
+def get_orders(image_path , boxes):
+    inputs = boxes2inputs(boxes)
+    inputs = prepare_inputs(inputs, finetuned_fully)
+    logits = finetuned_fully(**inputs).logits.cpu().squeeze(0)
+    predictions = parse_logits(logits, len(boxes))
+    return predictions
 
 model_dir = snapshot_download("omarelsayeed/DETR-ARABIC-DOCUMENT-LAYOUT-ANALYSIS") + "/rtdetr_1024_crops.pt"
 model = RTDETR(model_dir)