import os
os.system('pip install --upgrade gdown')
import gdown
gdown.download(id='1C3XPT8s-ONt88vlNykTkGt8c8frfB9U_', 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

config = Config('configs/train_iternet.yaml')
config.model_vision_checkpoint = None
model = get_model(config)
model = load(model, 'workdir/train-iternet/best-train-iternet.pth')
charset = CharsetMapper(filename=config.dataset_charset_path, max_length=config.dataset_max_length + 1)

def process_image(image):
  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 = "Made with IterNet"
description = "I hate captchas"

iface = gr.Interface(fn=process_image, 
                     inputs=gr.inputs.Image(type="pil"), 
                     outputs=gr.outputs.Textbox(),
                     title=title,
                     description=description,
                     examples=glob.glob('figs_captchas/*.jpg'))

iface.launch(debug=True)