110 lines
3.2 KiB
Python
110 lines
3.2 KiB
Python
# -*- coding:utf-8 -*-
|
|
"""InsightFace buffalo_l model wrapper."""
|
|
import os
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
_model = None
|
|
_lock = False
|
|
|
|
MODEL_NAME = "buffalo_l"
|
|
|
|
def get_model():
|
|
global _model, _lock
|
|
if _model is None and not _lock:
|
|
_lock = True
|
|
try:
|
|
import insightface
|
|
from insightface.app import FaceAnalysis
|
|
_model = FaceAnalysis(name=MODEL_NAME, providers=['CUDAExecutionProvider'])
|
|
_model.prepare(ctx_id=0, det_size=(640, 640))
|
|
except Exception as e:
|
|
_lock = False
|
|
raise RuntimeError(f"Failed to load InsightFace: {e}")
|
|
return _model
|
|
|
|
|
|
def load_image(image_input):
|
|
"""Load image from path, URL, or base64."""
|
|
if isinstance(image_input, str):
|
|
if image_input.startswith("/") and os.path.exists(image_input):
|
|
return Image.open(image_input).convert("RGB")
|
|
elif image_input.startswith("http"):
|
|
import urllib.request
|
|
import io
|
|
with urllib.request.urlopen(image_input, timeout=10) as resp:
|
|
return Image.open(io.BytesIO(resp.read())).convert("RGB")
|
|
raise ValueError(f"Cannot load image: {image_input[:50] if isinstance(image_input, str) else 'unknown'}")
|
|
|
|
|
|
def detect(image_input):
|
|
"""Detect faces in an image. Returns list of face info."""
|
|
model = get_model()
|
|
img = load_image(image_input)
|
|
img_np = np.array(img)
|
|
|
|
faces = model.get(img_np)
|
|
|
|
results = []
|
|
for i, face in enumerate(faces):
|
|
results.append({
|
|
"face_id": f"face_{i}",
|
|
"bbox": face.bbox.tolist(),
|
|
"det_score": round(float(face.det_score), 4),
|
|
"age": int(face.age) if face.age > 0 else None,
|
|
"gender": "M" if face.gender == 1 else "F",
|
|
"embedding_dim": len(face.embedding)
|
|
})
|
|
|
|
return {"faces": results, "count": len(results)}
|
|
|
|
|
|
def recognize(image_input):
|
|
"""Get face embeddings from an image."""
|
|
model = get_model()
|
|
img = load_image(image_input)
|
|
img_np = np.array(img)
|
|
|
|
faces = model.get(img_np)
|
|
|
|
results = []
|
|
for i, face in enumerate(faces):
|
|
# Normalize embedding to unit vector
|
|
embedding = face.embedding / np.linalg.norm(face.embedding)
|
|
results.append({
|
|
"face_id": f"face_{i}",
|
|
"bbox": face.bbox.tolist(),
|
|
"det_score": round(float(face.det_score), 4),
|
|
"embedding": embedding.tolist()
|
|
})
|
|
|
|
return {"faces": results, "count": len(results)}
|
|
|
|
|
|
def compare(embed1, embed2):
|
|
"""Compare two face embeddings. Returns cosine similarity."""
|
|
e1 = np.array(embed1)
|
|
e2 = np.array(embed2)
|
|
|
|
# Normalize
|
|
e1 = e1 / np.linalg.norm(e1)
|
|
e2 = e2 / np.linalg.norm(e2)
|
|
|
|
similarity = float(np.dot(e1, e2))
|
|
|
|
return {
|
|
"similarity": round(similarity, 6),
|
|
"is_same": similarity > 0.4,
|
|
"confidence": "high" if similarity > 0.6 else ("medium" if similarity > 0.4 else "low")
|
|
}
|
|
|
|
|
|
def health_check():
|
|
"""Check model status."""
|
|
model = get_model()
|
|
return {
|
|
"model": MODEL_NAME,
|
|
"loaded": model is not None,
|
|
"det_size": [640, 640]
|
|
}
|