xref: /OK3568_Linux_fs/external/rknn-toolkit2/examples/darknet/yolov3_416x416/test.py (revision 4882a59341e53eb6f0b4789bf948001014eff981) 
1import numpy as np
2import cv2
3from rknn.api import RKNN
4
5from yolov3_utils import yolov3_post_process, draw, download_yolov3_weight
6
7GRID0 = 13
8GRID1 = 26
9GRID2 = 52
10LISTSIZE = 85
11SPAN = 3
12
13if __name__ == '__main__':
14
15    MODEL_PATH = './yolov3.cfg'
16    WEIGHT_PATH = './yolov3.weights'
17    RKNN_MODEL_PATH = './yolov3_416.rknn'
18    im_file = './dog_bike_car_416x416.jpg'
19    DATASET = './dataset.txt'
20
21    # Download yolov3.weight
22    download_yolov3_weight(WEIGHT_PATH)
23
24    # Create RKNN object
25    rknn = RKNN(verbose=True)
26
27    # Pre-process config
28    print('--> Config model')
29    rknn.config(mean_values=[0, 0, 0], std_values=[255, 255, 255])
30    print('done')
31
32    # Load model
33    print('--> Loading model')
34    ret = rknn.load_darknet(model=MODEL_PATH, weight=WEIGHT_PATH)
35    if ret != 0:
36        print('Load model failed!')
37        exit(ret)
38    print('done')
39
40    # Build model
41    print('--> Building model')
42    ret = rknn.build(do_quantization=True, dataset=DATASET)
43    if ret != 0:
44        print('Build model failed!')
45        exit(ret)
46    print('done')
47
48    # Export rknn model
49    print('--> Export rknn model')
50    ret = rknn.export_rknn(RKNN_MODEL_PATH)
51    if ret != 0:
52        print('Export rknn model failed!')
53        exit(ret)
54    print('done')
55
56    # Set inputs
57    img = cv2.imread(im_file)
58    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
59
60    # Init runtime environment
61    print('--> Init runtime environment')
62    ret = rknn.init_runtime()
63    if ret != 0:
64        print('Init runtime environment failed!')
65        exit(ret)
66    print('done')
67
68    # Inference
69    print('--> Running model')
70    outputs = rknn.inference(inputs=[img])
71    print('done')
72
73    input0_data = outputs[0]
74    np.save('./darknet_yolov3_416x416_0.npy', input0_data)
75    input1_data = outputs[1]
76    np.save('./darknet_yolov3_416x416_1.npy', input1_data)
77    input2_data = outputs[2]
78    np.save('./darknet_yolov3_416x416_2.npy', input1_data)
79
80    input0_data = input0_data.reshape(SPAN, LISTSIZE, GRID0, GRID0)
81    input1_data = input1_data.reshape(SPAN, LISTSIZE, GRID1, GRID1)
82    input2_data = input2_data.reshape(SPAN, LISTSIZE, GRID2, GRID2)
83
84    input_data = []
85    input_data.append(np.transpose(input0_data, (2, 3, 0, 1)))
86    input_data.append(np.transpose(input1_data, (2, 3, 0, 1)))
87    input_data.append(np.transpose(input2_data, (2, 3, 0, 1)))
88
89    boxes, classes, scores = yolov3_post_process(input_data)
90
91    image = cv2.imread(im_file)
92    if boxes is not None:
93        draw(image, boxes, scores, classes)
94
95    print('Save results to results.jpg!')
96    cv2.imwrite('result.jpg', image)
97
98    rknn.release()
99