dschandra commited on
Commit
ac50f17
·
verified ·
1 Parent(s): dcfc73d

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +54 -0
app.py ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import requests
3
+ from PIL import Image
4
+ from io import BytesIO
5
+ import torch
6
+ from transformers import BlipProcessor, BlipForConditionalGeneration
7
+ import os
8
+
9
+ # Load BLIP model for caption generation
10
+ processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
11
+ model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
12
+
13
+ # Function to generate caption
14
+ def generate_caption(image):
15
+ inputs = processor(images=image, return_tensors="pt")
16
+ out = model.generate(**inputs)
17
+ caption = processor.decode(out[0], skip_special_tokens=True)
18
+ return caption
19
+
20
+ # Streamlit app
21
+ def main():
22
+ st.title("Image Caption Generator")
23
+
24
+ # Upload image
25
+ uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "png", "jpeg"])
26
+ if uploaded_file is not None:
27
+ # Open the uploaded image
28
+ image = Image.open(uploaded_file)
29
+ st.image(image, caption="Uploaded Image", use_column_width=True)
30
+
31
+ # Generate caption
32
+ if st.button("Generate Caption"):
33
+ caption = generate_caption(image)
34
+ st.write("Generated Caption: ", caption)
35
+
36
+ # Save image with caption
37
+ if st.button("Save Image with Caption"):
38
+ # Draw the caption on the image
39
+ draw = ImageDraw.Draw(image)
40
+ font = ImageFont.load_default()
41
+ draw.text((10, 10), caption, fill="white", font=font)
42
+
43
+ # Save the image to a file
44
+ save_path = os.path.join("saved_images", "image_with_caption.jpg")
45
+ image.save(save_path)
46
+
47
+ # Provide download button
48
+ with open(save_path, "rb") as f:
49
+ st.download_button("Download Image with Caption", f, file_name="image_with_caption.jpg")
50
+
51
+ if __name__ == "__main__":
52
+ if not os.path.exists("saved_images"):
53
+ os.makedirs("saved_images")
54
+ main()