initial commit

ExtractEmbedding.py

@@ -0,0 +1,59 @@
+import time
+import os
+import torch
+import numpy as np
+import torchvision
+import torch.nn.functional as F
+from torchvision.datasets import ImageFolder
+import torchvision.transforms as transforms
+from tqdm import tqdm
+import pickle
+import argparse
+from PIL import Image
+
+concat = lambda x: np.concatenate(x, axis=0)
+to_np = lambda x: x.data.to("cpu").numpy()
+
+
+class Wrapper(torch.nn.Module):
+    def __init__(self, model):
+        super(Wrapper, self).__init__()
+        self.model = model
+        self.avgpool_output = None
+        self.query = None
+        self.cossim_value = {}
+
+        def fw_hook(module, input, output):
+            self.avgpool_output = output.squeeze()
+
+        self.model.avgpool.register_forward_hook(fw_hook)
+
+    def forward(self, input):
+        _ = self.model(input)
+        return self.avgpool_output
+
+    def __repr__(self):
+        return "Wrappper"
+
+
+def QueryToEmbedding(query_pil):
+    dataset_transform = transforms.Compose(
+        [
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
+
+    model = torchvision.models.resnet50(pretrained=True)
+    model.eval()
+    myw = Wrapper(model)
+
+    # query_pil = Image.open(query_path)
+    query_pt = dataset_transform(query_pil)
+
+    with torch.no_grad():
+        embedding = to_np(myw(query_pt.unsqueeze(0)))
+
+    return np.asarray([embedding])

README.md

@@ -1,8 +1,8 @@
 ---
-title: CHM Corr
+title: CHM-Corr
 emoji: 🐨
 colorFrom: yellow
-colorTo: red
+colorTo: blue
 sdk: gradio
 sdk_version: 3.1.1
 app_file: app.py

SaveEmbedding.py

@@ -0,0 +1,100 @@
+import time
+import os
+import torch
+import numpy as np
+import torchvision
+import torch.nn.functional as F
+from torchvision.datasets import ImageFolder
+import torchvision.transforms as transforms
+from tqdm import tqdm
+import pickle
+import argparse
+
+
+concat = lambda x: np.concatenate(x, axis=0)
+to_np = lambda x: x.data.to("cpu").numpy()
+
+
+class Wrapper(torch.nn.Module):
+    def __init__(self, model):
+        super(Wrapper, self).__init__()
+        self.model = model
+        self.avgpool_output = None
+        self.query = None
+        self.cossim_value = {}
+
+        def fw_hook(module, input, output):
+            self.avgpool_output = output.squeeze()
+
+        self.model.avgpool.register_forward_hook(fw_hook)
+
+    def forward(self, input):
+        _ = self.model(input)
+        return self.avgpool_output
+
+    def __repr__(self):
+        return "Wrappper"
+
+
+def run(training_set_path):
+    # Standard ImageNet Transformer to apply imagenet's statistics to input batch
+    dataset_transform = transforms.Compose(
+        [
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
+
+    training_imagefolder = ImageFolder(
+        root=training_set_path, transform=dataset_transform
+    )
+    train_loader = torch.utils.data.DataLoader(
+        training_imagefolder,
+        batch_size=512,
+        shuffle=False,
+        num_workers=2,
+        pin_memory=True,
+    )
+    print(f"# of Training folder samples: {len(training_imagefolder)}")
+    ########################################################################################################################
+    model = torchvision.models.resnet50(pretrained=True)
+    model.eval()
+    myw = Wrapper(model)
+
+    training_embeddings = []
+    training_labels = []
+
+    with torch.no_grad():
+        for _, (data, target) in enumerate(tqdm(train_loader)):
+            embeddings = to_np(myw(data))
+            labels = to_np(target)
+
+            training_embeddings.append(embeddings)
+            training_labels.append(labels)
+
+    training_embeddings_concatted = concat(training_embeddings)
+    training_labels_concatted = concat(training_labels)
+    
+    print(training_embeddings_concatted.shape)
+    return training_embeddings_concatted, training_labels_concatted
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Saving Embeddings")
+    parser.add_argument("--train", help="Path to the Dataaset", type=str, required=True)
+    args = parser.parse_args()
+
+    embeddings, labels = run(args.train)
+
+    # Caluclate Accuracy
+    with open(f"embeddings.pickle", "wb") as f:
+        pickle.dump(embeddings, f)
+
+    with open(f"labels.pickle", "wb") as f:
+        pickle.dump(labels, f)
+
+
+if __name__ == "__main__":
+    main()

SimSearch.py

@@ -0,0 +1,66 @@
+import faiss
+import numpy as np
+
+
+class FaissNeighbors:
+    def __init__(self):
+        self.index = None
+        self.y = None
+
+    def fit(self, X, y):
+        self.index = faiss.IndexFlatL2(X.shape[1])
+        self.index.add(X.astype(np.float32))
+        self.y = y
+
+    def get_distances_and_indices(self, X, top_K=1000):
+        distances, indices = self.index.search(X.astype(np.float32), k=top_K)
+        return np.copy(distances), np.copy(indices), np.copy(self.y[indices])
+
+    def get_nearest_labels(self, X, top_K=1000):
+        distances, indices = self.index.search(X.astype(np.float32), k=top_K)
+        return np.copy(self.y[indices])
+
+
+class FaissCosineNeighbors:
+    def __init__(self):
+        self.cindex = None
+        self.y = None
+
+    def fit(self, X, y):
+        self.cindex = faiss.index_factory(
+            X.shape[1], "Flat", faiss.METRIC_INNER_PRODUCT
+        )
+        X = np.copy(X)
+        X = X.astype(np.float32)
+        faiss.normalize_L2(X)
+        self.cindex.add(X)
+        self.y = y
+
+    def get_distances_and_indices(self, Q, topK):
+        Q = np.copy(Q)
+        faiss.normalize_L2(Q)
+        distances, indices = self.cindex.search(Q.astype(np.float32), k=topK)
+        return np.copy(distances), np.copy(indices), np.copy(self.y[indices])
+
+    def get_nearest_labels(self, Q, topK=1000):
+        Q = np.copy(Q)
+        faiss.normalize_L2(Q)
+        distances, indices = self.cindex.search(Q.astype(np.float32), k=topK)
+        return np.copy(self.y[indices])
+
+
+class SearchableTrainingSet:
+    def __init__(self, embeddings, labels):
+        self.simsearcher = FaissCosineNeighbors()
+        self.X_train = embeddings
+        self.y_train = labels
+
+    def build_index(self):
+        self.simsearcher.fit(self.X_train, self.y_train)
+
+    def search(self, query, k=20):
+        nearest_data_points = self.simsearcher.get_distances_and_indices(
+            Q=query, topK=100
+        )
+        # topKs = [x[0] for x in Counter(nearest_data_points[0]).most_common(k)]
+        return nearest_data_points

app.py

@@ -0,0 +1,79 @@
+import pickle
+from collections import Counter
+import numpy as np
+import gradio as gr
+import gdown
+import torchvision
+from torchvision.datasets import ImageFolder
+
+from SimSearch import FaissCosineNeighbors, SearchableTrainingSet
+from ExtractEmbedding import QueryToEmbedding
+
+concat = lambda x: np.concatenate(x, axis=0)
+
+gdown.download(id="116CiA_cXciGSl72tbAUDoN-f1B9Frp89")
+gdown.download(id="1SDtq6ap7LPPpYfLbAxaMGGmj0EAV_m_e")
+
+# CUB training set
+gdown.cached_download(
+    url="https://drive.google.com/uc?id=1iR19j7532xqPefWYT-BdtcaKnsEokIqo",
+    path="./CUB_train.zip",
+    quiet=False,
+    md5="1bd99e73b2fea8e4c2ebcb0e7722f1b1",
+)
+
+# EXTRACT
+torchvision.datasets.utils.extract_archive(
+    from_path="CUB_train.zip",
+    to_path="Training/",
+    remove_finished=False,
+)
+
+
+# Caluclate Accuracy
+with open(f"./embeddings.pickle", "rb") as f:
+    Xtrain = pickle.load(f)
+# FIXME: re-run the code to get the embeddings in the right format
+with open(f"./labels.pickle", "rb") as f:
+    ytrain = pickle.load(f)
+
+searcher = SearchableTrainingSet(Xtrain, ytrain)
+searcher.build_index()
+
+# Extract label names
+training_folder = ImageFolder(root="./Training/train/")
+id_to_bird_name = {
+    x[1]: x[0].split("/")[-2].replace(".", " ") for x in training_folder.imgs
+}
+
+
+def search(query_imag, searcher=searcher):
+    query_embedding = QueryToEmbedding(query_imag)
+    indices, scores, labels = searcher.search(query_embedding, k=50)
+
+    result_ctr = Counter(labels[0][:20]).most_common(5)
+
+    top1_label = result_ctr[0][0]
+    top_indices = []
+
+    for a, b in zip(labels[0][:20], scores[0][:20]):
+        if a == top1_label:
+            top_indices.append(b)
+
+    gallery_images = [training_folder.imgs[int(X)][0] for X in top_indices[:5]]
+    predicted_labels = {id_to_bird_name[X[0]]: X[1] / 20.0 for X in result_ctr}
+
+    return predicted_labels, gallery_images
+
+
+demo = gr.Interface(
+    search,
+    gr.Image(type="pil"),
+    ["label", "gallery"],
+    examples=[["./examples/bird.jpg"]],
+    description="WIP - kNN on CUB dataset",
+    title="Work in Progress - CHM-Corr",
+)
+
+if __name__ == "__main__":
+    demo.launch()