Spaces:
Sleeping
Sleeping
| 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) | |