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


image_to_text = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")


st.title("Image Caption Generator")


image_upload = st.file_uploader("Upload Image (JPG or PNG):", type=["jpg", "jpeg", "png"])

if image_upload:
    image = Image.open(image_upload)
    st.image(image, caption="Uploaded Image.", use_column_width=True)


if st.button("Generate Caption") and image_upload:
    
    image_data = image_upload.read()
    caption = image_to_text(image_upload)[0]
    print(caption)