Update app.py

app.py

@@ -104,10 +104,12 @@ def pred_resume(pdf_path) -> dict:
                         if text.replace(" ","") !=  "":
                             bboxes.append(normalize_bbox([xmin, ymin, xmax, ymax], image.size))
                             words.append(decontracted(text))
+        text_reverse = {str(bboxes[i]): words[i] for i,_ in enumerate(words)}
         fake_label = ["O"] * len(words)
         encoding = processor(image, words, boxes=bboxes, word_labels=fake_label, truncation=True, stride=256,
                       padding="max_length", max_length=512, return_overflowing_tokens=True, return_offsets_mapping=True)
         labels = encoding["labels"]
+        key_box = encoding["bbox"]
         offset_mapping = encoding.pop('offset_mapping')
         overflow_to_sample_mapping = encoding.pop('overflow_to_sample_mapping')
         encoding = {k: torch.tensor(v) for k,v in encoding.items() if k != "labels"}
@@ -128,12 +130,16 @@ def pred_resume(pdf_path) -> dict:
                 if i>0:
                     labels[i] = labels[i][256:]
                     predictions[i] = predictions[i][256:]
+                    key_box[i] = key_box[i][256:]
             predictions = [j for i in predictions for j in i]
+        key_box = [j for i in key_box for j in i]
         labels = [j for i in labels for j in i]
         true_predictions = [id2label[pred] for pred, label in zip(predictions, labels) if label != -100]
-        for i, pred in enumerate(true_predictions):
+        key_box = [box for box, label in zip(key_box, labels) if label != -100]
+        for box, pred in zip(key_box, true_predictions):
             if pred in key_list:
-                result[pred].append(words[i])
+                result[pred].append(text_reverse[str(box)])
+    result = {k: list(set(v)) for k, v in result.items()}
     return str(result)
 def norm(result: str) -> str:
     result = ast.literal_eval(result)