Spaces:
Runtime error
Runtime error
| import matplotlib.pyplot as plt | |
| import nmslib | |
| import numpy as np | |
| import os | |
| import streamlit as st | |
| from PIL import Image | |
| from transformers import CLIPProcessor, FlaxCLIPModel | |
| import utils | |
| BASELINE_MODEL = "openai/clip-vit-base-patch32" | |
| MODEL_PATH = "flax-community/clip-rsicd-v2" | |
| IMAGE_VECTOR_FILE = "./vectors/test-bs128x8-lr5e-6-adam-ckpt-1.tsv" | |
| IMAGES_DIR = "./images" | |
| CAPTIONS_FILE = os.path.join(IMAGES_DIR, "test-captions.json") | |
| def app(): | |
| filenames, index = utils.load_index(IMAGE_VECTOR_FILE) | |
| model, processor = utils.load_model(MODEL_PATH, BASELINE_MODEL) | |
| image2caption = utils.load_captions(CAPTIONS_FILE) | |
| st.title("Retrieve Images given Text") | |
| st.markdown(""" | |
| This demo shows the image to text retrieval capabilities of this model, i.e., | |
| given a text query, we use our fine-tuned CLIP model to project the text query | |
| to the image/caption embedding space and search for nearby images (by | |
| cosine similarity) in this space. | |
| Our fine-tuned CLIP model was previously used to generate image vectors for | |
| our demo, and NMSLib was used for fast vector access. | |
| """) | |
| suggested_query = [ | |
| "ships", | |
| "school house", | |
| "military installation", | |
| "mountains", | |
| "beaches", | |
| "airports", | |
| "lakes" | |
| ] | |
| st.text("Some suggested queries to start you off with...") | |
| col0, col1, col2, col3, col4, col5, col6 = st.columns(7) | |
| # [1, 1.1, 1.3, 1.1, 1, 1, 1]) | |
| suggest_idx = -1 | |
| with col0: | |
| if st.button(suggested_query[0]): | |
| suggest_idx = 0 | |
| with col1: | |
| if st.button(suggested_query[1]): | |
| suggest_idx = 1 | |
| with col2: | |
| if st.button(suggested_query[2]): | |
| suggest_idx = 2 | |
| with col3: | |
| if st.button(suggested_query[3]): | |
| suggest_idx = 3 | |
| with col4: | |
| if st.button(suggested_query[4]): | |
| suggest_idx = 4 | |
| with col5: | |
| if st.button(suggested_query[5]): | |
| suggest_idx = 5 | |
| with col6: | |
| if st.button(suggested_query[6]): | |
| suggest_idx = 6 | |
| query = st.text_input("OR enter a text Query:") | |
| query = suggested_query[suggest_idx] if suggest_idx > -1 else query | |
| if st.button("Query") or suggest_idx > -1: | |
| inputs = processor(text=[query], images=None, return_tensors="jax", padding=True) | |
| query_vec = model.get_text_features(**inputs) | |
| query_vec = np.asarray(query_vec) | |
| ids, distances = index.knnQuery(query_vec, k=10) | |
| result_filenames = [filenames[id] for id in ids] | |
| for rank, (result_filename, score) in enumerate(zip(result_filenames, distances)): | |
| caption = "{:s} (score: {:.3f})".format(result_filename, 1.0 - score) | |
| col1, col2, col3 = st.columns([2, 10, 10]) | |
| col1.markdown("{:d}.".format(rank + 1)) | |
| col2.image(Image.open(os.path.join(IMAGES_DIR, result_filename)), | |
| caption=caption) | |
| caption_text = [] | |
| for caption in image2caption[result_filename]: | |
| caption_text.append("* {:s}\n".format(caption)) | |
| col3.markdown("".join(caption_text)) | |
| st.markdown("---") | |
| suggest_idx = -1 | |