import streamlit as st
from PIL import Image
import tensorflow as tf
import numpy as np
import pickle

# Load the pre-trained model
@st.cache(allow_output_mutation=True)
def load_model():
    model = tf.keras.models.load_model('./caption_model.h5')
    return model

# Load the tokenizer
@st.cache(allow_output_mutation=True)
def load_tokenizer():
    with open('tokenizer.pkl', 'rb') as handle:
        tokenizer = pickle.load(handle)
    return tokenizer

model = load_model()
tokenizer = load_tokenizer()

# Function to preprocess the image
def preprocess_image(image):
    image = image.resize((299, 299))  # Resize to the input size of the model
    image = np.array(image) / 255.0  # Normalize
    image = np.expand_dims(image, axis=0)  # Add batch dimension
    return image

# Function to generate caption
def generate_caption(image):
    image = preprocess_image(image)
    predictions = model.predict(image)
    predicted_caption = tokenizer.sequences_to_texts(predictions.argmax(axis=-1))
    return predicted_caption[0]

# Streamlit app
st.title("Image Captioning App")
st.write("Upload an image to generate a caption")

uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image.', use_column_width=True)
    st.write("")
    st.write("Generating caption...")
    caption = generate_caption(image)
    st.write(caption)