Skip to main content

Download notebook · View source · Open in Colab

Find Similar Faces in CelebA Kaggle Dataset

Introduction​

ApertureDB combines metadata filtering and vector similarity search over stored images. This example ingests a subset of the CelebA dataset from Kaggle, generates face embeddings during ingestion, and retrieves images through both kinds of query. ParallelLoader sends the ingestion requests concurrently.

Prerequisites.​

  • A running ApertureDB instance and configured client connection.
  • The ApertureDB SDK with the PyTorch and FaceNet dependencies installed by the package command below.
  • Kaggle API credentials, configured as described in the Kaggle API README.
  • Public network access and several GB of disk space for the full CelebA download and model weights. The example ingests 10,000 images after downloading the archive; generating their embeddings can take substantial time.

Run the cells in order. Each run uses a unique descriptor set and data-source marker. The images and vectors remain available after the searches; the final cell removes only this run's data.

Install pytorch​

%pip install --quiet "aperturedb[complete]"

Dataset and embedding components​

from pathlib import Path
from urllib.request import urlopen
from uuid import uuid4
import hashlib
import sys

# This helper is part of the SDK examples, not the installed package.
helper_dir = Path("celeba_example")
helper_dir.mkdir(exist_ok=True)
helper_path = helper_dir / "CelebADataKaggle.py"
helper_url = (
    "https://raw.githubusercontent.com/aperture-data/aperturedb-python/"
    "827399dc18baa0fbfca9f7c4abdecc990211f9e8/examples/CelebADataKaggle.py"
)
with urlopen(helper_url, timeout=30) as response:
    helper_bytes = response.read()
assert hashlib.sha256(helper_bytes).hexdigest() == "12bbc8b2ba1aab5238c2521fdb02f51fc9f8858850b6e7f613dcaf2f07ff0b10"
helper_path.write_bytes(helper_bytes)
sys.path.insert(0, str(helper_dir.resolve()))

from aperturedb.transformers.facenet import generate_embedding
from aperturedb.transformers.facenet_pytorch_embeddings import FacenetPyTorchEmbeddings
from aperturedb.transformers.common_properties import CommonProperties
from CelebADataKaggle import CelebADataKaggle

run_id = "celeba-" + uuid4().hex
search_set_name = "facenet_" + run_id

Data ingestion​

DescriptorSets and Descriptors​

A DescriptorSet holds vectors in a common search space. In this example, the stored vectors and the query vector come from the same FaceNet model, so their distances are comparable.

A Descriptor is an individual vector with optional metadata. The query language also refers to these vectors as embeddings or feature vectors.

Create the descriptor set​

The descriptor set stores the 512-dimensional FaceNet vectors with the Flat engine and L2 metric.

from aperturedb.CommonLibrary import create_connector
from aperturedb.Utils import Utils

# Connect to the ApertureDB instance.
con = create_connector()

utils = Utils(con)

# Create a new empty descriptor set.
utils.add_descriptorset(search_set_name, 512,
    metric=["L2"], engine="Flat")

Load Kaggle dataset into ApertureDB.​

The loader reads 10,000 CelebA images from Kaggle and generates a FaceNet descriptor for each image. It stores the images and descriptors, linking each descriptor to its source image.

ParallelLoader batches and sends the ingestion requests concurrently. Its implementation is available in the SDK repository.

from aperturedb.ParallelLoader import ParallelLoader

# KaggleData downloads the full archive before the subset is selected.
source = CelebADataKaggle()
sample_count = 10000
# Reserve the next image as a query that is outside this run's ingested subset.
query_image_name = source.collection[sample_count]["image_id"]
query_image_path = Path(source.workdir) / "img_align_celeba/img_align_celeba" / query_image_name
source.collection = source.collection[:sample_count]
dataset = CommonProperties(source, adb_data_source=run_id)
dataset = FacenetPyTorchEmbeddings(dataset, search_set_name=search_set_name, client=con)
print("Images selected for ingestion:", len(dataset))

loader = ParallelLoader(con)
loader.ingest(dataset, stats=True)

Ingest with the adb CLI​

The CLI can run the generator with additional transformers:

adb ingest from-generator CelebADataKaggle.py --transformer common_properties --transformer image_properties --transformer clip_pytorch_embeddings --transformer facenet_pytorch_embeddings --sample-count 10000

Run this command from the celeba_example directory created above, where the helper was downloaded. It stores Kaggle images, generates CLIP and FaceNet embeddings, and adds image metadata. This is a separate ingestion alternative: it uses the transformers' default descriptor-set names, so the following cells, which use this notebook's unique run marker and set, apply to the SDK ingestion path.

Query examples​

CelebA provides image attributes such as Arched_Eyebrows, Attractive, Bags_Under_Eyes, and Bald. These attributes become image properties in ApertureDB; the complete list describes the dataset columns.

This query selects up to five images from this run whose Bald property is 1, then displays the images and their metadata.

from aperturedb import Images, Utils
import pandas as pd
from aperturedb.Constraints import Constraints

# Connect to ApertureDB.
con = create_connector()
utils = Utils.Utils(con)

print("Dataset marker:", run_id)

# Find the first 5 images with the 'Bald' attribute
images = Images.Images(con)
images.search(limit=5, constraints=Constraints().equal("adb_data_source", run_id).equal("Bald", 1))
images.display(show_bboxes=False)

pd.json_normalize(images.get_properties(images.get_props_names()).values())

The query image is encoded with the same FaceNet model used during ingestion. FindDescriptor selects five nearby vectors, and FindImage follows their connections to retrieve the source images in the same query.

Similarity depends on the embedding model and the images present in the dataset. The result shows matches within the ingested subset.

from IPython.display import display
from PIL import Image
import numpy as np
import cv2
import matplotlib.pyplot as plt

# Pick a query image that is not already in the dataset.
image_name = query_image_path
pilImage = Image.open(image_name)
display(pilImage.resize((int(pilImage.width * 0.3), int(pilImage.height * 0.3))))

# Generate a descriptor (embedding) from our query image.
embedding = generate_embedding(Image.open(image_name).convert('RGB'))
if not embedding.any():
    raise ValueError("FaceNet could not detect a face in the query image; choose another image.")

# This ApertureDB query finds images that are similary to the query image.
q = [{
    # Find descriptors similar to the input descriptor.
    "FindDescriptor": {
        "set": search_set_name, # Search in our CelebA descriptor set.
        "k_neighbors": 5, # Return the 5 nearest neighbors.
        "_ref": 1 # Assign a reference ID to this result that we can use it below.
    }
}, {
    # Retrieve the images associated with the results above.
    "FindImage": {
        # Find images connected to the descriptors returned above.
        "is_connected_to": {
            "ref": 1,
        },
        # Return binary image data.
        "blobs": True
    }
}]

# Run the query.
# As additional input, include the descriptor data generated from our query image above.
result, image_bytes = con.query(q, [embedding.cpu().detach().numpy().tobytes()])

# display the returned images
for img in image_bytes:
    nparr = np.frombuffer(img, dtype=np.uint8)
    image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    fig1, ax1 = plt.subplots()
    plt.imshow(image), plt.axis("off")

Cleanup​

Delete the descriptor set, bounding boxes, and images created by this run. Downloaded source files remain on disk.

con.query([{"DeleteDescriptorSet": {"with_name": search_set_name}}])
for kind in ("BoundingBox", "Image"):
    con.query([{"Delete" + kind: {"constraints": ["$adb_data_source", "==", run_id]}}])
    con.print_last_response()