Spaces:
Running
Running
| import os | |
| import cloudinary | |
| import cloudinary.uploader | |
| from pinecone import Pinecone | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| class CloudDB: | |
| def __init__(self): | |
| cloudinary.config( | |
| cloud_name=os.getenv("CLOUDINARY_CLOUD_NAME"), | |
| api_key=os.getenv("CLOUDINARY_API_KEY"), | |
| api_secret=os.getenv("CLOUDINARY_API_SECRET") | |
| ) | |
| self.pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY")) | |
| # Connect to the TWO new indexes | |
| self.index_faces = self.pc.Index("enterprise-faces") | |
| self.index_objects = self.pc.Index("enterprise-objects") | |
| def upload_image(self, file_path, folder_name="visual_search"): | |
| response = cloudinary.uploader.upload(file_path, folder=folder_name) | |
| return response['secure_url'] | |
| def add_vector(self, data_dict, image_url, image_id): | |
| vector_list = data_dict["vector"].tolist() if hasattr(data_dict["vector"], 'tolist') else data_dict["vector"] | |
| payload = [{ | |
| "id": image_id, | |
| "values": vector_list, | |
| "metadata": {"image_url": image_url} | |
| }] | |
| if data_dict["type"] == "face": | |
| self.index_faces.upsert(vectors=payload) | |
| else: | |
| self.index_objects.upsert(vectors=payload) | |
| def search(self, query_dict, top_k=10, min_score=0.45): | |
| vector_list = query_dict["vector"].tolist() if hasattr(query_dict["vector"], 'tolist') else query_dict["vector"] | |
| results = [] | |
| if query_dict["type"] == "face": | |
| response = self.index_faces.query(vector=vector_list, top_k=top_k, include_metadata=True) | |
| RAW_THRESHOLD = 0.35 | |
| for match in response['matches']: | |
| raw_score = match['score'] | |
| if raw_score >= RAW_THRESHOLD: | |
| ui_score = 0.75 + ((raw_score - RAW_THRESHOLD) / (1.0 - RAW_THRESHOLD)) * 0.24 | |
| ui_score = min(0.99, ui_score) | |
| results.append({ | |
| "url": match['metadata']['image_url'], | |
| "score": ui_score, | |
| "caption": "👤 Verified Identity Match" | |
| }) | |
| else: | |
| response = self.index_objects.query(vector=vector_list, top_k=top_k, include_metadata=True) | |
| for match in response['matches']: | |
| if match['score'] >= min_score: | |
| results.append({ | |
| "url": match['metadata']['image_url'], | |
| "score": match['score'], | |
| "caption": "🎯 Visual & Semantic Match" | |
| }) | |
| return results |