from rknn.api import RKNN


def convert_rec_model():
    rknn = RKNN(verbose=True)
    
    print('--> Config model')
    rknn.config(mean_values=[[127.5, 127.5, 127.5]], 
                std_values=[[127.5, 127.5, 127.5]],
                # quant_img_RGB2BGR=True,
                target_platform='rk3588')
    
    print('--> Loading model')
    ret = rknn.load_onnx(model='/home/admin-root/haotian/康达瑞贝斯机器狗/rec_shape_20250815.onnx')
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    
    print('--> Building model')
    ret = rknn.build(do_quantization=False
                    #  , dataset='./rec_dataset.txt'
                     )
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    
    ret = rknn.export_rknn('./rec_model.rknn')
    if ret != 0:
        print('Export model failed!')
        exit(ret)
    
    rknn.release()

convert_rec_model()