Upload 11 files

app.py

@@ -0,0 +1,102 @@
+import os
+
+import gradio as gr
+from PIL import Image
+from sklearn.metrics.pairwise import cosine_similarity
+
+from clip_embedding import Clip
+from efficientnet_embedding import EfficientNet
+from vit_embedding import Vit
+from resnet_embedding import Resnet
+from dino_embedding import Dino
+from histogram_embedding import cosine, get_embedding
+from bovw_embedding import Bovw
+
+resnet = Resnet()
+vit = Vit()
+efficientnet = EfficientNet()
+bovw = Bovw()
+dino = Dino()
+clip = Clip()
+
+
+def get_image_embedding(image: Image.Image, name):
+    match name:
+        case "ResNet":
+            return resnet.get_embedding(image).cpu().numpy()
+        case "VIT":
+            return vit.get_embedding(image).cpu().numpy()
+        case "EfficientNet":
+            return efficientnet.get_embedding(image).cpu().numpy()
+        case "Histogram":
+            return get_embedding(image)
+        case "BOVW":
+            return bovw.get_embedding(image)
+        case "DINO":
+            return dino.get_embedding(image).cpu().numpy()
+        case _:
+            return clip.get_embedding(image).cpu().numpy()
+
+
+
+def compare_images(main_img, compare_imgs, name):
+    results = []
+    if name in ("Histogram", "ResNet", "BOVW"):
+        main_emb = get_image_embedding(main_img, name)
+        for img in compare_imgs:
+            emb = get_image_embedding(img, name)
+            results.append((img, round(cosine(main_emb, emb) * 100, 2)))
+    else:
+        main_embedding = get_image_embedding(main_img, name)
+        for img in compare_imgs:
+            emb = get_image_embedding(img, name)
+            score = cosine_similarity(main_embedding, emb)[0][0]
+            percentage = round(score * 100, 2)
+            results.append((img, percentage))
+
+
+    results.sort(key=lambda x: x[1], reverse=True)
+
+    return results
+
+
+model_list = ["CLIP", "VIT", "EfficientNet", "ResNet", "DINO", "Histogram", "BOVW"]
+
+
+with gr.Blocks() as demo:
+    gr.Tab("Image Embedding")
+    gr.Markdown("# Image Similarity Finder")
+    gr.Markdown(
+        "Upload a main image and compare it to others. Results show similarity percentages using embeddings.")
+
+    with gr.Row():
+        with gr.Column():
+            main_image = gr.Image(type="pil", label="Main Image")
+            compare_images_input = gr.File(file_count="multiple", file_types=["image"], label="Comparison Images")
+            modelName = gr.Dropdown(model_list, label="Model", value=model_list[0])
+            submit_btn = gr.Button("Compare")
+
+        with gr.Column():
+            gallery = gr.Gallery(label="Similarity Results")
+            similarity_text = gr.Textbox(label="Similarity Scores")
+
+
+    def process_comparison(main_img, compare_files, name):
+        compare_imgs = [Image.open(file.name) for file in compare_files]
+        results = compare_images(main_img, compare_imgs, name)
+
+        # Prepare outputs
+        images = [result[0] for result in results]
+        scores = [f"Image: {os.path.basename(result[0].filename)} -> Similarity: {result[1]:.2f}%" for result in
+                  results]
+
+        return images, "\n".join(scores)
+
+
+    submit_btn.click(
+        fn=process_comparison,
+        inputs=[main_image, compare_images_input, modelName],
+        outputs=[gallery, similarity_text]
+    )
+
+demo.launch()

bovw-codebook.pkl

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

bovw_embedding.py

@@ -0,0 +1,32 @@
+import numpy as np
+import cv2
+from PIL import Image
+import joblib
+from scipy.cluster.vq import vq
+
+
+class Bovw:
+    def __init__(self):
+        self.k, self.codebook = joblib.load("bovw-codebook.pkl")
+        self.idf = np.load("idf.npy")
+        self.sift = cv2.SIFT_create()
+
+    def get_embedding(self, pil_image: Image.Image) -> np.ndarray:
+        img_np = np.array(pil_image.convert("RGB"))
+        img_np = cv2.resize(img_np, (224, 224))
+        gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY).astype('uint8')
+
+        keypoints, descriptors = self.sift.detectAndCompute(gray, None)
+
+        if descriptors is None or len(descriptors) == 0:
+            return np.zeros(self.k)  # return zero-vector if no features found
+
+        visual_words, _ = vq(descriptors, self.codebook)
+
+        freq_vector = np.zeros(self.k)
+        for word in visual_words:
+            freq_vector[word] += 1
+
+        tfidf_vector = freq_vector * self.idf
+
+        return tfidf_vector

clip_embedding.py

@@ -0,0 +1,16 @@
+import torch
+from transformers import CLIPProcessor, CLIPModel
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+class Clip:
+    def __init__(self):
+        self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
+        self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+
+    def get_embedding(self, img):
+        inputs = self.processor(images=img, return_tensors="pt").to(device)
+        with torch.no_grad():
+            embeddings = self.model.get_image_features(**inputs)
+        return embeddings

dino_embedding.py

@@ -0,0 +1,21 @@
+import torch
+from torchvision import transforms
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+class Dino:
+    def __init__(self):
+        self.model = torch.hub.load('facebookresearch/dino:main', 'dino_vitb16').to(device)
+        self.model.eval()
+        self.transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5] * 3, std=[0.5] * 3),
+        ])
+
+    def get_embedding(self, image):
+        img_tensor = self.transform(image).unsqueeze(0)
+        with torch.no_grad():
+            embedding = self.model(img_tensor)
+        return embedding

efficientnet_embedding.py

@@ -0,0 +1,21 @@
+from torchvision import models, transforms
+import torch
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+class EfficientNet:
+    def __init__(self):
+        self.model = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.DEFAULT).to(device)
+        self.model.classifier = torch.nn.Identity()
+        self.model.eval()
+        self.transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+        ])
+
+    def get_embedding(self, image):
+        img_tensor = self.transform(image).unsqueeze(0)
+        with torch.no_grad():
+            embedding = self.model(img_tensor)
+        return embedding

histogram_embedding.py

@@ -0,0 +1,21 @@
+import cv2
+import numpy as np
+
+
+def cosine(a, b):
+    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
+
+
+def get_embedding(img, bins=32):
+    img = np.array(img)
+    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+
+    img = cv2.resize(img, (224, 224))
+
+    blue = cv2.calcHist([img], [0], None, [bins], [0, 256])
+    green = cv2.calcHist([img], [1], None, [bins], [0, 256])
+    red = cv2.calcHist([img], [2], None, [bins], [0, 256])
+    vector = np.concatenate([blue, green, red], axis=0)
+    vector = vector.flatten()  # Flatten to 1D array
+
+    return vector

idf.npy

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

resnet_embedding.py

@@ -0,0 +1,23 @@
+import torch
+import torchvision
+import torchvision.models as models
+import torchvision.transforms as transforms
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+class Resnet:
+    def __init__(self):
+        self.model = models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT).to(device)
+        self.model = torch.nn.Sequential(*list(self.model.children())[:-1])
+        self.model.eval()
+        self.transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+        ])
+
+    def get_embedding(self, image):
+        img_tensor = self.transform(image).unsqueeze(0)
+        with torch.no_grad():
+            embedding = self.model(img_tensor).squeeze()
+        return embedding

vit_embedding.py

@@ -0,0 +1,18 @@
+from transformers import ViTImageProcessor, ViTModel
+import torch
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+class Vit:
+    def __init__(self):
+        self.model = ViTModel.from_pretrained("google/vit-base-patch16-224-in21k").to(device)
+        self.processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
+        self.model.eval()
+
+    def get_embedding(self, image):
+        inputs = self.processor(images=image, return_tensors="pt").to(device)
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            embedding = outputs.last_hidden_state[:, 0, :]
+        return embedding