import streamlit as st
import requests
from PIL import Image, ImageDraw, ImageFont
import os
from transformers import BlipProcessor, BlipForConditionalGeneration
from random import choice

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

# Function to generate a detailed, creative caption
def generate_caption(image):
    # Prepare image for the model
    inputs = processor(images=image, return_tensors="pt")
    
    # Generate the caption
    out = model.generate(**inputs)
    
    # Decode the output to get a readable caption
    caption = processor.decode(out[0], skip_special_tokens=True)
    
    # Enhance the caption with more details (assuming the image context is about nature, people, etc.)
    detailed_caption = enhance_caption(caption)
    
    return detailed_caption

# Enhance caption with surrounding context and social media flavor
def enhance_caption(caption):
    # Possible enhancements for social media flavor
    emojis = ["🌅", "🌊", "🌻", "🌴", "🏖️", "🦋", "🌞", "🌇"]
    hashtags = ["#NatureLovers", "#TravelVibes", "#Sunset", "#BeachLife", "#AdventureAwaits"]
    
    # Select random emojis and hashtags to add a fun social media touch
    selected_emoji = choice(emojis)
    selected_hashtags = " ".join(choice(hashtags) for _ in range(2))  # Select two random hashtags
    
    # Creating a more detailed, fun caption for social media
    enhanced_caption = f"{caption} {selected_emoji} {selected_hashtags}"
    
    return enhanced_caption

# Streamlit app
def main():
    st.title("Creative Image Caption Generator")
    st.write("Upload an image, and let the AI generate a creative and descriptive caption for it, ready for social media sharing!")
    
    # 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)
        
        # Display the uploaded image
        st.image(image, caption="Uploaded Image", use_column_width=True)
        
        # Generate caption button
        if st.button("Generate Creative Caption"):
            caption = generate_caption(image)
            
            # Display the generated caption
            st.write("### Generated Creative Caption: ")
            st.write(caption)
            
            # Save image with caption overlay
            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 with the caption
                save_path = os.path.join("saved_images", "image_with_caption.jpg")
                image.save(save_path)

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

if __name__ == "__main__":
    # Create a directory to save the image
    if not os.path.exists("saved_images"):
        os.makedirs("saved_images")
        
    main()