import os
import glob
import gradio as gr
import yaml
from utils import load_config, get_args, download_model
from ocr import OCR
import PIL
import numpy as np

config = load_config()
download_model('inference.pdmodel', config['inference.pdmodel'])
download_model('inference.pdiparams', config['inference.pdiparams'])

model = OCR(get_args(config['svtr_nums']['params']))

def process_image(image):
  img = image.convert('RGB')
  img = np.array(img)
  val, conf = model(img)
  return val, conf

title = "Interactive demo: SVTR-Tiny"
description = "Demo for SVTR-Tiny. The model supports only digit recognition. To use it, simply upload a (single-text line) image or use one of the example images below and click 'submit'. Results will show up in a few seconds."
article = "<p style='text-align: center'><a href='https://arxiv.org/pdf/2205.00159.pdf'>SVTR: Scene Text Recognition with a Single Visual Model</a> | <a href='https://github.com/PaddlePaddle/PaddleOCR/blob/release/2.7/doc/doc_en/algorithm_rec_svtr_en.md'>Github Repo</a></p>"

demo = gr.Interface(fn=process_image,
                     inputs=gr.inputs.Image(type="pil"),
                     outputs=gr.outputs.Textbox(),
                     title=title,
                     description=description,
                     article=article,
                     examples=glob.glob('examples/*.jpeg'))

demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)