Update svm_vgg_preprocessor.py

svm_vgg_preprocessor.py

@@ -1,21 +1,23 @@
+# svm_vgg_preprocessor.py
 import torch
 import numpy as np
 from torchvision import transforms
 from torchvision.models import vgg16
 from PIL import Image
 
-# Initialize VGG model (matches your training setup)
 class FeatureExtractor:
     def __init__(self):
-        self.vgg_model = vgg16(pretrained=True)
-        self.vgg_model.eval()  # Set to evaluation mode
+        # Load pretrained VGG16
+        self.vgg = vgg16(weights='DEFAULT')
+        self.vgg.eval()
         
-        # For FC-CNN features (classifier-based)
-        self.fc_extractor = torch.nn.Sequential(
-            *list(self.vgg_model.classifier.children())[:-2]  # Remove last 2 layers
+        # Feature extractors
+        self.conv_features = self.vgg.features
+        self.fc_features = torch.nn.Sequential(
+            *list(self.vgg.classifier.children())[:-2]
         )
         
-        # Standard VGG preprocessing
+        # Preprocessing
         self.preprocess = transforms.Compose([
             transforms.Resize((224, 224)),
             transforms.ToTensor(),
@@ -26,22 +28,28 @@ class FeatureExtractor:
         ])
 
     def extract_fc_cnn_features(self, image_path):
-        """Extract fully-connected layer features"""
-        image = Image.open(image_path).convert('RGB')
-        image_tensor = self.preprocess(image).unsqueeze(0)
+        """Correct FC-CNN feature extraction"""
+        # Load and preprocess image
+        img = Image.open(image_path).convert('RGB')
+        img_tensor = self.preprocess(img).unsqueeze(0)
         
         with torch.no_grad():
-            features = self.fc_extractor(image_tensor)
+            # Get convolutional features
+            conv_out = self.conv_features(img_tensor)
+            # Flatten for FC layers
+            flattened = torch.flatten(conv_out, 1)
+            # Get FC features
+            fc_features = self.fc_features(flattened)
             
-        return features.squeeze().numpy().flatten()
+        return fc_features.squeeze().numpy().flatten()
 
     def extract_fv_cnn_features(self, image_path):
-        """Extract convolutional layer features"""
-        image = Image.open(image_path).convert('RGB')
-        image_tensor = self.preprocess(image).unsqueeze(0)
+        """FV-CNN feature extraction (unchanged)"""
+        img = Image.open(image_path).convert('RGB')
+        img_tensor = self.preprocess(img).unsqueeze(0)
         
         with torch.no_grad():
-            conv_features = self.vgg_model.features(image_tensor)
+            conv_features = self.conv_features(img_tensor)
             
         return conv_features.squeeze().numpy().flatten()