Update app.py

app.py

@@ -59,10 +59,18 @@ def iob_to_label(label):
     return label
 
 
+def print_content(image, predictions, token_boxes, ner_tags):
+    width, height = image.size
+    content_list = []
+    for prediction, box, tag in zip(predictions, token_boxes, ner_tags):
+        if not is_subword[idx]:
+            predicted_label = iob_to_label(prediction)
+            content = image.crop(box).copy().convert("RGB")
+            content_list.append((predicted_label, tag, content))
+    return content_list
 
-def process_image(image):
 
-    print(type(image))
+def process_image(image, ner_tags):
     width, height = image.size
 
     # encode
@@ -84,11 +92,10 @@ def process_image(image):
     # draw predictions over the image
     draw = ImageDraw.Draw(image)
     font = ImageFont.load_default()
-    for prediction, box in zip(true_predictions, true_boxes):
+    for prediction, box, ner_tag in zip(true_predictions, true_boxes, ner_tags):
         predicted_label = iob_to_label(prediction)
         draw.rectangle(box, outline=label2color[predicted_label])
-        draw.text((box[0]+10, box[1]-10), text=predicted_label, fill=label2color[predicted_label], font=font)
-        draw.text((box[0]+10, box[1]+10), text=dataset[0]["ner_tags"][id], fill='black', font=font)
+        draw.text((box[0]+10, box[1]+10), text=ner_tag, fill='black', font=font)
     
     return image
 
@@ -112,4 +119,4 @@ iface = gr.Interface(fn=process_image,
                      css=css,
                      analytics_enabled = True, enable_queue=True)
 
-iface.launch(inline=False, share=False, debug=False)
+iface.launch(inline=False, share=False, debug=False)