Update one.py

one.py

@@ -1,72 +1,50 @@
+import torch
+import torch.nn as nn
+import torchvision.models as models
+import torchvision.transforms as transforms
 import os
-import numpy as np
 from PIL import Image
-import tensorflow as tf
-from tensorflow.keras.applications.resnet50 import ResNet50
 
-# Set the path to the dataset
-data_path = '/content/lfw/'
-
-# Load the dataset
-images = []
-labels = []
-
-for folder_name in os.listdir(data_path):
-    folder_path = os.path.join(data_path, folder_name)
-    if not os.path.isdir(folder_path):
-        continue
-    for file_name in os.listdir(folder_path):
-        file_path = os.path.join(folder_path, file_name)
-        if not file_path.endswith('.jpg'):
-            continue
-        image = np.array(Image.open(file_path).convert('RGB'))
-        label = folder_name
-        images.append(image)
-        labels.append(label)
-
-# Convert to numpy arrays
-images = np.array(images)
-labels = np.array(labels)
-
-# Perform necessary preprocessing on the images
-preprocessed_images = tf.keras.applications.resnet50.preprocess_input(images)
-
-# Obtain a ResNet50 model pre-trained on ImageNet
-model = ResNet50(include_top=False, pooling='avg')
-
-# Extract features from the penultimate layer of the network
-features = model.predict(preprocessed_images)
-
-# Store the features in a dictionary
-features_dict = {}
-for i in range(len(labels)):
-    features_dict[labels[i]] = features[i]
-
-# Use a nearest neighbor algorithm to obtain the 10 most similar images to each query image
-from sklearn.neighbors import NearestNeighbors
-
-# Initialize the nearest neighbor algorithm with cosine distance
-nn = NearestNeighbors(n_neighbors=10, metric='cosine')
-
-# Fit the algorithm to the features
-nn.fit(list(features_dict.values()))
-
-# Define a function to retrieve the most similar images to a query image
-def retrieve_similar_images(query_image_path):
-    # Load the query image
-    query_image = np.array(Image.open(query_image_path).convert('RGB'))
-
-    # Perform necessary preprocessing on the query image
-    preprocessed_query_image = tf.keras.applications.resnet50.preprocess_input(np.array([query_image]))
-
-    # Extract features from the query image
-    query_features = model.predict(preprocessed_query_image)
-
-    # Use the nearest neighbor algorithm to retrieve the most similar images
-    distances, indices = nn.kneighbors(query_features)
-
-    # Display the most similar images
-    for i in range(len(indices[0])):
-        image_path = list(features_dict.keys())[list(features_dict.values()).index(features[indices[0][i]])]
-        image = Image.open(os.path.join(data_path, image_path)).convert('RGB')
-        image.show()
+# Define the ResNet-50 model
+model = models.resnet50(pretrained=True)
+
+# Remove the classification head (the fully connected layer)
+num_features = model.fc.in_features
+model.fc = nn.Identity()
+
+# Set the model to evaluation mode
+model.eval()
+
+# Define the preprocessing transforms
+preprocess = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    )
+])
+
+# Define the dictionary to store the feature vectors
+features = {}
+
+# Iterate over the images and extract the features
+image_dir = 'lfw'
+for root, dirs, files in os.walk(image_dir):
+    for file in files:
+        # Load the image
+        image_path = os.path.join(root, file)
+        image = Image.open(image_path).convert('RGB')
+
+        # Apply the preprocessing transforms
+        input_tensor = preprocess(image)
+        input_batch = input_tensor.unsqueeze(0)
+
+        # Extract the features from the penultimate layer
+        with torch.no_grad():
+            features_tensor = model(input_batch)
+            features_vector = torch.squeeze(features_tensor).numpy()
+
+        # Store the feature vector in the dictionary
+        features[file] = features_vector