from rknn.api import RKNN
import os

ONNX_MODEL = 'runs/train/YNDC_20251010_yolov5_optimize/weights/best.onnx'
RKNN_MODEL = 'runs/train/YNDC_20251010_yolov5_optimize/weights/best_rk3576.rknn'
DATASET_TXT = '/data6/zhangw/yolov5/data/images/lianghua.txt'

def main():
    print("==== 1. 创建 RKNN 对象 ====")
    rknn = RKNN(verbose=True)

    print("==== 2. RKNN 配置 ====")
    rknn.config(
        target_platform='rk3576',
        mean_values=[[0, 0, 0]],
        std_values=[[255, 255, 255]],
        quantized_dtype='w8a8',
        optimization_level=3
    )

    print("==== 3. 写入量化校准文件 ====")
    os.makedirs(os.path.dirname(DATASET_TXT), exist_ok=True)
    with open(DATASET_TXT, 'w') as f:
        f.write('/data6/zhangw/yolov5/data/images/bus.jpg\n')
        f.write('/data6/zhangw/yolov5/data/images/zidane.jpg\n')

    print("==== 4. 加载 ONNX ====")
    ret = rknn.load_onnx(model=ONNX_MODEL)
    if ret != 0:
        print("❌ 加载 ONNX 失败")
        exit(ret)

    print("==== 5. 构建 RKNN（量化） ====")
    ret = rknn.build(do_quantization=True, dataset=DATASET_TXT)
    if ret != 0:
        print("❌ 构建 RKNN 失败")
        exit(ret)

    print("==== 6. 导出 RKNN ====")
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print("❌ 导出失败")
        exit(ret)

    print(f"🎉 RKNN 模型导出成功: {RKNN_MODEL}")

    rknn.release()

if __name__ == '__main__':
    main()