Update one.py

one.py

@@ -1,50 +1,49 @@
 import torch
-import torch.nn as nn
-import torchvision.models as models
 import torchvision.transforms as transforms
-import os
 from PIL import Image
+from transformers import AutoModel, AutoTokenizer
 
-# 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()
+# Load the pre-trained ResNet50 model from Hugging Face
+model_name = 'pytorch/vision:v0.9.0'
+model = AutoModel.from_pretrained(model_name)
 
 # Set the model to evaluation mode
 model.eval()
 
-# Define the preprocessing transforms
+# Load the tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Define the preprocessing pipeline
 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]
-    )
+    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
+# Define a function to extract features from an image
+def extract_features(image_path, model, tokenizer, preprocess):
+    # Load the image
+    image = Image.open(image_path).convert('RGB')
+    # Apply the preprocessing pipeline
+    image = preprocess(image)
+    # Add a batch dimension to the image tensor
+    image = image.unsqueeze(0)
+    # Encode the image using the tokenizer
+    inputs = tokenizer(image, padding=True, truncation=True, return_tensors='pt')
+    # Pass the inputs through the model to get the features
+    outputs = model(**inputs)
+    # Return the features
+    return outputs.last_hidden_state.squeeze().detach().numpy()
+
+# Define a dictionary to store the features
+features_dict = {}
+
+# Loop over the images and extract the features
+for image_name in image_names:
+    # Extract the features for this image
+    image_path = os.path.join(images_folder, image_name)
+    features = extract_features(image_path, model, tokenizer, preprocess)
+    # Add the features to the dictionary
+    features_dict[image_name] = features