Spaces:
Running
Running
Download models/image_processor.py from cc1234/stashface_onnx: direct link, hf CLI and curl.
- Browser
- Download file 5.81 kB
-
https://huggingface.co/spaces/cc1234/stashface_onnx/resolve/main/models/image_processor.py
- Command line
-
hf download hf://spaces/cc1234/stashface_onnx/models/image_processor.py
-
curl -L -o image_processor.py https://huggingface.co/spaces/cc1234/stashface_onnx/resolve/main/models/image_processor.py
5.81 kB
| import re | |
| import numpy as np | |
| from models.face_recognition import AdaFaceRecognition, extract_faces, _face_to_base64 | |
| def _int8_similarity(v1, v2): | |
| dot = np.dot(v1.astype(np.float32), v2.astype(np.float32)) | |
| norm = max(np.linalg.norm(v1.astype(np.float32)), 1e-12) * max(np.linalg.norm(v2.astype(np.float32)), 1e-12) | |
| return dot / norm | |
| def _int8_score_to_confidence(score): | |
| return min(1.0, max(0.0, score / 12000.0)) | |
| def _normalize_name(name): | |
| name = name.strip().lower() | |
| name = re.sub(r'[-_\'.]+', ' ', name) | |
| return re.sub(r'\s+', ' ', name) | |
| def _consolidate(entries, sim_threshold=0.85): | |
| if not entries: | |
| return [] | |
| entries = sorted(entries, key=lambda e: e["score"], reverse=True) | |
| used = [False] * len(entries) | |
| groups = [] | |
| # Pass 1: group by exact lowercase name | |
| for i in range(len(entries)): | |
| if used[i]: | |
| continue | |
| name = _normalize_name(entries[i]["fields"].get("name", "")) | |
| if not name: | |
| continue | |
| group = [i] | |
| used[i] = True | |
| for j in range(i + 1, len(entries)): | |
| if used[j]: | |
| continue | |
| if _normalize_name(entries[j]["fields"].get("name", "")) == name: | |
| group.append(j) | |
| used[j] = True | |
| groups.append(group) | |
| # Pass 2: merge unnamed entries into existing groups by vector similarity | |
| for i in range(len(entries)): | |
| if used[i]: | |
| continue | |
| for g_idx, group in enumerate(groups): | |
| for m_idx in group: | |
| sim = _int8_similarity(entries[i]["vector"], entries[m_idx]["vector"]) | |
| if sim > sim_threshold: | |
| groups[g_idx].append(i) | |
| used[i] = True | |
| break | |
| if used[i]: | |
| break | |
| # Pass 3: remaining entries group among themselves | |
| for i in range(len(entries)): | |
| if used[i]: | |
| continue | |
| group = [i] | |
| used[i] = True | |
| for j in range(i + 1, len(entries)): | |
| if used[j]: | |
| continue | |
| sim = _int8_similarity(entries[i]["vector"], entries[j]["vector"]) | |
| if sim > sim_threshold: | |
| group.append(j) | |
| used[j] = True | |
| groups.append(group) | |
| result = [] | |
| for group in groups: | |
| members = [entries[idx] for idx in group] | |
| sources = {} | |
| source_urls = {} | |
| for m in members: | |
| src = m["fields"].get("source", "") | |
| if src: | |
| sources.setdefault(src, []).append(m) | |
| if src not in source_urls: | |
| url = m["fields"].get("url", "") | |
| if url: | |
| source_urls[src] = url | |
| all_images = [] | |
| for m in members: | |
| img = m["fields"].get("image", "") | |
| if img and img not in all_images: | |
| all_images.append(img) | |
| def completeness(m): | |
| name = m["fields"].get("name", "") | |
| image = m["fields"].get("image", "") | |
| return (1 if name and image else 0 if name else -1, m["score"]) | |
| best = max(members, key=completeness) | |
| result.append({ | |
| "id": best["id"], | |
| "name": best["fields"].get("name", ""), | |
| "score": best["score"], | |
| "confidence": _int8_score_to_confidence(best["score"]), | |
| "image": best["fields"].get("image", ""), | |
| "all_images": all_images, | |
| "country": best["fields"].get("country") or None, | |
| "gender": best["fields"].get("gender") or None, | |
| "url": best["fields"].get("url", ""), | |
| "sources": list(sources.keys()), | |
| "source_urls": source_urls, | |
| "duplicates": len(members) - 1, | |
| }) | |
| return result | |
| def image_search_performers(image, data_manager, results=3, detection_threshold=0.5): | |
| image_array = np.array(image) | |
| if image_array.ndim < 2 or image_array.size == 0: | |
| raise ValueError("No faces found: uploaded image is empty or corrupt") | |
| ensemble = AdaFaceRecognition() | |
| faces = extract_faces(image_array, min_confidence=detection_threshold) | |
| if not faces: | |
| raise ValueError("No faces found") | |
| response = [] | |
| for face_info in faces: | |
| face = face_info['face'] | |
| face_batch = np.expand_dims(face, axis=0) | |
| embeddings_batch = ensemble.get_adaface_embeddings_batch(face_batch) | |
| adaface_emb = embeddings_batch[0] | |
| raw_entries = data_manager.query_with_vectors(adaface_emb, max(results * 5, 50)) | |
| consolidated = _consolidate(raw_entries, sim_threshold=0.75) | |
| consolidated = consolidated[:results] | |
| im_b64 = _face_to_base64(face) | |
| performers = [] | |
| for entry in consolidated: | |
| performers.append({ | |
| 'id': entry['id'], | |
| 'name': entry['name'], | |
| 'confidence': int(_int8_score_to_confidence(entry['score']) * 100), | |
| 'image': entry['image'], | |
| 'all_images': entry.get('all_images', [entry['image']]), | |
| 'country': entry['country'], | |
| 'gender': entry.get('gender'), | |
| 'source': entry['sources'][0] if entry['sources'] else '', | |
| 'all_sources': entry['sources'], | |
| 'source_urls': entry['source_urls'], | |
| 'hits': 1 + entry['duplicates'], | |
| 'distance': entry['score'], | |
| 'performer_url': entry['url'], | |
| }) | |
| response.append({ | |
| 'image': im_b64, | |
| 'area': face_info['facial_area'], | |
| 'confidence': face_info['confidence'], | |
| 'landmarks': face_info['landmarks'], | |
| 'performers': performers, | |
| }) | |
| return response | |