Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import os | |
| from tensorflow.keras.applications.resnet50 import ResNet50, preprocess_input | |
| from tensorflow.keras.preprocessing import image | |
| from tensorflow.keras.models import Model | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| from extract_features import extract_features, download_and_extract_dataset | |
| # Ensure dataset is available | |
| download_and_extract_dataset() | |
| # Ensure features exist — otherwise extract them | |
| FEATURES_PATH = 'tmp_features/features.npy' | |
| FILENAMES_PATH = 'tmp_features/filenames.npy' | |
| if not (os.path.exists(FEATURES_PATH) and os.path.exists(FILENAMES_PATH)): | |
| print(" Features not found. Generating...") | |
| extract_features() | |
| else: | |
| print(" Found cached features. Loading...") | |
| # Load features | |
| features = np.load(FEATURES_PATH) | |
| filenames = np.load(FILENAMES_PATH) | |
| # Load model | |
| base_model = ResNet50(weights='imagenet', include_top=False, pooling='avg') | |
| model = Model(inputs=base_model.input, outputs=base_model.output) | |
| # Load saved features and filenames | |
| FEATURES_PATH = 'features/features.npy' | |
| FILENAMES_PATH = 'features/filenames.npy' | |
| features = np.load(FEATURES_PATH) | |
| filenames = np.load(FILENAMES_PATH) | |
| def preprocess_img(img_path): | |
| img = image.load_img(img_path, target_size=(224, 224)) | |
| x = image.img_to_array(img) | |
| x = np.expand_dims(x, axis=0) | |
| return preprocess_input(x) | |
| def get_image_embedding(img_path): | |
| img_tensor = preprocess_img(img_path) | |
| embedding = model.predict(img_tensor, verbose=0)[0] | |
| return embedding | |
| def recommend(query_img_path, top_k=5): | |
| query_embedding = get_image_embedding(query_img_path).reshape(1, -1) | |
| similarities = cosine_similarity(query_embedding, features)[0] | |
| # Load filenames | |
| filenames_arr = np.load(FILENAMES_PATH) | |
| # Exclude the image that's an exact match to the uploaded one | |
| similarity_list = [] | |
| for i, fname in enumerate(filenames_arr): | |
| if os.path.abspath(fname) != os.path.abspath(query_img_path): # avoid self-match | |
| similarity_list.append((fname, similarities[i])) | |
| # Sort and return top_k | |
| similarity_list = sorted(similarity_list, key=lambda x: x[1], reverse=True) | |
| return similarity_list[:top_k] | |
| # Example usage | |
| if __name__ == "__main__": | |
| query_path = 'data/sneakers/0ed14658bcd730081c79a5224bfb4100.jpg' | |
| results = recommend(query_path, top_k=5) | |
| for path, score in results: | |
| print(f"{path} - Similarity: {score:.4f}") | |