import gradio as gr
from tensorflow.keras.models import load_model
import cv2
import numpy as np
from tensorflow.keras.applications import ResNet50
from tensorflow.keras.models import Model
from tensorflow.keras.preprocessing.sequence import pad_sequences
import pickle

# Define the previous model
incept_model = ResNet50(include_top=True)
last = incept_model.layers[-2].output
modele = Model(inputs=incept_model.input, outputs=last)

model = load_model("model.h5")

pickle_open = open("count_words.pkl","rb")
count_words = pickle.load(pickle_open)
pickle_open.close()

pickle_open = open("inv_dict.pkl","rb")
inv_dict = pickle.load(pickle_open)
pickle_open.close()

def getImage(image):
    img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    img = cv2.resize(img, (224,224))
    img = np.reshape(img, (1,224,224,3))
    return img

def predict(image):
    test_feature = modele.predict(getImage(image)).reshape(1,2048)

    test_img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    test_img = cv2.cvtColor(test_img, cv2.COLOR_BGR2RGB)

    text_inp = ['startofseq']
    count = 0
    caption = ''

    while count < 25:
        count += 1

        encoded = [count_words.get(word, count_words.get('unk', 0)) for word in text_inp]
        encoded = [encoded]

        encoded = pad_sequences(encoded, padding='post', truncating='post', maxlen=34)

        prediction = np.argmax(model.predict([test_feature, encoded]))
        sampled_word = inv_dict[prediction]
        if sampled_word == 'endofseq':
            break
        caption = caption + ' ' + sampled_word

        

        text_inp.append(sampled_word)

    return caption

input_image = gr.inputs.Image(label="Upload an image")
output_text = gr.outputs.Textbox(label="Predictions")

gr.Interface(fn=predict, inputs=input_image, outputs=output_text).launch()