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():
# Load SLIP model from Facebook AI Research
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
# Load CLIP model from HuggingFace
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
# Load images' descriptions and embeddings
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 = (
"
"
)
html += "
"
html += ""
for url, title, link in url_list:
html2 = f"
"
if len(link) > 0:
html2 = f"" + html2 + ""
html = html + html2
html += ""
html += "
"
html += ""
for url, title, link in url_list_slip:
html2 = f"
"
if len(link) > 0:
html2 = f"" + html2 + ""
html = html + html2
html += " "
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(
"""
""",
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)