Update app.py

app.py

@@ -59,24 +59,27 @@ def infer(image):
     # Use this if you're loading images
     #image = Image.open(img_path).convert("RGB")
     #image = image.convert("RGB") # loading PDFs
-    encoding = processor(image, return_offsets_mapping=True, return_tensors="pt", truncation=True, max_length=514)#max_positional_embeddings
-    offset_mapping = encoding.pop('offset_mapping')
-    outputs = model(**encoding)
-    predictions = outputs.logits.argmax(-1).squeeze().tolist()
-    token_boxes = encoding.bbox.squeeze().tolist()
-    width, height = image.size
-    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 = ImageDraw.Draw(image)
-
-    font = ImageFont.load_default()
-
-    for prediction, box in zip(true_predictions, true_boxes):
-        predicted_label = iob_to_label(prediction).lower()
-        draw.rectangle(box, outline=label2color[predicted_label])
-        draw.text((box[0]+10, box[1]-10), text=predicted_label, fill=label2color[predicted_label], font=font)
+    try:
+        encoding = processor(image, return_offsets_mapping=True, return_tensors="pt", truncation=True, max_length=514)#max_positional_embeddings
+        offset_mapping = encoding.pop('offset_mapping')
+        outputs = model(**encoding)
+        predictions = outputs.logits.argmax(-1).squeeze().tolist()
+        token_boxes = encoding.bbox.squeeze().tolist()
+        width, height = image.size
+        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 = ImageDraw.Draw(image)
+    
+        font = ImageFont.load_default()
+    
+        for prediction, box in zip(true_predictions, true_boxes):
+            predicted_label = iob_to_label(prediction).lower()
+            draw.rectangle(box, outline=label2color[predicted_label])
+            draw.text((box[0]+10, box[1]-10), text=predicted_label, fill=label2color[predicted_label], font=font)
+    except Exception as e:
+        print(e)
 
     return image