stashface_onnx / web /interface.py
cc1234475
Add vector search API for raw 512-dim AdaFace embeddings
ec7950d
Raw History Blame
27 kB
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 "<p>No faces detected or no matches found.</p>"
if isinstance(json_results, dict) and "error" in json_results:
return f"""
<div class="performer-card">
<div class="face-info">
<h3 style="color: #ff6b6b;">Error</h3>
<p>{json_results['error']}</p>
</div>
</div>
"""
html_parts = []
html_parts.append("""
<style>
body, .gradio-container {
background-color: #1e1e1e !important;
color: #d4d4d4 !important;
}
.performer-card {
border: 1px solid #404040;
border-radius: 16px;
padding: 24px;
margin: 16px 0;
background: #2d2d2d;
box-shadow: 0 4px 12px rgba(0,0,0,0.3);
color: #d4d4d4;
}
.face-info {
background: #3c3c3c;
padding: 12px 16px;
border-radius: 10px;
margin-bottom: 24px;
border: 1px solid #4a4a4a;
display: flex;
align-items: center;
gap: 16px;
}
.face-info-content {
flex: 1;
}
.face-info h3 {
color: #ffffff;
margin: 0 0 4px 0;
font-size: 1.1em;
}
.face-info p {
color: #cccccc;
margin: 2px 0;
font-size: 0.9em;
}
.face-info p strong {
color: #9cdcfe;
}
.detected-face img {
width: 72px;
height: 72px;
border-radius: 8px;
object-fit: cover;
border: 2px solid #569cd6;
box-shadow: 0 2px 8px rgba(0,0,0,0.3);
flex-shrink: 0;
}
.performer-grid {
display: grid;
grid-template-columns: repeat(auto-fit, minmax(280px, 1fr));
gap: 24px;
margin-top: 8px;
}
.performer-item {
border: 1px solid #4a4a4a;
border-radius: 16px;
padding: 0;
background: #333333;
transition: all 0.3s ease;
box-shadow: 0 2px 8px rgba(0,0,0,0.2);
display: flex;
flex-direction: column;
overflow: hidden;
}
.performer-item:hover {
border-color: #569cd6;
box-shadow: 0 8px 24px rgba(0,0,0,0.5);
transform: translateY(-4px);
}
.performer-image-wrap {
width: 100%;
background: #2a2a2a;
display: flex;
justify-content: center;
align-items: center;
overflow: hidden;
}
.performer-image {
width: 100%;
max-height: 280px;
object-fit: cover;
object-position: top;
display: block;
border-bottom: 1px solid #4a4a4a;
transition: transform 0.4s ease;
}
.performer-image:hover {
transform: scale(1.04);
}
.performer-body {
padding: 16px;
display: flex;
flex-direction: column;
gap: 8px;
}
.performer-header {
display: flex;
align-items: center;
flex-wrap: wrap;
gap: 8px;
}
.performer-name {
color: #ffffff;
margin: 0;
font-size: 1.15em;
font-weight: 600;
line-height: 1.3;
}
.performer-name a {
color: #569cd6;
text-decoration: none;
transition: color 0.3s ease;
}
.performer-name a:hover {
color: #9cdcfe;
text-decoration: underline;
}
.performer-country {
color: #999999;
font-size: 0.85em;
margin-left: auto;
}
.country-flag {
font-size: 1.1em;
margin-right: 4px;
vertical-align: middle;
}
.confidence-bar {
background: #404040;
border-radius: 8px;
overflow: hidden;
height: 24px;
margin: 4px 0 0 0;
border: 1px solid #4a4a4a;
width: 100%;
}
.confidence-fill {
height: 100%;
transition: width 0.5s ease;
text-align: center;
line-height: 24px;
color: white;
font-size: 12px;
font-weight: bold;
text-shadow: 0 1px 2px rgba(0,0,0,0.5);
}
.high-confidence {
background: linear-gradient(135deg, #4caf50, #66bb6a);
}
.medium-confidence {
background: linear-gradient(135deg, #ff9800, #ffb74d);
}
.low-confidence {
background: linear-gradient(135deg, #f44336, #ef5350);
}
.source-badge {
display: inline-block;
width: 18px;
height: 18px;
border-radius: 50%;
text-align: center;
line-height: 18px;
font-size: 10px;
font-weight: 700;
color: #fff;
vertical-align: middle;
cursor: help;
}
.gender-badge {
display: inline-block;
padding: 2px 10px;
border-radius: 10px;
font-size: 0.72em;
font-weight: 600;
letter-spacing: 0.02em;
vertical-align: middle;
border: 1px solid transparent;
}
.gender-female {
background: rgba(220, 120, 200, 0.18);
color: #e8a8d8;
border-color: rgba(220, 120, 200, 0.35);
}
.gender-male {
background: rgba(100, 160, 230, 0.18);
color: #9cc2ee;
border-color: rgba(100, 160, 230, 0.35);
}
.gender-non-binary {
background: rgba(120, 210, 130, 0.18);
color: #a8e0b0;
border-color: rgba(120, 210, 130, 0.35);
}
.gender-transgender-female {
background: rgba(220, 140, 200, 0.12);
color: #d8a8c8;
border-color: rgba(220, 140, 200, 0.28);
}
.gender-transgender-male {
background: rgba(100, 160, 230, 0.12);
color: #a8c0e0;
border-color: rgba(100, 160, 230, 0.28);
}
.source-stashdb { background: #4a6fa5; }
.source-javstash { background: #c0392b; }
.source-fansdb { background: #27ae60; }
</style>
""")
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"""
<div class="performer-card">
<div class="face-info">
<div class="detected-face">
<img src="{face_image_b64}" alt="Detected Face {i+1}">
</div>
<div class="face-info-content">
<h3>Face {i+1}</h3>
<p><strong>Detection:</strong> {face_confidence:.1%} &nbsp;|&nbsp; <strong>Matches:</strong> {len(performers)}</p>
</div>
</div>
""")
if performers:
html_parts.append('<div class="performer-grid">')
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'<span class="source-badge source-{source}" title="{source}">{initial}</span>'
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'<span class="gender-badge gender-{gender_key}">{gender_display}</span>'
html_parts.append(f"""
<div class="performer-item">
<div class="performer-image-wrap">
<img src="{performer['image']}" alt="{performer['name']}" class="performer-image" onerror="this.style.display='none'">
</div>
<div class="performer-body">
<div class="performer-header">
<h4 class="performer-name"><a href="{performer['performer_url']}" target="_blank">{performer['name']}</a></h4>
{source_badge}
{gender_badge}
<span class="performer-country"><span class="country-flag">{country_flag}</span>{country_code if country_code else 'Unknown'}</span>
</div>
<div class="confidence-bar">
<div class="confidence-fill {confidence_class}" style="width: {performer['confidence']}%">
{performer['confidence']}%
</div>
</div>
</div>
</div>
""")
html_parts.append('</div>')
else:
html_parts.append('<p><em>No performer matches found for this face.</em></p>')
html_parts.append('</div>')
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"<div class='performer-card'><h3>Error</h3><p>{str(e)}</p></div>"
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="<p>Upload an image and click search to see results.</p>"
)
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)