Spaces:
Sleeping
Sleeping
File size: 2,108 Bytes
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 | 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 |