|
板凳
楼主 |
发表于 2019-1-30 15:57:47
|
只看该作者
- import numpy as np
- import re
- import math
- import random
- import cv2
- import copy
- import time
- from rknn.api import RKNN
- INPUT_SIZE = 300
- NUM_RESULTS = 1917
- NUM_CLASSES = 91
- Y_SCALE = 10.0
- X_SCALE = 10.0
- H_SCALE = 5.0
- W_SCALE = 5.0
- def expit(x):
- return 1. / (1. + math.exp(-x))
- def load_box_priors():
- box_priors_ = []
- fp = open('./box_priors.txt', 'r')
- ls = fp.readlines()
- for s in ls:
- aList = re.findall('([-+]?\d+(\.\d*)?|\.\d+)([eE][-+]?\d+)?', s)
- for ss in aList:
- aNum = float((ss[0]+ss[2]))
- box_priors_.append(aNum)
- fp.close()
- box_priors = np.array(box_priors_)
- box_priors = box_priors.reshape(4, NUM_RESULTS)
- return box_priors
- if __name__ == '__main__':
- # Create RKNN object
- rknn = RKNN()
- # Config for Model Input PreProcess
- rknn.config(channel_mean_value='128 128 128 128', reorder_channel='0 1 2')
- # Direct Load RKNN Model
- print('--> Loading model')
- rknn.load_rknn('./ssd_mobilenet_v1_coco.rknn')
- print('done')
- # Set inputs
- webcam = cv2.VideoCapture(0)
- # init runtime environment
- print('--> Init runtime environment')
- ret = rknn.init_runtime(host='rk3399pro')
- if ret != 0:
- print('Init runtime environment failed')
- exit(ret)
- print('done')
- while True:
- ref, frame = webcam.read()
- orig_img = frame
- img = copy.deepcopy(orig_img)
- img = cv2.resize(img, (INPUT_SIZE, INPUT_SIZE), interpolation=cv2.INTER_CUBIC)
- # Inference
- print('--> Running model')
- start = time.time()
- outputs = rknn.inference(inputs=[img])
- print("total time: ")
- print(time.time() - start)
- print('done')
- # print('inference result: ', outputs)
- cv2.imshow("test", orig_img)
- c= cv2.waitKey(30) & 0xff
- if c==27:
- webcam.release()
- break
- # Release RKNN Context
- rknn.release()
|
|