Added image retrieval

.gitattributes

@@ -35,3 +35,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 files/brand_bank.index filter=lfs diff=lfs merge=lfs -text
 files/caption_bank.index filter=lfs diff=lfs merge=lfs -text
 files/finetuned.pth filter=lfs diff=lfs merge=lfs -text
+files/combined.tar filter=lfs diff=lfs merge=lfs -text
+files/index/image.index filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,20 +1,20 @@
+import os
+
 import clip
 import faiss
 import torch
-import numpy as np
+import tarfile
 import gradio as gr
 import pandas as pd
+from PIL import Image
+from braceexpand import braceexpand
 
 
 # Load model
-#checkpoint_path = "files/finetuned_from_dopamine.pth"
 checkpoint_path = "ViT-B/16"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model, preprocess = clip.load(checkpoint_path, device=device, jit=False)
 
-bb_one = None
-bb_two = None
-
 
 def generate_caption(img):
     # Load caption bank
@@ -70,36 +70,56 @@ def estimate_price_and_usage(img):
     return "Estimated price: 50-100 SEK - Usage: Reuse - Saved C02: 4 kg"
 
 
-def select_handler(img, evt: gr.SelectData):
-    global bb_one, bb_two
-    line_width = 20
-    mask = np.zeros(img.shape[:2], dtype=np.uint8)
-
-    # Reset if creating a new bbox
-    if bb_one is not None and bb_two is not None:
-        bb_one = None
-        bb_two = None
+def retrieve(query):
+    index_folder = "files/index"
+    num_results = 3
 
-    if bb_one is not None:
-        bb_two = evt.index
+    # Read image metadata
+    metadata_df = pd.read_parquet(os.path.join(index_folder, "metadata.parquet"))
+    key_list = metadata_df["key"].tolist()
 
-        # Make sure the bbox is in the right order
-        if bb_one[0] > bb_two[0]:
-            bb_one[0], bb_two[0] = bb_two[0], bb_one[0]
-        if bb_one[1] > bb_two[1]:
-            bb_one[1], bb_two[1] = bb_two[1], bb_one[1]
+    # Load the index
+    index = faiss.read_index(os.path.join(index_folder, "image.index"))
 
-        # Fill in a square, then hollow it out to get a bbox
-        mask[bb_one[1]:bb_two[1], bb_one[0]:bb_two[0]] = 1
-        mask[bb_one[1]+line_width:bb_two[1]-line_width,
-             bb_one[0]+line_width:bb_two[0]-line_width] = 0
-        return (img, [(mask, "bbox")])
+    # Encode the query
+    if isinstance(query, str):
+        print("Query is a string")
+        text = clip.tokenize([query]).to(device)
+        query_features = model.encode_text(text)
     else:
-        bb_one = evt.index
-        # Make a small dot
-        mask[bb_one[1]-line_width:bb_one[1]+line_width,
-             bb_one[0]-line_width:bb_one[0]+line_width] = 1
-        return (img, [(mask, "bbox")])
+        print("Query is an image")
+        query_features = model.encode_image(preprocess(query).unsqueeze(0).to(device))
+    query_features = query_features / query_features.norm(dim=-1, keepdim=True)
+    query_features = query_features.cpu().detach().numpy().astype("float32")
+
+    d, i = index.search(query_features, num_results)
+    print(f"Found {num_results} items with query '{query}'")
+    indices = i[0]
+    similarities = d[0]
+
+    min_d = min(similarities)
+    max_d = max(similarities)
+    print(f"The minimum similarity is {min_d:.2f} and the maximum is {max_d:.2f}")
+
+    # Uncomment to generate combined.tar, combine the image_tars into a single tarfile
+    """
+    dataset_dir = "/fs/sefs1/circularfashion/sellpy/front_balanced"
+    image_tars = [os.path.join(dataset_dir, file) for file in sorted(braceexpand("{00000..00010}.tar"))]
+    with tarfile.open("files/combined.tar", "w") as combined_tar:
+        for tar in image_tars:
+            with tarfile.open(tar, "r") as tar_file:
+                for member in tar_file.getmembers():
+                    combined_tar.addfile(member, tar_file.extractfile(member))
+    """
+
+    images = []
+    for idx in indices:
+        image_name = key_list[idx]
+        with tarfile.open("files/combined.tar", "r") as tar_file:
+            image = tar_file.extractfile(f"{image_name}.jpg")
+            image = Image.open(image).copy()
+            images.append(image)
+    return images
 
 
 theme = gr.Theme.from_hub("JohnSmith9982/small_and_pretty")
@@ -108,25 +128,39 @@ with gr.Blocks(
     theme=theme,
     css="footer {visibility: hidden}",
 ) as demo:
-    with gr.Row(variant="compact"):
-        input_img = gr.Image(type="pil", show_label=False)
-        with gr.Column(min_width="80"):
-            btn_generate_caption = gr.Button("Create Description").style(size="sm")
-            generated_caption = gr.Textbox(label="Description", show_label=False)
-    with gr.Row(variant="compact"):
-        brand_img = gr.Image(type="pil", show_label=False)
-        with gr.Column(min_width="80"):
-            btn_predict_brand = gr.Button("Predict Brand").style(size="sm")
-            predicted_brand = gr.Textbox(label="Brand", show_label=False)
-
-    with gr.Column(variant="compact"):
-        btn_estimate = gr.Button("Estimate Price, Reuse, and Saved C02").style(size="sm")
-        text_box = gr.Textbox(label="Estimates:", show_label=False)
+    with gr.Tab("Captioning and Prediction"):
+        with gr.Row(variant="compact"):
+            input_img = gr.Image(type="pil", show_label=False)
+            with gr.Column(min_width="80"):
+                btn_generate_caption = gr.Button("Create Description").style(size="sm")
+                generated_caption = gr.Textbox(label="Description", show_label=False)
+        with gr.Row(variant="compact"):
+            brand_img = gr.Image(type="pil", show_label=False)
+            with gr.Column(min_width="80"):
+                btn_predict_brand = gr.Button("Predict Brand").style(size="sm")
+                predicted_brand = gr.Textbox(label="Brand", show_label=False)
+
+        with gr.Column(variant="compact"):
+            btn_estimate = gr.Button("Estimate Price, Reuse, and Saved C02").style(size="sm")
+            text_box = gr.Textbox(label="Estimates:", show_label=False)
+    with gr.Tab("Image Retrieval"):
+        with gr.Row(variant="compact"):
+            with gr.Column():
+                query_img = gr.Image(type="pil", label="Image Query")
+                btn_image_query = gr.Button("Retrieve Garments").style(size="sm")
+                img_query_gallery = gr.Gallery(show_label=False).style(rows=1, columns=3)
+        with gr.Row(variant="compact"):
+            with gr.Column():
+                query_text = gr.Textbox(label="Text Query", placeholder="Enter a description")
+                btn_text_query = gr.Button("Retrieve Garments").style(size="sm")
+                text_query_gallery = gr.Gallery(show_label=False).style(rows=1, columns=3)
 
     # Listeners
     btn_generate_caption.click(fn=generate_caption, inputs=input_img, outputs=generated_caption)
     btn_predict_brand.click(fn=predict_brand, inputs=brand_img, outputs=predicted_brand)
     btn_estimate.click(fn=estimate_price_and_usage, inputs=input_img, outputs=text_box)
+    btn_image_query.click(fn=retrieve, inputs=query_img, outputs=img_query_gallery)
+    btn_text_query.click(fn=retrieve, inputs=query_text, outputs=text_query_gallery)
 
 
 if __name__ == "__main__":

files/combined.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:12201a61246802fa690e35ba79ffaa310adf9a1a99ca28a227779e34ea4a7f5f
+size 723496960

files/index/image.index

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

files/index/metadata.parquet

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