import gradio as gr
import numpy as np
import cv2
import tensorflow as tf
import keras
from keras import layers, models

model = model = tf.keras.models.load_model('model/ocr_model.h5', compile=False)

def preprocessImage(img, shape):
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = cv2.resize(img, (shape))
    img = (img/255).astype(np.float32)
    img = img.T
    img = np.expand_dims(img, axis=-1)
    return img

label2char ={0: ' ',
 1: "'",
 2: '-',
 3: 'A',
 4: 'B',
 5: 'C',
 6: 'D',
 7: 'E',
 8: 'F',
 9: 'G',
 10: 'H',
 11: 'I',
 12: 'J',
 13: 'K',
 14: 'L',
 15: 'M',
 16: 'N',
 17: 'O',
 18: 'P',
 19: 'Q',
 20: 'R',
 21: 'S',
 22: 'T',
 23: 'U',
 24: 'V',
 25: 'W',
 26: 'X',
 27: 'Y',
 28: 'Z',
 29: '`'}

def getStringFromEncode(lst :list):
    return ''.join([label2char[i] if i in label2char else '' for i in lst])

def decode_batch_predictions(pred):
    pred = pred[:, :-2] # first two layers of ctc garbage
    input_len = np.ones(pred.shape[0])*pred.shape[1]
    
    results = keras.backend.ctc_decode(pred, 
                                        input_length=input_len,
                                        greedy=True)[0][0]
    
    output_text = []
    for res in results.numpy():
        outstr = getStringFromEncode(res)
        output_text.append(outstr)
    
    # return final text results
    return output_text

def predict(img):
    img = preprocessImage(img, (256,64))
    img = np.expand_dims(img, axis=0) # 1 image in batch

    preds = model.predict(img)
    pred_texts = decode_batch_predictions(preds)
    return pred_texts[0]

iface = gr.Interface(fn=predict, inputs="image", outputs="text")
iface.launch()