TrOCR-SIN-demo / app.py
kavg's picture
Update app.py
37c49d0 verified
raw
history blame contribute delete
876 Bytes
from transformers import VisionEncoderDecoderModel
from transformers import TrOCRProcessor, AutoTokenizer, ViTImageProcessor
import gradio as gr
from PIL import Image
def OCR(image):
model = VisionEncoderDecoderModel.from_pretrained("kavg/TrOCR-SIN")
tokenizer = AutoTokenizer.from_pretrained('NLPC-UOM/SinBERT-large')
feature_extractor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
processor = TrOCRProcessor(image_processor=feature_extractor, tokenizer=tokenizer)
pixel_values = processor(image, return_tensors="pt").pixel_values
generated_ids = model.generate(pixel_values)
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
return generated_text
demo = gr.Interface(fn=OCR, inputs=gr.Image(show_label=False, type="pil"),outputs=gr.Textbox())
demo.launch(debug=True)