| import faiss | |
| import numpy as np | |
| import pickle | |
| import os | |
| class SimilaritySearchEngine: | |
| def __init__(self): | |
| base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| embeddings_path = os.path.join(base_dir, 'data', 'embeddings.pkl') | |
| with open(embeddings_path, 'rb') as f: | |
| data = pickle.load(f) | |
| self.embeddings = data['embeddings'] | |
| # Convert Windows paths to Linux paths | |
| self.image_paths = [os.path.normpath(path).replace('\\', '/') | |
| for path in data['image_paths']] | |
| dimension = len(self.embeddings[0]) | |
| self.index = faiss.IndexFlatL2(dimension) | |
| self.index.add(np.array(self.embeddings)) | |
| def search_similar_images(self, query_embedding, top_k=5): | |
| distances, indices = self.index.search(np.array([query_embedding]), top_k) | |
| return [self.image_paths[idx] for idx in indices[0]], distances[0] |