192 lines
7.6 KiB
Python
192 lines
7.6 KiB
Python
|
||
import cv2
|
||
import numpy as np
|
||
from insightface.app import FaceAnalysis
|
||
import base64
|
||
import requests
|
||
from .database import database, users
|
||
from .milvus_helpers import collection as milvus_collection
|
||
|
||
class FaceRecognitionService:
|
||
"""封装所有核心人脸识别和数据库操作的业务逻辑"""
|
||
def __init__(self):
|
||
self.app = FaceAnalysis(name="buffalo_l", providers=['CPUExecutionProvider'])
|
||
self.app.prepare(ctx_id=0, det_size=(640, 640))
|
||
|
||
def _decode_image(self, image_source):
|
||
"""私有方法，从URL或Base64解码图片"""
|
||
img = None
|
||
if image_source.face_data:
|
||
try:
|
||
img_data = base64.b64decode(image_source.face_data)
|
||
nparr = np.frombuffer(img_data, np.uint8)
|
||
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
||
except Exception as e:
|
||
raise ValueError(f"Base64解码失败: {e}")
|
||
elif image_source.url:
|
||
try:
|
||
response = requests.get(image_source.url, timeout=10)
|
||
response.raise_for_status()
|
||
img_array = np.asarray(bytearray(response.content), dtype=np.uint8)
|
||
img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
|
||
except Exception as e:
|
||
raise ValueError(f"从URL获取图片失败: {e}")
|
||
|
||
if img is None:
|
||
raise ValueError("无法加载图片")
|
||
return img
|
||
|
||
async def register_new_face(self, user_id: int, name: str, image_source):
|
||
"""注册新的人脸，元数据存入PostgreSQL，向量存入Milvus"""
|
||
query = users.select().where(users.c.id == user_id)
|
||
user = await database.fetch_one(query)
|
||
|
||
if user and user.name != name:
|
||
raise ValueError(f"ID {user_id} 已被 '{user.name}' 注册，无法更改为 '{name}'。")
|
||
|
||
image = self._decode_image(image_source)
|
||
faces = self.app.get(image)
|
||
|
||
if not faces:
|
||
raise ValueError("图片中未检测到人脸。")
|
||
if len(faces) > 1:
|
||
raise ValueError("注册图片中只能包含一张人脸。")
|
||
|
||
embedding = faces[0].embedding
|
||
|
||
async with database.transaction():
|
||
if not user:
|
||
query = users.insert().values(id=user_id, name=name)
|
||
await database.execute(query)
|
||
|
||
data_to_insert = [[user_id], [embedding]]
|
||
milvus_collection.insert(data_to_insert)
|
||
milvus_collection.flush()
|
||
|
||
count_expr = f"user_id == {user_id}"
|
||
count = milvus_collection.query(expr=count_expr, output_fields=["user_id"])
|
||
|
||
return {"id": user_id, "name": name, "registered_faces_count": len(count)}
|
||
|
||
async def detect_faces(self, image_source):
|
||
"""在Milvus中进行1:N人脸识别"""
|
||
image = self._decode_image(image_source)
|
||
detected_faces = self.app.get(image)
|
||
|
||
if not detected_faces:
|
||
return []
|
||
|
||
search_vectors = [face.embedding for face in detected_faces]
|
||
search_params = {"metric_type": "IP", "params": {"nprobe": 10}}
|
||
|
||
results = milvus_collection.search(
|
||
data=search_vectors, anns_field="embedding", param=search_params,
|
||
limit=1, output_fields=["user_id"]
|
||
)
|
||
|
||
final_results = []
|
||
user_ids_to_fetch = {hits[0].entity.get('user_id') for hits in results if hits}
|
||
|
||
if not user_ids_to_fetch:
|
||
return []
|
||
|
||
# 批量查询用户信息，提高效率
|
||
user_query = users.select().where(users.c.id.in_(user_ids_to_fetch))
|
||
user_records = await database.fetch_all(user_query)
|
||
user_map = {user.id: user.name for user in user_records}
|
||
|
||
for i, hits in enumerate(results):
|
||
if not hits:
|
||
continue
|
||
best_hit = hits[0]
|
||
similarity = best_hit.distance
|
||
|
||
RECOGNITION_THRESHOLD = 0.5
|
||
if similarity > RECOGNITION_THRESHOLD:
|
||
matched_user_id = best_hit.entity.get('user_id')
|
||
user_name = user_map.get(matched_user_id)
|
||
|
||
if user_name:
|
||
count_expr = f"user_id == {matched_user_id}"
|
||
count_res = milvus_collection.query(expr=count_expr, output_fields=["user_id"])
|
||
x1, y1, x2, y2 = detected_faces[i].bbox.astype(int)
|
||
final_results.append({
|
||
"id": matched_user_id,
|
||
"name": user_name,
|
||
"registered_faces_count": len(count_res),
|
||
"confidence": float(similarity),
|
||
"location": {"x": x1, "y": y1, "width": x2 - x1, "height": y2 - y1}
|
||
})
|
||
return final_results
|
||
|
||
async def verify_face(self, user_id: int, image_source):
|
||
"""在Milvus中进行1:1人脸认证"""
|
||
image = self._decode_image(image_source)
|
||
detected_faces = self.app.get(image)
|
||
|
||
if not detected_faces:
|
||
raise ValueError("图片中未检测到人脸。")
|
||
if len(detected_faces) > 1:
|
||
raise ValueError("用于认证的图片中只能包含一张人脸。")
|
||
|
||
embedding = detected_faces[0].embedding
|
||
|
||
search_params = {"metric_type": "IP", "params": {"nprobe": 10}}
|
||
# 在指定user_id的范围内进行搜索
|
||
expr = f"user_id == {user_id}"
|
||
results = milvus_collection.search(
|
||
data=[embedding], anns_field="embedding", param=search_params,
|
||
limit=1, expr=expr, output_fields=["user_id"]
|
||
)
|
||
|
||
best_similarity = 0.0
|
||
if results and results[0]:
|
||
best_similarity = results[0][0].distance
|
||
|
||
VERIFICATION_THRESHOLD = 0.6
|
||
match = bool(best_similarity > VERIFICATION_THRESHOLD)
|
||
|
||
return {"match": match, "confidence": float(best_similarity)}
|
||
|
||
async def delete_user(self, user_id: int):
|
||
"""从PostgreSQL和Milvus中删除用户"""
|
||
query = users.select().where(users.c.id == user_id)
|
||
user = await database.fetch_one(query)
|
||
if not user:
|
||
raise ValueError(f"ID为 {user_id} 的用户不存在。")
|
||
|
||
async with database.transaction():
|
||
delete_query = users.delete().where(users.c.id == user_id)
|
||
await database.execute(delete_query)
|
||
|
||
# 从 Milvus 删除
|
||
expr = f"user_id == {user_id}"
|
||
milvus_collection.delete(expr)
|
||
|
||
return {"id": user.id, "name": user.name}
|
||
|
||
async def get_user(self, user_id: int):
|
||
"""获取单个用户的元数据和在Milvus中的人脸数量"""
|
||
query = users.select().where(users.c.id == user_id)
|
||
user = await database.fetch_one(query)
|
||
if not user:
|
||
return None
|
||
|
||
count_expr = f"user_id == {user_id}"
|
||
count_res = milvus_collection.query(expr=count_expr, output_fields=["user_id"])
|
||
return {"id": user.id, "name": user.name, "registered_faces_count": len(count_res)}
|
||
|
||
async def list_all_users(self, skip: int = 0, limit: int = 100):
|
||
"""列出所有用户，并从Milvus统计人脸数量"""
|
||
query = users.select().offset(skip).limit(limit)
|
||
user_records = await database.fetch_all(query)
|
||
|
||
results = []
|
||
for user in user_records:
|
||
count_expr = f"user_id == {user.id}"
|
||
count_res = milvus_collection.query(expr=count_expr, output_fields=["user_id"])
|
||
results.append({"id": user.id, "name": user.name, "registered_faces_count": len(count_res)})
|
||
return results
|
||
|
||
face_service = FaceRecognitionService()
|