forked from axinc-ai/ailia-models
-
Notifications
You must be signed in to change notification settings - Fork 0
/
vgg16.py
executable file
·134 lines (104 loc) · 3.48 KB
/
vgg16.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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
import sys
import time
import cv2
import ailia
import vgg16_labels
# import original modules
sys.path.append('../../util')
from utils import get_base_parser, update_parser # noqa: E402
from model_utils import check_and_download_models # noqa: E402
from image_utils import load_image # noqa: E402
from classifier_utils import plot_results, print_results # noqa: E402
import webcamera_utils # noqa: E402
# ======================
# Parameters
# ======================
WEIGHT_PATH = 'vgg16_pytorch.onnx'
MODEL_PATH = 'vgg16_pytorch.onnx.prototxt'
REMOTE_PATH = 'https://storage.googleapis.com/ailia-models/vgg16/'
IMAGE_PATH = 'pizza.jpg'
IMAGE_HEIGHT = 224
IMAGE_WIDTH = 224
SLEEP_TIME = 0
# TODO
# model_path = "VGG16.prototxt"
# weight_path = "VGG16.caffemodel"
# img_path = './pizza.jpg'
# ======================
# Arguemnt Parser Config
# ======================
parser = get_base_parser('Image classification model: VGG16', IMAGE_PATH, None)
args = update_parser(parser)
# ======================
# Main functions
# ======================
def recognize_from_image():
# prepare input data
input_data = load_image(
args.input,
(IMAGE_HEIGHT, IMAGE_WIDTH),
normalize_type='ImageNet',
gen_input_ailia=True
)
# net initialize
net = ailia.Net(MODEL_PATH, WEIGHT_PATH, env_id=args.env_id)
# inference
print('Start inference...')
if args.benchmark:
print('BENCHMARK mode')
for i in range(5):
start = int(round(time.time() * 1000))
preds_ailia = net.predict(input_data)
end = int(round(time.time() * 1000))
print(f'\tailia processing time {end - start} ms')
else:
preds_ailia = net.predict(input_data)
# postprocessing
print_results(preds_ailia, vgg16_labels.imagenet_category)
print('Script finished successfully.')
def recognize_from_video():
# net initialize
net = ailia.Net(MODEL_PATH, WEIGHT_PATH, env_id=args.env_id)
capture = webcamera_utils.get_capture(args.video)
# create video writer if savepath is specified as video format
if args.savepath is not None:
f_h = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
f_w = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
save_h, save_w = webcamera_utils.calc_adjust_fsize(
f_h, f_w, IMAGE_HEIGHT, IMAGE_WIDTH
)
writer = webcamera_utils.get_writer(args.savepath, save_h, save_w)
else:
writer = None
while(True):
ret, frame = capture.read()
if (cv2.waitKey(1) & 0xFF == ord('q')) or not ret:
break
input_image, input_data = webcamera_utils.preprocess_frame(
frame, IMAGE_HEIGHT, IMAGE_WIDTH, normalize_type='ImageNet'
)
# inference
preds_ailia = net.predict(input_data)
# postprocessing
plot_results(input_image, preds_ailia, vgg16_labels.imagenet_category)
cv2.imshow('frame', input_image)
time.sleep(SLEEP_TIME)
# save results
if writer is not None:
writer.write(input_image)
capture.release()
cv2.destroyAllWindows()
if writer is not None:
writer.release()
print('Script finished successfully.')
def main():
# model files check and download
check_and_download_models(WEIGHT_PATH, MODEL_PATH, REMOTE_PATH)
if args.video is not None:
# video mode
recognize_from_video()
else:
# image mode
recognize_from_image()
if __name__ == '__main__':
main()