xref: /OK3568_Linux_fs/external/rknn-toolkit2/examples/tensorflow/ssd_mobilenet_v1/test.py (revision 4882a59341e53eb6f0b4789bf948001014eff981) 
1import numpy as np
2
3import re
4import math
5import random
6import cv2
7
8from rknn.api import RKNN
9
10INPUT_SIZE = 300
11
12NUM_RESULTS = 1917
13NUM_CLASSES = 91
14
15Y_SCALE = 10.0
16X_SCALE = 10.0
17H_SCALE = 5.0
18W_SCALE = 5.0
19
20
21def expit(x):
22    return 1. / (1. + math.exp(-x))
23
24
25def unexpit(y):
26    return -1.0 * math.log((1.0 / y) - 1.0)
27
28
29def CalculateOverlap(xmin0, ymin0, xmax0, ymax0, xmin1, ymin1, xmax1, ymax1):
30    w = max(0.0, min(xmax0, xmax1) - max(xmin0, xmin1))
31    h = max(0.0, min(ymax0, ymax1) - max(ymin0, ymin1))
32    i = w * h
33    u = (xmax0 - xmin0) * (ymax0 - ymin0) + (xmax1 - xmin1) * (ymax1 - ymin1) - i
34
35    if u <= 0.0:
36        return 0.0
37
38    return i / u
39
40
41def load_box_priors():
42    box_priors_ = []
43    fp = open('./box_priors.txt', 'r')
44    ls = fp.readlines()
45    for s in ls:
46        aList = re.findall('([-+]?\d+(\.\d*)?|\.\d+)([eE][-+]?\d+)?', s)
47        for ss in aList:
48            aNum = float((ss[0]+ss[2]))
49            box_priors_.append(aNum)
50    fp.close()
51
52    box_priors = np.array(box_priors_)
53    box_priors = box_priors.reshape(4, NUM_RESULTS)
54
55    return box_priors
56
57
58if __name__ == '__main__':
59
60    # Create RKNN object
61    rknn = RKNN(verbose=True)
62
63    # Pre-process config
64    print('--> Config model')
65    rknn.config(mean_values=[127.5, 127.5, 127.5], std_values=[127.5, 127.5, 127.5])
66    print('done')
67
68    # Load model
69    print('--> Loading model')
70    ret = rknn.load_tensorflow(tf_pb='./ssd_mobilenet_v1_coco_2017_11_17.pb',
71                               inputs=['Preprocessor/sub'],
72                               outputs=['concat', 'concat_1'],
73                               input_size_list=[[1, INPUT_SIZE, INPUT_SIZE, 3]])
74    if ret != 0:
75        print('Load model failed!')
76        exit(ret)
77    print('done')
78
79    # Build Model
80    print('--> Building model')
81    ret = rknn.build(do_quantization=True, dataset='./dataset.txt')
82    if ret != 0:
83        print('Build model failed!')
84        exit(ret)
85    print('done')
86
87    # Export rknn model
88    print('--> Export rknn model')
89    ret = rknn.export_rknn('./ssd_mobilenet_v1_coco.rknn')
90    if ret != 0:
91        print('Export rknn model failed!')
92        exit(ret)
93    print('done')
94
95    # Set inputs
96    orig_img = cv2.imread('./road.bmp')
97    img = cv2.cvtColor(orig_img, cv2.COLOR_BGR2RGB)
98    img = cv2.resize(img, (INPUT_SIZE, INPUT_SIZE), interpolation=cv2.INTER_CUBIC)
99
100    # Init runtime environment
101    print('--> Init runtime environment')
102    ret = rknn.init_runtime()
103    if ret != 0:
104        print('Init runtime environment failed!')
105        exit(ret)
106    print('done')
107
108    # Inference
109    print('--> Running model')
110    outputs = rknn.inference(inputs=[img])
111    print('done')
112
113    predictions = outputs[0].reshape((1, NUM_RESULTS, 4))
114    np.save('./tensorflow_ssd_mobilenet_v1_0.npy', outputs[0])
115    outputClasses = outputs[1].reshape((1, NUM_RESULTS, NUM_CLASSES))
116    np.save('./tensorflow_ssd_mobilenet_v1_1.npy', outputs[0])
117    candidateBox = np.zeros([2, NUM_RESULTS], dtype=int)
118    classScore = [-1000.0] * NUM_RESULTS
119    vaildCnt = 0
120
121    box_priors = load_box_priors()
122
123    # Post Process
124    # got valid candidate box
125    for i in range(0, NUM_RESULTS):
126        topClassScore = -1000
127        topClassScoreIndex = -1
128
129        # Skip the first catch-all class.
130        for j in range(1, NUM_CLASSES):
131            score = expit(outputClasses[0][i][j])
132
133            if score > topClassScore:
134                topClassScoreIndex = j
135                topClassScore = score
136
137        if topClassScore > 0.4:
138            candidateBox[0][vaildCnt] = i
139            candidateBox[1][vaildCnt] = topClassScoreIndex
140            classScore[vaildCnt] = topClassScore
141            vaildCnt += 1
142
143    # calc position
144    for i in range(0, vaildCnt):
145        if candidateBox[0][i] == -1:
146            continue
147
148        n = candidateBox[0][i]
149        ycenter = predictions[0][n][0] / Y_SCALE * box_priors[2][n] + box_priors[0][n]
150        xcenter = predictions[0][n][1] / X_SCALE * box_priors[3][n] + box_priors[1][n]
151        h = math.exp(predictions[0][n][2] / H_SCALE) * box_priors[2][n]
152        w = math.exp(predictions[0][n][3] / W_SCALE) * box_priors[3][n]
153
154        ymin = ycenter - h / 2.
155        xmin = xcenter - w / 2.
156        ymax = ycenter + h / 2.
157        xmax = xcenter + w / 2.
158
159        predictions[0][n][0] = ymin
160        predictions[0][n][1] = xmin
161        predictions[0][n][2] = ymax
162        predictions[0][n][3] = xmax
163
164    # NMS
165    for i in range(0, vaildCnt):
166        if candidateBox[0][i] == -1:
167            continue
168
169        n = candidateBox[0][i]
170        xmin0 = predictions[0][n][1]
171        ymin0 = predictions[0][n][0]
172        xmax0 = predictions[0][n][3]
173        ymax0 = predictions[0][n][2]
174
175        for j in range(i+1, vaildCnt):
176            m = candidateBox[0][j]
177
178            if m == -1:
179                continue
180
181            xmin1 = predictions[0][m][1]
182            ymin1 = predictions[0][m][0]
183            xmax1 = predictions[0][m][3]
184            ymax1 = predictions[0][m][2]
185
186            iou = CalculateOverlap(xmin0, ymin0, xmax0, ymax0, xmin1, ymin1, xmax1, ymax1)
187
188            if iou >= 0.45:
189                candidateBox[0][j] = -1
190
191    # Draw result
192    for i in range(0, vaildCnt):
193        if candidateBox[0][i] == -1:
194            continue
195
196        n = candidateBox[0][i]
197
198        xmin = max(0.0, min(1.0, predictions[0][n][1])) * INPUT_SIZE
199        ymin = max(0.0, min(1.0, predictions[0][n][0])) * INPUT_SIZE
200        xmax = max(0.0, min(1.0, predictions[0][n][3])) * INPUT_SIZE
201        ymax = max(0.0, min(1.0, predictions[0][n][2])) * INPUT_SIZE
202
203        print("%d @ (%d, %d) (%d, %d) score=%f" % (candidateBox[1][i], xmin, ymin, xmax, ymax, classScore[i]))
204        cv2.rectangle(orig_img, (int(xmin), int(ymin)), (int(xmax), int(ymax)),
205                      (random.random()*255, random.random()*255, random.random()*255), 3)
206
207    cv2.imwrite("out.jpg", orig_img)
208
209    rknn.release()
210