import gradio as gr
import tensorflow as tf
from huggingface_hub import from_pretrained_keras
from tensorflow import keras

num_to_char = keras.layers.StringLookup(
    vocabulary=sorted(
        set("abcdefghijklmnpqrstuvwxyz123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ".upper())
    ),
    mask_token=None,
    invert=True,
)

model = from_pretrained_keras("wangxinhe/luogu-captcha-recognition", compile=False)

# Get the prediction model by extracting layers till the output layer
prediction_model = keras.models.Model(
    model.input[0], model.get_layer(name="dense2").output
)
prediction_model.summary()


def ocr(img):
    # Convert to float32 in [0, 1] range
    img = tf.image.convert_image_dtype(img, tf.float32)
    # Transpose the image because we want the time
    # dimension to correspond to the width of the image.
    img = tf.transpose(img, perm=[1, 0, 2])

    preds = prediction_model(tf.expand_dims(img, axis=0))

    # Use greedy search. For complex tasks, you can use beam search
    results = keras.backend.ctc_decode(
        preds, input_length=[preds.shape[1]], greedy=True
    )[0][0][:, :4]
    return tf.strings.reduce_join(num_to_char(results[0])).numpy().decode("ascii")


iface = gr.Interface(
    fn=ocr,
    inputs=gr.Image(
        shape=(90, 35),
        source="upload",
        label="CAPTCHA image",
    ),
    outputs="textbox",
)
iface.launch()