-
Notifications
You must be signed in to change notification settings - Fork 33
/
extract_example_rgb.py
46 lines (35 loc) · 1.26 KB
/
extract_example_rgb.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
import torch
from torch import nn
from pretrainedmodels import bninception
from torchvision import transforms
from glob import glob
from PIL import Image
import lmdb
from tqdm import tqdm
from os.path import basename
from argparse import ArgumentParser
env = lmdb.open('features/rgb', map_size=1099511627776)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = bninception(pretrained=None)
state_dict = torch.load('models/TSN-rgb.pth.tar')['state_dict']
state_dict = {k.replace('module.base_model.','') : v for k,v in state_dict.items()}
model.load_state_dict(state_dict, strict=False)
model.last_linear = nn.Identity()
model.global_pool = nn.AdaptiveAvgPool2d(1)
model.to(device)
transform = transforms.Compose([
transforms.Resize([256, 454]),
transforms.ToTensor(),
transforms.Lambda(lambda x: x[[2,1,0],...]*255), #to BGR
transforms.Normalize(mean=[104, 117, 128],
std=[1, 1, 1]),
])
imgs = sorted(glob('data/sample_rgb/*.jpg'))
model.eval()
for im in tqdm(imgs,'Extracting features'):
key = basename(im)
img = Image.open(im)
data = transform(img).unsqueeze(0).to(device)
feat = model(data).squeeze().detach().cpu().numpy()
with env.begin(write=True) as txn:
txn.put(key.encode(),feat.tobytes())