Update src/cloud_db.py

src/cloud_db.py

@@ -4,55 +4,65 @@ 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"))
+        # 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"):
-        """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
+    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"]
         
-        self.index.upsert(vectors=[{
+        payload = [{
             "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
-        )
-        
+        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 = []
-        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']
-                })
+        
+        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