AI-Powered Image Retrieval Using Vector Databases

---

Overview

A content-based image retrieval system that uses deep learning and vector similarity search to find visually similar images. Upload any image and the system retrieves the most similar images from a pre-indexed collection of 2,020 images across 101 categories — in milliseconds.

---

Features

• Upload any image and retrieve visually similar results instantly
• Fine-tuned ResNet18 model for 128-dimensional feature embeddings
• FAISS vector database for millisecond similarity search
• Interactive Streamlit web interface with adjustable settings
• 101 image categories from the Caltech-101 dataset
• GPU-accelerated inference (auto-detected, falls back to CPU)

---

🏗️ Architecture

Query Image → ResNet18 Feature Extractor → 128-dim Embedding → FAISS Inner Product Search → Top-K Similar Images + Scores

Step	Component	Detail
1	Input	Upload any JPG/PNG image
2	Preprocessing	Resize 128×128, Normalize (ImageNet stats)
3	Feature Extraction	ResNet18 backbone → 512-dim features
4	Embedding	Linear projection → 128-dim normalized vector
5	Vector Search	FAISS IndexFlatIP cosine similarity
6	Output	Top-K similar images with scores

Model Design

• Backbone: ResNet18 (pretrained on ImageNet)
• Fine-tuning: Last 2 blocks unfrozen for domain adaptation
• Embedding Layer: Linear(512→128) + BatchNorm + ReLU
• Similarity Metric: Cosine similarity via normalized embeddings
• Vector Database: FAISS IndexFlatIP (Inner Product)

📁 Project Structure

AI-Powered-Image-Retrieval-Vector-Database/
├── app.py                              # Main Streamlit application
├── precompute.sh                       # Feature extraction pipeline
├── train_val.json                      # Dataset metadata
├── web-app.png                         # App demo screenshot
├── src/
│   ├── __init__.py
│   └── model.py                        # ResNet18 transfer learning model
└── utils/
    ├── __init__.py
    ├── image_utils.py                  # Image preprocessing & feature extraction
    ├── faiss_utils.py                  # FAISS index build/load/search
    ├── display_utils.py                # Streamlit results display
    ├── data_utils.py                   # Data loading utilities
    └── precompute_features.py          # Gallery feature precomputation

Getting Started

Prerequisites

pip install torch torchvision streamlit faiss-cpu pillow numpy tqdm requests

Step 1 — Place Your Trained Model

mkdir -p weights data
# Place model.pth in the weights/ directory

Step 2 — Precompute Features & Build FAISS Index

# Default: 20 images per class
sh precompute.sh

# Custom number of images per class
sh precompute.sh --num-per-class 50

Step 3 — Launch the Web App

streamlit run app.py

---

App Features

Feature	Description
Image Upload	Supports JPG, JPEG, PNG formats
Number of Results	Adjustable slider (1–10 results)
Similarity Scores	Cosine similarity score shown per result
Category Display	Shows all 101 available categories
GPU Support	Auto-detects and uses GPU if available
Custom Paths	Override FAISS index and features paths

---

Tech Stack

Component	Technology
Deep Learning	PyTorch, TorchVision
Model	ResNet18 (Transfer Learning)
Vector Database	FAISS (Facebook AI Similarity Search)
Web Interface	Streamlit
Image Processing	PIL, torchvision transforms
Dataset	Caltech-101 (101 categories)

---

Performance

• Index Size: 2,020 images across 101 categories
• Embedding Dimension: 128
• Search Speed: Millisecond-level similarity search via FAISS
• Similarity Metric: Cosine similarity (normalized inner product)
• Top Results: Configurable (default: 5)

---

How It Works

1. Feature Extraction — ResNet18 backbone extracts 512-dim features, projected to 128-dim normalized embeddings
2. Indexing — All gallery image embeddings stored in FAISS flat index
3. Query — Uploaded image processed through same pipeline
4. Search — FAISS performs exact inner product search in milliseconds
5. Display — Top-K results shown with category labels and similarity scores