|
import os |
|
import urllib.request |
|
from collections import OrderedDict |
|
from html import escape |
|
|
|
import pandas as pd |
|
import numpy as np |
|
|
|
import torch |
|
|
|
from transformers import CLIPProcessor, CLIPModel |
|
|
|
import streamlit as st |
|
|
|
import models |
|
from tokenizer import SimpleTokenizer |
|
|
|
cuda_available = torch.cuda.is_available() |
|
|
|
model_url = "https://dl.fbaipublicfiles.com/slip/slip_large_100ep.pt" |
|
model_filename = "slip_large_100ep.pt" |
|
|
|
|
|
def get_model(model): |
|
if isinstance(model, torch.nn.DataParallel) or isinstance( |
|
model, torch.nn.parallel.DistributedDataParallel |
|
): |
|
return model.module |
|
else: |
|
return model |
|
|
|
|
|
@st.cache( |
|
show_spinner=False, |
|
hash_funcs={ |
|
CLIPModel: lambda _: None, |
|
CLIPProcessor: lambda _: None, |
|
dict: lambda _: None, |
|
}, |
|
) |
|
def load(): |
|
|
|
if model_filename not in os.listdir(): |
|
urllib.request.urlretrieve(model_url, model_filename) |
|
ckpt = torch.load("slip_large_100ep.pt", map_location="cpu") |
|
state_dict = OrderedDict() |
|
for k, v in ckpt["state_dict"].items(): |
|
state_dict[k.replace("module.", "")] = v |
|
old_args = ckpt["args"] |
|
slip_model = getattr(models, "SLIP_VITL16")( |
|
rand_embed=False, |
|
ssl_mlp_dim=old_args.ssl_mlp_dim, |
|
ssl_emb_dim=old_args.ssl_emb_dim, |
|
) |
|
if cuda_available: |
|
slip_model.cuda() |
|
slip_model.load_state_dict(state_dict, strict=True) |
|
slip_model = get_model(slip_model) |
|
tokenizer = SimpleTokenizer() |
|
del ckpt |
|
del state_dict |
|
|
|
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") |
|
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") |
|
|
|
df = {0: pd.read_csv("data.csv"), 1: pd.read_csv("data2.csv")} |
|
embeddings = {0: np.load("embeddings.npy"), 1: np.load("embeddings2.npy")} |
|
slip_embeddings = { |
|
0: np.load("embeddings_slip_large.npy"), |
|
1: np.load("embeddings2_slip_large.npy"), |
|
} |
|
for k in [0, 1]: |
|
embeddings[k] = np.divide( |
|
embeddings[k], np.sqrt(np.sum(embeddings[k] ** 2, axis=1, keepdims=True)) |
|
) |
|
return model, processor, slip_model, tokenizer, df, embeddings, slip_embeddings |
|
|
|
|
|
model, processor, slip_model, tokenizer, df, embeddings, slip_embeddings = load() |
|
|
|
source = {0: "\nSource: Unsplash", 1: "\nSource: The Movie Database (TMDB)"} |
|
|
|
|
|
def get_html(url_list, url_list_slip, height=150): |
|
html = ( |
|
"<div style='display: flex; flex-wrap: wrap; justify-content: space-evenly;'>" |
|
) |
|
html += "<span style='margin-top: 20px; max-width: 1200px; display: flex; align-content: flex-start; flex-wrap: wrap; justify-content: space-evenly; width: 50%'>" |
|
html += "<div style='width: 100%; text-align: center;'><b>CLIP</b> (<a href='https://arxiv.org/abs/2103.00020'>Arxiv</a>, <a href='https://github.com/openai/CLIP'>GitHub</a>) from OpenAI</div>" |
|
for url, title, link in url_list: |
|
html2 = f"<img title='{escape(title)}' style='height: {height}px; margin: 5px' src='{escape(url)}'>" |
|
if len(link) > 0: |
|
html2 = f"<a href='{escape(link)}' target='_blank'>" + html2 + "</a>" |
|
html = html + html2 |
|
html += "</span>" |
|
html += "<span style='margin-top: 20px; max-width: 1200px; display: flex; align-content: flex-start; flex-wrap: wrap; justify-content: space-evenly; width: 50%; border-left: solid; border-color: #ffc423; border-width: thin;'>" |
|
html += "<div style='width: 100%; text-align: center;'><b>SLIP</b> (<a href='https://arxiv.org/abs/2112.12750'>Arxiv</a>, <a href='https://github.com/facebookresearch/SLIP'>GitHub</a>) from Meta AI</div>" |
|
for url, title, link in url_list_slip: |
|
html2 = f"<img title='{escape(title)}' style='height: {height}px; margin: 5px' src='{escape(url)}'>" |
|
if len(link) > 0: |
|
html2 = f"<a href='{escape(link)}' target='_blank'>" + html2 + "</a>" |
|
html = html + html2 |
|
html += "</span></div>" |
|
return html |
|
|
|
|
|
def compute_text_embeddings(list_of_strings): |
|
inputs = processor(text=list_of_strings, return_tensors="pt", padding=True) |
|
return model.get_text_features(**inputs) |
|
|
|
|
|
def compute_text_embeddings_slip(list_of_strings): |
|
texts = tokenizer(list_of_strings) |
|
if cuda_available: |
|
texts = texts.cuda(non_blocking=True) |
|
texts = texts.view(-1, 77).contiguous() |
|
return slip_model.encode_text(texts) |
|
|
|
|
|
def image_search(query, corpus, n_results=24): |
|
text_embeddings = compute_text_embeddings([query]).detach().numpy() |
|
text_embeddings_slip = compute_text_embeddings_slip([query]).detach().numpy() |
|
k = 0 if corpus == "Unsplash" else 1 |
|
results = np.argsort((embeddings[k] @ text_embeddings.T)[:, 0])[ |
|
-1 : -n_results - 1 : -1 |
|
] |
|
results_slip = np.argsort((slip_embeddings[k] @ text_embeddings_slip.T)[:, 0])[ |
|
-1 : -n_results - 1 : -1 |
|
] |
|
return ( |
|
[ |
|
( |
|
df[k].iloc[i]["path"], |
|
df[k].iloc[i]["tooltip"] + source[k], |
|
df[k].iloc[i]["link"], |
|
) |
|
for i in results |
|
], |
|
[ |
|
( |
|
df[k].iloc[i]["path"], |
|
df[k].iloc[i]["tooltip"] + source[k], |
|
df[k].iloc[i]["link"], |
|
) |
|
for i in results_slip |
|
], |
|
) |
|
|
|
|
|
description = """ |
|
# Comparing CLIP and SLIP side by side |
|
|
|
**Enter your query and hit enter** |
|
|
|
CLIP and SLIP are ML models that encode images and texts as vectors so that the vectors of an image and its caption are similar. They can notably be used for zero-shot image classification, text-based image retrieval or image generation. |
|
|
|
Cf. this Twitter [thread](https://twitter.com/vivien000000/status/1475829936443334660) with some suprising differences between CLIP and SLIP. |
|
|
|
*Built with OpenAI's [CLIP](https://openai.com/blog/clip/) model, Meta AI's [SLIP](https://github.com/facebookresearch/SLIP) model, 🤗 Hugging Face's [transformers library](https://huggingface.co/transformers/), [Streamlit](https://streamlit.io/), 25k images from [Unsplash](https://unsplash.com/) and 8k images from [The Movie Database (TMDB)](https://www.themoviedb.org/)* |
|
""" |
|
|
|
|
|
st.markdown( |
|
""" |
|
<style> |
|
.block-container{ |
|
max-width: 1200px; |
|
} |
|
div.row-widget.stRadio > div{ |
|
flex-direction:row; |
|
display: flex; |
|
justify-content: center; |
|
} |
|
div.row-widget.stRadio > div > label{ |
|
margin-left: 5px; |
|
margin-right: 5px; |
|
} |
|
section.main>div:first-child { |
|
padding-top: 0px; |
|
} |
|
section:not(.main)>div:first-child { |
|
padding-top: 30px; |
|
} |
|
div.reportview-container > section:first-child{ |
|
max-width: 320px; |
|
} |
|
#MainMenu { |
|
visibility: hidden; |
|
} |
|
footer { |
|
visibility: hidden; |
|
} |
|
</style>""", |
|
unsafe_allow_html=True, |
|
) |
|
st.sidebar.markdown(description) |
|
_, c, _ = st.columns((1, 3, 1)) |
|
query = c.text_input("", value="clouds at sunset") |
|
corpus = st.radio("", ["Unsplash", "Movies"]) |
|
if len(query) > 0: |
|
results, results_slip = image_search(query, corpus) |
|
st.markdown(get_html(results, results_slip), unsafe_allow_html=True) |
|
|
