add files

app.py

@@ -111,12 +111,33 @@ def patch_resize_transform(patch_image_size=480, is_document=False):
     return _patch_resize_transform
 
 
+reader = ReaderLite()
+overrides={"eval_cider":False, "beam":8, "max_len_b":128, "patch_image_size":480,
+           "orig_patch_image_size":224, "no_repeat_ngram_size":0, "seed":7}
+models, cfg, task = checkpoint_utils.load_model_ensemble_and_task(
+    utils.split_paths('checkpoints/ocr.pt'),
+    arg_overrides=overrides
+)
+
+# Move models to GPU
+for model in models:
+    model.eval()
+    if use_fp16:
+        model.half()
+    if use_cuda and not cfg.distributed_training.pipeline_model_parallel:
+        model.cuda()
+    model.prepare_for_inference_(cfg)
+
+# Initialize generator
+generator = task.build_generator(models, cfg.generation)
+
+bos_item = torch.LongTensor([task.src_dict.bos()])
+eos_item = torch.LongTensor([task.src_dict.eos()])
+pad_idx = task.src_dict.pad()
+
+
 # Construct input for caption task
 def construct_sample(task, image: Image, patch_image_size=480):
-    bos_item = torch.LongTensor([task.src_dict.bos()])
-    eos_item = torch.LongTensor([task.src_dict.eos()])
-    pad_idx = task.src_dict.pad()
-
     patch_image = patch_resize_transform(patch_image_size)(image).unsqueeze(0)
     patch_mask = torch.tensor([True])
     src_text = encode_text(task, "图片上的文字是什么?", append_bos=True, append_eos=True).unsqueeze(0)
@@ -141,35 +162,11 @@ def apply_half(t):
     return t
 
 
-def ocr(ckpt, img, out_img):
-    reader = ReaderLite()
-    overrides={"eval_cider":False, "beam":8, "max_len_b":128, "patch_image_size":480, "orig_patch_image_size":224, "no_repeat_ngram_size":0, "seed":7}
-    models, cfg, task = checkpoint_utils.load_model_ensemble_and_task(
-        utils.split_paths(ckpt),
-        arg_overrides=overrides
-    )
-
-    # Move models to GPU
-    for model in models:
-        model.eval()
-        if use_fp16:
-            model.half()
-        if use_cuda and not cfg.distributed_training.pipeline_model_parallel:
-            model.cuda()
-        model.prepare_for_inference_(cfg)
-
-    # Initialize generator
-    generator = task.build_generator(models, cfg.generation)
-
-    bos_item = torch.LongTensor([task.src_dict.bos()])
-    eos_item = torch.LongTensor([task.src_dict.eos()])
-    pad_idx = task.src_dict.pad()
-
+def ocr(img):
     orig_image = Image.open(img)
     results = get_images(img, reader)
     box_list, image_list = zip(*results)
     draw_boxes(orig_image, box_list)
-    orig_image.save(out_img)
 
     ocr_result = []
     for box, image in zip(box_list, image_list):
@@ -183,7 +180,8 @@ def ocr(ckpt, img, out_img):
         ocr_result.append(result[0]['ocr'].replace(' ', ''))
 
     result = '\n'.join(ocr_result)
-    return result
+
+    return orig_image, result
 
 
 title = "OFA-OCR"
@@ -192,7 +190,8 @@ description = "Gradio Demo for OFA-OCR. Upload your own image or click any one o
 article = "<p style='text-align: center'><a href='https://github.com/OFA-Sys/OFA' target='_blank'>OFA Github " \
           "Repo</a></p> "
 examples = [['lihe.png'], ['chinese.jpg'], ['paibian.jpeg'], ['shupai.png'], ['zuowen.jpg']]
-io = gr.Interface(fn=ocr, inputs=gr.inputs.Image(type='pil'), outputs=gr.outputs.Textbox(label="Caption"),
+io = gr.Interface(fn=ocr, inputs=gr.inputs.Image(type='pil'),
+                  outputs=[gr.outputs.Image(type='pil'), gr.outputs.Textbox(label="OCR result")],
                   title=title, description=description, article=article, examples=examples,
                   allow_flagging=False, allow_screenshot=False)
 io.launch(cache_examples=True)