Aroy/LangPlayground
1
0
Fork 0
Code Activity Actions

Add: ask to load image index if it exists

Browse source
This commit is contained in:
Aroy-Art 2024-09-02 21:36:40 +02:00
parent 977be10b92
commit 43a979041e
Signed by: Aroy
GPG key ID: DB9689E9391DD156
1 changed files with 27 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
Python/similar-images-classification/main.py
View file

@ -58,6 +58,10 @@ image_paths = [
] ]
features = [] features = []
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: for image_path in image_paths:
img_feature = extract_features(image_path, model) img_feature = extract_features(image_path, model)
features.append(img_feature) features.append(img_feature)
@ -68,6 +72,22 @@ d = features.shape[1]
index = faiss.IndexFlatL2(d) index = faiss.IndexFlatL2(d)
index.add(features) 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): def find_similar_images(query_image_path, index, k=6):
query_feature = extract_features(query_image_path, model).reshape(1, -1) query_feature = extract_features(query_image_path, model).reshape(1, -1)