import onnxruntime as onr
import numpy as np
import glob
import gradio as gr
characters = ['z', 's', 'h', 'q', 'd', 'v', '2', '7', '8', 'x', 'y', '5', 'e', 'a', 'u', '4', 'k', 'n', 'm', 'c', 'p']
img_width = 130
img_height = 50
max_length = 7
Model = onr.InferenceSession('model.onnx')
ModelName = Model.get_inputs()[0].name
def solve_task(img):
    img = img.astype(np.float32) / 255.
    img = img.transpose([1, 0, 2])
    img = np.array([img])
    result_tensor = Model.run(None, {ModelName: img})[0]
    answer, accuracy = get_result(result_tensor)

    return answer
def get_result(pred):
    accuracy = 1
    last = None
    ans = []
    for item in pred[0]:
        char_ind = item.argmax()
        if char_ind != last and char_ind != 0 and char_ind != len(characters) + 1:
            ans.append(characters[char_ind - 1])
            accuracy *= item[char_ind]
        last = char_ind
    answ = "".join(ans)[:max_length]
    return answ, accuracy


title = "captcha solver"
description = "hate captcha"

iface = gr.Interface(fn=solve_task, 
                     inputs=gr.inputs.Image((img_width, img_height)), 
                     outputs=gr.outputs.Textbox(),
                     title=title,
                     examples=glob.glob('examples/*.jfif'),
                     description=description)
iface.launch()