Update image_search.py

image_search.py

@@ -264,6 +264,7 @@ from chromadb.utils.embedding_functions import EmbeddingFunction
 # Optional dataset downloader
 from download import download_flickr8k_dataset
 
+TOP_K = 5 
 
 # -------------------------
 # HuggingFace CLIP Wrapper
@@ -427,3 +428,11 @@ def update_custom_image_paths(folder_path):
     global custom_image_paths
     custom_image_paths = [os.path.join(folder_path, name) for name in os.listdir(folder_path) if name.lower().endswith((".jpg", ".jpeg", ".png"))]
 
+def query_by_image(image_path, nn_model, image_paths):
+    image = Image.open(image_path).convert("RGB")
+    tensor = clip_preprocess(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        image_features = clip_model.encode_image(tensor).cpu().numpy()
+    dists, indices = nn_model.kneighbors(image_features, n_neighbors=TOP_K)
+    return [image_paths[i] for i in indices[0]]
+