feature extracted with pretrained model by using avg pool in the last layer

feature_extraction.py

+import numpy as np
+import torchvision.models as models
+import torchvision as tv
+import torch
+import torch.nn as nn
+import os
+from PIL import Image
+import argparse
+
+
+class FeatureExtractor(nn.Module):
+    def __init__(self, arch=tv.models.resnet50):
+        super(FeatureExtractor, self).__init__()
+        self.tranform = tv.transforms.Compose([
+            tv.transforms.Resize([224, 224]), tv.transforms.ToTensor(),
+            tv.transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                    std=[0.229, 0.224, 0.225])])
+        self.model = nn.Sequential(*(list(arch.children())[:-2] + [nn.MaxPool2d(4, 1)]))
+
+    def forward(self, frame):
+        features = self.model(frame)
+        features = features.reshape((features.shape[0], -1))
+        return features
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Features Extraction')
+    parser.add_argument('--frames', metavar='Frames-dir',
+                        default='./frames',
+                        help='path to input frames')
+    parser.add_argument('--model', default='resnet50',
+                        help='pretrained model architecture e.g. resnet50 or alexnet')
+    args = parser.parse_args()
+    frames_dir = args.frames
+    if args.model =='alexnet':
+        model_arch = models.alexnet(pretrained=True)
+    else:
+        model_arch = models.resnet50(pretrained=True)
+    isCuda = torch.cuda.is_available()
+    if isCuda:
+        model = FeatureExtractor(model_arch).cuda()
+    else:
+        model = FeatureExtractor(model_arch)
+    feature = dict()
+
+    for i, video in enumerate(os.listdir(frames_dir)):
+        video = os.path.join(frames_dir, video)
+        print(video)
+        feature[i] = []
+        for frame in os.listdir(video):
+            frame = os.path.join(video, frame)
+            print(frame)
+            img = Image.open(frame)
+            img = model.tranform(img)
+            img = img.view((1,) + img.shape)
+            feat = model(img)
+            print(feat.shape)
+            feature[i].append(feat.cpu().detach().numpy()[0])