FaceNet / pineconevectors.py
mdirshad09's picture
Rename pinecone.py to pineconevectors.py
f4150c0
import uuid
import pinecone
class Pinecone:
def __init__(self, api_key, environment, index_name):
self.api_key = api_key
self.environment = environment
self.index_name = index_name
self.index = None
def initialize_index(self):
pinecone.init(api_key=self.api_key, environment=self.environment)
self.index = pinecone.Index(self.index_name)
return self.index
def upsert_data(self, img_id, embeddings, path, img_url, page_url, face_coordinates, single_face):
vec_id = str(uuid.uuid4())
data = []
embedding_as_list = embeddings.tolist()
if face_coordinates is not None:
coordinates_1d = [str(coord) for coord in face_coordinates]
metadata = {'Image id': img_id, 'directory path': path, 'Image URL': img_url, 'Page URL': page_url, 'Face Coordinates': coordinates_1d, 'Single Face': single_face}
data.append((vec_id, embedding_as_list, metadata))
self.index.upsert(data)
def search_data(self, query_embedding):
matches = self.index.query(
vector=query_embedding,
top_k=10,
include_values=True,
include_metadata = True
)
return matches