Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 16 additions & 11 deletions examples/face_recognition/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,10 @@

@dataclasses.dataclass
class ImageRect:
top: int
left: int
bottom: int
right: int
min_x: int
min_y: int
max_x: int
max_y: int


@dataclasses.dataclass
Expand All @@ -28,7 +28,12 @@ class FaceBase:
MAX_IMAGE_WIDTH = 1280


@cocoindex.op.function(cache=True, behavior_version=1, gpu=True)
@cocoindex.op.function(
cache=True,
behavior_version=1,
gpu=True,
arg_relationship=(cocoindex.op.ArgRelationship.RECTS_BASE_IMAGE, "content"),
)
def extract_faces(content: bytes) -> list[FaceBase]:
"""Extract the first pages of a PDF."""
orig_img = Image.open(io.BytesIO(content)).convert("RGB")
Expand All @@ -48,17 +53,17 @@ def extract_faces(content: bytes) -> list[FaceBase]:
locs = face_recognition.face_locations(np.array(img), model="cnn")

faces: list[FaceBase] = []
for top, right, bottom, left in locs:
for min_y, max_x, max_y, min_x in locs:
rect = ImageRect(
left=int(left * ratio),
top=int(top * ratio),
right=int(right * ratio),
bottom=int(bottom * ratio),
min_x=int(min_x * ratio),
min_y=int(min_y * ratio),
max_x=int(max_x * ratio),
max_y=int(max_y * ratio),
)

# Crop the face and save it as a PNG.
buf = io.BytesIO()
orig_img.crop((rect.left, rect.top, rect.right, rect.bottom)).save(
orig_img.crop((rect.min_x, rect.min_y, rect.max_x, rect.max_y)).save(
buf, format="PNG"
)
face = buf.getvalue()
Expand Down
2 changes: 1 addition & 1 deletion examples/face_recognition/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ version = "0.1.0"
description = "Build index for papers with both metadata and content embeddings"
requires-python = ">=3.11"
dependencies = [
"cocoindex>=0.1.67",
"cocoindex>=0.1.71",
"face-recognition>=1.3.0",
"pillow>=10.0.0",
"numpy>=1.26.0",
Expand Down