Update app.py

app.py

@@ -1,58 +1,45 @@
-# install requirements
-import requests
+import streamlit as st
 from PIL import Image
+import requests
 import torch
 from torchvision import transforms
 from torchvision.transforms.functional import InterpolationMode
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from models.blip import blip_decoder
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-def load_demo_image(image_size, device, img_url):
-    raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')   
+def load_image(image_url, image_size, device):
+    raw_image = Image.open(requests.get(image_url, stream=True).raw).convert('RGB')
 
-    w, h = raw_image.size
+    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_captions(image, model):
-    # beam search
-    #captions = model.generate(image, sample=False, num_beams=3, max_length=20, min_length=5, num_return_sequences=3) 
-    # nucleus sampling
-    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):
-        print(f'caption {i+1}: {caption}')
-
-model_name = 'EleutherAI/gpt-neo-1.3B'
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
-model.eval()
-
-# Streamlit app code
-import streamlit as st
-
-st.title('Image Caption Generator')
-
-# Get user input
-img_url = st.text_input('Enter image URL', 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg')
-
-if img_url:
-    # Load image
+def generate_caption(image_url):
     image_size = 384
-    image = load_demo_image(image_size, device, img_url)
-    
-    # Generate captions
-    input_text = tokenizer.decode(tokenizer(image.tolist()[0]).input_ids)
-    print(f'Input text: {input_text}')
-    generate_captions(image, model)
+    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
+
+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}')