feature_extraction.py

import torch
from torchvision import models, transforms
from PIL import Image
import os
import numpy as np

# Load pre-trained ResNet model
model = models.resnet50(weights="IMAGENET1K_V1")  # Updated to match the new API
model = torch.nn.Sequential(*list(model.children())[:-1])  # Remove classification layer
model.eval()

# Define image transformations
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Function to extract features from an image
def extract_features(image_path):
    image = Image.open(image_path).convert('RGB')
    image = transform(image).unsqueeze(0)  # Add batch dimension
    with torch.no_grad():
        features = model(image).squeeze().numpy()  # Flatten the output
    return features

# Extract features from all .jpg files in a folder
def extract_features_from_folder(folder_path):
    feature_list = []
    image_paths = []
    
    for root, _, files in os.walk(folder_path):
        for file in files:
            if file.lower().endswith('.jpg'):
                image_path = os.path.join(root, file)
                features = extract_features(image_path)
                print(f"Processed: {image_path}, Feature shape: {features.shape}")  # Debugging print
                feature_list.append(features)
                image_paths.append(image_path)
    
    # Convert to numpy array for clustering
    features_array = np.array(feature_list)
    print(f"Total number of images processed: {len(feature_list)}")
    print(f"Feature array shape: {features_array.shape}")
    
    return features_array, image_paths

if __name__ == "__main__":
    folder_path = "jay1-others"  # Change this to your folder path
    features, paths = extract_features_from_folder(folder_path)
    
    if features.size == 0:
        print("No features extracted. Please check your image paths.")
    else:
        np.save("features.npy", features)  # Save features for clustering
        np.save("paths.npy", paths)        # Save paths for later reference
        print("Feature extraction complete.")