Version 0.0

app.py

@@ -0,0 +1,112 @@
+import streamlit as st
+from sentence_transformers import SentenceTransformer, util
+from pathlib import Path
+import pickle
+import requests
+from PIL import Image
+from io import BytesIO
+import pandas as pd
+from loguru import logger
+import torch
+
+T2I = "Text 2 Image"
+I2I = "Image 2 Image"
+def get_match(model, query, img_embs):
+    query_emb = model.encode([query], convert_to_tensor=True)
+    cosine_sim = util.pytorch_cos_sim(query_emb, img_embs)
+    return cosine_sim
+def text_2_image(model, img_emb, img_names, img_urls, n_top_k_images):
+    st.title("Text to Image")
+    st.write("This is the text to image mode. Enter a text to be converted to an image")
+    text = st.text_input("Enter the text to be converted to an image")
+    if text:
+        if st.button("Convert"):
+            st.write("The image with the most similar embedding is:")
+            cosine_sim = get_match(model, text, img_emb)
+            logger.info(cosine_sim.shape)
+            top_k_images_indices = torch.topk(cosine_sim, n_top_k_images, 1).indices
+            logger.info(top_k_images_indices.squeeze().tolist())
+            images_found = [img_names[top_k_best_image] for top_k_best_image in top_k_images_indices.squeeze().tolist()]
+            cols = st.columns(n_top_k_images)
+            for i, image_found in enumerate(images_found):
+                logger.success(f"Image match found: {image_found}")
+                img_url_best_match = img_urls.loc[img_urls["photo_id"] == image_found]
+                logger.info(img_url_best_match.photo_url)
+                if len(img_url_best_match) >= 1:
+                    response = requests.get(img_url_best_match.iloc[0]["photo_image_url"] + "?w=320")
+                    image = Image.open(BytesIO(response.content))
+                    with cols[i]:
+                        st.image(image, caption=f"{i+1}/{n_top_k_images} most similar")
+                else:
+                    st.error("No image found")
+
+
+def image_2_image(model, img_emb, img_names, img_urls,n_top_k_images):
+    st.title("Image to Image")
+    st.write("This is the image to image mode. Enter an image to be converted to an image")
+    image = st.file_uploader("Upload an image to be converted to an image", type=["jpg", "png", "jpeg"])
+    if image is not None:
+        image = Image.open(BytesIO(image.getvalue()))
+        st.image(image, caption="Uploaded image")
+        if st.button("Convert"):
+            st.write("The image with the most similar embedding is:")
+            cosine_sim = get_match(model, image.convert("RGB"), img_emb)
+            logger.info(cosine_sim.shape)
+            top_k_images_indices = torch.topk(cosine_sim, n_top_k_images, 1).indices
+            logger.info(top_k_images_indices.squeeze().tolist())
+            images_found = [img_names[top_k_best_image] for top_k_best_image in top_k_images_indices.squeeze().tolist()]
+            cols = st.columns(n_top_k_images)
+            for i, image_found in enumerate(images_found):
+                logger.success(f"Image match found: {image_found}")
+                img_url_best_match = img_urls.loc[img_urls["photo_id"] == image_found]
+                logger.info(img_url_best_match.photo_url)
+                if len(img_url_best_match) >= 1:
+                    response = requests.get(img_url_best_match.iloc[0]["photo_image_url"] + "?w=320")
+                    image = Image.open(BytesIO(response.content))
+                    with cols[i]:
+                        st.image(image, caption=f"{i+1}/{n_top_k_images} most similar")
+                else:
+                    st.error("No image found")
+
+@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+def load_model(name):
+    # st.sidebar.info("Loading model")
+    model = SentenceTransformer(name)
+    # st.sidebar.success(f"Model {name} loaded")
+    return model
+
+@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+def load_embeddings(filename):
+    st.sidebar.info("Loading Unsplash-Lite image embeddings")
+    with open(filename, "rb") as fIn:
+        img_names, img_emb = pickle.load(fIn)
+    st.sidebar.success("Images embeddings loaded")
+    return img_names, img_emb
+
+@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+def load_image_url_list(filename):
+    url_list = pd.read_csv(filename, sep='\t', header=0)
+    return url_list
+
+def main():
+    st.title("CLIP Image Search")
+    model = load_model("clip-ViT-B-32")
+    st.write("Select the mode to search for a match in Unsplash (thumbnail size) dataset. text2image mode needs a text as input and outputs the image with the most similar embedding (following cosine similarity). The Image to image mode is similar, but an input image is used instead of a text query")
+    emb_filename = Path("./unsplash-25k-photos-embeddings.pkl")
+    urls_file = "./photos.tsv000"
+    img_urls = load_image_url_list(urls_file)
+    img_names, img_emb = load_embeddings(emb_filename)
+    # Convert list of image names to a dict matching image IDs and their embedding index
+    img_names = {img_number: img_name.split('.')[0] for img_number, img_name in enumerate(img_names)}
+    st.sidebar.title("Settings")
+    app_mode = st.sidebar.selectbox("Choose the app mode",
+        [T2I, I2I])
+    n_images_to_search = st.sidebar.number_input("Select the number of images to search", min_value=1, max_value=6)
+    if app_mode == T2I:
+        st.sidebar.info("Text to image mode")
+        text_2_image(model, img_emb, img_names, img_urls,n_images_to_search)
+    elif app_mode == I2I:
+        st.sidebar.info("Image to image mode")
+        image_2_image(model, img_emb, img_names, img_urls, n_images_to_search)
+if __name__ == "__main__":
+    main()

unsplash-25k-photos-embeddings.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e60d49bd6334c29bc15fc2bc18c30b6d047a5584ad67c793eba376e95eaef8e
+size 51816207