initial commit

app.py

@@ -0,0 +1,86 @@
+import pickle
+import numpy as np
+import gradio as gr
+import clip
+import torch
+import annoy
+
+
+data = pickle.load(open("embeddings.pkl", "rb"))
+embeddings = data["embeddings"]
+image_features = torch.Tensor(embeddings)
+image_features /= image_features.norm(dim=-1, keepdim=True)
+
+
+n, d = embeddings.shape
+
+filenames = data["filenames"]
+thumbs = data["thumbs"]
+
+base_url = "https://static.renyi.hu/ai-shared/daniel/sameenergy/index/"
+urls = [base_url + filename for filename in filenames]
+
+
+model, preprocess = clip.load('RN50')
+
+
+def embed_text(text):
+    tokens = clip.tokenize([text])
+    with torch.no_grad():
+        text_features = model.encode_text(tokens)
+    assert text_features.shape == (1, d)
+    return text_features
+
+
+def similarities(text_features, topk=20):
+    text_features /= text_features.norm(dim=-1, keepdim=True)
+    # the softmax rounds up everything to 1, so does not distinguish between good fits.
+    similarity = (100.0 * image_features @ text_features.T) # .softmax(dim=-1)
+    values, indices = similarity[:, 0].topk(topk)
+    return values, indices
+
+
+def image_retrieval(text):
+    values, indices = similarities(embed_text(text), topk=20)
+    top_urls = np.array(urls)[indices]
+    return top_urls.tolist(), indices.numpy().tolist()
+
+
+def on_select(evt):
+    print("event:", evt)
+    return str(evt)
+    return f"You selected {evt.value} at {evt.index} from {evt.target}"
+
+
+def empty_gallery():
+    return [], []
+
+
+with gr.Blocks(css="footer {visibility: hidden}") as demo:
+    state = gr.State()
+
+    with gr.Row(variant="compact"):
+        text = gr.Textbox(
+            label="Enter your prompt",
+            show_label=False,
+            max_lines=1,
+            placeholder="Enter your prompt",
+        ).style(container=False)
+        btn = gr.Button("Search").style(full_width=False)
+
+
+    gallery = gr.Gallery(label="Images", show_label=False, elem_id="gallery"
+        ).style(columns=4, container=False)
+
+    demo.load(empty_gallery, None, [gallery, state])
+
+    selected = gr.Textbox(placeholder="Selected", show_label=False)
+
+    btn.click(image_retrieval, text, [gallery, state])
+
+    # does not work, function is called with None instead of event:
+    gallery.select(on_select, None, selected)
+
+
+if __name__ == "__main__":
+    demo.launch(height=2000)

create_embeddings.py

@@ -0,0 +1,57 @@
+import os
+import numpy as np
+import torch
+from PIL import Image
+import clip
+import pickle
+
+
+model, preprocess = clip.load('RN50')
+# model, preprocess = clip.load('ViT-L/14@336px')
+
+limit = 1e9
+batch_size = 100
+
+def do_batch(batch, embeddings):
+    image_batch = torch.tensor(np.stack(batch))
+    with torch.no_grad():
+        image_features = model.encode_image(image_batch).float()
+        embeddings += image_features.numpy().tolist()
+        print(f"{len(embeddings)} done")
+
+
+workdir = "./index"
+indx = os.listdir(workdir)
+embeddings = []
+filenames = []
+thumbs = []
+print("starting processing")
+batch = []
+for filename in indx:
+    if filename.lower().endswith("jpg"):
+        full_filename = os.path.join(workdir, filename)
+        rgb = Image.open(full_filename).convert("RGB")
+        img = preprocess(rgb)
+        rgb.thumbnail((128, 128))
+        thumb = np.array(rgb)
+        batch.append(img)
+        if len(batch) >= batch_size:
+            do_batch(batch, embeddings)
+            batch = []
+        filenames.append(filename)
+        thumbs.append(thumb)
+        if len(filenames) >= limit:
+            break
+
+# remaining
+if len(batch) > 0:
+    do_batch(batch, embeddings)
+
+embeddings = np.array(embeddings)
+assert len(embeddings) == len(filenames) == len(thumbs)
+print(f"processed {len(embeddings)} images")
+
+data = {"embeddings": embeddings, "filenames": filenames, "thumbs": thumbs}
+
+with open("embeddings.pkl", "wb") as f:
+    pickle.dump(data, f)

embeddings_nothumb.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59a6fce40441f2a5b61901f959dcee9836c5caa5813ef482e94c58a652a7c578
+size 2105705

visualize_embeddings.py

@@ -0,0 +1,21 @@
+import numpy as np
+from sklearn.manifold import TSNE
+import matplotlib.pyplot as plt
+import pickle
+
+data = pickle.load(open("embeddings.pkl", "rb"))
+embeddings = data["embeddings"]
+filenames = data["filenames"]
+thumbs = data["thumbs"]
+
+tsne = TSNE(n_components=2)
+reduced = tsne.fit_transform(embeddings)
+
+fig, ax = plt.subplots()
+# ax.scatter(reduced[:, 0], reduced[:, 1])
+delta = 0.5
+for i, txt in enumerate(filenames):
+    # ax.annotate(txt, (reduced[i, 0], reduced[i, 1]))
+    x, y = reduced[i]
+    ax.imshow(thumbs[i], extent=[x-delta, x+delta, y-delta, y+delta])
+plt.show()