import os
import gradio as gr
import omegaconf
import torch
from vietocr.model.transformerocr import VietOCR
from vietocr.model.vocab import Vocab
from vietocr.translate import translate, process_input

examples_data = os.listdir('examples')
examples_data = [os.path.join('examples', line.split('\t')[0]) for line in examples_data]

config = omegaconf.OmegaConf.load("vgg-seq2seq.yaml")
config = omegaconf.OmegaConf.to_container(config, resolve=True)

vocab = Vocab(config['vocab'])
model = VietOCR(len(vocab),
        config['backbone'],
        config['cnn'], 
        config['transformer'],
        config['seq_modeling'])
model.load_state_dict(torch.load('train_old.pth', map_location=torch.device('cpu')))

def predict(inp):
    img = process_input(inp, config['dataset']['image_height'], 
                    config['dataset']['image_min_width'], config['dataset']['image_max_width'])
    out = translate(img, model)[0].tolist()
    out = vocab.decode(out)
    return out

gr.Interface(fn=predict,
             title='Vietnamese Handwriting Recognition',
             inputs=gr.Image(type='pil'),
             outputs=gr.Text(),
             examples=examples_data,
).launch()