Spaces:
Running
Running
| import os | |
| from concurrent.futures import ThreadPoolExecutor, as_completed | |
| import numpy as np | |
| from models.face_recognition import ( | |
| AdaFaceRecognition, | |
| extract_faces, | |
| _get_face_analyzer, | |
| _dist_to_confidence, | |
| _face_to_base64, | |
| ) | |
| from utils.vtt_parser import parse_sprite_vtt | |
| HIT_CONF_THRESH = 0.15 | |
| CAST_SCORE_THRESH = 0.9 | |
| CAST_HIT_THRESH = 2 | |
| RECLAIM_SIM_THRESH = 0.35 | |
| GAP_TOLERANCE = 2 | |
| TOP_K_CANDIDATES = 10 | |
| CONFIDENCE_FLOOR = 0.05 | |
| # Core-aware parallelism. On Hugging Face Spaces (low vCPU), this auto-scales down. | |
| MAX_WORKERS = max(1, os.cpu_count() or 1) | |
| # Shared buffalo_sc 640x640 analyzer — ONNX sessions are thread-safe. | |
| _sprite_ana = None | |
| def _get_sprite_analyzer(): | |
| global _sprite_ana | |
| if _sprite_ana is None: | |
| _sprite_ana = _get_face_analyzer(det_size=(640, 640)) | |
| return _sprite_ana | |
| def _detect_one(frame_args): | |
| """Detect faces in a single frame. Called by ThreadPoolExecutor.""" | |
| idx, frame, mc, ana = frame_args | |
| results = [] | |
| for d in extract_faces(frame, min_confidence=mc, analyzer=ana): | |
| results.append((d['face'], idx)) | |
| return results | |
| def _detect_faces(frames, min_confidence): | |
| """Per-frame face detection with buffalo_sc at 640x640. | |
| Single-core: sequential per-frame. | |
| Multi-core: parallel per-frame via ThreadPoolExecutor. | |
| """ | |
| n_frames = len(frames) | |
| if n_frames == 0: | |
| return [] | |
| ana = _get_sprite_analyzer() | |
| if MAX_WORKERS <= 1: | |
| results = [] | |
| for idx, frame in enumerate(frames): | |
| for d in extract_faces(frame, min_confidence=min_confidence, analyzer=ana): | |
| results.append((d['face'], idx)) | |
| return results | |
| face_crops = [] | |
| with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool: | |
| args = [(idx, f, min_confidence, ana) for idx, f in enumerate(frames)] | |
| for fut in as_completed(pool.submit(_detect_one, a) for a in args): | |
| face_crops.extend(fut.result()) | |
| face_crops.sort(key=lambda x: x[1]) | |
| return face_crops | |
| def identify_sprite_performers(sprite_image, vtt_text, data_manager, | |
| detection_threshold=0.5): | |
| yield (0.0, "Parsing VTT cue sheet") | |
| cues = parse_sprite_vtt(vtt_text) | |
| if not cues: | |
| raise ValueError("No cues found in VTT file") | |
| yield (0.05, "Extracting sprite frames") | |
| sprite = np.array(sprite_image) | |
| frames = [sprite[c['y']:c['y'] + c['h'], c['x']:c['x'] + c['w']] for c in cues] | |
| yield (0.1, f"Detecting faces across {len(frames)} frames\u2026") | |
| face_crops = _detect_faces(frames, min_confidence=detection_threshold) | |
| yield (0.6, f"Found faces in {len(set(fc[1] for fc in face_crops))} frames") | |
| if not face_crops: | |
| raise ValueError("Face detection found no faces in any sprite frame.") | |
| yield (0.6, f"Generating embeddings for {len(face_crops)} faces\u2026") | |
| face_embs = AdaFaceRecognition().get_adaface_embeddings_batch( | |
| np.stack([fc[0] for fc in face_crops], axis=0) | |
| ) | |
| yield (0.7, f"Matching {len(face_embs)} faces against performer index\u2026") | |
| performer_data = {} | |
| face_results = [] | |
| for emb in face_embs: | |
| ids, dists = data_manager.query_adaface_index(emb, TOP_K_CANDIDATES) | |
| confident = [] | |
| all_hits = [] | |
| for pid, dist in zip(ids, dists): | |
| conf = _dist_to_confidence(dist) | |
| if conf < CONFIDENCE_FLOOR: | |
| break | |
| info = data_manager.get_performer_info(pid, conf, dist) | |
| if info and info.get('name'): | |
| p = performer_data.setdefault(pid, {"score": 0.0, "hits": 0, "conf_hits": 0}) | |
| p["score"] += conf | |
| p["hits"] += 1 | |
| all_hits.append((pid, conf)) | |
| if conf >= HIT_CONF_THRESH: | |
| p["conf_hits"] += 1 | |
| confident.append(pid) | |
| face_results.append((confident, emb, all_hits)) | |
| yield (0.85, "Clustering results by performer") | |
| cast = [pid for pid, p in performer_data.items() | |
| if p["score"] >= CAST_SCORE_THRESH and p["conf_hits"] >= CAST_HIT_THRESH] | |
| if not cast and face_crops: | |
| raise ValueError( | |
| f"Faces detected in {len(set(fc[1] for fc in face_crops))} frames, but " | |
| "no performer could be identified with confidence. " | |
| "Individual frame matches were too weak to cluster into a reliable result." | |
| ) | |
| centroids = {} | |
| for pid in cast: | |
| embs = [face_results[i][1] for i in range(len(face_results)) if pid in face_results[i][0]] | |
| if embs: | |
| centroids[pid] = np.mean(embs, axis=0) | |
| cue_sets = {pid: set() for pid in cast} | |
| face_indices = {pid: [] for pid in cast} | |
| for i, (confident, _, _) in enumerate(face_results): | |
| for pid in confident: | |
| if pid not in cue_sets: | |
| continue | |
| cue_sets[pid].add(face_crops[i][1]) | |
| face_indices[pid].append(i) | |
| for i, (confident, emb, _) in enumerate(face_results): | |
| if not confident and centroids: | |
| best = max(((pid, float(np.dot(emb, c))) for pid, c in centroids.items()), key=lambda x: x[1], default=(None, 0)) | |
| if best[0] is not None and best[1] >= RECLAIM_SIM_THRESH: | |
| cue_sets[best[0]].add(face_crops[i][1]) | |
| face_indices[best[0]].append(i) | |
| yield (0.95, "Building response") | |
| response = [] | |
| for pid in cast: | |
| info = data_manager.get_performer_info(pid, 0, 0) | |
| if not info: | |
| continue | |
| centroid = centroids[pid] | |
| best_idx = max(face_indices[pid], key=lambda i: float(np.dot(face_results[i][1], centroid))) if face_indices[pid] else None | |
| thumb = _face_to_base64(face_crops[best_idx][0]) if best_idx is not None else info.get('image', '') | |
| shot_groups = _group_adjacent_indices(cue_sets[pid]) | |
| shot_confs = [] | |
| for group in shot_groups: | |
| max_in_group = 0.0 | |
| for idx in face_indices[pid]: | |
| if face_crops[idx][1] not in group: | |
| continue | |
| for hit_pid, conf in face_results[idx][2]: | |
| if hit_pid == pid and conf > max_in_group: | |
| max_in_group = conf | |
| shot_confs.append(max_in_group) | |
| response.append({ | |
| "id": pid, | |
| "name": info['name'], | |
| "confidence": int((1.0 - np.prod([1.0 - c for c in shot_confs])) * 100), | |
| "thumbnail": thumb, | |
| "time_ranges": _merge_time_ranges(cue_sets[pid], cues), | |
| "hit_count": performer_data[pid]["hits"], | |
| "shot_count": len(shot_confs), | |
| "performer_url": info.get('performer_url', ''), | |
| "source": info.get('source', ''), | |
| "gender": info.get('gender'), | |
| }) | |
| response.sort(key=lambda x: x['hit_count'], reverse=True) | |
| yield {"performers": response} | |
| def _group_adjacent_indices(indices, gap_tolerance=GAP_TOLERANCE): | |
| if not indices: | |
| return [] | |
| sorted_idx = sorted(set(indices)) | |
| groups = [[sorted_idx[0]]] | |
| for i in range(1, len(sorted_idx)): | |
| if sorted_idx[i] - sorted_idx[i - 1] <= gap_tolerance + 1: | |
| groups[-1].append(sorted_idx[i]) | |
| else: | |
| groups.append([sorted_idx[i]]) | |
| return groups | |
| def _merge_time_ranges(cue_indices, cues, gap_tolerance=GAP_TOLERANCE): | |
| groups = _group_adjacent_indices(cue_indices, gap_tolerance) | |
| return [{"start": cues[g[0]]["start"], "end": cues[g[-1]]["end"]} for g in groups] | |