Update app.py

app.py

@@ -1,54 +1,67 @@
 import streamlit as st
 import requests
-from PIL import Image
-from io import BytesIO
-import torch
-from transformers import BlipProcessor, BlipForConditionalGeneration
+from PIL import Image, ImageDraw, ImageFont
 import os
+from transformers import BlipProcessor, BlipForConditionalGeneration
 
-# Load BLIP model for caption generation
+# 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 caption
+# 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("Image Caption Generator")
-
+    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)
-        st.image(image, caption="Uploaded Image", use_container_width=True)
         
-        # Generate caption
-        if st.button("Generate Caption"):
+        # 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)
-            st.write("Generated Caption: ", caption)
             
-            # Save image with caption
+            # 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 to a file
+                # Save the image with the caption
                 save_path = os.path.join("saved_images", "image_with_caption.jpg")
                 image.save(save_path)
 
-                # Provide download button
+                # 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()