import hashlib from collections import OrderedDict import gradio as gr import numpy as np from models.data_manager import DataManager from models.face_recognition import _dist_to_confidence from models.image_processor import image_search_performers from models.sprite_processor import identify_sprite_performers class _LRUCache: def __init__(self, maxsize=3): self.maxsize = maxsize self._cache = OrderedDict() @staticmethod def _hash_args(*args): h = hashlib.md5(usedforsecurity=False) for arg in args: if arg is None: h.update(b'\x00') elif isinstance(arg, bytes): h.update(arg) elif isinstance(arg, str): h.update(arg.encode('utf-8')) elif isinstance(arg, int): h.update(str(arg).encode()) elif hasattr(arg, 'tobytes'): h.update(arg.tobytes()) else: h.update(str(arg).encode()) return h.digest() def get(self, *args): key = self._hash_args(*args) if key in self._cache: self._cache.move_to_end(key) return self._cache[key] return None def put(self, value, *args): key = self._hash_args(*args) self._cache[key] = value self._cache.move_to_end(key) while len(self._cache) > self.maxsize: self._cache.popitem(last=False) _COUNTRY_FLAGS = { 'AD': '🇦🇩', 'AE': '🇦🇪', 'AF': '🇦🇫', 'AG': '🇦🇬', 'AI': '🇦🇮', 'AL': '🇦🇱', 'AM': '🇦🇲', 'AO': '🇦🇴', 'AQ': '🇦🇶', 'AR': '🇦🇷', 'AS': '🇦🇸', 'AT': '🇦🇹', 'AU': '🇦🇺', 'AW': '🇦🇼', 'AX': '🇦🇽', 'AZ': '🇦🇿', 'BA': '🇧🇦', 'BB': '🇧🇧', 'BD': '🇧🇩', 'BE': '🇧🇪', 'BF': '🇧🇫', 'BG': '🇧🇬', 'BH': '🇧🇭', 'BI': '🇧🇮', 'BJ': '🇧🇯', 'BL': '🇧🇱', 'BM': '🇧🇲', 'BN': '🇧🇳', 'BO': '🇧🇴', 'BQ': '🇧🇶', 'BR': '🇧🇷', 'BS': '🇧🇸', 'BT': '🇧🇹', 'BV': '🇧🇻', 'BW': '🇧🇼', 'BY': '🇧🇾', 'BZ': '🇧🇿', 'CA': '🇨🇦', 'CC': '🇨🇨', 'CD': '🇨🇩', 'CF': '🇨🇫', 'CG': '🇨🇬', 'CH': '🇨🇭', 'CI': '🇨🇮', 'CK': '🇨🇰', 'CL': '🇨🇱', 'CM': '🇨🇲', 'CN': '🇨🇳', 'CO': '🇨🇴', 'CR': '🇨🇷', 'CU': '🇨🇺', 'CV': '🇨🇻', 'CW': '🇨🇼', 'CX': '🇨🇽', 'CY': '🇨🇾', 'CZ': '🇨🇿', 'DE': '🇩🇪', 'DJ': '🇩🇯', 'DK': '🇩🇰', 'DM': '🇩🇲', 'DO': '🇩🇴', 'DZ': '🇩🇿', 'EC': '🇪🇨', 'EE': '🇪🇪', 'EG': '🇪🇬', 'EH': '🇪🇭', 'ER': '🇪🇷', 'ES': '🇪🇸', 'ET': '🇪🇹', 'FI': '🇫🇮', 'FJ': '🇫🇯', 'FK': '🇫🇰', 'FM': '🇫🇲', 'FO': '🇫🇴', 'FR': '🇫🇷', 'GA': '🇬🇦', 'GB': '🇬🇧', 'GD': '🇬🇩', 'GE': '🇬🇪', 'GF': '🇬🇫', 'GG': '🇬🇬', 'GH': '🇬🇭', 'GI': '🇬🇮', 'GL': '🇬🇱', 'GM': '🇬🇲', 'GN': '🇬🇳', 'GP': '🇬🇵', 'GQ': '🇬🇶', 'GR': '🇬🇷', 'GS': '🇬🇸', 'GT': '🇬🇹', 'GU': '🇬🇺', 'GW': '🇬🇼', 'GY': '🇬🇾', 'HK': '🇭🇰', 'HM': '🇭🇲', 'HN': '🇭🇳', 'HR': '🇭🇷', 'HT': '🇭🇹', 'HU': '🇭🇺', 'ID': '🇮🇩', 'IE': '🇮🇪', 'IL': '🇮🇱', 'IM': '🇮🇲', 'IN': '🇮🇳', 'IO': '🇮🇴', 'IQ': '🇮🇶', 'IR': '🇮🇷', 'IS': '🇮🇸', 'IT': '🇮🇹', 'JE': '🇯🇪', 'JM': '🇯🇲', 'JO': '🇯🇴', 'JP': '🇯🇵', 'KE': '🇰🇪', 'KG': '🇰🇬', 'KH': '🇰🇭', 'KI': '🇰🇮', 'KM': '🇰🇲', 'KN': '🇰🇳', 'KP': '🇰🇵', 'KR': '🇰🇷', 'KW': '🇰🇼', 'KY': '🇰🇾', 'KZ': '🇰🇿', 'LA': '🇱🇦', 'LB': '🇱🇧', 'LC': '🇱🇨', 'LI': '🇱🇮', 'LK': '🇱🇰', 'LR': '🇱🇷', 'LS': '🇱🇸', 'LT': '🇱🇹', 'LU': '🇱🇺', 'LV': '🇱🇻', 'LY': '🇱🇾', 'MA': '🇲🇦', 'MC': '🇲🇨', 'MD': '🇲🇩', 'ME': '🇲🇪', 'MF': '🇲🇫', 'MG': '🇲🇬', 'MH': '🇲🇭', 'MK': '🇲🇰', 'ML': '🇲🇱', 'MM': '🇲🇲', 'MN': '🇲🇳', 'MO': '🇲🇴', 'MP': '🇲🇵', 'MQ': '🇲🇶', 'MR': '🇲🇷', 'MS': '🇲🇸', 'MT': '🇲🇹', 'MU': '🇲🇺', 'MV': '🇲🇻', 'MW': '🇲🇼', 'MX': '🇲🇽', 'MY': '🇲🇾', 'MZ': '🇲🇿', 'NA': '🇳🇦', 'NC': '🇳🇨', 'NE': '🇳🇪', 'NF': '🇳🇫', 'NG': '🇳🇬', 'NI': '🇳🇮', 'NL': '🇳🇱', 'NO': '🇳🇴', 'NP': '🇳🇵', 'NR': '🇳🇷', 'NU': '🇳🇺', 'NZ': '🇳🇿', 'OM': '🇴🇲', 'PA': '🇵🇦', 'PE': '🇵🇪', 'PF': '🇵🇫', 'PG': '🇵🇬', 'PH': '🇵🇭', 'PK': '🇵🇰', 'PL': '🇵🇱', 'PM': '🇵🇲', 'PN': '🇵🇳', 'PR': '🇵🇷', 'PS': '🇵🇸', 'PT': '🇵🇹', 'PW': '🇵🇼', 'PY': '🇵🇾', 'QA': '🇶🇦', 'RE': '🇷🇪', 'RO': '🇷🇴', 'RS': '🇷🇸', 'RU': '🇷🇺', 'RW': '🇷🇼', 'SA': '🇸🇦', 'SB': '🇸🇧', 'SC': '🇸🇨', 'SD': '🇸🇩', 'SE': '🇸🇪', 'SG': '🇸🇬', 'SH': '🇸🇭', 'SI': '🇸🇮', 'SJ': '🇸🇯', 'SK': '🇸🇰', 'SL': '🇸🇱', 'SM': '🇸🇲', 'SN': '🇸🇳', 'SO': '🇸🇴', 'SR': '🇸🇷', 'SS': '🇸🇸', 'ST': '🇸🇹', 'SV': '🇸🇻', 'SX': '🇸🇽', 'SY': '🇸🇾', 'SZ': '🇸🇿', 'TC': '🇹🇨', 'TD': '🇹🇩', 'TF': '🇹🇫', 'TG': '🇹🇬', 'TH': '🇹🇭', 'TJ': '🇹🇯', 'TK': '🇹🇰', 'TL': '🇹🇱', 'TM': '🇹🇲', 'TN': '🇹🇳', 'TO': '🇹🇴', 'TR': '🇹🇷', 'TT': '🇹🇹', 'TV': '🇹🇻', 'TW': '🇹🇼', 'TZ': '🇹🇿', 'UA': '🇺🇦', 'UG': '🇺🇬', 'UM': '🇺🇲', 'US': '🇺🇸', 'UY': '🇺🇾', 'UZ': '🇺🇿', 'VA': '🇻🇦', 'VC': '🇻🇨', 'VE': '🇻🇪', 'VG': '🇻🇬', 'VI': '🇻🇮', 'VN': '🇻🇳', 'VU': '🇻🇺', 'WF': '🇼🇫', 'WS': '🇼🇸', 'YE': '🇾🇪', 'YT': '🇾🇹', 'ZA': '🇿🇦', 'ZM': '🇿🇲', 'ZW': '🇿🇼', } class WebInterface: def __init__(self, data_manager: DataManager): self.data_manager = data_manager self._img_cache = _LRUCache(maxsize=3) self._sprite_cache = _LRUCache(maxsize=3) def get_country_flag(self, country_code): if not country_code or len(country_code) != 2: return "" return _COUNTRY_FLAGS.get(country_code.upper(), "") def multiple_image_search(self, img, results): cached = self._img_cache.get(img, results) if cached is not None: return cached try: result = image_search_performers(img, self.data_manager, results) self._img_cache.put(result, img, results) return result except ValueError as e: if "No faces found" in str(e): return {"error": "No faces detected in the uploaded image. Please try uploading an image with visible faces."} else: raise e def format_results_for_visual_display(self, json_results): if not json_results: return "

No faces detected or no matches found.

" if isinstance(json_results, dict) and "error" in json_results: return f"""

Error

{json_results['error']}

""" html_parts = [] html_parts.append(""" """) for i, face_result in enumerate(json_results): face_confidence = face_result['confidence'] performers = face_result['performers'] face_image_b64 = f"data:image/jpeg;base64,{face_result['image']}" html_parts.append(f"""
Detected Face {i+1}

Face {i+1}

Detection: {face_confidence:.1%}  |  Matches: {len(performers)}

""") if performers: html_parts.append('
') for performer in performers: confidence_class = "high-confidence" if performer['confidence'] >= 80 else "medium-confidence" if performer['confidence'] >= 60 else "low-confidence" country_code = performer.get('country', '') country_flag = self.get_country_flag(country_code) country_display = f"{country_flag} {country_code}" if country_flag else (country_code if country_code else 'Unknown') source = performer.get('source', '') source_badge = '' if source: initial = source[0].upper() source_badge = f'{initial}' gender_raw = performer.get('gender') or '' gender_badge = '' if gender_raw: gender_key = gender_raw.lower().replace('_', '-') gender_display = gender_raw.lower().replace('_', ' ').title() if gender_key == 'non-binary': gender_display = 'Non-binary' gender_badge = f'{gender_display}' html_parts.append(f"""
{performer['name']}

{performer['name']}

{source_badge} {gender_badge} {country_flag}{country_code if country_code else 'Unknown'}
{performer['confidence']}%
""") html_parts.append('
') else: html_parts.append('

No performer matches found for this face.

') html_parts.append('
') return ''.join(html_parts) def multiple_image_search_with_visual(self, img, results): """Run face search and return HTML for visual display.""" try: json_results = self.multiple_image_search(img, results) return self.format_results_for_visual_display(json_results) except Exception as e: return f"

Error

{str(e)}

" def search_by_vector(self, payload): """Search the AdaFace index directly with a raw 512-dim embedding. Expects payload = {"embedding": [512 floats], "top_k": 5}. Skips detection and embedding generation entirely. """ if not isinstance(payload, dict): return {"error": "Payload must be a JSON object: {\"embedding\": [512 floats], \"top_k\": 5}"} embedding = payload.get("embedding") if embedding is None: return {"error": "Missing 'embedding' field."} try: emb = np.asarray(embedding, dtype=np.float32) except (TypeError, ValueError): return {"error": "'embedding' must be an array of 512 numbers."} if emb.ndim != 1 or emb.shape[0] != 512: return {"error": f"'embedding' must have exactly 512 dimensions, got {emb.size}."} norm = float(np.linalg.norm(emb)) if norm < 1e-12: return {"error": "'embedding' must be a non-zero vector."} emb = emb / norm try: top_k = int(payload.get("top_k", 5)) except (TypeError, ValueError): return {"error": "'top_k' must be an integer."} top_k = max(1, min(top_k, 50)) ids, dists = self.data_manager.query_adaface_index(emb, top_k) performers = [] for pid, dist in zip(ids, dists): info = self.data_manager.get_performer_info(pid, _dist_to_confidence(dist), dist) if info is not None: performers.append(info) return {"performers": performers, "count": len(performers)} def _create_vector_search_interface(self): """Create the raw-vector JSON API search interface""" with gr.Blocks() as interface: gr.Markdown("# Vector Search API") gr.Markdown("Search with a raw 512-dim AdaFace embedding (no image upload).") with gr.Row(): with gr.Column(): payload_input = gr.JSON( label="Request", value={"embedding": [0.0] * 512, "top_k": 5}, ) search_btn = gr.Button("Search") with gr.Column(): json_output = gr.JSON(label="JSON Results") search_btn.click( fn=self.search_by_vector, inputs=[payload_input], outputs=json_output, api_name="search_by_vector" ) return interface def _create_json_search_interface(self): """Create the JSON API search interface""" with gr.Blocks() as interface: gr.Markdown("# Face Recognition API") gr.Markdown("Upload an image and get JSON results - perfect for API integration.") with gr.Row(): with gr.Column(): img_input = gr.Image(type="pil") results_count = gr.Slider( label="Top matches", minimum=3, maximum=50, value=5, step=1 ) search_btn = gr.Button("Search") with gr.Column(): json_output = gr.JSON(label="JSON Results") search_btn.click( fn=self.multiple_image_search, inputs=[img_input, results_count], outputs=json_output, api_name="multiple_image_search" ) return interface def identify_sprite_performers_wrapper(self, sprite_img, vtt_file, threshold): if sprite_img is None: yield {"error": "Please upload a sprite image."} return if vtt_file is None: yield {"error": "Please upload a VTT file."} return try: if hasattr(vtt_file, 'read'): vtt_text = vtt_file.read().decode('utf-8') else: with open(vtt_file, 'r', encoding='utf-8') as f: vtt_text = f.read() except Exception as e: yield {"error": f"Failed to read VTT file: {e}"} return cached = self._sprite_cache.get(sprite_img, vtt_text, threshold) if cached is not None: yield cached return gen = identify_sprite_performers( sprite_img, vtt_text, self.data_manager, detection_threshold=threshold / 100.0, ) try: final = None for item in gen: if isinstance(item, tuple) and len(item) == 2: progress_val, desc = item yield {"progress": progress_val, "desc": desc} else: final = item yield item return except ValueError as e: yield {"error": str(e)} finally: if final is not None: self._sprite_cache.put(final, sprite_img, vtt_text, threshold) def _create_sprite_search_interface(self): with gr.Blocks() as interface: gr.Markdown("# Sprite/VTT Batch Face Identification") gr.Markdown("Upload a scene sprite image and its VTT file to identify all performers across the entire scene.") with gr.Row(): with gr.Column(): sprite_input = gr.Image(type="pil") vtt_input = gr.File(label="VTT File", file_types=[".vtt"]) threshold = gr.Slider( label="Detection threshold", minimum=1, maximum=100, value=50, step=1, ) search_btn = gr.Button("Identify Performers") with gr.Column(): json_output = gr.JSON(label="Identification Results") search_btn.click( fn=self.identify_sprite_performers_wrapper, inputs=[sprite_input, vtt_input, threshold], outputs=json_output, api_name="identify_sprite_performers", ) return interface def _create_visual_search_interface(self): with gr.Blocks() as interface: gr.Markdown("# Who is in the photo?") gr.Markdown("Upload an image of a person(s) and we'll show you who it is with photos and details.") with gr.Row(): with gr.Column(): img_input = gr.Image(type="pil") results_count = gr.Slider( label="Top matches", minimum=3, maximum=50, value=5, step=1 ) search_btn = gr.Button("Search") with gr.Column(): performer_info = gr.HTML( label="Performer Information", value="

Upload an image and click search to see results.

" ) search_btn.click( fn=self.multiple_image_search_with_visual, inputs=[img_input, results_count], outputs=[performer_info], api_name="multiple_image_search_with_visual" ) return interface def launch(self, server_name="0.0.0.0", server_port=7860, share=True): with gr.Blocks( css=""" .gradio-container { background-color: #1e1e1e !important; color: #d4d4d4 !important; } .dark { --background-fill-primary: #2d2d2d; --background-fill-secondary: #3c3c3c; --border-color-primary: #404040; --block-title-text-color: #ffffff; --body-text-color: #d4d4d4; } """ ) as demo: with gr.Tabs(): with gr.TabItem("Visual Search"): self._create_visual_search_interface() with gr.TabItem("JSON API"): self._create_json_search_interface() with gr.TabItem("Sprite Search"): self._create_sprite_search_interface() with gr.TabItem("Vector Search"): self._create_vector_search_interface() @demo.app.middleware("http") async def add_cors(request, call_next): response = await call_next(request) response.headers["Access-Control-Allow-Origin"] = "*" response.headers["Access-Control-Allow-Methods"] = "*" response.headers["Access-Control-Allow-Headers"] = "*" response.headers["Access-Control-Allow-Credentials"] = "true" return response demo.queue().launch(server_name=server_name, server_port=server_port, share=share, ssr_mode=False)