Reverb/Embrace-Vision

1import torch2from torchvision import models, transforms3from PIL import Image4import pickle5import os6from tqdm import tqdm  # Import tqdm for the progress bar7 8# Load a pretrained ResNet model9model = models.resnet50(pretrained=True)10model = model.eval()11 12# Define preprocessing transforms13preprocess = transforms.Compose([14    transforms.Resize((224, 224)),15    transforms.ToTensor(),16    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),17])18 19# Function to extract features from an image20def extract_features(image_path):21    image = Image.open(image_path).convert('RGB')22    input_tensor = preprocess(image)23    input_batch = input_tensor.unsqueeze(0)24 25    with torch.no_grad():26        output = model(input_batch)27 28    return output.squeeze().numpy()29 30# Directory containing your images31images_directory = "photos/"32 33# Process each image and save features34image_features = {}35for filename in tqdm(os.listdir(images_directory), desc="Processing Images"):36    if filename.endswith(".jpg") or filename.endswith(".png"):37        image_path = os.path.join(images_directory, filename)38        features = extract_features(image_path)39        image_features[filename] = features40 41# Save the features to a pickle file42output_file = "unsplash-25k-embeddings.pkl"43with open(output_file, 'wb') as f:44    pickle.dump(image_features, f)