Gateway / src /services /db_client.py
Adidasa108's picture
Upload 3 files
24fe247 verified
Raw
History Blame Contribute Delete
7.26 kB
import time
from typing import Any, Dict, List
import cloudinary
import cloudinary.uploader
import cloudinary.api
from pinecone import Pinecone, ServerlessSpec
from src.core.config import IDX_FACES, IDX_OBJECTS, FACE_MATCH_THRESHOLD
class PineconePool:
def __init__(self):
self._clients = {}
def get(self, api_key: str) -> Pinecone:
if api_key not in self._clients:
self._clients[api_key] = Pinecone(api_key=api_key)
return self._clients[api_key]
pinecone_pool = PineconePool()
def _set_cld_config(creds: dict):
cloudinary.config(
cloud_name=creds.get("cloud_name"),
api_key=creds.get("api_key"),
api_secret=creds.get("api_secret"),
secure=True
)
def cld_ping(creds: dict):
_set_cld_config(creds)
cloudinary.api.ping()
def cld_upload(file_obj, folder: str, creds: dict) -> dict:
_set_cld_config(creds)
return cloudinary.uploader.upload(file_obj, folder=folder)
def cld_root_folders(creds: dict) -> dict:
_set_cld_config(creds)
return cloudinary.api.root_folders()
def cld_list_folder_images(folder: str, creds: dict, cursor: str = None, page_size: int = 100) -> dict:
_set_cld_config(creds)
kwargs = {"type": "upload", "prefix": f"{folder}/", "max_results": page_size}
if cursor:
kwargs["next_cursor"] = cursor
return cloudinary.api.resources(**kwargs)
def cld_delete_resource(public_id: str, creds: dict):
_set_cld_config(creds)
cloudinary.uploader.destroy(public_id)
def cld_delete_folder_resources(folder: str, creds: dict):
_set_cld_config(creds)
cloudinary.api.delete_resources_by_prefix(f"{folder}/")
def cld_remove_folder(folder: str, creds: dict):
_set_cld_config(creds)
try:
cloudinary.api.delete_folder(folder)
except Exception:
pass
def cld_delete_all_paginated(creds: dict) -> int:
_set_cld_config(creds)
deleted = 0
cursor = None
while True:
kwargs = {"type": "upload", "max_results": 500}
if cursor:
kwargs["next_cursor"] = cursor
res = cloudinary.api.resources(**kwargs)
resources = res.get("resources", [])
if not resources:
break
pids = [r["public_id"] for r in resources]
cloudinary.api.delete_resources(pids)
deleted += len(pids)
cursor = res.get("next_cursor")
if not cursor:
break
return deleted
def ensure_indexes(pc: Pinecone) -> List[str]:
created = []
existing = [idx.name for idx in pc.list_indexes()]
for name in [IDX_FACES, IDX_OBJECTS]:
if name not in existing:
pc.create_index(
name=name,
dimension=1024 if name == IDX_FACES else 1536,
metric="cosine",
spec=ServerlessSpec(cloud="aws", region="us-east-1")
)
created.append(name)
return created
def delete_and_recreate_indexes(pc: Pinecone):
existing = [idx.name for idx in pc.list_indexes()]
for name in [IDX_FACES, IDX_OBJECTS]:
if name in existing:
pc.delete_index(name)
time.sleep(5)
ensure_indexes(pc)
def search_faces(idx, vec: List[float], det_score: float, filter_dict: dict = None) -> Dict[str, Any]:
query_kwargs = {"vector": vec, "top_k": 50, "include_metadata": True}
if filter_dict:
query_kwargs["filter"] = filter_dict
res = idx.query(**query_kwargs)
image_map = {}
for match in res.get("matches", []):
raw_score = match.get("score", 0)
# Drop strangers immediately
if raw_score < FACE_MATCH_THRESHOLD:
continue
meta = match.get("metadata", {})
url = meta.get("url")
if not url:
continue
if url not in image_map or image_map[url]["raw_score"] < raw_score:
image_map[url] = {
"raw_score": raw_score,
"face_crop": meta.get("face_crop", ""),
"folder": meta.get("folder", "uncategorized")
}
return image_map
def search_objects(idx, vec: List[float]) -> List[Dict[str, Any]]:
res = idx.query(vector=vec, top_k=50, include_metadata=True)
results = []
for match in res.get("matches", []):
meta = match.get("metadata", {})
results.append({
"url": meta.get("url", ""),
# NEW UPDATED CODE: Removed the * 100
"score": round(match.get("score", 0), 4),
"raw_score": match.get("score", 0),
"folder": meta.get("folder", "uncategorized")
})
return results
def merge_face_results(groups: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
merged = {}
for group in groups:
for match in group.get("matches", []):
url = match["url"]
if url not in merged or merged[url]["score"] < match["score"]:
merged[url] = match
return sorted(merged.values(), key=lambda x: x["score"], reverse=True)
def merge_object_results(nested_results: List[List[Dict[str, Any]]]) -> List[Dict[str, Any]]:
merged = {}
for res_list in nested_results:
for match in res_list:
url = match["url"]
if url not in merged or merged[url]["score"] < match["score"]:
merged[url] = match
return sorted(merged.values(), key=lambda x: x["score"], reverse=True)
def intersect_face_results(groups: List[Dict[str, Any]], min_faces_required: int = 1) -> List[Dict[str, Any]]:
"""
Returns photos appearing in >= min_faces_required face groups.
min_faces_required=1 → union (all photos matching any face).
min_faces_required=len(groups) → intersection (photos where ALL searched faces appear together).
Score = average across matched groups, normalized by total query face count.
"""
from collections import defaultdict
url_scores: Dict[str, list] = defaultdict(list)
url_folder: Dict[str, str] = {}
url_face_crop: Dict[str, str] = {}
for group in groups:
seen_urls_this_group = set()
for match in group.get("matches", []):
url = match["url"]
if url in seen_urls_this_group:
continue
seen_urls_this_group.add(url)
url_scores[url].append(match.get("raw_score", match.get("score", 0)))
url_folder[url] = match.get("folder", "uncategorized")
if not url_face_crop.get(url):
url_face_crop[url] = match.get("face_crop", "")
results = []
n_groups = len(groups)
for url, scores in url_scores.items():
if len(scores) >= min_faces_required:
results.append({
"url": url,
"score": round(sum(scores) / n_groups, 4),
"raw_score": round(sum(scores) / n_groups, 4),
"matched_faces": len(scores),
"total_query_faces": n_groups,
"face_crop": url_face_crop.get(url, ""),
"folder": url_folder.get(url, "uncategorized"),
"caption": f"👥 {len(scores)}/{n_groups} faces matched",
})
return sorted(results, key=lambda x: x["score"], reverse=True)