import cv2
import numpy as np
import os 
import tempfile
import pickle
import tensorflow as tf
import gradio as gr
from PIL import Image

# 
# Transformer 
from library.prediction import evaluate_single_image
from  library.transformer import Transformer
from library.customSchedule import learning_rate

top_k = 25000
num_layer = 4
d_model = 512
dff = 2048
num_heads = 8
row_size = 8
col_size = 8
target_vocab_size = top_k + 1
dropout_rate = 0.1


loaded_transformer = Transformer(num_layer, d_model, num_heads, dff, row_size, col_size,
                                 target_vocab_size, max_pos_encoding=target_vocab_size,
                                 rate=dropout_rate)

# Load the weights into the model
loaded_transformer.load_weights('model/')
# Use the loaded custom objects
loaded_transformer.compile(optimizer=tf.keras.optimizers.Adam(learning_rate))
print("Trasformer model loaded successfully")
global tokenizer
with open('transformer/tokenizer.pickle', 'rb') as handle:
    tokenizer = pickle.load(handle)
    tokenizer.word_index['<pad>'] = 0
    tokenizer.index_word[0] = '<pad>'


print("Tokenizer  loaded successfully")


def transformer(image_pil):
    print(image_pil)
    temp_file, temp_file_path = tempfile.mkstemp(suffix=".jpg")
    os.close(temp_file)
    image_pil.save(temp_file_path)
    caption = evaluate_single_image(temp_file_path, tokenizer, loaded_transformer)
    return caption
    
interface = gr.Interface(fn=transformer, inputs=gr.Image(type="pil"), outputs="text", api_name="predict")
interface.launch(debug=True)