Duplicate from s-l-s/cbir-image-similarity

.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: Cbir Image Similarity
+emoji: 🏢
+colorFrom: green
+colorTo: purple
+sdk: gradio
+sdk_version: 3.19.1
+app_file: src/app.py
+pinned: false
+license: openrail
+duplicated_from: s-l-s/cbir-image-similarity
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

requirements.txt

@@ -0,0 +1,6 @@
+gradio
+numpy==1.21.6
+opencv-python==4.5.4.60
+faiss-cpu
+transformers
+torch

src/.gitattributes

@@ -0,0 +1,4 @@
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text

src/CLIP.py

@@ -0,0 +1,13 @@
+from transformers import AutoProcessor, CLIPModel
+
+
+class CLIPImageEncoder:
+    def __init__(self, device="cpu"):
+        self.device = device
+        self.model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
+        self.processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32")
+
+    def encode_image(self, image_pil):
+        input = self.processor(images=image_pil, return_tensors="pt")
+        image_features = self.model.get_image_features(**input)
+        return image_features.cpu().detach().numpy()[0]

src/FaRL.py

@@ -0,0 +1,15 @@
+#import torch
+#import clip
+
+
+#class CLIPImageEncoder:
+#    def __init__(self, device="cpu"):
+#        self.device = device
+#        self.model, self.preprocess = clip.load("ViT-B/16", device=device)
+#
+#    def encode_image(self, image_pil):
+#        print("Encoding image with CLIP")
+#        with torch.no_grad():
+#            image_preprocessed = self.preprocess(image_pil).unsqueeze(0).to(self.device)
+#            image_features = self.model.encode_image(image_preprocessed)
+#            return image_features.cpu().numpy()[0]

src/app.py

@@ -0,0 +1,81 @@
+from colordescriptor import ColorDescriptor
+from CLIP import CLIPImageEncoder
+import gradio as gr
+import cv2
+import numpy as np
+from datasets import *
+
+dataset = load_dataset("huggan/CelebA-faces")
+candidate_subset = dataset["train"].select(range(1000)) # This is a small CBIR app! :D
+
+def emb_dataset(dataset):
+    # This function might need to be split up, to reduce start-up time of app
+    # It could also use batches to increase speed
+    # If indexes are saved in files, this is all not really necessary
+
+    ## Color Embeddings
+    cd = ColorDescriptor((8, 12, 3))
+    dataset_with_embeddings = dataset.map(lambda row: {'color_embeddings': cd.describe(row["image"])}) # we assume that dataset has a column 'image'
+    dataset_with_embeddings.add_faiss_index(column='color_embeddings')
+
+    ## CLIP Embeddings
+    clip_model = CLIPImageEncoder()
+    dataset_with_embeddings = dataset.map(lambda row: {'clip_embeddings': clip_model.encode_image(row["image"])})
+    dataset_with_embeddings.add_faiss_index(column='clip_embeddings')
+    dataset_with_embeddings # Just to check, if okay
+
+    return dataset_with_embeddings
+
+dataset_with_embeddings = emb_dataset(candidate_subset)
+
+# Main function, to find similar images
+# TODO: allow different descriptor/embedding functions
+# TODO: implement different distance measures
+
+def get_neighbors(query_image, selected_descriptor, top_k=5):
+    """Returns the top k nearest examples to the query image.
+
+    Args:
+        query_image: A PIL object representing the query image.
+        top_k: An integer representing the number of nearest examples to return.
+
+    Returns:
+        A list of the top_k most similar images as PIL objects.
+    """
+    if  "Color Descriptor" in selected_descriptor:
+        cd = ColorDescriptor((8, 12, 3))
+        qi_embedding = cd.describe(query_image)
+        qi_np = np.array(qi_embedding)
+        scores, retrieved_examples = dataset_with_embeddings.get_nearest_examples(
+            'color_embeddings', qi_np, k=top_k)
+        images = retrieved_examples['image'] #retrieved images is a dict, with images and embeddings
+        return images
+    if "CLIP" in selected_descriptor:
+        clip_model = CLIPImageEncoder()
+        qi_embedding = clip_model.encode_image(query_image)
+        scores, retrieved_examples = dataset_with_embeddings.get_nearest_examples(
+            'clip_embeddings', qi_embedding, k=top_k)
+        images = retrieved_examples['image']
+        return images
+    else:
+        print("This descriptor is not yet supported :(")
+        return []
+
+
+# Define the Gradio Interface
+
+
+iface = gr.Interface(
+    fn=get_neighbors,
+    inputs=[
+        gr.Image(type="pil", label="Your Image"),
+        gr.CheckboxGroup(["Color Descriptor", "LBP", "CLIP"], label="Descriptor method?"),
+    ],
+    outputs=gr.Gallery(),
+    title="Image Similarity Gallery",
+    description="Upload an image and get similar images",
+    allow_flagging="never"
+)
+
+# Launch the Gradio interface
+iface.launch()

src/colordescriptor.py

@@ -0,0 +1,60 @@
+import numpy as np
+import cv2
+from helper import pil_cv2_image_converter
+
+class ColorDescriptor:
+    def __init__(self, bins):
+        # store the number of bins for the 3D histogram
+        self.bins = bins
+
+    def histogram(self, image, mask):
+        # extract a 3D color histogram from the masked region of the
+        # image, using the supplied number of bins per channel
+        hist = cv2.calcHist([image], [0, 1, 2], mask, self.bins,
+                            [0, 180, 0, 256, 0, 256])
+
+        hist = cv2.normalize(hist, hist).flatten()
+        # return the histogram
+        return hist
+
+    def describe(self, image):
+        # first, convert image to cv2 from pil
+        # TODO: Add check, if already cv2 image
+        image = pil_cv2_image_converter(image)
+        # convert the image to the HSV color space and initialize
+        # the features used to quantify the image
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
+        features = []
+        # grab the dimensions and compute the center of the image
+        (h, w) = image.shape[:2]
+        (cX, cY) = (int(w * 0.5), int(h * 0.5))
+
+        # divide the image into four rectangles/segments (top-left,
+        # top-right, bottom-right, bottom-left)
+        segments = [(0, cX, 0, cY), (cX, w, 0, cY), (cX, w, cY, h),
+                    (0, cX, cY, h)]
+        # construct an elliptical mask representing the center of the
+        # image
+        (axesX, axesY) = (int(w * 0.75) // 2, int(h * 0.75) // 2)
+        ellipMask = np.zeros(image.shape[:2], dtype = "uint8")
+        cv2.ellipse(ellipMask, (cX, cY), (axesX, axesY), 0, 0, 360, 255, -1)
+
+        # loop over the segments
+        for (startX, endX, startY, endY) in segments:
+            # construct a mask for each corner of the image, subtracting
+            # the elliptical center from it
+            cornerMask = np.zeros(image.shape[:2], dtype = "uint8")
+            cv2.rectangle(cornerMask, (startX, startY), (endX, endY), 255, -1)
+            cornerMask = cv2.subtract(cornerMask, ellipMask)
+            # extract a color histogram from the image, then update the
+            # feature vector
+            hist = self.histogram(image, cornerMask)
+            features.extend(hist)
+
+        # extract a color histogram from the elliptical region and
+        # update the feature vector
+        hist = self.histogram(image, ellipMask)
+        features.extend(hist)
+
+        # return the feature vector
+        return features

src/helper.py

@@ -0,0 +1,9 @@
+import numpy as np
+import cv2
+
+def pil_cv2_image_converter(image):
+    numpy_image = np.array(image)
+    # Convert to an OpenCV image; notice the COLOR_RGB2BGR flag, which means
+    # that the color is converted from RGB to BGR format.
+    opencv_image = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
+    return opencv_image

src/index.py

@@ -0,0 +1,32 @@
+# # import the necessary packages
+# from colordescriptor import ColorDescriptor
+# import glob
+# import cv2
+
+
+# class Indexer:
+#         def __init__(self, indexPath):
+#                 # store our index path
+#                 self.indexPath = indexPath
+
+#         def index(self):
+#                 # initialize the color descriptor
+#                 cd = ColorDescriptor((8, 12, 3))
+
+#                 # open the output index file for writing
+#                 output = open(self.indexPath, "w")
+
+#                 # use glob to grab the image paths and loop over them
+#                 for imagePath in glob.glob("../static/images/" + "/*.png"):
+#                         # extract the image ID (i.e. the unique filename) from the image
+#                         # path and load the image itself
+#                         imageID = imagePath[imagePath.rfind("/") + 1:]
+#                         image = cv2.imread(imagePath)
+#                         # describe the image
+#                         features = cd.describe(image)
+#                         # write the features to file
+#                         features = [str(f) for f in features]
+#                         output.write("%s,%s\n" % (imageID, ",".join(features)))
+
+#                 # close the index file
+#                 output.close()

src/search.py

@@ -0,0 +1,31 @@
+# from colordescriptor import ColorDescriptor
+# from searcher import Searcher
+# import argparse
+# import cv2
+
+# # construct the argument parser and parse the arguments
+# ap = argparse.ArgumentParser()
+# ap.add_argument("-i", "--index", required = True,
+#         help = "Path to where the computed index will be stored")
+# ap.add_argument("-q", "--query", required = True,
+#         help = "Path to the query image")
+# ap.add_argument("-r", "--result-path", required = True,
+#         help = "Path to the result path")
+# args = vars(ap.parse_args())
+
+# # initialize the image descriptor
+# cd = ColorDescriptor((8, 12, 3))
+# # load the query image and describe it
+# query = cv2.imread(args["query"])
+# features = cd.describe(query)
+# # perform the search
+# searcher = Searcher(args["index"])
+# results = searcher.search(features)
+# # display the query
+# cv2.imshow("Query", query)
+# # loop over the results
+# for (score, resultID) in results:
+#         # load the result image and display it
+#         result = cv2.imread(args["result_path"] + "/" + resultID)
+#         cv2.imshow("Result", result)
+#         cv2.waitKey(0)

src/searcher.py

@@ -0,0 +1,47 @@
+# import numpy as np
+# import csv
+
+# class Searcher:
+#     def __init__(self, indexPath):
+#         # store our index path
+#         self.indexPath = indexPath
+
+
+#     def chi2_distance(self, histA, histB, eps = 1e-10):
+#         # compute the chi-squared distance
+#         d = 0.5 * np.sum([((a - b) ** 2) / (a + b + eps)
+#             for (a, b) in zip(histA, histB)])
+#         # return the chi-squared distance
+#         return d
+
+#     def search(self, queryFeatures, limit = 3):
+#         # initialize our dictionary of results
+#         results = {}
+#         # open the index file for reading
+#         with open(self.indexPath) as f:
+#             # initialize the CSV reader
+#             reader = csv.reader(f)
+#             # loop over the rows in the index
+#             for row in reader:
+#                 # parse out the image ID and features, then compute the
+#                 # chi-squared distance between the features in our index
+#                 # and our query features
+#                 features = [float(x) for x in row[1:]]
+#                 d = self.chi2_distance(features, queryFeatures)
+#                 # now that we have the distance between the two feature
+#                 # vectors, we can udpate the results dictionary -- the
+#                 # key is the current image ID in the index and the
+#                 # value is the distance we just computed, representing
+#                 # how 'similar' the image in the index is to our query
+#                 results[row[0]] = d
+
+#             # close the reader
+#             f.close()
+
+#         # sort our results, so that the smaller distances (i.e. the
+#         # more relevant images are at the front of the list)
+#         path = "home/user/app/static/images/"
+#         results = sorted([(v, f"{path}{k}") for (k, v) in results.items()])
+
+#         # return our (limited) results
+#         return results[:limit]