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

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

# Load model
print('--> Loading model')
ret = rknn.load_tensorflow(tf_pb='./vgg16.pb',
                            inputs=['image_input'],
                            outputs=['dense_2/Softmax'],
                            input_size_list=[[1, INPUT_SIZE, INPUT_SIZE, 3]])
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('./vgg16.rknn')
if ret != 0:
    print('Export rknn model failed!')
    exit(ret)
print('done')