Update app.py

app.py

@@ -102,33 +102,20 @@ def draw_boxes(image, boxes, predictions):
 def process_image(image):
     width, height = image.size
 
-    # encode
-    encoding = processor(
-        image, truncation=True, return_offsets_mapping=True, return_tensors="pt"
-    )
-    offset_mapping = encoding.pop("offset_mapping")
-
-    # forward pass
+    # create model input
+    encoding = processor(image, return_tensors="pt")
+    del encoding["pixel_values"]
+    # run inference
     outputs = model(**encoding)
-
-    # get predictions
     predictions = outputs.logits.argmax(-1).squeeze().tolist()
-    token_boxes = encoding.bbox.squeeze().tolist()
-
-    # only keep non-subword predictions
-    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]
-    ]
-    true_boxes = [
-        unnormalize_box(box, width, height)
-        for idx, box in enumerate(token_boxes)
-        if not is_subword[idx]
-    ]
-
-    draw_boxes(image, true_boxes, true_predictions)
+    # get labels
+    labels = [model.config.id2label[prediction] for prediction in predictions]
+    if output_image:
+        return draw_boxes(image, encoding["bbox"][0], labels)
+    else:
+        return labels
     
-    return image
+ 
 
 
 title = "OCR Document Parser : Information Extraction - Fine Tuned LiLT Language-independent Layout Transformer Model"