import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from huggingface_hub import snapshot_download

import numpy as np
import gradio as gr

max_length = 5
img_width = 200
img_height = 50

# Download the model from Hugging Face
model_path = snapshot_download(repo_id="keras-io/ocr-for-captcha")

# Load the TensorFlow SavedModel using TFSMLayer
prediction_layer = keras.layers.TFSMLayer(model_path, call_endpoint='serving_default')

with open("vocab.txt", "r") as f:
    vocab = f.read().splitlines()

# Mapping integers back to original characters
num_to_char = layers.StringLookup(
    vocabulary=vocab, mask_token=None, invert=True
)

def decode_batch_predictions(pred):
    input_len = np.ones(pred.shape[0]) * pred.shape[1]
    # Use greedy search. For complex tasks, you can use beam search
    results = keras.backend.ctc_decode(pred, input_length=input_len, greedy=True)[0][0][
        :, :max_length
    ]
    # Iterate over the results and get back the text
    output_text = []
    for res in results:
        res = tf.strings.reduce_join(num_to_char(res)).numpy().decode("utf-8")
        output_text.append(res)
    return output_text

def classify_image(img_path):
    # 1. Read image
    img = tf.io.read_file(img_path)
    # 2. Decode and convert to grayscale
    img = tf.io.decode_png(img, channels=1)
    # 3. Convert to float32 in [0, 1] range
    img = tf.image.convert_image_dtype(img, tf.float32)
    # 4. Resize to the desired size
    img = tf.image.resize(img, [img_height, img_width])
    # 5. 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])
    img = tf.expand_dims(img, axis=0)
    preds = prediction_layer(img)  # Use the TFSMLayer for prediction
    pred_text = decode_batch_predictions(preds)
    return pred_text[0]

image = gr.inputs.Image(type='filepath')
text = gr.outputs.Textbox()

iface = gr.Interface(classify_image, image, text,
  title="CGIP CAPTCHA RECOGNITION OCR",
  description="Keras Implementation of OCR model for reading captcha 🤖🦹🏻",
  examples=["dd764.png", "3p4nn.png", "ydd3g.png", "268g2.png", "36nx4.png", "3bnyf.png", "5p8fm.png", "8y6b3.png", "mnef5.png", "yxd7m.png"]
)

iface.launch()