-
Notifications
You must be signed in to change notification settings - Fork 10
/
predict.py
55 lines (41 loc) · 1.56 KB
/
predict.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
# -*- coding: utf-8 -*-
# @File : predict.py
# @Author: Runist
# @Time : 2021/12/14 16:15
# @Software: PyCharm
# @Brief:
import os
import torch
from PIL import Image
from dataloader import data_transform
from utils import create_model, model_parallel
from config import args
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# load image
image_path = "/home/rhdai/workspace/code/image_classification/dataset/daisy.jpg"
assert os.path.exists(image_path), "file: '{}' dose not exist.".format(image_path)
image = Image.open(image_path)
# [N, C, H, W]
image = data_transform["val"](image)
# expand batch dimension
image = torch.unsqueeze(image, dim=0)
# create model
model = create_model(args)
model = model_parallel(args, model).to(device)
# load model weights
model_weight_path = "{}/weights/epoch=20_val_acc=0.9643.pth".format(args.summary_dir)
model.load_state_dict(torch.load(model_weight_path, map_location=device))
model.eval()
with torch.no_grad():
# predict class
output = torch.squeeze(model(image.to(device))).cpu()
predict = torch.softmax(output, dim=0)
index = torch.argmax(predict).numpy()
print("prediction: {} prob: {:.3}\n".format(args.label_name[index],
predict[index].numpy()))
for i in range(len(predict)):
print("class: {} prob: {:.3}".format(args.label_name[i],
predict[i].numpy()))
if __name__ == '__main__':
main()