quocanh944 commited on
Commit
0034728
1 Parent(s): ff83535

Add application file

Browse files
Files changed (1) hide show
  1. app.py +43 -0
app.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import TrOCRProcessor, VisionEncoderDecoderModel
3
+ import requests
4
+ from PIL import Image
5
+
6
+ processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
7
+ model = VisionEncoderDecoderModel.from_pretrained("quocanh944/tr-ocr")
8
+
9
+ # load image examples from the IAM database
10
+ urls = ['./images/a01-000u-00.png',
11
+ './images/a01-000x-04.png',
12
+ './images/a01-003-10.png']
13
+
14
+ for idx, url in enumerate(urls):
15
+ image = Image.open(url)
16
+ image.save(f"image_{idx}.png")
17
+
18
+ def process_image(image):
19
+ # prepare image
20
+ pixel_values = processor(image, return_tensors="pt").pixel_values
21
+
22
+ # generate (no beam search)
23
+ generated_ids = model.generate(pixel_values)
24
+
25
+ # decode
26
+ generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
27
+
28
+ return generated_text
29
+
30
+ title = "Interactive demo: TrOCR"
31
+ description = "Demo for Microsoft's TrOCR, an encoder-decoder model consisting of an image Transformer encoder and a text Transformer decoder for state-of-the-art optical character recognition (OCR) on single-text line images. This particular model is fine-tuned on IAM, a dataset of annotated handwritten images. To use it, simply upload an image or use the example image below and click 'submit'. Results will show up in a few seconds."
32
+ article = "TrOCR: Transformer-based Optical Character Recognition with Pre-trained Models | Github Repo"
33
+ examples =[["image_0.png"], ["image_1.png"], ["image_2.png"]]
34
+
35
+ iface = gr.Interface(fn=process_image,
36
+ inputs=gr.inputs.Image(type="pil"),
37
+ outputs=gr.outputs.Textbox(),
38
+ title=title,
39
+ description=description,
40
+ article=article,
41
+ examples=examples)
42
+
43
+ iface.launch()