from rknn.api import RKNN

def convert_det_model():
    # 创建RKNN对象
    rknn = RKNN(verbose=True)
    
    # 配置模型
    print('--> Config model')
    rknn.config(
                mean_values=[[123.675, 116.28, 103.53]], 
                std_values=[[58.395, 57.12, 57.375]],
                # mean_values=[[0, 0, 0]], 
                # std_values=[[255, 255, 255]],
                # quant_img_RGB2BGR=True,
                target_platform='rk3588')  # 根据您的芯片型号调整
    
    # 加载ONNX模型
    print('--> Loading model')
    ret = rknn.load_onnx(model='/home/admin-root/haotian/康达瑞贝斯机器狗/det_shape_20250814.onnx')
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    
    # 构建模型
    print('--> Building model')
    ret = rknn.build(do_quantization=False
                    #  , dataset='./det_dataset.txt'
                     )
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    
    # 导出RKNN模型
    ret = rknn.export_rknn('./det_shape_255.rknn')
    if ret != 0:
        print('Export model failed!')
        exit(ret)
    
    rknn.release()

convert_det_model()