xref: /OK3568_Linux_fs/external/rknn-toolkit2/examples/functions/dynamic_input/test.py (revision 4882a59341e53eb6f0b4789bf948001014eff981)
1import numpy as np
2import cv2
3from rknn.api import RKNN
4
5
6def show_outputs(outputs):
7    output_ = outputs[0].reshape((-1, 1000))
8    for output in output_:
9        output_sorted = sorted(output, reverse=True)
10        top5_str = 'mobilenet_v1\n-----TOP 5-----\n'
11        for i in range(5):
12            value = output_sorted[i]
13            index = np.where(output == value)
14            for j in range(len(index)):
15                if (i + j) >= 5:
16                    break
17                if value > 0:
18                    topi = '{}: {}\n'.format(index[j], value)
19                else:
20                    topi = '-1: 0.0\n'
21                top5_str += topi
22        print(top5_str)
23
24
25def show_perfs(perfs):
26    perfs = 'perfs: {}\n'.format(outputs)
27    print(perfs)
28
29
30if __name__ == '__main__':
31
32    # Create RKNN object
33    rknn = RKNN(verbose=True)
34
35    # The multiple sets of input shapes specified by the user, to simulate the function of dynamic input.
36    # Please make sure the model can be dynamic when enable 'config.dynamic_input', and shape in dynamic_input are correctly!
37    dynamic_input = [
38        [[1,3,224,224]],    # set 0: [input0_224]
39        [[1,3,192,192]],    # set 1: [input0_192]
40        [[1,3,160,160]],    # set 2: [input0_160]
41    ]
42
43    # Pre-process config
44    print('--> Config model')
45    rknn.config(mean_values=[103.94, 116.78, 123.68], std_values=[58.82, 58.82, 58.82], quant_img_RGB2BGR=True, dynamic_input=dynamic_input)
46    print('done')
47
48    # Load model
49    print('--> Loading model')
50    ret = rknn.load_caffe(model='../../caffe/mobilenet_v2/mobilenet_v2.prototxt',
51                          blobs='../../caffe/mobilenet_v2/mobilenet_v2.caffemodel')
52    if ret != 0:
53        print('Load model failed!')
54        exit(ret)
55    print('done')
56
57    # Build model
58    print('--> Building model')
59    ret = rknn.build(do_quantization=True, dataset='../../caffe/mobilenet_v2/dataset.txt')
60    if ret != 0:
61        print('Build model failed!')
62        exit(ret)
63    print('done')
64
65    # Export rknn model
66    print('--> Export rknn model')
67    ret = rknn.export_rknn('./mobilenet_v2.rknn')
68    if ret != 0:
69        print('Export rknn model failed!')
70        exit(ret)
71    print('done')
72
73    # Init runtime environment
74    print('--> Init runtime environment')
75    ret = rknn.init_runtime()
76    if ret != 0:
77        print('Init runtime environment failed!')
78        exit(ret)
79    print('done')
80
81    # Set inputs
82    img = cv2.imread('./dog_224x224.jpg')
83
84    # Inference
85    print('--> Running model')
86    img2 = cv2.resize(img, (224,224))
87    img2 = np.expand_dims(img2, 0)
88    img2 = np.transpose(img2, (0,3,1,2))    # [1,3,224,224]
89    outputs = rknn.inference(inputs=[img2], data_format=['nchw'])
90    np.save('./functions_dynamic_input_0.npy', outputs[0])
91    show_outputs(outputs)
92
93    img3 = cv2.resize(img, (160,160))
94    img3 = np.expand_dims(img3, 0)
95    img3 = np.transpose(img3, (0,3,1,2))    # [1,3,160,160]
96    outputs = rknn.inference(inputs=[img3], data_format=['nchw'])
97    np.save('./functions_dynamic_input_1.npy', outputs[0])
98    show_outputs(outputs)
99    print('done')
100
101    rknn.release()
102