TrOCR-SIN-demo / app.py
kavg's picture
Update app.py
1683728 verified
raw
history blame
873 Bytes
from transformers import VisionEncoderDecoderModel
from transformers import TrOCRProcessor, AutoTokenizer, ViTImageProcessor
import gradio as gr
def OCR(image):
# Load the image
img = Image.open(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(img, 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(),outputs=gr.Textbox())
demo.launch(debug=True)