Build An Image Search Engine Using Python | Towards Data Science

We use search engine regularly. When we have queries, we can use the search engine like Google to retrieve the most relevant answer.
Most of the queries format is text-based. But not most of the time, the text is quite useful to find relevant answers.
For example, you want to search for a product on the internet, in this case, a t-shirt, but you don’t know the name of it. How could you find them? Well, you can write the description of that shirt.
The problem for using descriptions is that you will get wide varieties of products. And what makes it worse, they will be not similar with the product that you want to search, so you need a better way to retrieve them.
To solve it, we can use the image of the product, extracts its features, and use those features to retrieve similar products. We call this concept as content-based image retrieval.
In this article, I will show you on how to build an image search engine using Python.
Right before I explain to you on how to build an image retrieval using Python, let me explain to you the concept of content-based image retrieval.
Content-based image retrieval (CBIR) is a system for retrieving relevant images based on a given image. The system consists of an image query and an image database.
The process of the system will begin by extracting features on all images, whether it’s the query or the image database by using a feature extraction algorithm. Then, the system will calculate similarities between the query with all images on the database. At the end, the system will retrieve all the images that have a great similarity with the query.
Figure 1 describes the flow of the CBIR.
For extracting features, there are lots of options to choose. In this case, we can use a convolutional neural network (CNN) for extracting those features.
This model has a great capability to capture patterns than any other algorithms thanks to its convolutional layer that captures the neighbor for each instance of the data.
Convolutional Neural Network consists of layers, such as convolutional layer for feature extraction, pooling layer for sampling the features, and fully-connected layer for doing prediction.
For this case, we ignore most of the fully-connected layer and focus only to the feature extraction result only.
Figure 2 shows the example of the convolutional neural network architecture, which is the LeNet-5.
After we extract features, we calculate the distance between the query and all images. For doing this, we can use the Euclidean distance or l2 norm to measure it. If the number is getting smaller, the pair of image is similar to each other. The formula looks like this,
Where:
q = the query
img = the image
n = the number of feature vector element
i = the position of the vector.
Right after we know the concepts, now we can implement the system. To implement the CBIR, we will use Python as the programming language, Tensorflow for feature extraction, and Numpy for calculating the distance. If we summarize, here are the steps that we will do,
Download the dataset
Extract features from image database
Insert the query image and extract its features
Calculate the similarities with all images
Retrieve the most similar result
For the dataset, we will use a dataset from Kaggle called CBIR_dataset. You can access the dataset through the link here.
After we retrieve all of the images, now we can extract the features from all images using CNN and save those features on .npy format for later use. To specify the architecture, we will use VGG-16 architecture and pretrained weight from the ImageNet.
Side note, we can utilize GPU for extracting features on the images. But in this case, we use CPU only.
The code looks like this,
# Import the librariesfrom tensorflow.keras.preprocessing import imagefrom tensorflow.keras.applications.vgg16 import VGG16, preprocess_inputfrom tensorflow.keras.models import Modelfrom pathlib import Pathfrom PIL import Imageclass FeatureExtractor:    def __init__(self):        # Use VGG-16 as the architecture and ImageNet for the weight        base_model = VGG16(weights='imagenet')        # Customize the model to return features from fully-connected layer        self.model = Model(inputs=base_model.input, outputs=base_model.get_layer('fc1').output)    def extract(self, img):        # Resize the image        img = img.resize((224, 224))        # Convert the image color space        img = img.convert('RGB')        # Reformat the image        x = image.img_to_array(img)        x = np.expand_dims(x, axis=0)        x = preprocess_input(x)        # Extract Features        feature = self.model.predict(x)[0]        return feature / np.linalg.norm(feature)# Iterate through images (Change the path based on your image location)for img_path in sorted("<IMAGE DATABASE PATH LIST HERE>"):    print(img_path)    # Extract Features    feature = fe.extract(img=Image.open(img_path))    # Save the Numpy array (.npy) on designated path    feature_path = "<IMAGE FEATURE PATH HERE>.npy"    np.save(feature_path, feature)
After we extract features from all images, now we can try to retrieve similar images using a query. In this case, I will put the image of a lion that looks like this,
To retrieve similar images, we use code that looks like this,
# Import the librariesimport matplotlib.pyplot as pltimport numpy as np# Insert the image queryimg = Image.open("<IMAGE QUERY PATH HERE>")# Extract its featuresquery = fe.extract(img)# Calculate the similarity (distance) between imagesdists = np.linalg.norm(features - query, axis=1)# Extract 30 images that have lowest distanceids = np.argsort(dists)[:30]scores = [(dists[id], img_paths[id]) for id in ids]# Visualize the resultaxes=[]fig=plt.figure(figsize=(8,8))for a in range(5*6):    score = scores[a]    axes.append(fig.add_subplot(5, 6, a+1))    subplot_title=str(score[0])    axes[-1].set_title(subplot_title)      plt.axis('off')    plt.imshow(Image.open(score[1]))fig.tight_layout()plt.show()
Here is the final result,
Congratulations! You have made your own image search engine. Well, this is not pretty similar to Google, but at least you know the concept on how content-based image retrieval works.
If you have any knowledge on web development, maybe you can create a web application using a framework like Flask or Django to build your own search engine.
I hope this article increases your knowledge about image retrieval. If you are interested in this article, you can follow me on Medium. Also, you can connect with me on LinkedIn.
Thank you, and have a great day!
[1] Matsui Y. Simple Image Search Engine. https://github.com/matsui528/sis[2] Alkhawlani M, Elmogy M, El Bakry H. 2015. Text-based, Content-based, and Semantic-based Image Retrievals: A Survey. International Journal of Computer and Information Technology.[3] LeCun Y, Bottou L, Bengio Y, Haffner P. 1998. Gradient Based Learning Applied to Document Recognition. Proc. Of The IEEE.