#!/usr/bin/env python
# coding: utf-8

import datetime
from rknn.api import RKNN
from sys import exit
import argparse


ONNX_MODEL="model.onnx"
detailed_performance_log = True

# 添加命令行参数解析
parser = argparse.ArgumentParser(description='将ONNX模型转换为RKNN模型')
parser.add_argument('--quantize', '-q', action='store_true',
                   help='是否对模型进行量化')
parser.add_argument('--dataset', '-d', 
                   default="/home/zt/rk3588-nn/rknn_model_zoo/datasets/COCO/coco_subset_20.txt",
                   help='量化校准数据集的路径')
args = parser.parse_args()

DATASET=args.dataset  # 使用命令行参数指定数据集路径
QUANTIZE=args.quantize

# 根据是否量化决定输出文件名
if args.quantize:
    RKNN_MODEL=ONNX_MODEL.replace(".onnx","_int8.rknn")
else:
    RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")

timedate_iso = datetime.datetime.now().isoformat()

rknn = RKNN(verbose=True)
rknn.config(
    # mean_values=[x * 255 for x in [0.485, 0.456, 0.406]],
    # std_values=[x * 255 for x in [0.229, 0.224, 0.225]],
    quantized_dtype='w8a8',
    quantized_algorithm='kl_divergence',
    quantized_method='channel',
    quantized_hybrid_level=0,
    target_platform='rk3588',
    quant_img_RGB2BGR = False,
    float_dtype='float16',
    optimization_level=3,
    custom_string=f"converted by: qq: 232004040, email: 2302004040@qq.com at {timedate_iso}",
    remove_weight=False,
    compress_weight=False,
    inputs_yuv_fmt=None,
    single_core_mode=False,
    dynamic_input=None,
    model_pruning=False,
    op_target=None,
    quantize_weight=False,
    remove_reshape=False,
    sparse_infer=False,
    enable_flash_attention=False,
    )

ret = rknn.load_onnx(model=ONNX_MODEL, inputs=["/Transpose_output_0"], input_size_list=[[1,3,448,448]])
ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
ret = rknn.export_rknn(RKNN_MODEL)

# ret = rknn.init_runtime(target='rk3588',device_id='cbb956772bf5dac9',core_mask=RKNN.NPU_CORE_0,perf_debug=detailed_performance_log)
# rknn.eval_perf()
ret = rknn.accuracy_analysis(inputs=['img.npy'], target='rk3588')