main.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Image recognition with Python, OpenCV, OpenAI CLIP model and PostgreSQL `pgvector` \n",
    "\n",
    "This repository contains the working code for the example in the [blog post](https://aiven.io/developer/find-faces-with-pgvector)\n",
    "\n",
    "The below is the overall flow:\n",
    "\n",
    "![Overall flow](entire_flow.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 0: Install requirements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Face recognition\n",
    "\n",
    "Detect the faces from the [test-image](test-image.png) picture and store them under the `stored-faces` folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# import the OpenCV library - it's called cv2\n",
    "import cv2\n",
    "\n",
    "# load the Haar Cascade algorithm from the XML file into OpenCV\n",
    "haar_cascade = cv2.CascadeClassifier(\"haarcascade_frontalface_default.xml\")\n",
    "\n",
    "# read the test image as grayscale\n",
    "gray_img = cv2.imread(\"test-image.png\", cv2.IMREAD_GRAYSCALE)\n",
    "\n",
    "# find the faces in that image\n",
    "# this gives back an array of the x,y location of each face, and its width and height\n",
    "faces = haar_cascade.detectMultiScale(\n",
    "    gray_img, scaleFactor=1.05, minNeighbors=2, minSize=(100, 100)\n",
    ")\n",
    "\n",
    "# make sure the directory we're going to write to actually exists\n",
    "os.makedirs('stored-faces', exist_ok=True)\n",
    "\n",
    "i = 0\n",
    "# write all the faces out to files\n",
    "# for each face we found:\n",
    "for x, y, w, h in faces:\n",
    "    # crop the image to select only the face\n",
    "    cropped_image = gray_img[y : y + h, x : x + w]\n",
    "    # make up a filename for that face - we're just going to number them\n",
    "    target_file_name = f'stored-faces/{i}.jpg'\n",
    "    # report each file so we can tell we're doing something\n",
    "    print(target_file_name)\n",
    "    # and write the cropped face to the file\n",
    "    cv2.imwrite(\n",
    "        target_file_name,\n",
    "        cropped_image,\n",
    "    )\n",
    "    i = i + 1;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Embeddings Calculation\n",
    "\n",
    "Calculate embeddings from the faces and pushing to PostgreSQL\n",
    "\n",
    "Remember you need to have enabled pgvector in PostgreSQL:\n",
    "```sql\n",
    "CREATE EXTENSION vector;\n",
    "```\n",
    "and created the table to write to\n",
    "```sql\n",
    "CREATE TABLE pictures (picture text PRIMARY KEY, embedding vector(768));\n",
    "```\n",
    "\n",
    "In the following code, you'll need to change the `<SERVICE_URI>` string to the PostgreSQL Service URI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import the other required libraries\n",
    "import numpy as np\n",
    "from imgbeddings import imgbeddings\n",
    "from PIL import Image\n",
    "import psycopg2\n",
    "\n",
    "# connect to the database - replace <SERVICE URI> with the actual service URI\n",
    "# remember it needs to have the table defined already\n",
    "conn = psycopg2.connect(\"<SERVICE_URI>\")\n",
    "\n",
    "# load `imgbeddings` so we can calculate embeddings\n",
    "ibed = imgbeddings()\n",
    "\n",
    "for filename in os.listdir(\"stored-faces\"):\n",
    "    # read the image\n",
    "    img = Image.open(\"stored-faces/\" + filename)\n",
    "    # calculate the embedding for this face\n",
    "    embedding = ibed.to_embeddings(img)[0]\n",
    "    # and write it to PostgreSQL\n",
    "    cur = conn.cursor()\n",
    "    cur.execute(\"INSERT INTO pictures values (%s,%s)\", (filename, embedding.tolist()))\n",
    "    print(filename)\n",
    "conn.commit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Calculate embeddings on a new picture\n",
    "\n",
    "Find the face and calculate the embeddings on the picture `solo-image.png` used for research"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read the image\n",
    "img = Image.open(\"solo-image.png\")\n",
    "# calculating the embedding for that face\n",
    "embedding = ibed.to_embeddings(img)[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Find similar images by querying the Postgresql database using pgvector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image, display\n",
    "\n",
    "cur = conn.cursor()\n",
    "string_representation = \"[\"+ \",\".join(str(x) for x in embedding.tolist()) +\"]\"\n",
    "cur.execute(\"SELECT * FROM pictures ORDER BY embedding <-> %s LIMIT 1;\", (string_representation,))\n",
    "rows = cur.fetchall()\n",
    "for row in rows:\n",
    "    display(Image(filename=\"stored-faces/\"+row[0]))\n",
    "cur.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}