Spaces:
Sleeping
Sleeping
| import os | |
| import cloudinary | |
| import cloudinary.uploader | |
| from pinecone import Pinecone | |
| from dotenv import load_dotenv | |
| # Load keys from the .env file | |
| load_dotenv() | |
| class CloudDB: | |
| def __init__(self): | |
| # 1. Connect to Cloudinary | |
| cloudinary.config( | |
| cloud_name=os.getenv("CLOUDINARY_CLOUD_NAME"), | |
| api_key=os.getenv("CLOUDINARY_API_KEY"), | |
| api_secret=os.getenv("CLOUDINARY_API_SECRET") | |
| ) | |
| # 2. Connect to Pinecone | |
| self.pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY")) | |
| self.index = self.pc.Index(os.getenv("PINECONE_INDEX_NAME")) | |
| def upload_image(self, file_path, folder_name="visual_search"): | |
| """Uploads an image to Cloudinary and returns the public URL.""" | |
| response = cloudinary.uploader.upload(file_path, folder=folder_name) | |
| return response['secure_url'] | |
| def add_vector(self, vector, image_url, image_id): | |
| """Saves the vector and the image URL to Pinecone.""" | |
| # Convert numpy array to list for Pinecone | |
| vector_list = vector.tolist() if hasattr(vector, 'tolist') else vector | |
| self.index.upsert(vectors=[{ | |
| "id": image_id, | |
| "values": vector_list, | |
| "metadata": {"image_url": image_url} | |
| }]) | |
| def search(self, query_vector, top_k=10, min_score=0.60): # <-- CHANGED baseline to 0.60 | |
| """Searches Pinecone and filters out baseline 'random noise' matches.""" | |
| vector_list = query_vector.tolist() if hasattr(query_vector, 'tolist') else query_vector | |
| response = self.index.query( | |
| vector=vector_list, | |
| top_k=top_k, | |
| include_metadata=True | |
| ) | |
| results = [] | |
| for match in response['matches']: | |
| # Only keep the image if it's an ACTUAL mathematical match (60% or higher) | |
| if match['score'] >= min_score: | |
| results.append({ | |
| "url": match['metadata']['image_url'], | |
| "score": match['score'] | |
| }) | |
| return results |