Face Similarity Search

In this example we'll use PostgreSQL + pgvectors similarity search using the vecs library to identify the celebrities a person looks most similar to.

We'll start by loading a dataset of celebrity faces. Then we'll create embeddings for the faces using python's face_recognition library and store them in PostgreSQL with vecs. Finally we'll query the database with a user defined face to see which celebrities they look most like.

Setup GPU on Google Colab

If you are running this in Google Colab, you must first ensure that the GPU is enabled. You can set this by navigating to Runtime > Change runtime type, then choose one of the available GPUs under Hardware Accelerator.

Install Dependencies

python

!pip install -qU vecs datasets face_recognition flupy tqdm numpy matplotlib

Load the Dataset

First, we load a dataset of celebrity faces.

python

from datasets import load_dataset

people = load_dataset("ashraq/tmdb-people-image", split='train')
people

python

# Look at an example record from the dataset
person = people[15]
person

python

person['image'].resize((210, 315))

Embedding Model

Next, we can use face_recognition to produce a face embedding for each row (person) in the dataset.

python

import numpy as np
import face_recognition

# Display an example embedding
face_recognition.face_encodings(np.array(person['image']))[0]

Initialize the Vecs Collection

The vecs library wraps a pythonic interface around PostgreSQL and pgvector. A collection in vecs maps 1:1 with a PostgreSQL table.

First you will need to establish a connection to your database. You can find the Postgres connection string in the project connect page of your Supabase project.

Note: SQLAlchemy requires the connection string to start with postgresql:// (instead of postgres://). Don't forget to rename this after copying the string from the dashboard.

Note: You must use the "connection pooling" string (domain ending in *.pooler.supabase.com) with Google Colab since Colab does not support IPv6.

This will also work with any other Postgres provider that supports pgvector.

python

import vecs
DB_CONNECTION = "postgresql://postgres:password@localhost:5611/vecs_db"

# create vector store client
vx = vecs.create_client(DB_CONNECTION)

# create a PostgreSQL/pgvector table named "faces" to contain the face embeddings
faces = vx.get_or_create_collection(name="faces", dimension=128)

Create Embeddings for Each Face

Now we can iterate over the dataset, producing embeddings for the faces.

Note that it could take a few hours to produce all of the embeddings. If you're just testing it out, feel free to interrupt the loop after a few hundred iterations and continue with the next step.

python

from typing import List, Dict, Tuple
from PIL import Image
from flupy import flu
import numpy as np
from tqdm import tqdm

# Records we'll insert into the database
records: List[Tuple[str, np.ndarray, Dict]] = []

# Iterate over the dataset in chunks
for ix, person in tqdm(enumerate(people)):

    # Extract the person's image
    person_image = person['image']

    # Some of the images are grayscale with a single image channel
    # We'll normalize the image set by converting those to 3 channel RBG format
    if person_image.mode == 'L':
        # Extract the available channel
        L_channel = np.array(person_image)

        # Repeat that channel 3 times for R G B
        person_image = Image.fromarray(
            np.moveaxis(np.stack([L_channel, L_channel, L_channel]), 0, -1)
        )

    # Create embeddings for current chunk
    embeddings = face_recognition.face_encodings(np.array(person_image))

    # In some cases the face is too obscured to be detectable and no embedding
    # is produced. We'll skip those cases
    if len(embeddings) == 1:
        embedding = embeddings[0]
        records.append((
            f"{ix}",
            embedding,
            {k: v for k, v in person.items() if k != 'image'}
        ))

Insert the Embeddings into Postgres

python

faces.upsert(records)

Index the Collection

Indexing the collection creates an index on the vector column in Postgres that significantly improves performance of similarity queries.

python

faces.create_index()

Search for Similar Faces

Finally we can load a user defined face and search the database for other similar faces to find their look alikes. For simplicity, we'll grab a random face from the dataset as our query but it can be substituted for your own image.

Example Results

We'll create helper functions to display search results and try it out on several celebrities. Since our query faces are also in the dataset, the query face is the first in the result output.

python

from IPython.core.display import HTML
from PIL import Image, ImageDraw, ImageFont
import matplotlib.font_manager as fm
from typing import Dict, Any

def add_label(img, label_text, label_height=30):
    # Set the font and size
    font_path = fm.findfont(fm.FontProperties(family='Arial'))
    font = ImageFont.truetype(font_path, 15)
    
    # Create a new image with a white background
    label_img = Image.new('RGB', (img.width, label_height), color = (255, 255, 255))
    d = ImageDraw.Draw(label_img)

    # Calculate the width and height of the text to center it
    text_bbox = d.textbbox((0, 0), label_text, font)
    text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
    text_x = (label_img.width - text_width) // 2
    text_y = (label_img.height - text_height) // 2

    # Add the text to the label image
    d.text((text_x, text_y), label_text, fill=(0,0,0), font=font)

    # Concatenate the original image with the label image
    img_with_label = Image.new('RGB', (img.width, img.height + label_height))
    img_with_label.paste(img, (0, 0))
    img_with_label.paste(label_img, (0, img.height))

    return img_with_label

def resize_for_output(person_image):
    return person_image.resize((150, 220))
    
def render_similar_faces(person_image: Image) -> Image:
    # create query face embedding
    face_embedding = face_recognition.face_encodings(np.array(person_image))[0]
    
    # query database for similar results
    result = faces.query(face_embedding, limit=5, include_metadata=True)   

    captioned_images = [
        add_label(
            resize_for_output(person_image),
            "Query Image"
        ),
        Image.fromarray(255*np.ones((250,30,3), np.uint8))
    ]

    for person_id, person_metadata in result:
        result_person = people[int(person_id)]
        result_image = result_person['image']
        captioned_images.append(add_label(resize_for_output(result_person['image']), person_metadata["name"]))

    return Image.fromarray(np.hstack(captioned_images))

python

render_similar_faces(
    person_image=people[1014]['image']
)

python

render_similar_faces(person_image=people[15]['image'])

python

render_similar_faces(person_image=people[188]['image'])