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
|
import paddle
import paddlehub as hub
import paddlehub.vision.transforms as T
import cv2 as cv
import numpy as np
def preprocess(image_path):
''' preprocess input image file to np.ndarray
Args:
image_path(str): Path of input image file
Returns:
ProcessedImage(numpy.ndarray): A numpy.ndarray
variable which shape is (1, 3, 224, 224)
'''
transforms = T.Compose([
T.Resize((256, 256)),
T.CenterCrop(224),
T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])],
to_rgb=True)
return np.expand_dims(transforms(image_path), axis=0)
def export_onnx_resnet50(save_path):
''' export PaddlePaddle model to ONNX format
Args:
save_path(str): Path to save exported ONNX model
Returns:
None
'''
model = hub.Module(name="resnet50_vd_imagenet_ssld")
input_spec = paddle.static.InputSpec(
[1, 3, 224, 224], "float32", "image")
paddle.onnx.export(model, save_path,
input_spec=[input_spec],
opset_version=10)
if __name__ == '__main__':
save_path = './resnet50'
image_file = './data/cat.jpg'
labels = open('./data/labels.txt').read().strip().split('\n')
model = export_onnx_resnet50(save_path)
# load resnet50 use cv.dnn
net = cv.dnn.readNetFromONNX(save_path + '.onnx')
# read and preprocess image file
im = preprocess(image_file)
# inference
net.setInput(im)
result = net.forward(['save_infer_model/scale_0.tmp_0'])
# post process
class_id = np.argmax(result[0])
label = labels[class_id]
print("Image: {}".format(image_file))
print("Predict Category: {}".format(label))
|
