-
Notifications
You must be signed in to change notification settings - Fork 7
/
inference.py
74 lines (58 loc) · 2.08 KB
/
inference.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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
import sys
import time
import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
from skimage import io
sys.path.append('../')
from models.nets import ComboNet
class FacialBeautyPredictor:
"""
Facial Beauty Predictor
"""
def __init__(self, pretrained_model_path):
model = ComboNet(num_out=5, backbone_net_name='SEResNeXt50')
model = model.float()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
# model.load_state_dict(torch.load(pretrained_model_path))
if torch.cuda.device_count() > 1:
print("We are running on", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
model.load_state_dict(torch.load(pretrained_model_path))
else:
state_dict = torch.load(pretrained_model_path)
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
model.to(device)
model.eval()
self.device = device
self.model = model
def infer(self, img_file):
tik = time.time()
img = io.imread(img_file)
img = Image.fromarray(img.astype(np.uint8))
preprocess = transforms.Compose([
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
img = preprocess(img)
img.unsqueeze_(0)
img = img.to(self.device)
score, cls = self.model(img)
tok = time.time()
return {
'beauty': float(score.to('cpu').detach().item()),
'elapse': tok - tik
}
if __name__ == '__main__':
fbp = FacialBeautyPredictor(pretrained_model_path='/home/xulu/ModelZoo/ComboNet.pth')
print(fbp.infer('../test.jpg'))