Spaces:
Runtime error
Runtime error
File size: 2,532 Bytes
5462c07 716b1d9 fd16bc7 716b1d9 fd16bc7 5462c07 716b1d9 fd16bc7 5462c07 716b1d9 fd16bc7 716b1d9 fd16bc7 5462c07 fd16bc7 |
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 56 57 58 59 60 61 |
import streamlit as st
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')
@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
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(raw_image).unsqueeze(0).to(device)
return image, raw_image.resize((w//5,h//5))
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()
|