Near-duplicate image search

<!DOCTYPE html> <html lang="en"> <head> <meta charset="utf-8"> <meta name="viewport" content="width=device-width, initial-scale=1"> <meta name="description" content="Keras documentation"> <meta name="author" content="Keras Team"> <link rel="shortcut icon" href="https://keras.io/img/favicon.ico"> <link rel="canonical" href="https://keras.io/examples/vision/near_dup_search/" />  <meta property="og:title" content="Keras documentation: Near-duplicate image search"> <meta property="og:image" content="https://keras.io/img/logo-k-keras-wb.png"> <meta name="twitter:title" content="Keras documentation: Near-duplicate image search"> <meta name="twitter:image" content="https://keras.io/img/k-keras-social.png"> <meta name="twitter:card" content="summary"> <title>Near-duplicate image search</title>  <link href="/css/bootstrap.min.css" rel="stylesheet">  <link href="https://fonts.googleapis.com/css2?family=Open+Sans:wght@400;600;700;800&display=swap" rel="stylesheet">  <link href="/css/docs.css" rel="stylesheet"> <link href="/css/monokai.css" rel="stylesheet">  <script>(function(w,d,s,l,i){w[l]=w[l]||[];w[l].push({'gtm.start': new Date().getTime(),event:'gtm.js'});var f=d.getElementsByTagName(s)[0], j=d.createElement(s),dl=l!='dataLayer'?'&l='+l:'';j.async=true;j.src= 'https://www.googletagmanager.com/gtm.js?id='+i+dl;f.parentNode.insertBefore(j,f); })(window,document,'script','dataLayer','GTM-5DNGF4N'); </script> <script> (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){ (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o), m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m) })(window,document,'script','https://www.google-analytics.com/analytics.js','ga'); ga('create', 'UA-175165319-128', 'auto'); ga('send', 'pageview'); </script>  <script async defer src="https://buttons.github.io/buttons.js"></script> </head> <body>  <noscript><iframe src="https://www.googletagmanager.com/ns.html?id=GTM-5DNGF4N" height="0" width="0" style="display:none;visibility:hidden"></iframe></noscript>  <div class='k-page'> <div class="k-nav" id="nav-menu"> <a href='/'><img src='/img/logo-small.png' class='logo-small' /></a> <div class="nav flex-column nav-pills" role="tablist" aria-orientation="vertical"> <a class="nav-link" href="/about/" role="tab" aria-selected="">About Keras</a> <a class="nav-link" href="/getting_started/" role="tab" aria-selected="">Getting started</a> <a class="nav-link" href="/guides/" role="tab" aria-selected="">Developer guides</a> <a class="nav-link active" href="/examples/" role="tab" aria-selected="">Code examples</a> <a class="nav-sublink active" href="/examples/vision/">Computer Vision</a> <a class="nav-sublink2" href="/examples/vision/image_classification_from_scratch/">Image classification from scratch</a> <a class="nav-sublink2" href="/examples/vision/mnist_convnet/">Simple MNIST convnet</a> <a class="nav-sublink2" href="/examples/vision/image_classification_efficientnet_fine_tuning/">Image classification via fine-tuning with EfficientNet</a> <a class="nav-sublink2" href="/examples/vision/image_classification_with_vision_transformer/">Image classification with Vision Transformer</a> <a class="nav-sublink2" href="/examples/vision/attention_mil_classification/">Classification using Attention-based Deep Multiple Instance Learning</a> <a class="nav-sublink2" href="/examples/vision/mlp_image_classification/">Image classification with modern MLP models</a> <a class="nav-sublink2" href="/examples/vision/mobilevit/">A mobile-friendly Transformer-based model for image classification</a> <a class="nav-sublink2" href="/examples/vision/xray_classification_with_tpus/">Pneumonia Classification on TPU</a> <a class="nav-sublink2" href="/examples/vision/cct/">Compact Convolutional Transformers</a> <a class="nav-sublink2" href="/examples/vision/convmixer/">Image classification with ConvMixer</a> <a class="nav-sublink2" href="/examples/vision/eanet/">Image classification with EANet (External Attention Transformer)</a> <a class="nav-sublink2" href="/examples/vision/involution/">Involutional neural networks</a> <a class="nav-sublink2" href="/examples/vision/perceiver_image_classification/">Image classification with Perceiver</a> <a class="nav-sublink2" href="/examples/vision/reptile/">Few-Shot learning with Reptile</a> <a class="nav-sublink2" href="/examples/vision/semisupervised_simclr/">Semi-supervised image classification using contrastive pretraining with SimCLR</a> <a class="nav-sublink2" href="/examples/vision/swin_transformers/">Image classification with Swin Transformers</a> <a class="nav-sublink2" href="/examples/vision/vit_small_ds/">Train a Vision Transformer on small datasets</a> <a class="nav-sublink2" href="/examples/vision/shiftvit/">A Vision Transformer without Attention</a> <a class="nav-sublink2" href="/examples/vision/image_classification_using_global_context_vision_transformer/">Image Classification using Global Context Vision Transformer</a> <a class="nav-sublink2" href="/examples/vision/oxford_pets_image_segmentation/">Image segmentation with a U-Net-like architecture</a> <a class="nav-sublink2" href="/examples/vision/deeplabv3_plus/">Multiclass semantic segmentation using DeepLabV3+</a> <a class="nav-sublink2" href="/examples/vision/basnet_segmentation/">Highly accurate boundaries segmentation using BASNet</a> <a class="nav-sublink2" href="/examples/vision/fully_convolutional_network/">Image Segmentation using Composable Fully-Convolutional Networks</a> <a class="nav-sublink2" href="/examples/vision/retinanet/">Object Detection with RetinaNet</a> <a class="nav-sublink2" href="/examples/vision/keypoint_detection/">Keypoint Detection with Transfer Learning</a> <a class="nav-sublink2" href="/examples/vision/object_detection_using_vision_transformer/">Object detection with Vision Transformers</a> <a class="nav-sublink2" href="/examples/vision/3D_image_classification/">3D image classification from CT scans</a> <a class="nav-sublink2" href="/examples/vision/depth_estimation/">Monocular depth estimation</a> <a class="nav-sublink2" href="/examples/vision/nerf/">3D volumetric rendering with NeRF</a> <a class="nav-sublink2" href="/examples/vision/pointnet_segmentation/">Point cloud segmentation with PointNet</a> <a class="nav-sublink2" href="/examples/vision/pointnet/">Point cloud classification</a> <a class="nav-sublink2" href="/examples/vision/captcha_ocr/">OCR model for reading Captchas</a> <a class="nav-sublink2" href="/examples/vision/handwriting_recognition/">Handwriting recognition</a> <a class="nav-sublink2" href="/examples/vision/autoencoder/">Convolutional autoencoder for image denoising</a> <a class="nav-sublink2" href="/examples/vision/mirnet/">Low-light image enhancement using MIRNet</a> <a class="nav-sublink2" href="/examples/vision/super_resolution_sub_pixel/">Image Super-Resolution using an Efficient Sub-Pixel CNN</a> <a class="nav-sublink2" href="/examples/vision/edsr/">Enhanced Deep Residual Networks for single-image super-resolution</a> <a class="nav-sublink2" href="/examples/vision/zero_dce/">Zero-DCE for low-light image enhancement</a> <a class="nav-sublink2" href="/examples/vision/cutmix/">CutMix data augmentation for image classification</a> <a class="nav-sublink2" href="/examples/vision/mixup/">MixUp augmentation for image classification</a> <a class="nav-sublink2" href="/examples/vision/randaugment/">RandAugment for Image Classification for Improved Robustness</a> <a class="nav-sublink2" href="/examples/vision/image_captioning/">Image captioning</a> <a class="nav-sublink2" href="/examples/vision/nl_image_search/">Natural language image search with a Dual Encoder</a> <a class="nav-sublink2" href="/examples/vision/visualizing_what_convnets_learn/">Visualizing what convnets learn</a> <a class="nav-sublink2" href="/examples/vision/integrated_gradients/">Model interpretability with Integrated Gradients</a> <a class="nav-sublink2" href="/examples/vision/probing_vits/">Investigating Vision Transformer representations</a> <a class="nav-sublink2" href="/examples/vision/grad_cam/">Grad-CAM class activation visualization</a> <a class="nav-sublink2 active" href="/examples/vision/near_dup_search/">Near-duplicate image search</a> <a class="nav-sublink2" href="/examples/vision/semantic_image_clustering/">Semantic Image Clustering</a> <a class="nav-sublink2" href="/examples/vision/siamese_contrastive/">Image similarity estimation using a Siamese Network with a contrastive loss</a> <a class="nav-sublink2" href="/examples/vision/siamese_network/">Image similarity estimation using a Siamese Network with a triplet loss</a> <a class="nav-sublink2" href="/examples/vision/metric_learning/">Metric learning for image similarity search</a> <a class="nav-sublink2" href="/examples/vision/metric_learning_tf_similarity/">Metric learning for image similarity search using TensorFlow Similarity</a> <a class="nav-sublink2" href="/examples/vision/nnclr/">Self-supervised contrastive learning with NNCLR</a> <a class="nav-sublink2" href="/examples/vision/video_classification/">Video Classification with a CNN-RNN Architecture</a> <a class="nav-sublink2" href="/examples/vision/conv_lstm/">Next-Frame Video Prediction with Convolutional LSTMs</a> <a class="nav-sublink2" href="/examples/vision/video_transformers/">Video Classification with Transformers</a> <a class="nav-sublink2" href="/examples/vision/vivit/">Video Vision Transformer</a> <a class="nav-sublink2" href="/examples/vision/bit/">Image Classification using BigTransfer (BiT)</a> <a class="nav-sublink2" href="/examples/vision/gradient_centralization/">Gradient Centralization for Better Training Performance</a> <a class="nav-sublink2" href="/examples/vision/token_learner/">Learning to tokenize in Vision Transformers</a> <a class="nav-sublink2" href="/examples/vision/knowledge_distillation/">Knowledge Distillation</a> <a class="nav-sublink2" href="/examples/vision/fixres/">FixRes: Fixing train-test resolution discrepancy</a> <a class="nav-sublink2" href="/examples/vision/cait/">Class Attention Image Transformers with LayerScale</a> <a class="nav-sublink2" href="/examples/vision/patch_convnet/">Augmenting convnets with aggregated attention</a> <a class="nav-sublink2" href="/examples/vision/learnable_resizer/">Learning to Resize</a> <a class="nav-sublink2" href="/examples/vision/adamatch/">Semi-supervision and domain adaptation with AdaMatch</a> <a class="nav-sublink2" href="/examples/vision/barlow_twins/">Barlow Twins for Contrastive SSL</a> <a class="nav-sublink2" href="/examples/vision/consistency_training/">Consistency training with supervision</a> <a class="nav-sublink2" href="/examples/vision/deit/">Distilling Vision Transformers</a> <a class="nav-sublink2" href="/examples/vision/focal_modulation_network/">Focal Modulation: A replacement for Self-Attention</a> <a class="nav-sublink2" href="/examples/vision/forwardforward/">Using the Forward-Forward Algorithm for Image Classification</a> <a class="nav-sublink2" href="/examples/vision/masked_image_modeling/">Masked image modeling with Autoencoders</a> <a class="nav-sublink2" href="/examples/vision/sam/">Segment Anything Model with 🤗Transformers</a> <a class="nav-sublink2" href="/examples/vision/segformer/">Semantic segmentation with SegFormer and Hugging Face Transformers</a> <a class="nav-sublink2" href="/examples/vision/simsiam/">Self-supervised contrastive learning with SimSiam</a> <a class="nav-sublink2" href="/examples/vision/supervised-contrastive-learning/">Supervised Contrastive Learning</a> <a class="nav-sublink2" href="/examples/vision/temporal_latent_bottleneck/">When Recurrence meets Transformers</a> <a class="nav-sublink2" href="/examples/vision/yolov8/">Efficient Object Detection with YOLOV8 and KerasCV</a> <a class="nav-sublink" href="/examples/nlp/">Natural Language Processing</a> <a class="nav-sublink" href="/examples/structured_data/">Structured Data</a> <a class="nav-sublink" href="/examples/timeseries/">Timeseries</a> <a class="nav-sublink" href="/examples/generative/">Generative Deep Learning</a> <a class="nav-sublink" href="/examples/audio/">Audio Data</a> <a class="nav-sublink" href="/examples/rl/">Reinforcement Learning</a> <a class="nav-sublink" href="/examples/graph/">Graph Data</a> <a class="nav-sublink" href="/examples/keras_recipes/">Quick Keras Recipes</a> <a class="nav-link" href="/api/" role="tab" aria-selected="">Keras 3 API documentation</a> <a class="nav-link" href="/2.18/api/" role="tab" aria-selected="">Keras 2 API documentation</a> <a class="nav-link" href="/keras_tuner/" role="tab" aria-selected="">KerasTuner: Hyperparam Tuning</a> <a class="nav-link" href="/keras_hub/" role="tab" aria-selected="">KerasHub: Pretrained Models</a> </div> </div> <div class='k-main'> <div class='k-main-top'> <script> function displayDropdownMenu() { e = document.getElementById("nav-menu"); if (e.style.display == "block") { e.style.display = "none"; } else { e.style.display = "block"; document.getElementById("dropdown-nav").style.display = "block"; } } function resetMobileUI() { if (window.innerWidth <= 840) { document.getElementById("nav-menu").style.display = "none"; document.getElementById("dropdown-nav").style.display = "block"; } else { document.getElementById("nav-menu").style.display = "block"; document.getElementById("dropdown-nav").style.display = "none"; } var navmenu = document.getElementById("nav-menu"); var menuheight = navmenu.clientHeight; var kmain = document.getElementById("k-main-id"); kmain.style.minHeight = (menuheight + 100) + 'px'; } window.onresize = resetMobileUI; window.addEventListener("load", (event) => { resetMobileUI() }); </script> <div id='dropdown-nav' onclick="displayDropdownMenu();"> <svg viewBox="-20 -20 120 120" width="60" height="60"> <rect width="100" height="20"></rect> <rect y="30" width="100" height="20"></rect> <rect y="60" width="100" height="20"></rect> </svg> </div> <form class="bd-search d-flex align-items-center k-search-form" id="search-form"> <input type="search" class="k-search-input" id="search-input" placeholder="Search Keras documentation..." aria-label="Search Keras documentation..." autocomplete="off"> <button class="k-search-btn"> <svg width="13" height="13" viewBox="0 0 13 13"><title>search</title><path d="m4.8495 7.8226c0.82666 0 1.5262-0.29146 2.0985-0.87438 0.57232-0.58292 0.86378-1.2877 0.87438-2.1144 0.010599-0.82666-0.28086-1.5262-0.87438-2.0985-0.59352-0.57232-1.293-0.86378-2.0985-0.87438-0.8055-0.010599-1.5103 0.28086-2.1144 0.87438-0.60414 0.59352-0.8956 1.293-0.87438 2.0985 0.021197 0.8055 0.31266 1.5103 0.87438 2.1144 0.56172 0.60414 1.2665 0.8956 2.1144 0.87438zm4.4695 0.2115 3.681 3.6819-1.259 1.284-3.6817-3.7 0.0019784-0.69479-0.090043-0.098846c-0.87973 0.76087-1.92 1.1413-3.1207 1.1413-1.3553 0-2.5025-0.46363-3.4417-1.3909s-1.4088-2.0686-1.4088-3.4239c0-1.3553 0.4696-2.4966 1.4088-3.4239 0.9392-0.92727 2.0864-1.3969 3.4417-1.4088 1.3553-0.011889 2.4906 0.45771 3.406 1.4088 0.9154 0.95107 1.379 2.0924 1.3909 3.4239 0 1.2126-0.38043 2.2588-1.1413 3.1385l0.098834 0.090049z"></path></svg> </button> </form> <script> var form = document.getElementById('search-form'); form.onsubmit = function(e) { e.preventDefault(); var query = document.getElementById('search-input').value; window.location.href = '/search.html?query=' + query; return False } </script> </div> <div class='k-main-inner' id='k-main-id'> <div class='k-location-slug'> ► <a href='/examples/'>Code examples</a> / <a href='/examples/vision/'>Computer Vision</a> / Near-duplicate image search </div> <div class='k-content'> <h1 id="nearduplicate-image-search">Near-duplicate image search</h1> Author: <a href="https://twitter.com/RisingSayak">Sayak Paul</a> Date created: 2021/09/10 Last modified: 2023/08/30 <div class='example_version_banner keras_2'>ⓘ This example uses Keras 2</div> <img class="k-inline-icon" src="https://colab.research.google.com/img/colab_favicon.ico"/> <a href="https://colab.research.google.com/github/keras-team/keras-io/blob/master/examples/vision/ipynb/near_dup_search.ipynb">View in Colab</a> •<img class="k-inline-icon" src="https://github.com/favicon.ico"/> <a href="https://github.com/keras-team/keras-io/blob/master/examples/vision/near_dup_search.py">GitHub source</a> Description: Building a near-duplicate image search utility using deep learning and locality-sensitive hashing. <hr /> <h2 id="introduction">Introduction</h2> Fetching similar images in (near) real time is an important use case of information retrieval systems. Some popular products utilizing it include Pinterest, Google Image Search, etc. In this example, we will build a similar image search utility using <a href="https://towardsdatascience.com/understanding-locality-sensitive-hashing-49f6d1f6134">Locality Sensitive Hashing</a> (LSH) and <a href="https://en.wikipedia.org/wiki/Random_projection">random projection</a> on top of the image representations computed by a pretrained image classifier. This kind of search engine is also known as a near-duplicate (or near-dup) image detector. We will also look into optimizing the inference performance of our search utility on GPU using <a href="https://developer.nvidia.com/tensorrt">TensorRT</a>. There are other examples under <a href="https://keras.io/examples/vision">keras.io/examples/vision</a> that are worth checking out in this regard: <ul> <li><a href="https://keras.io/examples/vision/metric_learning">Metric learning for image similarity search</a></li> <li><a href="https://keras.io/examples/vision/siamese_network">Image similarity estimation using a Siamese Network with a triplet loss</a></li> </ul> Finally, this example uses the following resource as a reference and as such reuses some of its code: <a href="https://towardsdatascience.com/locality-sensitive-hashing-for-music-search-f2f1940ace23">Locality Sensitive Hashing for Similar Item Search</a>. Note that in order to optimize the performance of our parser, you should have a GPU runtime available. <hr /> <h2 id="setup">Setup</h2> <div class="codehilite"><pre><code>!pip install tensorrt </code></pre></div> <hr /> <h2 id="imports">Imports</h2> <div class="codehilite"><pre><code>import matplotlib.pyplot as plt import tensorflow as tf import tensorrt import numpy as np import time import tensorflow_datasets as tfds tfds.disable_progress_bar() </code></pre></div> <hr /> <h2 id="load-the-dataset-and-create-a-training-set-of-1000-images">Load the dataset and create a training set of 1,000 images</h2> To keep the run time of the example short, we will be using a subset of 1,000 images from the <code>tf_flowers</code> dataset (available through <a href="https://www.tensorflow.org/datasets/catalog/tf_flowers">TensorFlow Datasets</a>) to build our vocabulary. <div class="codehilite"><pre><code>train_ds, validation_ds = tfds.load( "tf_flowers", split=["train[:85%]", "train[85%:]"], as_supervised=True ) IMAGE_SIZE = 224 NUM_IMAGES = 1000 images = [] labels = [] for (image, label) in train_ds.take(NUM_IMAGES): image = tf.image.resize(image, (IMAGE_SIZE, IMAGE_SIZE)) images.append(image.numpy()) labels.append(label.numpy()) images = np.array(images) labels = np.array(labels) </code></pre></div> <hr /> <h2 id="load-a-pretrained-model">Load a pre-trained model</h2> In this section, we load an image classification model that was trained on the <code>tf_flowers</code> dataset. 85% of the total images were used to build the training set. For more details on the training, refer to <a href="https://github.com/sayakpaul/near-dup-parser/blob/main/bit-supervised-training.ipynb">this notebook</a>. The underlying model is a BiT-ResNet (proposed in <a href="https://arxiv.org/abs/1912.11370">Big Transfer (BiT): General Visual Representation Learning</a>). The BiT-ResNet family of models is known to provide excellent transfer performance across a wide variety of different downstream tasks. <div class="codehilite"><pre><code>!wget -q https://github.com/sayakpaul/near-dup-parser/releases/download/v0.1.0/flower_model_bit_0.96875.zip !unzip -qq flower_model_bit_0.96875.zip </code></pre></div> <div class="codehilite"><pre><code>bit_model = tf.keras.models.load_model("flower_model_bit_0.96875") bit_model.count_params() </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>23510597 </code></pre></div> </div> <hr /> <h2 id="create-an-embedding-model">Create an embedding model</h2> To retrieve similar images given a query image, we need to first generate vector representations of all the images involved. We do this via an embedding model that extracts output features from our pretrained classifier and normalizes the resulting feature vectors. <div class="codehilite"><pre><code>embedding_model = tf.keras.Sequential( [ tf.keras.layers.Input((IMAGE_SIZE, IMAGE_SIZE, 3)), tf.keras.layers.Rescaling(scale=1.0 / 255), bit_model.layers[1], tf.keras.layers.Normalization(mean=0, variance=1), ], name="embedding_model", ) embedding_model.summary() </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Model: "embedding_model" _________________________________________________________________ Layer (type) Output Shape Param # ================================================================= rescaling (Rescaling) (None, 224, 224, 3) 0 _________________________________________________________________ keras_layer (KerasLayer) (None, 2048) 23500352 _________________________________________________________________ normalization (Normalization (None, 2048) 0 ================================================================= Total params: 23,500,352 Trainable params: 23,500,352 Non-trainable params: 0 _________________________________________________________________ </code></pre></div> </div> Take note of the normalization layer inside the model. It is used to project the representation vectors to the space of unit-spheres. <hr /> <h2 id="hashing-utilities">Hashing utilities</h2> <div class="codehilite"><pre><code>def hash_func(embedding, random_vectors): embedding = np.array(embedding) # 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 </code></pre></div> The shape of the vectors coming out of <code>embedding_model</code> is <code>(2048,)</code>, and considering practical aspects (storage, retrieval performance, etc.) it is quite large. So, there arises a need to reduce the dimensionality of the embedding vectors without reducing their information content. This is where random projection comes into the picture. It is based on the principle that if the distance between a group of points on a given plane is approximately preserved, the dimensionality of that plane can further be reduced. Inside <code>hash_func()</code>, we first reduce the dimensionality of the embedding vectors. Then we compute the bitwise hash values of the images to determine their hash buckets. Images having same hash values are likely to go into the same hash bucket. From a deployment perspective, bitwise hash values are cheaper to store and operate on. <hr /> <h2 id="query-utilities">Query utilities</h2> The <code>Table</code> class is responsible for building a single hash table. Each entry in the hash table is a mapping between the reduced embedding of an image from our dataset and a unique identifier. Because our dimensionality reduction technique involves randomness, it can so happen that similar images are not mapped to the same hash bucket everytime the process run. To reduce this effect, we will take results from multiple tables into consideration – the number of tables and the reduction dimensionality are the key hyperparameters here. Crucially, you wouldn't reimplement locality-sensitive hashing yourself when working with real world applications. Instead, you'd likely use one of the following popular libraries: <ul> <li><a href="https://github.com/google-research/google-research/tree/master/scann">ScaNN</a></li> <li><a href="https://github.com/spotify/annoy">Annoy</a></li> <li><a href="https://github.com/vdaas/vald">Vald</a></li> </ul> <div class="codehilite"><pre><code>class Table: def __init__(self, hash_size, dim): self.table = {} self.hash_size = hash_size self.random_vectors = np.random.randn(hash_size, dim).T def add(self, id, vectors, label): # Create a unique indentifier. entry = {"id_label": str(id) + "_" + str(label)} # Compute the hash values. hashes = hash_func(vectors, self.random_vectors) # 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, vectors): # Compute hash value for the query vector. hashes = hash_func(vectors, self.random_vectors) 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 </code></pre></div> In the following <code>LSH</code> class we will pack the utilities to have multiple hash tables. <div class="codehilite"><pre><code>class LSH: def __init__(self, hash_size, dim, num_tables): self.num_tables = num_tables self.tables = [] for i in range(self.num_tables): self.tables.append(Table(hash_size, dim)) def add(self, id, vectors, label): for table in self.tables: table.add(id, vectors, label) def query(self, vectors): results = [] for table in self.tables: results.extend(table.query(vectors)) return results </code></pre></div> Now we can encapsulate the logic for building and operating with the master LSH table (a collection of many tables) inside a class. It has two methods: <ul> <li><code>train()</code>: Responsible for building the final LSH table.</li> <li><code>query()</code>: Computes the number of matches given a query image and also quantifies the similarity score.</li> </ul> <div class="codehilite"><pre><code>class BuildLSHTable: def __init__( self, prediction_model, concrete_function=False, hash_size=8, dim=2048, num_tables=10, ): self.hash_size = hash_size self.dim = dim self.num_tables = num_tables self.lsh = LSH(self.hash_size, self.dim, self.num_tables) self.prediction_model = prediction_model self.concrete_function = concrete_function def train(self, training_files): for id, training_file in enumerate(training_files): # Unpack the data. image, label = training_file if len(image.shape) < 4: image = image[None, ...] # Compute embeddings and update the LSH tables. # More on `self.concrete_function()` later. if self.concrete_function: features = self.prediction_model(tf.constant(image))[ "normalization" ].numpy() else: features = self.prediction_model.predict(image) self.lsh.add(id, features, label) def query(self, image, verbose=True): # Compute the embeddings of the query image and fetch the results. if len(image.shape) < 4: image = image[None, ...] if self.concrete_function: features = self.prediction_model(tf.constant(image))[ "normalization" ].numpy() else: features = self.prediction_model.predict(image) results = self.lsh.query(features) 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 </code></pre></div> <hr /> <h2 id="create-lsh-tables">Create LSH tables</h2> With our helper utilities and classes implemented, we can now build our LSH table. Since we will be benchmarking performance between optimized and unoptimized embedding models, we will also warm up our GPU to avoid any unfair comparison. <div class="codehilite"><pre><code># Utility to warm up the GPU. def warmup(): dummy_sample = tf.ones((1, IMAGE_SIZE, IMAGE_SIZE, 3)) for _ in range(100): _ = embedding_model.predict(dummy_sample) </code></pre></div> Now we can first do the GPU wam-up and proceed to build the master LSH table with <code>embedding_model</code>. <div class="codehilite"><pre><code>warmup() training_files = zip(images, labels) lsh_builder = BuildLSHTable(embedding_model) lsh_builder.train(training_files) </code></pre></div> At the time of writing, the wall time was 54.1 seconds on a Tesla T4 GPU. This timing may vary based on the GPU you are using. <hr /> <h2 id="optimize-the-model-with-tensorrt">Optimize the model with TensorRT</h2> For NVIDIA-based GPUs, the <a href="https://docs.nvidia.com/deeplearning/frameworks/tf-trt-user-guide/index.html">TensorRT framework</a> can be used to dramatically enhance the inference latency by using various model optimization techniques like pruning, constant folding, layer fusion, and so on. Here we will use the <a href="https://www.tensorflow.org/api_docs/python/tf/experimental/tensorrt"><code>tf.experimental.tensorrt</code></a> module to optimize our embedding model. <div class="codehilite"><pre><code># First serialize the embedding model as a SavedModel. embedding_model.save("embedding_model") # Initialize the conversion parameters. params = tf.experimental.tensorrt.ConversionParams( precision_mode="FP16", maximum_cached_engines=16 ) # Run the conversion. converter = tf.experimental.tensorrt.Converter( input_saved_model_dir="embedding_model", conversion_params=params ) converter.convert() converter.save("tensorrt_embedding_model") </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model. WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model. INFO:tensorflow:Assets written to: embedding_model/assets INFO:tensorflow:Assets written to: embedding_model/assets INFO:tensorflow:Linked TensorRT version: (0, 0, 0) INFO:tensorflow:Linked TensorRT version: (0, 0, 0) INFO:tensorflow:Loaded TensorRT version: (0, 0, 0) INFO:tensorflow:Loaded TensorRT version: (0, 0, 0) INFO:tensorflow:Assets written to: tensorrt_embedding_model/assets INFO:tensorflow:Assets written to: tensorrt_embedding_model/assets </code></pre></div> </div> Notes on the parameters inside of <code>tf.experimental.tensorrt.ConversionParams()</code>: <ul> <li><code>precision_mode</code> defines the numerical precision of the operations in the to-be-converted model.</li> <li><code>maximum_cached_engines</code> specifies the maximum number of TRT engines that will be cached to handle dynamic operations (operations with unknown shapes).</li> </ul> To learn more about the other options, refer to the <a href="https://www.tensorflow.org/api_docs/python/tf/experimental/tensorrt/ConversionParams">official documentation</a>. You can also explore the different quantization options provided by the <a href="https://www.tensorflow.org/api_docs/python/tf/experimental/tensorrt"><code>tf.experimental.tensorrt</code></a> module. <div class="codehilite"><pre><code># Load the converted model. root = tf.saved_model.load("tensorrt_embedding_model") trt_model_function = root.signatures["serving_default"] </code></pre></div> <hr /> <h2 id="build-lsh-tables-with-optimized-model">Build LSH tables with optimized model</h2> <div class="codehilite"><pre><code>warmup() training_files = zip(images, labels) lsh_builder_trt = BuildLSHTable(trt_model_function, concrete_function=True) lsh_builder_trt.train(training_files) </code></pre></div> Notice the difference in the wall time which is 13.1 seconds. Earlier, with the unoptimized model it was 54.1 seconds. We can take a closer look into one of the hash tables and get an idea of how they are represented. <div class="codehilite"><pre><code>idx = 0 for hash, entry in lsh_builder_trt.lsh.tables[0].table.items(): if idx == 5: break if len(entry) < 5: print(hash, entry) idx += 1 </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>145 [{'id_label': '3_4'}, {'id_label': '727_3'}] 5 [{'id_label': '12_4'}] 128 [{'id_label': '30_2'}, {'id_label': '480_2'}] 208 [{'id_label': '34_2'}, {'id_label': '132_2'}, {'id_label': '984_2'}] 188 [{'id_label': '42_0'}, {'id_label': '135_3'}, {'id_label': '436_3'}, {'id_label': '670_3'}] </code></pre></div> </div> <hr /> <h2 id="visualize-results-on-validation-images">Visualize results on validation images</h2> In this section we will first writing a couple of utility functions to visualize the similar image parsing process. Then we will benchmark the query performance of the models with and without optimization. First, we take 100 images from the validation set for testing purposes. <div class="codehilite"><pre><code>validation_images = [] validation_labels = [] for image, label in validation_ds.take(100): image = tf.image.resize(image, (224, 224)) validation_images.append(image.numpy()) validation_labels.append(label.numpy()) validation_images = np.array(validation_images) validation_labels = np.array(validation_labels) validation_images.shape, validation_labels.shape </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>((100, 224, 224, 3), (100,)) </code></pre></div> </div> Now we write our visualization utilities. <div class="codehilite"><pre><code>def plot_images(images, labels): plt.figure(figsize=(20, 10)) columns = 5 for (i, image) in enumerate(images): ax = plt.subplot(len(images) // columns + 1, columns, i + 1) if i == 0: ax.set_title("Query Image\n" + "Label: {}".format(labels[i])) else: ax.set_title("Similar Image # " + str(i) + "\nLabel: {}".format(labels[i])) plt.imshow(image.astype("int")) plt.axis("off") def visualize_lsh(lsh_class): idx = np.random.choice(len(validation_images)) image = validation_images[idx] label = validation_labels[idx] results = lsh_class.query(image) candidates = [] labels = [] overlaps = [] for idx, r in enumerate(sorted(results, key=results.get, reverse=True)): if idx == 4: break image_id, label = r.split("_")[0], r.split("_")[1] candidates.append(images[int(image_id)]) labels.append(label) overlaps.append(results[r]) candidates.insert(0, image) labels.insert(0, label) plot_images(candidates, labels) </code></pre></div> <h3 id="nontrt-model">Non-TRT model</h3> <div class="codehilite"><pre><code>for _ in range(5): visualize_lsh(lsh_builder) visualize_lsh(lsh_builder) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Matches: 507 Matches: 554 Matches: 438 Matches: 370 Matches: 407 Matches: 306 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/near_dup_search/near_dup_search_41_1.png" /> <img alt="png" src="/img/examples/vision/near_dup_search/near_dup_search_41_2.png" /> <img alt="png" src="/img/examples/vision/near_dup_search/near_dup_search_41_3.png" /> <img alt="png" src="/img/examples/vision/near_dup_search/near_dup_search_41_4.png" /> <img alt="png" src="/img/examples/vision/near_dup_search/near_dup_search_41_5.png" /> <img alt="png" src="/img/examples/vision/near_dup_search/near_dup_search_41_6.png" /> <h3 id="trt-model">TRT model</h3> <div class="codehilite"><pre><code>for _ in range(5): visualize_lsh(lsh_builder_trt) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Matches: 458 Matches: 181 Matches: 280 Matches: 280 Matches: 503 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/near_dup_search/near_dup_search_43_1.png" /> <img alt="png" src="/img/examples/vision/near_dup_search/near_dup_search_43_2.png" /> <img alt="png" src="/img/examples/vision/near_dup_search/near_dup_search_43_3.png" /> <img alt="png" src="/img/examples/vision/near_dup_search/near_dup_search_43_4.png" /> <img alt="png" src="/img/examples/vision/near_dup_search/near_dup_search_43_5.png" /> As you may have noticed, there are a couple of incorrect results. This can be mitigated in a few ways: <ul> <li>Better models for generating the initial embeddings especially for noisy samples. We can use techniques like <a href="https://arxiv.org/abs/1801.07698">ArcFace</a>, <a href="https://arxiv.org/abs/2004.11362">Supervised Contrastive Learning</a>, etc. that implicitly encourage better learning of representations for retrieval purposes.</li> <li>The trade-off between the number of tables and the reduction dimensionality is crucial and helps set the right recall required for your application.</li> </ul> <hr /> <h2 id="benchmarking-query-performance">Benchmarking query performance</h2> <div class="codehilite"><pre><code>def benchmark(lsh_class): warmup() start_time = time.time() for _ in range(1000): image = np.ones((1, 224, 224, 3)).astype("float32") _ = lsh_class.query(image, verbose=False) end_time = time.time() - start_time print(f"Time taken: {end_time:.3f}") benchmark(lsh_builder) benchmark(lsh_builder_trt) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Time taken: 54.359 Time taken: 13.963 </code></pre></div> </div> We can immediately notice a stark difference between the query performance of the two models. <hr /> <h2 id="final-remarks">Final remarks</h2> In this example, we explored the TensorRT framework from NVIDIA for optimizing our model. It's best suited for GPU-based inference servers. There are other choices for such frameworks that cater to different hardware platforms: <ul> <li><a href="https://www.tensorflow.org/lite">TensorFlow Lite</a> for mobile and edge devices.</li> <li><a href="hhttps://onnx.ai/">ONNX</a> for commodity CPU-based servers.</li> <li><a href="https://tvm.apache.org/">Apache TVM</a>, compiler for machine learning models covering various platforms.</li> </ul> Here are a few resources you might want to check out to learn more about applications based on vector similary search in general: <ul> <li><a href="http://ann-benchmarks.com/">ANN Benchmarks</a></li> <li><a href="https://arxiv.org/abs/1908.10396">Accelerating Large-Scale Inference with Anisotropic Vector Quantization(ScaNN)</a></li> <li><a href="https://arxiv.org/abs/1806.03198">Spreading vectors for similarity search</a></li> <li><a href="https://cloud.google.com/architecture/building-real-time-embeddings-similarity-matching-system">Building a real-time embeddings similarity matching system</a></li> </ul> </div> <div class='k-outline'> <div class='k-outline-depth-1'> <a href='#nearduplicate-image-search'>Near-duplicate image search</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#introduction'>Introduction</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#setup'>Setup</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#imports'>Imports</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#load-the-dataset-and-create-a-training-set-of-1000-images'>Load the dataset and create a training set of 1,000 images</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#load-a-pretrained-model'>Load a pre-trained model</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#create-an-embedding-model'>Create an embedding model</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#hashing-utilities'>Hashing utilities</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#query-utilities'>Query utilities</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#create-lsh-tables'>Create LSH tables</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#optimize-the-model-with-tensorrt'>Optimize the model with TensorRT</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#build-lsh-tables-with-optimized-model'>Build LSH tables with optimized model</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#visualize-results-on-validation-images'>Visualize results on validation images</a> </div> <div class='k-outline-depth-3'> <a href='#nontrt-model'>Non-TRT model</a> </div> <div class='k-outline-depth-3'> <a href='#trt-model'>TRT model</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#benchmarking-query-performance'>Benchmarking query performance</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#final-remarks'>Final remarks</a> </div> </div> </div> </div> </div> </body> <footer style="float: left; width: 100%; padding: 1em; border-top: solid 1px #bbb;"> <a href="https://policies.google.com/terms">Terms</a> | <a href="https://policies.google.com/privacy">Privacy</a> </footer> </html>

CINXE.COM

Near-duplicate image search