import streamlit as st
import requests
from PIL import Image
from io import BytesIO
import torch
from transformers import BlipProcessor, BlipForConditionalGeneration
import os

# Load BLIP model for caption generation
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

# Function to generate caption
def generate_caption(image):
    inputs = processor(images=image, return_tensors="pt")
    out = model.generate(**inputs)
    caption = processor.decode(out[0], skip_special_tokens=True)
    return caption

# Streamlit app
def main():
    st.title("Image Caption Generator")

    # Upload image
    uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "png", "jpeg"])
    if uploaded_file is not None:
        # Open the uploaded image
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image", use_container_width=True)
        
        # Generate caption
        if st.button("Generate Caption"):
            caption = generate_caption(image)
            st.write("Generated Caption: ", caption)
            
            # Save image with caption
            if st.button("Save Image with Caption"):
                # Draw the caption on the image
                draw = ImageDraw.Draw(image)
                font = ImageFont.load_default()
                draw.text((10, 10), caption, fill="white", font=font)
                
                # Save the image to a file
                save_path = os.path.join("saved_images", "image_with_caption.jpg")
                image.save(save_path)

                # Provide download button
                with open(save_path, "rb") as f:
                    st.download_button("Download Image with Caption", f, file_name="image_with_caption.jpg")

if __name__ == "__main__":
    if not os.path.exists("saved_images"):
        os.makedirs("saved_images")
    main()