import streamlit as st
from transformers import pipeline
from PIL import Image

# Load the image-to-text model pipeline
@st.cache_resource
def load_pipeline():
    return pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")

pipe = load_pipeline()

# Streamlit app title
st.title("🖼️ Image Caption Generator using AI")
st.write("Upload an image, and the model will generate a descriptive caption.")

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

if uploaded_file:
    # Display the uploaded image
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image', use_column_width=True)

    # Generate a caption for the image
    st.write("Generating caption...")
    caption = pipe(image)

    # Display the generated caption
    if caption:
        st.subheader("Generated Caption:")
        st.write(caption[0]['generated_text'])
    else:
        st.write("Couldn't generate a caption. Please try a different image.")

# Footer information
st.markdown("---")
st.write("Created with 🤗 Transformers and Streamlit")