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