Spaces:
Build error
Build error
| 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() | |