Create app.py

app.py

@@ -0,0 +1,67 @@
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+import os
+from sklearn.neighbors import NearestNeighbors
+import numpy as np
+
+
+# Load pre-trained ResNet-50 model
+model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet50', weights=None)
+model.eval()
+
+# Define image transformation
+transform = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.CenterCrop((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+# Directory containing images
+images_dir = "picture/"
+
+# List all image files in directory
+image_files = [f for f in os.listdir(images_dir) if f.endswith('.jpg') or f.endswith('.png')]
+
+if not image_files:
+    print("No images found in directory")
+else:
+    # Dictionary to store feature vectors
+    feature_dict = {}
+
+    # Loop through images in the directory
+    for filename in image_files:
+        # Load image
+        image_path = os.path.join(images_dir, filename)
+        with Image.open(image_path) as img:
+            img = transform(img).unsqueeze(0)
+
+        # Extract features from penultimate layer
+        with torch.no_grad():
+            features = model(img)
+            features = torch.squeeze(features).detach().numpy()
+
+        feature_dict[filename] = features
+
+    # Convert dictionary of feature vectors to array
+    feature_array = np.array(list(feature_dict.values()))
+
+    if len(feature_array) == 0:
+        print("No feature vectors extracted")
+    else:
+        # Fit nearest neighbor model on feature vectors
+        nbrs = NearestNeighbors(n_neighbors=10, algorithm='auto').fit(feature_array)
+
+        # Loop through images again to query nearest neighbors
+        for query_filename in image_files:
+            query_feature = feature_dict[query_filename]
+            distances, indices = nbrs.kneighbors(query_feature.reshape(1, -1))
+
+            print("Query image:", query_filename)
+            print("Most similar images:")
+            for i, idx in enumerate(indices[0]):
+                if i == 0:
+                    continue  # Skip first index, as it will always be the query image itself
+                print(image_files[idx])
+            print("-----")