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 = "gjuggler/swin-tiny-patch4-window7-224-finetuned-birds" 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(224), 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