# svm_vgg_preprocessor.py
import torch
import numpy as np
from torchvision import transforms
from torchvision.models import vgg16
from PIL import Image

class FeatureExtractor:
    def __init__(self):
        # Load VGG16 with original architecture
        self.vgg = vgg16(weights='DEFAULT')
        self.vgg.eval()
        
        # Use only convolutional features (matches original code)
        self.conv_extractor = self.vgg.features
        
        # Original preprocessing
        self.preprocess = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])

    def extract_fc_cnn_features(self, image_path):
        """Matches original FC-CNN feature extraction (conv features)"""
        img = Image.open(image_path).convert('RGB')
        img_tensor = self.preprocess(img).unsqueeze(0)
        
        with torch.no_grad():
            features = self.conv_extractor(img_tensor)
            
        return features.squeeze().numpy().flatten()

    def extract_fv_cnn_features(self, image_path):
        """Matches original FV-CNN feature extraction (conv features)"""
        img = Image.open(image_path).convert('RGB')
        img_tensor = self.preprocess(img).unsqueeze(0)
        
        with torch.no_grad():
            features = self.conv_extractor(img_tensor)
            
        return features.squeeze().numpy().flatten()

    def extract_combined_features(self, image_path):
        """Original concatenation of identical features"""
        fc = self.extract_fc_cnn_features(image_path)
        fv = self.extract_fv_cnn_features(image_path)
        return np.concatenate((fc, fv))