Content-Based Image Retrieval

Module: Computer Vision | Difficulty: Intermediate

Similarity Metrics

Cosine Similarity

Euclidean Distance

Triplet Loss

Approximate Nearest Neighbor (ANN)

• LSH: Locality-sensitive hashing
• IVF: Inverted file index
• HNSW: Hierarchical navigable small world

Mean Average Precision (mAP)

import torch
import torch.nn as nn
import numpy as np

class ImageRetrieval:
    def __init__(self, model, device='cuda'):
        self.model = model.eval()
        self.device = device
        self.database = []
        self.labels = []
    
    def build_index(self, dataloader):
        for images, labels in dataloader:
            with torch.no_grad():
                feats = self.model(images.to(self.device))
                feats = nn.functional.normalize(feats, dim=1)
                self.database.append(feats.cpu())
                self.labels.extend(labels)
        self.database = torch.cat(self.database)
    
    def query(self, q_img, top_k=10):
        with torch.no_grad():
            q_feat = self.model(q_img.to(self.device))
            q_feat = nn.functional.normalize(q_feat, dim=1)
        dists = torch.cdist(q_feat.cpu(), self.database)
        indices = dists.topk(top_k, dim=1, largest=False).indices
        return [(self.labels[i.item()], dists[0, i].item()) for i in indices[0]]

Key Takeaways

• Metric learning optimizes embedding space for similarity search
• ANN indexes enable billion-scale retrieval in milliseconds
• Evaluation uses mAP across query sets