supabase/examples/ai/face_similarity.ipynb

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

In [1]:

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

Load the Dataset

First, we load a dataset of celebrity faces.

In [ ]:

from datasets import load_dataset

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

In [3]:

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

Out [3]:

{'adult': False,
 'also_known_as': "['Morgan Porterfield Freeman Jr.', 'Morgan J. Freeman', 'مورغان فريمان', '모건 프리먼', 'モーガン・フリーマン', 'Морган Фриман', 'Морган Фримен', 'มอร์แกน ฟรีแมน', '摩根·弗里曼', 'Μόργκαν Φρίμαν', 'مورگان فریمن', 'Морґан Фрімен', 'Μόργκαν Πόρτερφιλντ Φρίμαν Τζούνιορ']",
 'biography': "Morgan Freeman (born June 1, 1937) is an American actor, director, and narrator. Noted for his distinctive deep voice, Freeman is known for his various roles in a wide variety of film genres. Throughout his career spanning over five decades, he has received multiple accolades, including an Academy Award, a Screen Actors Guild Award, and a Golden Globe Award.\n\nBorn in Memphis, Tennessee, Freeman was raised in Mississippi where he began acting in school plays. He studied theatre arts in Los Angeles and appeared in stage productions in his early career. He rose to fame in the 1970s for his role in the children's television series The Electric Company. Freeman then appeared in the Shakespearean plays Coriolanus and Julius Caesar, the former of which earned him an Obie Award. His breakout role was in Street Smart (1987), playing a hustler, which earned him an Academy Award nomination for Best Supporting Actor. He achieved further stardom in Glory, the biographical drama Lean on Me, and comedy-drama Driving Miss Daisy (all 1989), the latter of which garnered him his first Academy Award nomination for Best Actor.\n\nIn 1992, Freeman starred alongside Clint Eastwood in the western revenge film Unforgiven; this would be the first of several collaborations with Eastwood. In 1994, he starred in the prison drama The Shawshank Redemption for which he received another Academy Award nomination. Freeman also starred in David Fincher's crime thriller Se7en (1995), and Steven Spielberg's historical drama Amistad (1997). Freeman won the Academy Award for Best Supporting Actor for his performance in Clint Eastwood's 2004 sports drama Million Dollar Baby. In 2009, he received his fifth Oscar nomination for playing former South African President Nelson Mandela in Eastwood's Invictus. Freeman is also known for his performance as Lucius Fox in Christopher Nolan's The Dark Knight Trilogy (2005–2012).\n\nIn addition to acting, Freeman has directed the drama Bopha! (1993). He also founded film production company Revelations Entertainment with business partner Lori McCreary. He is the recipient of the Kennedy Center Honor, the AFI Life Achievement Award, the Cecil B. DeMille Award, and the Screen Actors Guild Life Achievement Award. For his performances in theatrical productions, he has won three Obie Awards, one of the most prestigious honors for recognizing excellence in theatre.\n\nDescription above from the Wikipedia article Morgan Freeman, licensed under CC-BY-SA, full list of contributors on Wikipedia.",
 'birthday': '1937-06-01',
 'deathday': None,
 'gender': 2,
 'homepage': None,
 'id': 192,
 'imdb_id': 'nm0000151',
 'known_for_department': 'Acting',
 'name': 'Morgan Freeman',
 'place_of_birth': 'Memphis, Tennessee, USA',
 'popularity': 95.033,
 'profile_path': '905k0RFzH0Kd6gx8oSxRdnr6FL.jpg',
 'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=421x632>}

In [4]:

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

Out [4]:

Embedding Model

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

In [5]:

import numpy as np
import face_recognition

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

Out [5]:

array([-0.13341002,  0.13275896,  0.16850984,  0.06553721,  0.01849797,
       -0.10046144,  0.08926678, -0.10535249,  0.09548245,  0.01443161,
        0.29246074, -0.05078556, -0.09225237, -0.20251125,  0.06319536,
        0.13064152, -0.13121371, -0.14068669, -0.14239085, -0.07584596,
       -0.00425112,  0.05749821,  0.03404564,  0.00690284, -0.05748112,
       -0.30998233, -0.05940549, -0.12016021,  0.13422301, -0.09362517,
        0.00096729,  0.00121274, -0.28880239, -0.10776889, -0.02401398,
        0.02188838,  0.04223322, -0.05020416,  0.13429314, -0.06062891,
       -0.15768167,  0.04524709,  0.09769595,  0.18990684,  0.16210739,
       -0.04965826, -0.03221413, -0.03109819,  0.01347625, -0.18865313,
        0.04968366,  0.08835649,  0.12916593,  0.08170474, -0.01250134,
       -0.10153808, -0.04734591,  0.05344258, -0.15603815,  0.05757158,
        0.03810269, -0.07773421, -0.11656785,  0.00838896,  0.11370895,
        0.10145531, -0.01232528, -0.19529839,  0.07950203, -0.16160172,
       -0.0103126 ,  0.10836543, -0.07068133, -0.08057274, -0.2873418 ,
        0.11391003,  0.38112539,  0.02767274, -0.24450374,  0.02405314,
       -0.17997386,  0.0244493 ,  0.01259282,  0.02163479, -0.0256572 ,
        0.00239573, -0.15525642,  0.02548696,  0.13868049, -0.01777178,
       -0.00228858,  0.19361472,  0.00259691, -0.00303981,  0.00783446,
       -0.01679439, -0.0077012 , -0.02348143, -0.09466386, -0.06156742,
        0.02911264, -0.00982514, -0.04479714,  0.08873156, -0.18456687,
        0.09431744,  0.06609484, -0.02497767,  0.09436633,  0.0584843 ,
        0.02720976, -0.12375437,  0.07636171, -0.17666884,  0.1997328 ,
        0.16992269, -0.05255131,  0.13182855,  0.00585927,  0.1533282 ,
       -0.03610092,  0.01495183, -0.07237726, -0.01772323,  0.03974599,
       -0.00466948, -0.03418808,  0.02540907])

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.

In [ ]:

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.

In [ ]:

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

In [15]:

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.

In [16]:

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.

In [17]:

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))

In [18]:

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

Out [18]:

In [19]:

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

Out [19]:

In [20]:

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

Out [20]: