Update app.py

app.py

@@ -1,25 +1,12 @@
-from cord_inference import prediction as cord_prediction
-from sroie_inference import prediction as sroie_prediction
+from inference import prediction as cord_sroie_prediction
 import gradio as gr
 import json
 
 def prediction(image):
 
-    # first use the model fine-tuned on sroie (for now it is Theivaprakasham/layoutlmv3-finetuned-sroie)
-    # on the image, which returns a JSON with some info and an image with the corresponding boxes blurred
-    j1, image_blurred = sroie_prediction(image)
-    
-    # then use the model fine-tuned on cord on the blurred image
-    img = image_blurred.copy()
-    j2, image_final = cord_prediction(img)
+    j2, image_final = cord_sroie_prediction(img)
 
-    # link the two json files
-    if len(j1) == 0:
-        j3 = j2
-    else:
-        j3 = json.dumps(j1).split('}')[0] + ', ' + json.dumps(j2).split('{')[1]
-
-    return j1, image_blurred, j2, image_final, j3
+    return j2, image_final
 
 
 title = "Interactive demo: LayoutLMv3 for receipts"
@@ -32,10 +19,7 @@ css = """.output_image, .input_image {height: 600px !important}"""
 # for now it shows also the intermediate steps
 iface = gr.Interface(fn=prediction,
                      inputs=gr.Image(type="pil"),
-                     outputs=[gr.JSON(label="sroie parsing"),
-                              gr.Image(type="pil", label="blurred image"),
-                              gr.JSON(label="cord parsing"),
-                              gr.Image(type="pil", label="annotated image"),
+                     outputs=[gr.Image(type="pil", label="annotated image"),
                               gr.JSON(label="final output")],
                      title=title,
                      description=description,