Spaces:
Sleeping
Sleeping
File size: 8,729 Bytes
4b61829 e7e57d0 4b61829 e7e57d0 4b61829 | 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 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | import asyncio
import time
import traceback
from fastapi import APIRouter, File, Form, HTTPException, Request, UploadFile, Depends
from src.core.config import (
DEFAULT_PINECONE_KEY, IDX_FACES, IDX_OBJECTS,
FACE_MATCH_THRESHOLD, REDIS_TTL_SEARCH, REDIS_TTL_INFERENCE,
FACE_SPACE_URL, OBJECT_SPACE_URL,
FACE_QUALITY_GATE_SEARCH,
)
from src.core.security import get_verified_keys
from src.services.db_client import (
intersect_face_results, merge_face_results, merge_object_results,
pinecone_pool, search_faces, search_objects,
)
from src.services.cache import cache_get, cache_set, make_face_search_key, make_obj_search_key, make_inference_key
from src.core.logging import log
from src.common.utils import face_ui_score, get_ip, is_default_key, to_list
router = APIRouter()
@router.post("/api/search")
async def search_database(
request: Request,
file: UploadFile = File(...),
detect_faces: bool = Form(True),
user_id: str = Form(""),
keys: dict = Depends(get_verified_keys)
):
ip = get_ip(request)
start = time.perf_counter()
mode = "guest" if is_default_key(keys["pinecone_key"], DEFAULT_PINECONE_KEY) else "personal"
log("INFO", "search.start", user_id=user_id or "anonymous", ip=ip, mode=mode,
filename=file.filename, detect_faces=detect_faces)
try:
file_bytes = await file.read()
ai_manager = request.app.state.ai
sem = request.app.state.ai_semaphore
if FACE_SPACE_URL and OBJECT_SPACE_URL:
# Distributed mode: call remote Spaces for embedding, check Redis cache first
print(f"[DEBUG] Calling face space at: {FACE_SPACE_URL}") # add this
print(f"[DEBUG] Calling object space at: {OBJECT_SPACE_URL}") # add this
from src.services.space_client import embed_face, embed_object
inference_cache_key = make_inference_key(file_bytes, detect_faces, "search")
vectors = await cache_get(inference_cache_key)
if vectors is None:
face_task = embed_face(file_bytes, quality_gate=FACE_QUALITY_GATE_SEARCH) if detect_faces else asyncio.sleep(0)
obj_task = embed_object(file_bytes)
face_vecs, obj_vecs = await asyncio.gather(face_task, obj_task, return_exceptions=True)
if isinstance(face_vecs, Exception) or face_vecs is None:
log("WARNING", "search.face_space_unavailable", user_id=user_id or "anonymous", ip=ip)
face_vecs = []
if isinstance(obj_vecs, Exception) or obj_vecs is None:
raise HTTPException(503, "Object AI service is waking up (free tier). Please retry in 60 seconds.")
vectors = [*(face_vecs or []), *(obj_vecs or [])]
await cache_set(inference_cache_key, vectors, ttl=REDIS_TTL_INFERENCE)
else:
# Monolith mode: local inference
async with sem:
vectors = await ai_manager.process_image_bytes_async(file_bytes, detect_faces=detect_faces, mode="search")
inference_ms = round((time.perf_counter() - start) * 1000)
face_vectors = [v for v in vectors if v["type"] == "face"]
object_vectors = [v for v in vectors if v["type"] == "object"]
lanes_used = list({v["type"] for v in vectors})
log("INFO", "search.inference_done", user_id=user_id or "anonymous", ip=ip, mode=mode,
face_vecs=len(face_vectors), obj_vecs=len(object_vectors), inference_ms=inference_ms)
pc = pinecone_pool.get(keys["pinecone_key"])
idx_obj = pc.Index(IDX_OBJECTS)
idx_face = pc.Index(IDX_FACES)
if detect_faces and face_vectors:
return await _run_face_search(face_vectors, object_vectors, idx_face, idx_obj, start, user_id, ip, mode, lanes_used)
else:
return await _run_object_search(object_vectors, idx_obj, start, user_id, ip, mode, lanes_used)
except HTTPException:
raise
except Exception as e:
log("ERROR", "search.error", user_id=user_id or "anonymous", ip=ip, mode=mode,
error=str(e), traceback=traceback.format_exc()[-800:])
raise HTTPException(500, str(e))
async def _run_face_search(face_vectors, object_vectors, idx_face, idx_obj, start, user_id, ip, mode, lanes_used) -> dict:
async def _query_face(fv: dict) -> dict:
vec = to_list(fv["vector"])
det_score = fv.get("det_score", 1.0)
try:
ck = make_face_search_key(vec, FACE_MATCH_THRESHOLD)
image_map = await cache_get(ck)
if image_map is None:
image_map = await asyncio.to_thread(search_faces, idx_face, vec, det_score)
await cache_set(ck, image_map, ttl=REDIS_TTL_SEARCH)
except Exception as e:
if "404" in str(e):
raise HTTPException(404, "Pinecone index not found. Go to Settings → Verify & Save.")
raise
return {
"query_face_idx": fv.get("face_idx", 0),
"query_face_crop": fv.get("face_crop", ""),
"query_bbox": fv.get("bbox", []),
"det_score": det_score,
"face_width_px": fv.get("face_width_px", 0),
"matches": sorted(
[
{
"url": url,
"score": face_ui_score(d["raw_score"]),
"raw_score": round(d["raw_score"], 4),
"face_crop": d["face_crop"],
"folder": d["folder"],
"caption": "👤 Verified Identity",
}
for url, d in image_map.items()
],
key=lambda x: x["score"], reverse=True,
)[:50],
}
async def _query_obj_single(ov: dict) -> list:
vec = to_list(ov["vector"])
try:
ck = make_obj_search_key(vec)
result = await cache_get(ck)
if result is None:
result = await asyncio.to_thread(search_objects, idx_obj, vec)
await cache_set(ck, result, ttl=REDIS_TTL_SEARCH)
return result
except Exception as e:
if "404" in str(e):
raise HTTPException(404, "Pinecone index not found.")
raise
face_tasks = [_query_face(fv) for fv in face_vectors]
obj_tasks = [_query_obj_single(ov) for ov in object_vectors]
all_results = await asyncio.gather(*face_tasks, *obj_tasks)
raw_groups = list(all_results[:len(face_tasks)])
obj_nested = list(all_results[len(face_tasks):])
merged_face = merge_face_results(raw_groups)
merged_objects = merge_object_results(obj_nested)
face_groups = [g for g in raw_groups if g.get("matches")]
# Photos where ALL searched faces appear together (only meaningful for multi-face queries)
group_results = intersect_face_results(face_groups, min_faces_required=len(face_groups)) if len(face_groups) > 1 else []
duration_ms = round((time.perf_counter() - start) * 1000)
log("INFO", "search.complete", user_id=user_id or "anonymous", ip=ip, mode=mode,
lanes=["face", "object"], face_groups=len(face_groups), face_results=len(merged_face),
group_results=len(group_results), object_results=len(merged_objects), duration_ms=duration_ms)
return {
"mode": "face",
"face_groups": face_groups,
"results": merged_face,
"group_results": group_results,
"object_results": merged_objects,
}
async def _run_object_search(object_vectors, idx_obj, start, user_id, ip, mode, lanes_used) -> dict:
if not object_vectors:
return {"mode": "object", "results": [], "face_groups": []}
async def _query_obj(ov: dict) -> list:
vec = to_list(ov["vector"])
try:
ck = make_obj_search_key(vec)
result = await cache_get(ck)
if result is None:
result = await asyncio.to_thread(search_objects, idx_obj, vec)
await cache_set(ck, result, ttl=REDIS_TTL_SEARCH)
return result
except Exception as e:
if "404" in str(e):
raise HTTPException(404, "Pinecone index not found.")
raise
nested = await asyncio.gather(*[_query_obj(ov) for ov in object_vectors])
final = merge_object_results(nested)
duration_ms = round((time.perf_counter() - start) * 1000)
log("INFO", "search.complete", user_id=user_id or "anonymous", ip=ip, mode=mode,
lanes=lanes_used, results=len(final), duration_ms=duration_ms)
return {"mode": "object", "results": final, "face_groups": []} |