Update app.py

app.py

@@ -1,45 +1,60 @@
 import streamlit as st
-from PIL import Image
-import requests
 import torch
+import requests
+from PIL import Image
 from torchvision import transforms
 from torchvision.transforms.functional import InterpolationMode
+
 from models.blip import blip_decoder
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-def load_image(image_url, image_size, device):
-    raw_image = Image.open(requests.get(image_url, stream=True).raw).convert('RGB')
-
+@st.cache(show_spinner=False)
+def load_demo_image(image_size, device):
+    img_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg'
+    raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
     w,h = raw_image.size
-    display(raw_image.resize((w//5,h//5)))
-    
     transform = transforms.Compose([
-        transforms.Resize((image_size,image_size),interpolation=InterpolationMode.BICUBIC),
+        transforms.Resize((image_size, image_size), interpolation=InterpolationMode.BICUBIC),
         transforms.ToTensor(),
         transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
-        ]) 
-    image = transform(raw_image).unsqueeze(0).to(device)   
-    return image
-
-def generate_caption(image_url):
-    image_size = 384
-    model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_capfilt_large.pth'    
-    model = blip_decoder(pretrained=model_url, image_size=image_size, vit='base')
-    model.eval()
-    model = model.to(device)
-    image = load_image(image_url, image_size, device)
-    with torch.no_grad():
-        captions = []
-        for i in range(3):
-            caption = model.generate(image, sample=True, top_p=0.9, max_length=20, min_length=5)
-            captions.append(caption[0])
-        return captions
+    ])
+    image = transform(raw_image).unsqueeze(0).to(device)
+    return image, raw_image.resize((w//5,h//5))
 
-st.title("Image Caption Generator")
-
-image_url = st.text_input("Enter the image URL:")
-if image_url:
-    captions = generate_caption(image_url)
-    for i, caption in enumerate(captions):
-        st.write(f'caption {i+1}: {caption}')
+def main():
+    st.set_page_config(page_title="Image Captioning App")
+    st.title("Image Captioning App")
+    st.write("This app generates captions for images using a pre-trained model.")
+    
+    # Load image
+    image_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+    if image_file is not None:
+        image = Image.open(image_file)
+        image_size = 384
+        transform = transforms.Compose([
+            transforms.Resize((image_size, image_size), interpolation=InterpolationMode.BICUBIC),
+            transforms.ToTensor(),
+            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
+        ])
+        image = transform(image).unsqueeze(0).to(device)
+        
+        # Generate captions
+        with torch.no_grad():
+            model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_capfilt_large.pth'
+            model = blip_decoder(pretrained=model_url, image_size=image_size, vit='base')
+            model.eval()
+            model = model.to(device)
+            num_captions = 3
+            captions = []
+            for i in range(num_captions):
+                caption = model.generate(image, sample=True, top_p=0.9, max_length=20, min_length=5)
+                captions.append(caption[0])
+            for i, caption in enumerate(captions):
+                st.write(f'Caption {i+1}: {caption}')
+            
+        # Display uploaded image
+        st.image(image_file, caption='Uploaded image', use_column_width=True)
+        
+if __name__ == "__main__":
+    main()