import gradio as gr
from PIL import Image
from vietocr.tool.predictor import Predictor
from vietocr.tool.config import Cfg


config = Cfg.load_config_from_name('vgg_transformer')
config['weights'] = 'ocr.pth'
config['device'] = 'cpu'
config['cnn']['pretrained']=False
config['predictor']['beamsearch']=False

detector = Predictor(config)


def run(source):
    global detector
    img = Image.open(source) 
    result = detector.predict(img)
    return result
    

gr.Interface(
    run,
    inputs=gr.Image(label="Upload image", type="filepath"),
    outputs="label",
    title="Transformer Ocr",
).launch()