Spaces:
Running
Running
File size: 2,717 Bytes
9282c17 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | 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
|