Image Retrieval Using Pre-trained ResNet in 5 Steps

Got an image and wanna find similar images from your image folder or database? Image retrieval can be very useful in a lot applications. For example, a search engine with an image search function, or helping people find an image they are looking for but not quite sure where to start from.

In this article, we’ll use ResNet50 and K-Nearest Neighbours to develop a K-NN classifier. The dataset we use is Cifar-100 collected by Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. (Learning Multiple Layers of Features from Tiny Images, Alex Krizhevsky, 2009).

For the demo, we created a smaller dataset from the original dataset by randomly selecting 1 image from each class in the test set and 5 images from each class in the training set. Therefore, we have a dataset of 500 training images as candidates and 100 test images. The folder structure looks as follows:


This folder structure facilitates data loading with the torchvision ImageFolder API.

Now let’s get started with the exciting part!

  1. Load data
from torchvision import datasets, transforms

2. Load pre-trained ResNet Model

import torch
from torchvision import models

3. Calculate the features for the candidates

import numpy as np

4. Fit the candidates

from sklearn.neighbors import NearestNeighbors

5. Start predicting!

def predict(test_loader):
res = []
for i, data in enumerate(test_loader):
image, label = data
image =
output = net(image)
output = output.detach().cpu().numpy().reshape(-1,[1:]))

_, neighbors_idx = knn.kneighbors(output)
res.append((test_dataset.imgs[i][0], neighbors_idx[0],))

To visualise the results, you can use this helper function:

import matplotlib.pyplot as plt
from PIL import Image

And then plot some results!

test_image_1, candidates_idx = res[5]
plot_result(test_image_1, candidates_idx, candidates_path)

Let’s take a look at some results:

Example 1: camel

Top 10 nearest neighbours for query image “camel”.

Example 2: armchair

Top 10 nearest neighbours for query image “armchair”.

Example 3:

Top 10 nearest neighbours for query image “bowl”.

Example 4:

Top 10 nearest neighbours for query image “bed”.

Voilà! The results look quite good, as the nearest neighbours are actually quite similar to the query image. Of course, in some cases it works less well, such as the example with the pink bed. But at least it got a couch as a similar neighbour, which is kind of similar to a bed 😂

If you like my articles, don’t forget to give them a clap 👏 and share them with your friends and colleagues! Cheers! 😉



Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store