172 lines
4.6 KiB
Python
172 lines
4.6 KiB
Python
"""
|
|
Face Recognition Service
|
|
|
|
Uses InsightFace with ONNX Runtime for:
|
|
- Face detection (SCRFD)
|
|
- Face embedding generation (ArcFace)
|
|
- Gender and age estimation
|
|
|
|
Provides REST API endpoints for the digest pipeline.
|
|
"""
|
|
|
|
from fastapi import FastAPI, HTTPException
|
|
from pydantic import BaseModel
|
|
import insightface
|
|
import numpy as np
|
|
import base64
|
|
import cv2
|
|
import logging
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
app = FastAPI(
|
|
title="Face Recognition Service",
|
|
description="Face detection and embedding generation using InsightFace",
|
|
version="1.0.0"
|
|
)
|
|
|
|
# Global face analyzer - loaded on startup
|
|
face_analyzer = None
|
|
|
|
|
|
@app.on_event("startup")
|
|
async def load_models():
|
|
"""Load InsightFace models on startup."""
|
|
global face_analyzer
|
|
logger.info("Loading InsightFace models...")
|
|
|
|
try:
|
|
# buffalo_l includes SCRFD (detection) + ArcFace (recognition)
|
|
face_analyzer = insightface.app.FaceAnalysis(
|
|
name='buffalo_l',
|
|
providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
|
|
)
|
|
# Prepare with detection size
|
|
face_analyzer.prepare(ctx_id=0, det_size=(640, 640))
|
|
logger.info("InsightFace models loaded successfully")
|
|
except Exception as e:
|
|
logger.error(f"Failed to load InsightFace models: {e}")
|
|
raise
|
|
|
|
|
|
class AnalyzeRequest(BaseModel):
|
|
"""Request body for face analysis."""
|
|
image_base64: str
|
|
|
|
|
|
class Face(BaseModel):
|
|
"""Detected face with embedding and attributes."""
|
|
bbox: list[float] # [x1, y1, x2, y2]
|
|
embedding: list[float] # 512-D vector
|
|
gender: int # 0=female, 1=male
|
|
age: int
|
|
det_score: float # Detection confidence
|
|
|
|
|
|
class AnalyzeResponse(BaseModel):
|
|
"""Response from face analysis."""
|
|
faces: list[Face]
|
|
count: int
|
|
women_count: int
|
|
men_count: int
|
|
|
|
|
|
class HealthResponse(BaseModel):
|
|
"""Health check response."""
|
|
status: str
|
|
model_loaded: bool
|
|
|
|
|
|
@app.post("/analyze", response_model=AnalyzeResponse)
|
|
async def analyze_faces(request: AnalyzeRequest):
|
|
"""
|
|
Detect faces in an image and return embeddings with attributes.
|
|
|
|
- **image_base64**: Base64 encoded image (JPEG or PNG)
|
|
|
|
Returns detected faces with:
|
|
- Bounding boxes
|
|
- 512-D face embeddings (for recognition)
|
|
- Gender (0=female, 1=male)
|
|
- Age estimate
|
|
- Detection confidence score
|
|
"""
|
|
if face_analyzer is None:
|
|
raise HTTPException(503, "Face analyzer not initialized")
|
|
|
|
try:
|
|
# Decode base64 image
|
|
img_data = base64.b64decode(request.image_base64)
|
|
img_array = np.frombuffer(img_data, np.uint8)
|
|
img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
|
|
|
|
if img is None:
|
|
raise HTTPException(400, "Invalid image data - could not decode")
|
|
|
|
# Detect and analyze faces
|
|
faces = face_analyzer.get(img)
|
|
|
|
# Convert to response format
|
|
face_results = []
|
|
women_count = 0
|
|
men_count = 0
|
|
|
|
for face in faces:
|
|
gender = int(face.gender) if hasattr(face, 'gender') else -1
|
|
if gender == 0:
|
|
women_count += 1
|
|
elif gender == 1:
|
|
men_count += 1
|
|
|
|
face_results.append(Face(
|
|
bbox=face.bbox.tolist(),
|
|
embedding=face.embedding.tolist() if hasattr(face, 'embedding') else [],
|
|
gender=gender,
|
|
age=int(face.age) if hasattr(face, 'age') else -1,
|
|
det_score=float(face.det_score) if hasattr(face, 'det_score') else 0.0
|
|
))
|
|
|
|
logger.info(f"Detected {len(faces)} faces ({women_count} women, {men_count} men)")
|
|
|
|
return AnalyzeResponse(
|
|
faces=face_results,
|
|
count=len(faces),
|
|
women_count=women_count,
|
|
men_count=men_count
|
|
)
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Face analysis failed: {e}")
|
|
raise HTTPException(500, f"Face analysis failed: {str(e)}")
|
|
|
|
|
|
@app.get("/health", response_model=HealthResponse)
|
|
async def health():
|
|
"""Health check endpoint."""
|
|
return HealthResponse(
|
|
status="ok",
|
|
model_loaded=face_analyzer is not None
|
|
)
|
|
|
|
|
|
@app.get("/")
|
|
async def root():
|
|
"""Root endpoint with service info."""
|
|
return {
|
|
"service": "Face Recognition",
|
|
"version": "1.0.0",
|
|
"endpoints": {
|
|
"/analyze": "POST - Analyze faces in image",
|
|
"/health": "GET - Health check"
|
|
}
|
|
}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=5001)
|