File size: 1,985 Bytes
ac50f17 e552807 ac50f17 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
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()
|