Duplicate from sayakpaul/fetch-similar-images

Co-authored-by: Sayak Paul <[email protected]>

.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,14 @@
+---
+title: Fetch Similar Beans 🪴
+emoji: ⚡
+colorFrom: yellow
+colorTo: blue
+sdk: gradio
+sdk_version: 3.12.0
+app_file: app.py
+pinned: false
+license: apache-2.0
+duplicated_from: sayakpaul/fetch-similar-images
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,73 @@
+"""
+Thanks to Freddy Boulton (https://github.com/freddyaboulton) for helping with this.
+"""
+
+
+import pickle
+
+import gradio as gr
+from datasets import load_dataset
+from transformers import AutoModel
+
+# `LSH` and `Table` imports are necessary in order for the
+# `lsh.pickle` file to load successfully.
+from similarity_utils import LSH, BuildLSHTable, Table
+
+seed = 42
+
+# Only runs once when the script is first run.
+with open("lsh.pickle", "rb") as handle:
+    loaded_lsh = pickle.load(handle)
+
+# Load model for computing embeddings.
+model_ckpt = "nateraw/vit-base-beans"
+model = AutoModel.from_pretrained(model_ckpt)
+lsh_builder = BuildLSHTable(model)
+lsh_builder.lsh = loaded_lsh
+
+# Candidate images.
+dataset = load_dataset("beans")
+candidate_dataset = dataset["train"].shuffle(seed=seed)
+
+
+def query(image, top_k):
+    results = lsh_builder.query(image)
+
+    # Should be a list of string file paths for gr.Gallery to work
+    images = []
+    # List of labels for each image in the gallery
+    labels = []
+
+    candidates = []
+
+    for idx, r in enumerate(sorted(results, key=results.get, reverse=True)):
+        if idx == top_k:
+            break
+        image_id, label = r.split("_")[0], r.split("_")[1]
+        candidates.append(candidate_dataset[int(image_id)]["image"])
+        labels.append(f"Label: {label}")
+
+    for i, candidate in enumerate(candidates):
+        filename = f"similar_{i}.png"
+        candidate.save(filename)
+        images.append(filename)
+
+    # The gallery component can be a list of tuples, where the first element is a path to a file
+    # and the second element is an optional caption for that image
+    return list(zip(images, labels))
+
+
+title = "Fetch Similar Beans 🪴"
+description = "This Space demos an image similarity system. You can refer to [this notebook](TODO) to know the details of the system. You can pick any image from the available samples below. On the right hand side, you'll find the similar images returned by the system. The example images have been named with their corresponding integer class labels for easier identification. The fetched images will also have their integer labels tagged so that you can validate the correctness of the results."
+
+# You can set the type of gr.Image to be PIL, numpy or str (filepath)
+# Not sure what the best for this demo is.
+gr.Interface(
+    query,
+    inputs=[gr.Image(type="pil"), gr.Slider(value=5, minimum=1, maximum=10, step=1)],
+    outputs=gr.Gallery().style(grid=[3], height="auto"),
+    # Filenames denote the integer labels. Know here: https://hf.co/datasets/beans
+    title=title,
+    description=description,
+    examples=[["0.png", 5], ["1.png", 5], ["2.png", 5]],
+).launch()

lsh.pickle

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

requirements.txt

@@ -0,0 +1,5 @@
+transformers==4.25.1
+datasets==2.7.1
+numpy==1.21.6
+torch==1.12.1
+torchvision

similarity_utils.py

@@ -0,0 +1,175 @@
+from typing import List, Union
+
+import datasets
+import numpy as np
+import torch
+import torchvision.transforms as T
+from PIL import Image
+from tqdm.auto import tqdm
+from transformers import AutoFeatureExtractor, AutoModel
+
+seed = 42
+hash_size = 8
+hidden_dim = 768  # ViT-base
+np.random.seed(seed)
+
+
+# Device.
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load model for computing embeddings..
+model_ckpt = "nateraw/vit-base-beans"
+extractor = AutoFeatureExtractor.from_pretrained(model_ckpt)
+
+# Data transformation chain.
+transformation_chain = T.Compose(
+    [
+        # We first resize the input image to 256x256 and then we take center crop.
+        T.Resize(int((256 / 224) * extractor.size["height"])),
+        T.CenterCrop(extractor.size["height"]),
+        T.ToTensor(),
+        T.Normalize(mean=extractor.image_mean, std=extractor.image_std),
+    ]
+)
+
+
+# Define random vectors to project with.
+random_vectors = np.random.randn(hash_size, hidden_dim).T
+
+
+def hash_func(embedding, random_vectors=random_vectors):
+    """Randomly projects the embeddings and then computes bit-wise hashes."""
+    if not isinstance(embedding, np.ndarray):
+        embedding = np.array(embedding)
+    if len(embedding.shape) < 2:
+        embedding = np.expand_dims(embedding, 0)
+
+    # Random projection.
+    bools = np.dot(embedding, random_vectors) > 0
+    return [bool2int(bool_vec) for bool_vec in bools]
+
+
+def bool2int(x):
+    y = 0
+    for i, j in enumerate(x):
+        if j:
+            y += 1 << i
+    return y
+
+
+def compute_hash(model: Union[torch.nn.Module, str]):
+    """Computes hash on a given dataset."""
+    device = model.device
+
+    def pp(example_batch):
+        # Prepare the input images for the model.
+        image_batch = example_batch["image"]
+        image_batch_transformed = torch.stack(
+            [transformation_chain(image) for image in image_batch]
+        )
+        new_batch = {"pixel_values": image_batch_transformed.to(device)}
+
+        # Compute embeddings and pool them i.e., take the representations from the [CLS]
+        # token.
+        with torch.no_grad():
+            embeddings = model(**new_batch).last_hidden_state[:, 0].cpu().numpy()
+
+        # Compute hashes for the batch of images.
+        hashes = [hash_func(embeddings[i]) for i in range(len(embeddings))]
+        example_batch["hashes"] = hashes
+        return example_batch
+
+    return pp
+
+
+class Table:
+    def __init__(self, hash_size: int):
+        self.table = {}
+        self.hash_size = hash_size
+
+    def add(self, id: int, hashes: List[int], label: int):
+        # Create a unique indentifier.
+        entry = {"id_label": str(id) + "_" + str(label)}
+
+        # Add the hash values to the current table.
+        for h in hashes:
+            if h in self.table:
+                self.table[h].append(entry)
+            else:
+                self.table[h] = [entry]
+
+    def query(self, hashes: List[int]):
+        results = []
+
+        # Loop over the query hashes and determine if they exist in
+        # the current table.
+        for h in hashes:
+            if h in self.table:
+                results.extend(self.table[h])
+        return results
+
+
+class LSH:
+    def __init__(self, hash_size, num_tables):
+        self.num_tables = num_tables
+        self.tables = []
+        for i in range(self.num_tables):
+            self.tables.append(Table(hash_size))
+
+    def add(self, id: int, hash: List[int], label: int):
+        for table in self.tables:
+            table.add(id, hash, label)
+
+    def query(self, hashes: List[int]):
+        results = []
+        for table in self.tables:
+            results.extend(table.query(hashes))
+        return results
+
+
+class BuildLSHTable:
+    def __init__(
+        self,
+        model: Union[torch.nn.Module, None],
+        batch_size: int = 48,
+        hash_size: int = hash_size,
+        dim: int = hidden_dim,
+        num_tables: int = 10,
+    ):
+        self.hash_size = hash_size
+        self.dim = dim
+        self.num_tables = num_tables
+        self.lsh = LSH(self.hash_size, self.num_tables)
+
+        self.batch_size = batch_size
+        self.hash_fn = compute_hash(model.to(device))
+
+    def build(self, ds: datasets.DatasetDict):
+        dataset_hashed = ds.map(self.hash_fn, batched=True, batch_size=self.batch_size)
+
+        for id in tqdm(range(len(dataset_hashed))):
+            hash, label = dataset_hashed[id]["hashes"], dataset_hashed[id]["labels"]
+            self.lsh.add(id, hash, label)
+
+    def query(self, image, verbose=True):
+        if isinstance(image, str):
+            image = Image.open(image).convert("RGB")
+
+        # Compute the hashes of the query image and fetch the results.
+        example_batch = dict(image=[image])
+        hashes = self.hash_fn(example_batch)["hashes"][0]
+
+        results = self.lsh.query(hashes)
+        if verbose:
+            print("Matches:", len(results))
+
+        # Calculate Jaccard index to quantify the similarity.
+        counts = {}
+        for r in results:
+            if r["id_label"] in counts:
+                counts[r["id_label"]] += 1
+            else:
+                counts[r["id_label"]] = 1
+        for k in counts:
+            counts[k] = float(counts[k]) / self.dim
+        return counts