Prathamesh1420's picture
Upload 11 files
66c277d verified
raw
history blame
1.78 kB
import pickle
import tensorflow
import numpy as np
from numpy.linalg import norm
from tensorflow.keras.preprocessing import image
from tensorflow.keras.layers import GlobalMaxPooling2D
from tensorflow.keras.applications.resnet50 import ResNet50, preprocess_input
from sklearn.neighbors import NearestNeighbors
import cv2
# Load the precomputed feature vectors and filenames from pickle files
feature_list = np.array(pickle.load(open('embeddings.pkl', 'rb')))
filenames = pickle.load(open('filenames.pkl', 'rb'))
# Load ResNet50 model without the top layer for feature extraction
model = ResNet50(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
model.trainable = False
# Create a Sequential model with ResNet50 and GlobalMaxPooling2D layers
model = tensorflow.keras.Sequential([
model,
GlobalMaxPooling2D()
])
# Load and preprocess the query image
img = image.load_img('sample/khade.jpg', target_size=(224, 224))
img_array = image.img_to_array(img)
expanded_img_array = np.expand_dims(img_array, axis=0)
preprocessed_img = preprocess_input(expanded_img_array)
# Extract features from the query image and normalize
result = model.predict(preprocessed_img).flatten()
normalized_result = result / norm(result)
# Initialize NearestNeighbors model and fit with the feature vectors
neighbors = NearestNeighbors(n_neighbors=6, algorithm='brute', metric='euclidean')
neighbors.fit(feature_list)
# Find the nearest neighbors (excluding itself)
distances, indices = neighbors.kneighbors([normalized_result])
print(indices) # Print the indices of nearest neighbors
# Display the nearest neighbor images
for file in indices[0][0:5]:
temp_img = cv2.imread(filenames[file])
cv2.imshow('output', cv2.resize(temp_img, (312, 312)))
cv2.waitKey(0)