1*4882a593Smuzhiyunimport numpy as np 2*4882a593Smuzhiyunimport cv2 3*4882a593Smuzhiyunfrom rknn.api import RKNN 4*4882a593Smuzhiyun 5*4882a593Smuzhiyunif __name__ == '__main__': 6*4882a593Smuzhiyun 7*4882a593Smuzhiyun # Create RKNN object 8*4882a593Smuzhiyun rknn = RKNN(verbose=True) 9*4882a593Smuzhiyun 10*4882a593Smuzhiyun # Pre-process config 11*4882a593Smuzhiyun print('--> Config model') 12*4882a593Smuzhiyun rknn.config(mean_values=[[127.5, 127.5, 127.5], [0, 0, 0], [0, 0, 0], [127.5]], 13*4882a593Smuzhiyun std_values=[[128, 128, 128], [1, 1, 1], [1, 1, 1], [128]]) 14*4882a593Smuzhiyun print('done') 15*4882a593Smuzhiyun 16*4882a593Smuzhiyun # Load model 17*4882a593Smuzhiyun print('--> Loading model') 18*4882a593Smuzhiyun ret = rknn.load_tensorflow(tf_pb='./conv_128.pb', 19*4882a593Smuzhiyun inputs=['input1', 'input2', 'input3', 'input4'], 20*4882a593Smuzhiyun outputs=['output'], 21*4882a593Smuzhiyun input_size_list=[[1, 128, 128, 3], [1, 128, 128, 3], [1, 128, 128, 3], [1, 128, 128, 1]]) 22*4882a593Smuzhiyun if ret != 0: 23*4882a593Smuzhiyun print('Load model failed!') 24*4882a593Smuzhiyun exit(ret) 25*4882a593Smuzhiyun print('done') 26*4882a593Smuzhiyun 27*4882a593Smuzhiyun # Build model 28*4882a593Smuzhiyun print('--> Building model') 29*4882a593Smuzhiyun ret = rknn.build(do_quantization=True, dataset='./dataset.txt') 30*4882a593Smuzhiyun if ret != 0: 31*4882a593Smuzhiyun print('Build model failed!') 32*4882a593Smuzhiyun exit(ret) 33*4882a593Smuzhiyun print('done') 34*4882a593Smuzhiyun 35*4882a593Smuzhiyun # Export rknn model 36*4882a593Smuzhiyun print('--> Export rknn model') 37*4882a593Smuzhiyun ret = rknn.export_rknn('./conv_128.rknn') 38*4882a593Smuzhiyun if ret != 0: 39*4882a593Smuzhiyun print('Export rknn model failed!') 40*4882a593Smuzhiyun exit(ret) 41*4882a593Smuzhiyun print('done') 42*4882a593Smuzhiyun 43*4882a593Smuzhiyun # Init runtime environment 44*4882a593Smuzhiyun print('--> Init runtime environment') 45*4882a593Smuzhiyun ret = rknn.init_runtime() 46*4882a593Smuzhiyun if ret != 0: 47*4882a593Smuzhiyun print('Init runtime environment failed!') 48*4882a593Smuzhiyun exit(ret) 49*4882a593Smuzhiyun print('done') 50*4882a593Smuzhiyun 51*4882a593Smuzhiyun # Set inputs 52*4882a593Smuzhiyun img = cv2.imread('./dog_128x128.jpg') 53*4882a593Smuzhiyun img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # nhwc 54*4882a593Smuzhiyun 55*4882a593Smuzhiyun img_gray = cv2.imread('./dog_128x128_gray.png', cv2.IMREAD_GRAYSCALE) 56*4882a593Smuzhiyun img_gray = np.expand_dims(img_gray, -1) # nhwc 57*4882a593Smuzhiyun 58*4882a593Smuzhiyun input2 = np.load('input2.npy').astype('float32') # nchw 59*4882a593Smuzhiyun 60*4882a593Smuzhiyun input3 = np.load('input3.npy').astype('float32') # nchw 61*4882a593Smuzhiyun 62*4882a593Smuzhiyun # Inference 63*4882a593Smuzhiyun print('--> Running model') 64*4882a593Smuzhiyun outputs = rknn.inference(inputs=[img, input2, input3, img_gray], data_format=['nhwc', 'nchw', 'nchw', 'nhwc']) 65*4882a593Smuzhiyun np.save('./functions_multi_input_test_0.npy', outputs[0]) 66*4882a593Smuzhiyun print('done') 67*4882a593Smuzhiyun outputs[0] = outputs[0].reshape((1, -1)) 68*4882a593Smuzhiyun print('inference result: ', outputs) 69*4882a593Smuzhiyun 70*4882a593Smuzhiyun rknn.release() 71