dschandra's picture
Update app.py
2d43de5 verified
raw
history blame
2.51 kB
import streamlit as st
import requests
from PIL import Image, ImageDraw, ImageFont
import os
from transformers import BlipProcessor, BlipForConditionalGeneration
# 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 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)
# Return the creative caption
return 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.")
# 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()