import tensorflow as tf
from matplotlib import pyplot as plt
from skimage.transform import rescale, resize
import pickle as pkl
import numpy as np
import os
import cv2
from PIL import Image,ImageFont, ImageDraw
import CALTextModel
import gradio as gr
import data


#### training setup parameters ####


lambda_val=1e-4
gamma_val=1

os.environ['CUDA_VISIBLE_DEVICES'] = '-1'



################################### Utility functions###################################



# Create an instance of the model
CALTEXT = CALTextModel.CALTEXT_Model(training=False)
CALTEXT.load_weights('final_caltextModel/cp-0037.ckpt')
test_loss = tf.keras.metrics.Mean(name='test_loss')





examples = [
        ['sample_test_images/59-11.png'],
		['sample_test_images/59-21.png'],
		['sample_test_images/59-32.png'],
		['sample_test_images/59-37.png'],
		['sample_test_images/91-47.png'],
		['sample_test_images/91-49.png'],
    ]




def recognize_text(input_image):
    x, x_mask=data.preprocess_img(input_image)
    output_str, gifImage=CALTextModel.predict(x, x_mask)
    return output_str,gifImage


title = "CALText Demo"
description = "<p style='text-align: center'>Gradio demo for an CALText model architecture <a href='https://github.com/nazar-khan/CALText'>[GitHub Code]</a> trained on the <a href='http://faculty.pucit.edu.pk/nazarkhan/work/urdu_ohtr/pucit_ohul_dataset.html'>PUCIT-OHUL</a> dataset. To use it, simply add your image, or click one of the examples to load them.  </p>"
article = "<p style='text-align: center'></p>"
css = "#0 {object-fit: contain;} #1 {object-fit: contain;}"
inputs = gr.inputs.Image(label="Input Image")

demo = gr.Interface(fn=recognize_text,
                    inputs=inputs,
                    outputs=[gr.Textbox(label="Output"), 
                    gr.Image(label="Attended Regions")],
                    examples=examples,
                    title=title,
                    description=description,
                    article=article,allow_flagging='never')

demo.launch()