Spaces:
Running
Running
| import io | |
| import base64 | |
| import numpy as np | |
| from PIL import Image | |
| from models.face_recognition import EnsembleFaceRecognition, extract_faces | |
| def get_face_predictions(face, ensemble, data_manager, results): | |
| """ | |
| Get predictions for a single face | |
| Parameters: | |
| face: Face image array | |
| ensemble: EnsembleFaceRecognition instance | |
| data_manager: DataManager instance | |
| results: Number of results to return | |
| Returns: | |
| List of (name, confidence, distance) tuples | |
| """ | |
| face_batch = np.stack([face, face[:, ::-1, :]], axis=0) | |
| embeddings_batch = ensemble.get_face_embeddings_batch(face_batch) | |
| facenet = np.mean(embeddings_batch['facenet'], axis=0) | |
| arc = np.mean(embeddings_batch['arc'], axis=0) | |
| ids, scores = data_manager.query_multi(facenet, arc, max(results, 50)) | |
| if not ids: | |
| return [] | |
| # WeightedReRanker normalizes each field to [0,1] then sums; with two | |
| # equal-weight fields the max combined score is 2.0. Dividing by 2.0 | |
| # gives an absolute average similarity, and the sqrt curve matches the | |
| # old distance_to_score shape. | |
| return [(name, (score / 2.0) ** 0.5, 0.0) for name, score in zip(ids, scores)] | |
| def image_search_performers(image, data_manager, threshold=0.5, results=3): | |
| """ | |
| Search for multiple performers in an image | |
| Parameters: | |
| image: PIL Image object | |
| data_manager: DataManager instance | |
| threshold: Confidence threshold | |
| results: Number of results to return | |
| Returns: | |
| List of dictionaries with face image and performer information | |
| """ | |
| image_array = np.array(image) | |
| ensemble = EnsembleFaceRecognition({"facenet": 0.5, "arc": 0.5}) | |
| faces = extract_faces(image_array) | |
| if not faces: | |
| raise ValueError("No faces found") | |
| response = [] | |
| for face in faces: | |
| predictions = get_face_predictions(face['face'], ensemble, data_manager, results) | |
| cimage = Image.fromarray(face['face']) | |
| buf = io.BytesIO() | |
| cimage.save(buf, format='JPEG') | |
| im_b64 = base64.b64encode(buf.getvalue()).decode('ascii') | |
| performers = [] | |
| for name, confidence, distance in predictions: | |
| if confidence < threshold: | |
| break | |
| if len(performers) >= results: | |
| break | |
| performer_info = data_manager.get_performer_info(name, confidence, distance) | |
| if performer_info: | |
| performers.append(performer_info) | |
| response.append({ | |
| 'image': im_b64, | |
| 'area': face['facial_area'], | |
| 'confidence': face['confidence'], | |
| 'landmarks': face['landmarks'], | |
| 'performers': performers | |
| }) | |
| return response | |