File size: 2,703 Bytes
3e805ab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0df58f0
 
 
3e805ab
 
 
 
 
0df58f0
 
3e805ab
0df58f0
3e805ab
 
 
0df58f0
3e805ab
0df58f0
 
 
 
 
 
 
3e805ab
0df58f0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3e805ab
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
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