generate_image_features.py (9949B)
1 # Copyright (c) Facebook, Inc. and its affiliates. 2 3 """ 4 Run with for example: 5 python3 mmf/tools/scripts/features/frcnn/extract_features_vgg.py \ 6 --image_dir ./example_images --output_folder ./output_features 7 """ 8 9 import argparse 10 import copy 11 import logging 12 import os 13 import torchvision 14 import numpy as np 15 import torch 16 import glob 17 import math 18 import cv2 19 import os 20 21 22 def get_image_files( 23 image_dir, 24 exclude_list=None, 25 partition=None, 26 max_partition=None, 27 start_index=0, 28 end_index=None, 29 output_folder=None, 30 ): 31 files = glob.glob(os.path.join(image_dir, "*.png")) 32 files.extend(glob.glob(os.path.join(image_dir, "*.jpg"))) 33 files.extend(glob.glob(os.path.join(image_dir, "*.jpeg"))) 34 35 files = set(files) 36 exclude = set() 37 38 if os.path.exists(exclude_list): 39 with open(exclude_list) as f: 40 lines = f.readlines() 41 for line in lines: 42 exclude.add(line.strip("\n").split(os.path.sep)[-1].split(".")[0]) 43 output_ignore = set() 44 if output_folder is not None: 45 output_files = glob.glob(os.path.join(output_folder, "*.npy")) 46 for f in output_files: 47 file_name = f.split(os.path.sep)[-1].split(".")[0] 48 output_ignore.add(file_name) 49 50 for f in list(files): 51 file_name = f.split(os.path.sep)[-1].split(".")[0] 52 if file_name in exclude or file_name in output_ignore: 53 files.remove(f) 54 55 files = list(files) 56 files = sorted(files) 57 58 if partition is not None and max_partition is not None: 59 interval = math.floor(len(files) / max_partition) 60 if partition == max_partition: 61 files = files[partition * interval :] 62 else: 63 files = files[partition * interval : (partition + 1) * interval] 64 65 if end_index is None: 66 end_index = len(files) 67 68 files = files[start_index:end_index] 69 70 return files 71 72 73 def chunks(array, chunk_size): 74 for i in range(0, len(array), chunk_size): 75 yield array[i : i + chunk_size], i 76 77 78 class ModelFeatureExtractor(torch.nn.Module): 79 def __init__(self, model): 80 super(ModelFeatureExtractor, self).__init__() 81 # Extract VGG-16 Feature Layers 82 self.features = list(model.features) 83 self.features = torch.nn.Sequential(*self.features) 84 85 # Extract VGG-16 Average Pooling Layer 86 self.pooling = model.avgpool 87 # Convert the image into one-dimensional vector 88 self.flatten = torch.nn.Flatten() 89 # Extract the first part of fully-connected layer from VGG16 90 #self.fc = model.classifier[0] 91 92 def forward(self, x): 93 # It will take the input 'x' until it returns the feature vector called 'out' 94 feature = self.features(x) 95 pooled = self.pooling(feature) 96 pooled = self.flatten(pooled) 97 #out = self.fc(out) 98 return feature, pooled 99 100 101 class FeatureExtractor: 102 def __init__(self): 103 self.args = self.get_parser().parse_args() 104 self.args.output_folder = self.args.output_folder + '_' + self.args.modelname + '/' 105 if not os.path.exists(self.args.output_folder): 106 os.mkdir(self.args.output_folder) 107 if self.args.modelname == 'vgg': 108 model_vgg_pretrained = torchvision.models.vgg16(pretrained=True) 109 self.model = ModelFeatureExtractor(model_vgg_pretrained) 110 elif self.args.modelname == 'b7': 111 model_b7_pretrained = torchvision.models.efficientnet_b7(pretrained=True) 112 self.model = ModelFeatureExtractor(model_b7_pretrained) 113 else: 114 print("Model unknown") 115 116 # Change the device to GPU 117 self.device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu") 118 self.model = self.model.to(self.device) 119 120 # Transform the image, so it becomes readable with the model 121 self.transform_center = torchvision.transforms.Compose([ 122 torchvision.transforms.ToPILImage(), 123 torchvision.transforms.CenterCrop(512), 124 torchvision.transforms.Resize(448), 125 torchvision.transforms.ToTensor(), 126 torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], 127 std=[0.229, 0.224, 0.225]) 128 ]) 129 130 def get_parser(self): 131 parser = argparse.ArgumentParser() 132 parser.add_argument( 133 "--start_index", default=0, type=int, help="Index to start from " 134 ) 135 parser.add_argument("--end_index", default=None, type=int, help="") 136 parser.add_argument("--batch_size", type=int, default=1, help="Batch size") 137 parser.add_argument( 138 "--num_features", 139 type=int, 140 default=100, 141 help="Number of features to extract.", 142 ) 143 parser.add_argument( 144 "--output_folder", type=str, default="/data/almanach/user/smontari/scratch/constraint_challenge/mmsrl/all/features", help="Output folder" 145 ) 146 parser.add_argument("--image_dir", type=str, help="Image directory or file", default='/data/almanach/user/smontari/scratch/constraint_challenge/mmsrl/all/images/') 147 parser.add_argument( 148 "--modelname", 149 type=str, 150 help="The name of the feature to extract", 151 default="vgg", 152 ) 153 parser.add_argument( 154 "--exclude_list", 155 type=str, 156 help="List of images to be excluded from feature conversion. " 157 + "Each image on a new line", 158 default="./list", 159 ) 160 parser.add_argument( 161 "--confidence_threshold", 162 type=float, 163 default=0, 164 help="Threshold of detection confidence above which boxes will be selected", 165 ) 166 # TODO finish background flag 167 parser.add_argument( 168 "--background", 169 action="store_true", 170 help="The model will output predictions for the background class when set", 171 ) 172 parser.add_argument( 173 "--padding", 174 type=str, 175 default=None, 176 help="You can set your padding, i.e. 'max_detections'", 177 ) 178 parser.add_argument( 179 "--visualize", 180 type=bool, 181 default=False, 182 help="Add this flag to save the extra file used for visualization", 183 ) 184 parser.add_argument( 185 "--partition", 186 type=int, 187 default=None, 188 help="Add this flag to save the extra file used for visualization", 189 ) 190 parser.add_argument( 191 "--max_partition", 192 type=int, 193 default=None, 194 help="Add this flag to save the extra file used for visualization", 195 ) 196 return parser 197 198 def _save_feature(self, file_name, feature, pooled): 199 file_base_name = os.path.basename(file_name) 200 file_base_name = file_base_name.split(".")[0] 201 feat_list_base_name = file_base_name + ".npy" 202 np.save( 203 os.path.join(self.args.output_folder, file_base_name + ".npy"), 204 feature.cpu().numpy(), 205 ) 206 np.save( 207 os.path.join(self.args.output_folder, file_base_name + "_pooled.npy"), 208 pooled.cpu().numpy(), 209 ) 210 211 def extract_features(self): 212 image_dir = self.args.image_dir 213 if os.path.isfile(image_dir): 214 feature, pooled = self.get_image_center([image_dir]) 215 self._save_feature(image_dir, feature, pooled) 216 else: 217 files = get_image_files( 218 self.args.image_dir, 219 exclude_list=self.args.exclude_list, 220 partition=self.args.partition, 221 max_partition=self.args.max_partition, 222 start_index=self.args.start_index, 223 end_index=self.args.end_index, 224 ) 225 226 finished = 0 227 total = len(files) 228 failed = 0 229 failedNames = [] 230 231 file_names = copy.deepcopy(files) 232 233 for chunk, begin_idx in chunks(files, self.args.batch_size): 234 try: 235 feature, pooled = self.get_image_center(chunk) 236 for idx, file_name in enumerate(chunk): 237 self._save_feature( 238 file_names[begin_idx + idx], 239 feature, pooled, 240 ) 241 finished += len(chunk) 242 243 if finished % 200 == 0: 244 print(f"Processed {finished}/{total}") 245 except Exception: 246 failed += len(chunk) 247 for idx, file_name in enumerate(chunk): 248 failedNames.append(file_names[begin_idx + idx]) 249 logging.exception("message") 250 if self.args.partition is not None: 251 print("Partition " + str(self.args.partition) + " done.") 252 print("Failed: " + str(failed)) 253 print("Failed Names: " + str(failedNames)) 254 255 def get_image_center(self, chunk): 256 # Set the image path 257 imgs: torch.Tensor = torch.empty((len(chunk), 3, 448, 448), dtype=torch.float32) 258 for i, path in enumerate(chunk): 259 img = cv2.imread(path) 260 img = self.transform_center(img) 261 # Reshape the image. PyTorch model reads 4-dimensional tensor 262 # [batch_size, channels, width, height] 263 img = img.reshape(3, 448, 448) 264 imgs[i] = img 265 266 imgs = imgs.to(self.device) 267 # We only extract features, so we don't need gradient 268 with torch.no_grad(): 269 # Extract the feature from the image 270 #feature = model_vgg(imgs).squeeze() 271 feature, pooled = self.model(imgs) 272 273 return feature, pooled 274 275 276 277 if __name__ == "__main__": 278 279 feature_extractor = FeatureExtractor() 280 feature_extractor.extract_features()