extract_image_features.py (2219B)
1 import numpy as np 2 import os 3 import torch 4 5 class ImageFeatureExtractor(): 6 def __init__(self, config): 7 self.config = config 8 9 def process_image(self, images): 10 """ 11 return pooled or full vecteur, for mlp or attention pooling respectively. 12 If all features are used, return pooled features concatenated (only for mlp model) 13 """ 14 15 if self.config.features_path == "all": 16 if self.config.pooling == 'mlp': 17 b7_pooled = np.load(str(images).replace('.png', '_pooled.npy').replace('images', "features_b7")).squeeze(0) 18 vgg_pooled = np.load(str(images).replace('.png', '_pooled.npy').replace('images', "features_vgg")).squeeze(0) 19 frcnn = np.load(str(images).replace('.png', '.npy').replace('images', "features_frcnn")).mean(0) 20 all_features = np.concatenate((vgg_pooled, b7_pooled, frcnn), axis=0)[None,:] 21 else: 22 b7 = np.load(str(images).replace('.png', '.npy').replace('images', "features_b7")).squeeze(0).reshape((2560,196)).transpose() 23 vgg = np.load(str(images).replace('.png', '.npy').replace('images', "features_vgg")).squeeze(0).reshape((512,196)).transpose() 24 frcnn = np.load(str(images).replace('.png', '.npy').replace('images', "features_frcnn")) # 2048, 36 25 all_features = [vgg, b7, frcnn] 26 raise NotImplementedError() 27 return all_features 28 elif self.config.features_path in "features_frcnn": 29 feature = np.load(str(images).replace('.png', '.npy').replace('images', self.config.features_path)) 30 return feature.mean(0)[None,:] if self.config.pooling == 'mlp' else feature 31 elif self.config.features_path in ["features_b7", "features_vgg"]: 32 extension = '_pooled.npy' if self.config.pooling == 'mlp' else '.npy' 33 vec = np.load(str(images).replace('.png', extension).replace('images', self.config.features_path)) 34 return vec if self.config.pooling == 'mlp' else vec.reshape((vec.shape[1],196)).transpose() 35 else: 36 raise RuntimeError(f"Unknown config value for feature path: {self.config.features_path}")