Build an Image Search Engine With Python & MongoDB

1
pip install -r requirements.txt

1
# Set the value below to your cluster:
2
export MONGODB_URI="mongodb+srv://image_search_demo:my_password_not_yours@sandbox.abcde.mongodb.net/image_search_demo?retryWrites=true&w=majority"

1
jupyter notebook "Image Search.ipynb"

1
client = MongoClient(MONGODB_URI)
2
db = client.get_database(DATABASE_NAME)
3
4
# Ensure the collection exists, because otherwise you can't add a search index to it.
5
try:
6
    db.create_collection(IMAGE_COLLECTION_NAME)
7
except CollectionInvalid:
8
    # This is raised when the collection already exists.
9
    print("Images collection already exists")
10
11
# Add a search index (if it doesn't already exist):
12
collection = db.get_collection(IMAGE_COLLECTION_NAME)
13
if len(list(collection.list_search_indexes(name="default"))) == 0:
14
    print("Creating search index...")
15
    collection.create_search_index(
16
        SearchIndexModel(
17
            {
18
                "mappings": {
19
                    "dynamic": True,
20
                    "fields": {
21
                        "embedding": {
22
                            "dimensions": 768,
23
                            "similarity": "cosine",
24
                            "type": "knnVector",
25
                        }
26
                    },
27
                }
28
            },
29
            name="default",
30
        )
31
    )
32
    print("Done.")
33
else:
34
    print("Vector search index already exists")

1
{
2
    "mappings": {
3
        "dynamic": True,
4
        "fields": {
5
            "embedding": {
6
                "dimensions": 768,
7
                "similarity": "cosine",
8
                "type": "knnVector",
9
            }
10
        },
11
    }
12
}

1
# Load CLIP model.
2
# This may print out warnings, which can be ignored.
3
model = SentenceTransformer("clip-ViT-L-14")

1
emb = model.encode(Image.open(path))

1
collection.insert_one(
2
    {
3
        "_id": re.sub("images/", "", path),
4
        "embedding": emb.tolist(),
5
    }
6
)

1
def image_search(search_phrase):
2
    """
3
    Use MongoDB Vector Search to search for a matching image.
4
5
    The search_phrase is first converted to a vector embedding using
6
    the model loaded earlier in the Jupyter notebook. The vector is then used
7
    to search MongoDB for matching images.
8
    """
9
    emb = model.encode(search_phrase)
10
    cursor = collection.aggregate(
11
        [
12
            {
13
                "$vectorSearch": {
14
                    "index": "default",
15
                    "path": "embedding",
16
                    "queryVector": emb.tolist(),
17
                    "numCandidates": 100,
18
                    "limit": 9,
19
                }
20
            },
21
            {"$project": {"_id": 1, "score": {"$meta": "vectorSearchScore"}}},
22
        ]
23
    )
24
25
    return list(cursor)

1
display_images(image_search("sharks in the water"))

1
display_images(image_search("corgi in the snow"))

1
display_images(image_search("childhood"))

1
display_images(image_search("ennui"))