85 lines
2.7 KiB
Python
85 lines
2.7 KiB
Python
#!/opt/homebrew/bin/python3.11
|
|
"""
|
|
Extract face embedding from an image using InsightFace + CoreML FaceNet.
|
|
|
|
Usage:
|
|
python3 scripts/extract_face_embedding.py <image_path>
|
|
|
|
Output: JSON with "embedding" key (512 floats) or "error" key.
|
|
Exit code: 0 on success, 1 on failure.
|
|
"""
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
# Prefer venv if it exists (has insightface + coremltools installed)
|
|
VENV_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "venv")
|
|
VENV_SITE = os.path.join(VENV_PATH, "lib", "python3.11", "site-packages")
|
|
if os.path.isdir(VENV_SITE):
|
|
sys.path.insert(0, VENV_SITE)
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
MODELS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "models")
|
|
FACENET_PATH = os.path.join(MODELS_DIR, "facenet512.mlpackage")
|
|
|
|
|
|
def extract_embedding(image_path: str):
|
|
import io
|
|
import warnings
|
|
warnings.filterwarnings("ignore")
|
|
|
|
# Suppress InsightFace verbose stdout during model loading
|
|
old_stdout = sys.stdout
|
|
sys.stdout = io.StringIO()
|
|
try:
|
|
import insightface
|
|
from insightface.app import FaceAnalysis
|
|
import coremltools as ct
|
|
|
|
app = FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"])
|
|
app.prepare(ctx_id=0, det_thresh=0.5)
|
|
coreml_model = ct.models.MLModel(FACENET_PATH)
|
|
finally:
|
|
sys.stdout = old_stdout
|
|
|
|
img_bytes = open(image_path, "rb").read()
|
|
nparr = np.frombuffer(img_bytes, np.uint8)
|
|
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
|
if img is None:
|
|
print(json.dumps({"error": "Failed to decode image"}))
|
|
sys.exit(1)
|
|
|
|
# Detect faces
|
|
faces = app.get(img)
|
|
if not faces:
|
|
print(json.dumps({"error": "No face detected"}))
|
|
sys.exit(1)
|
|
|
|
largest = max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
|
|
x1, y1, x2, y2 = [int(v) for v in largest.bbox]
|
|
x1, y1 = max(0, x1), max(0, y1)
|
|
x2, y2 = min(img.shape[1], x2), min(img.shape[0], y2)
|
|
if x2 <= x1 or y2 <= y1:
|
|
print(json.dumps({"error": "Invalid face bbox"}))
|
|
sys.exit(1)
|
|
|
|
face_img = img[y1:y2, x1:x2]
|
|
face_img = cv2.resize(face_img, (160, 160))
|
|
normalized = (face_img.astype(np.float32) / 127.5) - 1.0
|
|
normalized = np.transpose(normalized, (2, 0, 1))
|
|
input_array = np.expand_dims(normalized, axis=0)
|
|
|
|
result = coreml_model.predict({"input": input_array})
|
|
emb_key = [k for k in result.keys() if k.startswith("var_")][0]
|
|
embedding = result[emb_key].flatten().tolist()
|
|
print(json.dumps({"embedding": embedding}))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 2:
|
|
print(json.dumps({"error": "Usage: extract_face_embedding.py <image_path>"}))
|
|
sys.exit(1)
|
|
extract_embedding(sys.argv[1])
|