File size: 2,816 Bytes
39ae7cb 53a4b0d 39ae7cb 53a4b0d 39ae7cb 53a4b0d 39ae7cb 53a4b0d 39ae7cb 53a4b0d 39ae7cb 53a4b0d 39ae7cb 53a4b0d 4785fbb 53a4b0d 39ae7cb 53a4b0d 39ae7cb 53a4b0d 39ae7cb 53a4b0d 39ae7cb |
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 86 87 88 |
import io
import base64
import numpy as np
from models.face_recognition import ArcFaceRecognition, extract_faces
def get_face_predictions(face, arcface_model, data_manager, results):
"""
Get predictions for a single face
Parameters:
face: Face image array
arcface_model: ArcFaceRecognition instance
data_manager: DataManager instance
results: Number of results to return
Returns:
List of (name, confidence) tuples from vector search
"""
# Create batch with original and flipped images
face_batch = np.stack([face, face[:, ::-1, :]], axis=0)
# Get embeddings for both orientations in one batch call
embeddings_batch = arcface_model.get_face_embeddings_batch(face_batch)
arc = np.mean(embeddings_batch, axis=0)
# Get predictions from ArcFace model
query_limit = max(results, 50)
arc_results = data_manager.query_arc_index(arc, query_limit)
# Convert distances to confidence scores and return top results
names, distances = arc_results
confidences = 1.0 - np.array(distances) # Cosine distance to cosine similarity
# Sort by confidence and return top results
sorted_indices = np.argsort(confidences)[::-1]
return [(names[i], confidences[i]) for i in sorted_indices[:results]]
def image_search_performers(image, data_manager, threshold=0.5, results=4):
"""
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)
arcface_model = ArcFaceRecognition()
try:
faces = extract_faces(image_array)
except ValueError:
raise ValueError("No faces found")
response = []
for face in faces:
predictions = get_face_predictions(face['face'], arcface_model, data_manager, results)
# Crop and encode face image
area = face['facial_area']
cimage = image.crop((area['x'], area['y'], area['x'] + area['w'], area['y'] + area['h']))
buf = io.BytesIO()
cimage.save(buf, format='JPEG')
im_b64 = base64.b64encode(buf.getvalue()).decode('ascii')
# Get performer information
performers = []
for name, confidence in predictions:
performer_info = data_manager.get_performer_info(data_manager.faces[name], confidence)
if performer_info:
performers.append(performer_info)
response.append({
'image': im_b64,
'area': area,
'confidence': face['confidence'],
'performers': performers
})
return response
|