Skip to content

Commit b8c957d

Browse files
authored
example: add the example of face recognition (#780)
* example: add the example of face recognition * docs: fix typo
1 parent fdf5c2e commit b8c957d

File tree

9 files changed

+193
-10
lines changed

9 files changed

+193
-10
lines changed

README.md

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -27,18 +27,18 @@ Ultra performant data transformation framework for AI, with core engine written
2727
⭐ Drop a star to help us grow!
2828

2929
<div align="center">
30-
30+
3131
<!-- Keep these links. Translations will automatically update with the README. -->
32-
[Deutsch](https://readme-i18n.com/cocoindex-io/cocoindex?lang=de) |
33-
[English](https://readme-i18n.com/cocoindex-io/cocoindex?lang=en) |
34-
[Español](https://readme-i18n.com/cocoindex-io/cocoindex?lang=es) |
35-
[français](https://readme-i18n.com/cocoindex-io/cocoindex?lang=fr) |
36-
[日本語](https://readme-i18n.com/cocoindex-io/cocoindex?lang=ja) |
37-
[한국어](https://readme-i18n.com/cocoindex-io/cocoindex?lang=ko) |
38-
[Português](https://readme-i18n.com/cocoindex-io/cocoindex?lang=pt) |
39-
[Русский](https://readme-i18n.com/cocoindex-io/cocoindex?lang=ru) |
32+
[Deutsch](https://readme-i18n.com/cocoindex-io/cocoindex?lang=de) |
33+
[English](https://readme-i18n.com/cocoindex-io/cocoindex?lang=en) |
34+
[Español](https://readme-i18n.com/cocoindex-io/cocoindex?lang=es) |
35+
[français](https://readme-i18n.com/cocoindex-io/cocoindex?lang=fr) |
36+
[日本語](https://readme-i18n.com/cocoindex-io/cocoindex?lang=ja) |
37+
[한국어](https://readme-i18n.com/cocoindex-io/cocoindex?lang=ko) |
38+
[Português](https://readme-i18n.com/cocoindex-io/cocoindex?lang=pt) |
39+
[Русский](https://readme-i18n.com/cocoindex-io/cocoindex?lang=ru) |
4040
[中文](https://readme-i18n.com/cocoindex-io/cocoindex?lang=zh)
41-
41+
4242
</div>
4343

4444
</br>
@@ -183,6 +183,7 @@ It defines an index flow like this:
183183
| [FastAPI Server with Docker](examples/fastapi_server_docker) | Run the semantic search server in a Dockerized FastAPI setup |
184184
| [Product Recommendation](examples/product_recommendation) | Build real-time product recommendations with LLM and graph database|
185185
| [Image Search with Vision API](examples/image_search) | Generates detailed captions for images using a vision model, embeds them, enables live-updating semantic search via FastAPI and served on a React frontend|
186+
| [Face Recognition](examples/face_recognition) | Recognize faces in images and build embedding index |
186187
| [Paper Metadata](examples/paper_metadata) | Index papers in PDF files, and build metadata tables for each paper |
187188

188189
More coming and stay tuned 👀!

examples/face_recognition/.env

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
# Postgres database address for cocoindex
2+
COCOINDEX_DATABASE_URL=postgres://cocoindex:cocoindex@localhost/cocoindex
Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
# Recognize faces in images and build embedding index
2+
[![GitHub](https://img.shields.io/github/stars/cocoindex-io/cocoindex?color=5B5BD6)](https://github.com/cocoindex-io/cocoindex)
3+
4+
5+
In this example, we will recognize faces in images and build embedding index.
6+
7+
We appreciate a star ⭐ at [CocoIndex Github](https://github.com/cocoindex-io/cocoindex) if this is helpful.
8+
9+
## Steps
10+
### Indexing Flow
11+
12+
1. We will ingest a list of images.
13+
2. For each image, we:
14+
- Extract faces from the image.
15+
- Compute embeddings for each face.
16+
3. We will export to the following tables in Postgres with PGVector:
17+
- Filename, rect, embedding for each face.
18+
19+
20+
## Prerequisite
21+
22+
1. [Install Postgres](https://cocoindex.io/docs/getting_started/installation#-install-postgres) if you don't have one.
23+
24+
2. dependencies:
25+
26+
```bash
27+
pip install -e .
28+
```
29+
30+
## Run
31+
32+
Update index, which will also setup the tables at the first time:
33+
34+
```bash
35+
cocoindex update --setup main.py
36+
```
37+
38+
You can also run the command with `-L`, which will watch for file changes and update the index automatically.
39+
40+
```bash
41+
cocoindex update --setup -L main.py
42+
```
43+
44+
## CocoInsight
45+
I used CocoInsight (Free beta now) to troubleshoot the index generation and understand the data lineage of the pipeline. It just connects to your local CocoIndex server, with zero pipeline data retention. Run following command to start CocoInsight:
46+
47+
```
48+
cocoindex server -ci main.py
49+
```
50+
51+
Then open the CocoInsight UI at [https://cocoindex.io/cocoinsight](https://cocoindex.io/cocoinsight).
146 KB
Loading
786 KB
Loading
173 KB
Loading
36.9 KB
Loading

examples/face_recognition/main.py

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
import cocoindex
2+
import io
3+
import dataclasses
4+
import datetime
5+
import typing
6+
7+
import face_recognition
8+
from PIL import Image
9+
import numpy as np
10+
11+
12+
@dataclasses.dataclass
13+
class ImageRect:
14+
top: int
15+
left: int
16+
bottom: int
17+
right: int
18+
19+
20+
@dataclasses.dataclass
21+
class FaceBase:
22+
"""A face in an image."""
23+
24+
rect: ImageRect
25+
image: bytes
26+
27+
28+
MAX_IMAGE_WIDTH = 1280
29+
30+
31+
@cocoindex.op.function(cache=True, behavior_version=1, gpu=True)
32+
def extract_faces(content: bytes) -> list[FaceBase]:
33+
"""Extract the first pages of a PDF."""
34+
orig_img = Image.open(io.BytesIO(content)).convert("RGB")
35+
36+
# The model is too slow on large images, so we resize them if too large.
37+
if orig_img.width > MAX_IMAGE_WIDTH:
38+
ratio = orig_img.width * 1.0 / MAX_IMAGE_WIDTH
39+
img = orig_img.resize(
40+
(MAX_IMAGE_WIDTH, int(orig_img.height / ratio)),
41+
resample=Image.Resampling.BICUBIC,
42+
)
43+
else:
44+
ratio = 1.0
45+
img = orig_img
46+
47+
# Extract face locations.
48+
locs = face_recognition.face_locations(np.array(img), model="cnn")
49+
50+
faces: list[FaceBase] = []
51+
for top, right, bottom, left in locs:
52+
rect = ImageRect(
53+
left=int(left * ratio),
54+
top=int(top * ratio),
55+
right=int(right * ratio),
56+
bottom=int(bottom * ratio),
57+
)
58+
59+
# Crop the face and save it as a PNG.
60+
buf = io.BytesIO()
61+
orig_img.crop((rect.left, rect.top, rect.right, rect.bottom)).save(
62+
buf, format="PNG"
63+
)
64+
face = buf.getvalue()
65+
faces.append(FaceBase(rect, face))
66+
67+
return faces
68+
69+
70+
@cocoindex.op.function(cache=True, behavior_version=1, gpu=True)
71+
def extract_face_embedding(
72+
face: bytes,
73+
) -> cocoindex.Vector[cocoindex.Float32, typing.Literal[128]]:
74+
"""Extract the embedding of a face."""
75+
img = Image.open(io.BytesIO(face)).convert("RGB")
76+
embedding = face_recognition.face_encodings(
77+
np.array(img),
78+
known_face_locations=[(0, img.width - 1, img.height - 1, 0)],
79+
)[0]
80+
return embedding
81+
82+
83+
@cocoindex.flow_def(name="FaceRecognition")
84+
def face_recognition_flow(
85+
flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope
86+
) -> None:
87+
"""
88+
Define an example flow that embeds files into a vector database.
89+
"""
90+
data_scope["images"] = flow_builder.add_source(
91+
cocoindex.sources.LocalFile(path="images", binary=True),
92+
refresh_interval=datetime.timedelta(seconds=10),
93+
)
94+
95+
face_embeddings = data_scope.add_collector()
96+
97+
with data_scope["images"].row() as image:
98+
# Extract faces
99+
image["faces"] = image["content"].transform(extract_faces)
100+
101+
with image["faces"].row() as face:
102+
face["embedding"] = face["image"].transform(extract_face_embedding)
103+
104+
# Collect embeddings
105+
face_embeddings.collect(
106+
filename=image["filename"],
107+
rect=face["rect"],
108+
embedding=face["embedding"],
109+
)
110+
111+
face_embeddings.export(
112+
"face_embeddings",
113+
cocoindex.targets.Postgres(),
114+
primary_key_fields=["filename", "rect"],
115+
)
Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
[project]
2+
name = "cocoindex-face-recognition-example"
3+
version = "0.1.0"
4+
description = "Build index for papers with both metadata and content embeddings"
5+
requires-python = ">=3.11"
6+
dependencies = [
7+
"cocoindex>=0.1.67",
8+
"face-recognition>=1.3.0",
9+
"pillow>=10.0.0",
10+
"numpy>=1.26.0",
11+
]
12+
13+
[tool.setuptools]
14+
packages = []

0 commit comments

Comments
 (0)