From 43a979041e30e50a8c01e436e8583304d74123cb Mon Sep 17 00:00:00 2001 From: Aroy-Art Date: Mon, 2 Sep 2024 21:36:40 +0200 Subject: [PATCH] Add: ask to load image index if it exists --- Python/similar-images-classification/main.py | 34 ++++++++++++++++---- 1 file changed, 27 insertions(+), 7 deletions(-) diff --git a/Python/similar-images-classification/main.py b/Python/similar-images-classification/main.py index efc2dbc..a0c1406 100644 --- a/Python/similar-images-classification/main.py +++ b/Python/similar-images-classification/main.py @@ -58,15 +58,35 @@ image_paths = [ ] features = [] -for image_path in image_paths: - img_feature = extract_features(image_path, model) - features.append(img_feature) +if os.path.exists("image_index.bin"): + if query_yes_no("Load the index?", default="yes"): + index = faiss.read_index("image_index.bin") + else: + for image_path in image_paths: + img_feature = extract_features(image_path, model) + features.append(img_feature) -features = np.array(features) + features = np.array(features) -d = features.shape[1] -index = faiss.IndexFlatL2(d) -index.add(features) + d = features.shape[1] + index = faiss.IndexFlatL2(d) + index.add(features) + + if query_yes_no("Save the index?", default="yes"): + faiss.write_index(index, "image_index.bin") +else: + for image_path in image_paths: + img_feature = extract_features(image_path, model) + features.append(img_feature) + + features = np.array(features) + + d = features.shape[1] + index = faiss.IndexFlatL2(d) + index.add(features) + + if query_yes_no("Save the index?", default="yes"): + faiss.write_index(index, "image_index.bin") def find_similar_images(query_image_path, index, k=6):