Skip to main content

Extract embeddings using CLIP model and Index in ApertureDB

Install prerequisites​

!pip install wheel ftfy regex tqdm
!pip install git+https://github.com/openai/CLIP.git

Common Imports​

from aperturedb.Connector import Connector
from aperturedb.Images import Images
from aperturedb.Query import Query
from aperturedb.Constraints import Constraints
from aperturedb.Utils import Utils
from aperturedb.ParallelQuery import execute_batch

from IPython.display import display
from tqdm import tqdm
from PIL import Image

import numpy as np
import clip
import torch
import cv2

Choose the relevant CLIPS model​

# Choose the model to be used.
descriptor_set = "ViT-B/16"
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load(descriptor_set, device=device)
clip.available_models()

Connect with ApertureDB​

# We will need this to query images and update embeddings (descriptors)
db = Connector(host="aperturedb", port=55555, user="admin", password="admin")

Enrich the existing images by connecting embeddings generated by CLIP​

This is to exemplify the enrichment loop and how ApertureDB's query mechanism can be used to arbitrarily decorate already ingested data

Warning This will delete existing descriptors

utils = Utils(db)  # ApertureDB util package for common functions like remove_descriptorset()

# !! This will remove the existing descriptor set along with the indexed descriptors
utils.remove_descriptorset(descriptor_set)
vector_length = 512
utils.add_descriptorset(descriptor_set, vector_length,
                           metric="CS", engine="FaissFlat")

# Query images, extract descriptor using CLIPS, index it back into the DB
query = [{
    "FindImage": {
        "limit": 10,
        
        "results": {
            "count": True,
            "list": ["_uniqueid"]
        },
        "blobs": True
    }
}]
r, _ = db.query(query, [])
count = r[0]["FindImage"]["count"]
print(count)

params_t = query[0]["FindImage"]
for i in tqdm(range(0, count, 200)):
    params_t["offset"] = i
    q = [{"FindImage": params_t}]
    r, blobs = db.query(q, [])
    update_q = [ 
        {
            "FindImage": {
                "_ref": i+1,
                "constraints": {
                    "_uniqueid": ["==", e["_uniqueid"]]
                }
            }
        } for i, e in enumerate(r[0]["FindImage"]["entities"])]
    
    descs = []
    for img in update_q:
        descs.append(
            {
                "AddDescriptor": {
                    "set": descriptor_set,
                    "connect":{
                        "ref": img["FindImage"]["_ref"]
                    }
                }
            })
    update_q.extend(descs)
    
    desc_blobs = []
    for b in blobs:
        nparr = np.fromstring(b, np.uint8)
        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = preprocess(Image.fromarray(image)).unsqueeze(0).to(device)
        image_features = model.encode_image(image)
        
        if device == "cuda":
            image_features = image_features.float()
            desc_blobs.append(image_features.detach().cpu().numpy().tobytes())
        else:
            desc_blobs.append(image_features.detach().numpy().tobytes())

    _, r, b = execute_batch(db=db, q=update_q, blobs=desc_blobs)