from rknn.api import RKNN

onnx_path = '/data6/zhangw/yolov5/runs/train/YNDC_20251010_yolov5_optimize/weights/best.onnx'
rknn_path = '/data6/zhangw/Projects/output_rknn/yolov5_best.rknn'

print('==== 1. 创建 RKNN 对象 ====')
rknn = RKNN(verbose=True)

print('==== 2. 加载 ONNX ====')
ret = rknn.load_onnx(model=onnx_path)
if ret != 0:
    print('load onnx failed!')
    exit(ret)

print('==== 3. 模型构建 ====')
ret = rknn.build(do_quantization=False)
if ret != 0:
    print('build failed!')
    exit(ret)

print('==== 4. 导出 RKNN ====')
ret = rknn.export_rknn(rknn_path)
if ret != 0:
    print('export failed!')
    exit(ret)

print('==== 完成！RKNN 模型已生成 ====')