| import os |
| import pickle |
| import sys |
| import keras |
| import numpy as np |
| from keras.preprocessing import image |
| from keras.layers import GlobalMaxPooling2D |
| from keras.applications.resnet50 import ResNet50, preprocess_input |
| from sklearn.neighbors import NearestNeighbors |
| from numpy.linalg import norm |
|
|
| model = None |
| feature_list = None |
| filenames = None |
|
|
| def load_model(): |
| global model, feature_list, filenames |
| if model is None: |
| model = ResNet50(weights='imagenet', include_top=False, input_shape=(224, 224, 3)) |
| model.trainable = False |
| model = keras.Sequential([ |
| model, |
| GlobalMaxPooling2D() |
| ]) |
| |
| script_dir = os.path.dirname(os.path.abspath(__file__)) |
| embeddings_path = os.path.join(script_dir, 'res_vector_embeddings.pkl') |
| filenames_path = os.path.join(script_dir, 'res_filenames.pkl') |
|
|
| try: |
| with open(embeddings_path, 'rb') as emb_file, open(filenames_path, 'rb') as name_file: |
| feature_list = pickle.load(emb_file) |
| filenames = pickle.load(name_file) |
| except FileNotFoundError as e: |
| print(f"Error: {e}. Check if the required files exist in the specified path.") |
| sys.exit(1) |
| except Exception as e: |
| print(f"Error loading pickle files: {e}") |
| sys.exit(1) |
|
|
| def find_similar_images(image_path): |
| if model is None or feature_list is None or filenames is None: |
| load_model() |
|
|
| try: |
| query_img = image.load_img(image_path, target_size=(224, 224)) |
| query_img_array = image.img_to_array(query_img) |
| expanded_query_img_array = np.expand_dims(query_img_array, axis=0) |
| preprocessed_query_img = preprocess_input(expanded_query_img_array) |
| query_result = model.predict(preprocessed_query_img).flatten() |
| normalized_query_result = query_result / norm(query_result) |
|
|
| neighbors = NearestNeighbors(n_neighbors=100, algorithm='brute', metric='euclidean') |
| neighbors.fit(feature_list) |
|
|
| distances, indices = neighbors.kneighbors([normalized_query_result]) |
|
|
| similar_image_paths = [filenames[idx] for idx in indices[0][1:]] |
| return similar_image_paths |
| except FileNotFoundError as e: |
| print(f"Error: {e}. Check if the specified image file exists.") |
| return [] |
| except Exception as e: |
| print(f"An error occurred: {e}") |
| return [] |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| def extract_filenames(paths): |
| return [path.split("\\")[-1] for path in paths] |
