#!/usr/bin/env python3
"""将 YOLOv5-pose 权重转换为 RKNN 模型(RK3568/RK3576/RK3588)。

依赖:
  pip install rknn-toolkit2 ultralytics onnx onnxsim
用法:
  python3 scripts/export_pose.py --weights yolov5s-pose.pt \
      --rknn models/yolov5s-pose_rk3588_int8.rknn --chip rk3588
"""
import sys
import os
import argparse

def export_onnx(weights, onnx_path, input_size=640):
    from ultralytics import YOLO
    model = YOLO(weights)
    model.export(format='onnx', imgsz=input_size, simplify=True,
                 opset=12, dynamic=False)
    default = os.path.splitext(weights)[0] + '.onnx'
    if os.path.exists(default):
        os.replace(default, onnx_path)
    print(f"[ONNX] -> {onnx_path}")

def convert(onnx_path, rknn_path, chip, input_size=640, fp=False):
    from rknn.api import RKNN
    rknn = RKNN(verbose=False)
    if rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]],
                   target_platform=f'rk{chip}',
                   optimization_level=3) != 0:
        raise RuntimeError("rknn.config 失败")
    if rknn.load_onnx(model=onnx_path, outputs=None) != 0:
        raise RuntimeError("load_onnx 失败")
    if fp:
        if rknn.build(do_quantization=False) != 0:
            raise RuntimeError("build(fp) 失败")
    else:
        if rknn.build(do_quantization=True, dataset='scripts/dataset.txt') != 0:
            raise RuntimeError("build(int8) 失败")
    if rknn.export_rknn(rknn_path) != 0:
        raise RuntimeError("export 失败")
    print(f"[RKNN] -> {rknn_path}")
    rknn.release()

if __name__ == '__main__':
    ap = argparse.ArgumentParser()
    ap.add_argument('--weights', required=True)
    ap.add_argument('--rknn', default='models/yolov5s-pose_rk3588_int8.rknn')
    ap.add_argument('--chip', default='rk3588', help='rk3568/rk3576/rk3588')
    ap.add_argument('--input-size', type=int, default=640)
    ap.add_argument('--fp16', action='store_true', help='导出fp16而非int8')
    args = ap.parse_args()

    onnx = os.path.splitext(args.weights)[0] + '.onnx'
    export_onnx(args.weights, onnx, args.input_size)
    convert(onnx, args.rknn, args.chip, args.input_size, fp=args.fp16)