import os
os.system('pip install --upgrade gdown')
import gdown
gdown.download(id='1mYM_26qHUom_5NU7iutHneB_KHlLjL5y', output='workdir.zip')
os.system('unzip workdir.zip')

import glob
import gradio as gr
from demo import get_model, preprocess, postprocess, load
from utils import Config, Logger, CharsetMapper

def process_image(image):
  config = Config('configs/train_abinet.yaml')
  config.model_vision_checkpoint = None
  model = get_model(config)
  model = load(model, 'workdir/train-abinet/best-train-abinet.pth')
  charset = CharsetMapper(filename=config.dataset_charset_path, max_length=config.dataset_max_length + 1)

  img = image.convert('RGB')
  img = preprocess(img, config.dataset_image_width, config.dataset_image_height)
  res = model(img)
  return postprocess(res, charset, 'alignment')[0][0]

title = "张博强毕设中期展示（文本识别部分）"
description = "西北工业大学航海学院张博强毕设，目前识别部分进度为复现abinet，本网页为abinet复现的可视化web端展示"
#article = "<p style='text-align: center'><a href='https://arxiv.org/pdf/2103.06495.pdf'>Read Like Humans: Autonomous, Bidirectional and Iterative Language Modeling for Scene Text Recognition</a> | <a href='https://github.com/FangShancheng/ABINet'>Github Repo</a></p>"

iface = gr.Interface(fn=process_image, 
                     inputs=[gr.inputs.Image(type="pil")], 
                     outputs=[gr.outputs.Textbox()],
                     title=title,
                     description=description,)
                     #examples=glob.glob('figs/test/*.png'))
iface.launch(enable_queue=True)