from rknn.api import RKNN
import cv2
import numpy as np

# Create RKNN object
rknn = RKNN(verbose=True)

# Pre-process config
print('--> Config model')
rknn.config(mean_values=[128, 128, 128], std_values=[128, 128, 128], target_platform='rk3566')
print('done')

# Load model (from https://www.tensorflow.org/lite/guide/hosted_models?hl=zh-cn)
print('--> Loading model')
ret = rknn.load_tflite(model='mobilenet_v1_1.0_224.tflite')
if ret != 0:
    print('Load model failed!')
    exit(ret)
print('done')

# Build model
print('--> Building model')
ret = rknn.build(do_quantization=True, dataset='./dataset.txt')
if ret != 0:
    print('Build model failed!')
    exit(ret)
print('done')

# Export rknn model
print('--> Export rknn model')
ret = rknn.export_rknn('./mobilenet_v1.rknn')
if ret != 0:
    print('Export rknn model failed!')
    exit(ret)
print('done')

# Set inputs
img = cv2.imread('./dog_224x224.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.expand_dims(img, 0)

# Init runtime environment
print('--> Init runtime environment')
ret = rknn.init_runtime()
if ret != 0:
    print('Init runtime environment failed!')
    exit(ret)
print('done')

# Inference
print('--> Running model')
outputs = rknn.inference(inputs=[img], data_format=['nhwc'])
np.save('./tflite_mobilenet_v1_0.npy', outputs[0])
# show_outputs(outputs)
print('done')

rknn.release()