| | """ |
| | # Copyright (c) 2022, salesforce.com, inc. |
| | # All rights reserved. |
| | # SPDX-License-Identifier: BSD-3-Clause |
| | # For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause |
| | """ |
| |
|
| | import streamlit as st |
| | from app import device, load_demo_image |
| | from app.utils import load_model_cache |
| | from lavis.processors import load_processor |
| | from PIL import Image |
| |
|
| |
|
| | def app(): |
| | |
| | model_type = st.sidebar.selectbox("Model:", ["BLIP_base", "BLIP_large"]) |
| |
|
| | sampling_method = st.sidebar.selectbox( |
| | "Sampling method:", ["Beam search", "Nucleus sampling"] |
| | ) |
| |
|
| | st.markdown( |
| | "<h1 style='text-align: center;'>Image Description Generation</h1>", |
| | unsafe_allow_html=True, |
| | ) |
| |
|
| | instructions = """Try the provided image or upload your own:""" |
| | file = st.file_uploader(instructions) |
| |
|
| | use_beam = sampling_method == "Beam search" |
| |
|
| | col1, col2 = st.columns(2) |
| |
|
| | if file: |
| | raw_img = Image.open(file).convert("RGB") |
| | else: |
| | raw_img = load_demo_image() |
| |
|
| | col1.header("Image") |
| |
|
| | w, h = raw_img.size |
| | scaling_factor = 720 / w |
| | resized_image = raw_img.resize((int(w * scaling_factor), int(h * scaling_factor))) |
| |
|
| | col1.image(resized_image, use_column_width=True) |
| | col2.header("Description") |
| |
|
| | cap_button = st.button("Generate") |
| |
|
| | |
| | vis_processor = load_processor("blip_image_eval").build(image_size=384) |
| |
|
| | if cap_button: |
| | if model_type.startswith("BLIP"): |
| | blip_type = model_type.split("_")[1].lower() |
| | model = load_model_cache( |
| | "blip_caption", |
| | model_type=f"{blip_type}_coco", |
| | is_eval=True, |
| | device=device, |
| | ) |
| |
|
| | img = vis_processor(raw_img).unsqueeze(0).to(device) |
| | captions = generate_caption( |
| | model=model, image=img, use_nucleus_sampling=not use_beam |
| | ) |
| |
|
| | col2.write("\n\n".join(captions), use_column_width=True) |
| |
|
| |
|
| | def generate_caption( |
| | model, image, use_nucleus_sampling=False, num_beams=3, max_length=40, min_length=5 |
| | ): |
| | samples = {"image": image} |
| |
|
| | captions = [] |
| | if use_nucleus_sampling: |
| | for _ in range(5): |
| | caption = model.generate( |
| | samples, |
| | use_nucleus_sampling=True, |
| | max_length=max_length, |
| | min_length=min_length, |
| | top_p=0.9, |
| | ) |
| | captions.append(caption[0]) |
| | else: |
| | caption = model.generate( |
| | samples, |
| | use_nucleus_sampling=False, |
| | num_beams=num_beams, |
| | max_length=max_length, |
| | min_length=min_length, |
| | ) |
| | captions.append(caption[0]) |
| |
|
| | return captions |
| |
|
