|
|
|
import torch |
|
import numpy as np |
|
from torchvision import transforms |
|
from torchvision.models import vgg16 |
|
from PIL import Image |
|
|
|
class FeatureExtractor: |
|
def __init__(self): |
|
|
|
self.vgg = vgg16(weights='DEFAULT') |
|
self.vgg.eval() |
|
|
|
|
|
self.conv_extractor = self.vgg.features |
|
|
|
|
|
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)) |