Update search

app.py

@@ -2,34 +2,52 @@ import gradio as gr
 import random
 from datasets import load_dataset
 from sentence_transformers import SentenceTransformer, util
+import logging
+from PIL import Image
+# Create a custom logger
+logger = logging.getLogger(__name__)
 
-model = SentenceTransformer('clip-ViT-B-32')
-
-def fake_gan():
-    images = [
-        (random.choice(
-            [
-                "https://upload.wikimedia.org/wikipedia/commons/6/69/NASA-HS201427a-HubbleUltraDeepField2014-20140603.jpg",
-                "https://upload.wikimedia.org/wikipedia/commons/7/73/Cycliste_%C3%A0_place_d%27Italie-Paris.jpg",
-                "https://upload.wikimedia.org/wikipedia/commons/3/31/Great_white_shark_south_africa.jpg",
-            ]
-        ), f"label {i}" if i != 0 else "label" * 50)
-        for i in range(3)
-    ]
-    return images
-
-def search_images_from_text(text):
-    emb = model.encode(text)
-    return fake_gan()
-
-def search_images_from_image(image):
-    image_emb = model.encode(image)
-    return fake_gan()
+# Set the level of this logger. INFO means that it will log all INFO, WARNING, ERROR, and CRITICAL messages.
+logger.setLevel(logging.INFO)
+
+# Create handlers
+c_handler = logging.StreamHandler()
+c_handler.setLevel(logging.INFO)
+
+# Create formatters and add it to handlers
+c_format = logging.Formatter('%(name)s - %(levelname)s - %(message)s')
+c_handler.setFormatter(c_format)
+
+# Add handlers to the logger
+logger.addHandler(c_handler)
+
+class SearchEngine:
+    def __init__(self):
+        self.model = SentenceTransformer('clip-ViT-B-32')
+        self.embedding_dataset = load_dataset("JLD/unsplash25k-image-embeddings", trust_remote_code=True, split="train").with_format("torch", device="cuda:0")
+        image_dataset = load_dataset("jamescalam/unsplash-25k-photos", trust_remote_code=True, revision="refs/pr/3")
+        self.image_dataset = {image["photo_id"]: image["photo_image_url"] for image in image_dataset["train"]}
+
+    def get_candidates(self, query_embedding, top_k=5):
+        logger.info("Getting candidates")
+        candidates = util.semantic_search(query_embeddings=query_embedding.unsqueeze(0), corpus_embeddings=self.embedding_dataset["image_embedding"].squeeze(1), top_k=top_k)[0]
+        return [self.image_dataset.get(self.embedding_dataset[candidate["corpus_id"]]["image_id"], "https://upload.wikimedia.org/wikipedia/commons/6/69/NASA-HS201427a-HubbleUltraDeepField2014-20140603.jpg") for candidate in candidates]
+
+    def search_images_from_text(self, text):
+        logger.info("Searching images from text")
+        emb = self.model.encode(text, convert_to_tensor=True, device="cuda:0")
+        return self.get_candidates(query_embedding=emb)
+
+    def search_images_from_image(self, image):
+        logger.info("Searching images from image")
+        emb = self.model.encode(Image.fromarray(image), convert_to_tensor=True, device="cuda:0")
+        return self.get_candidates(query_embedding=emb)
 
 def main():
-    dataset = load_dataset("JLD/unsplash25k-image-embeddings", trust_remote_code=True, split="train").with_format("torch", device="cuda:0")
-    text_to_image_iface = gr.Interface(fn=search_images_from_text, inputs="text", outputs="gallery")
-    image_to_image_iface = gr.Interface(fn=search_images_from_image, inputs="image", outputs="gallery")
+    logger.info("Loading dataset")
+    search_engine = SearchEngine()
+    text_to_image_iface = gr.Interface(fn=search_engine.search_images_from_text, inputs="text", outputs="gallery")
+    image_to_image_iface = gr.Interface(fn=search_engine.search_images_from_image, inputs="image", outputs="gallery")
     demo = gr.TabbedInterface([text_to_image_iface, image_to_image_iface], ["Text query", "Image query"])
     demo.launch()