diff --git a/docs/命令.md b/docs/命令.md index da22303..dc7fd88 100644 --- a/docs/命令.md +++ b/docs/命令.md @@ -120,4 +120,10 @@ ls -l ./rk3588-agent_linux_arm64 -.\minio.exe server E:\minio\data --address ":9000" --console-address ":9001" \ No newline at end of file +.\minio.exe server E:\minio\data --address ":9000" --console-address ":9001" + + +python scripts/pt2rknn_yolov8.py ` + --pt yolov8n.pt ` + --out models/yolov8n-640.rknn ` + --imgsz 640 diff --git a/scripts/pt2rknn_yolov8.py b/scripts/pt2rknn_yolov8.py new file mode 100644 index 0000000..40fb536 --- /dev/null +++ b/scripts/pt2rknn_yolov8.py @@ -0,0 +1,229 @@ +#!/usr/bin/env python3 +import argparse +import os +import subprocess +import sys +from pathlib import Path + + +def run_cmd(cmd): + print("[CMD]", " ".join(cmd)) + ret = subprocess.run(cmd, check=False) + if ret.returncode != 0: + raise RuntimeError(f"command failed ({ret.returncode}): {' '.join(cmd)}") + + +def export_pt_to_onnx(pt_path: Path, onnx_path: Path, imgsz: int, opset: int): + try: + from ultralytics import YOLO # type: ignore + + print("[INFO] exporting via ultralytics python API") + model = YOLO(str(pt_path)) + exported = model.export( + format="onnx", + imgsz=imgsz, + opset=opset, + simplify=True, + dynamic=False, + nms=False, + ) + exported_path = Path(str(exported)) + if exported_path.exists() and exported_path.resolve() != onnx_path.resolve(): + onnx_path.write_bytes(exported_path.read_bytes()) + except Exception as e: + print(f"[WARN] ultralytics API export failed: {e}") + print("[INFO] fallback to CLI: yolo export ...") + run_cmd( + [ + "yolo", + "export", + f"model={pt_path}", + "format=onnx", + f"imgsz={imgsz}", + f"opset={opset}", + "simplify=True", + "dynamic=False", + "nms=False", + ] + ) + default_onnx = pt_path.with_suffix(".onnx") + if not default_onnx.exists(): + raise RuntimeError(f"ONNX not found after export: {default_onnx}") + if default_onnx.resolve() != onnx_path.resolve(): + onnx_path.write_bytes(default_onnx.read_bytes()) + + if not onnx_path.exists(): + raise RuntimeError(f"ONNX export failed, file missing: {onnx_path}") + print(f"[OK] ONNX: {onnx_path}") + + +def _value_info_shape(value_info): + if value_info is None: + return [] + dims = [] + tt = value_info.type.tensor_type + for d in tt.shape.dim: + if d.HasField("dim_value"): + dims.append(int(d.dim_value)) + else: + dims.append(-1) + return dims + + +def _find_value_info(model, name): + for vi in list(model.graph.value_info) + list(model.graph.output) + list(model.graph.input): + if vi.name == name: + return vi + return None + + +def ensure_v8_output_layout_cxn(in_onnx: Path, out_onnx: Path): + try: + import onnx + from onnx import TensorProto, helper, shape_inference + except Exception as e: + raise RuntimeError(f"missing dependency 'onnx': {e}") + + model = onnx.load(str(in_onnx)) + model = shape_inference.infer_shapes(model) + if len(model.graph.output) != 1: + raise RuntimeError( + f"expect single output for this media-server v8 path, got {len(model.graph.output)}" + ) + + out = model.graph.output[0] + out_name = out.name + shape = _value_info_shape(_find_value_info(model, out_name)) + print(f"[INFO] ONNX output shape: {shape if shape else 'unknown'}") + + # Expected CxN style: [1, 4+num_classes, num_boxes], e.g. [1,84,8400] + if len(shape) == 3 and shape[1] > 4 and shape[2] > shape[1]: + # already [1,84,8400] like + onnx.save(model, str(out_onnx)) + print("[OK] output layout already CxN") + return + + # Common Ultralytics style: [1, num_boxes, 4+num_classes], e.g. [1,8400,84] + if len(shape) == 3 and shape[2] > 4 and shape[1] > shape[2]: + transposed = out_name + "_cxn" + node = helper.make_node( + "Transpose", + inputs=[out_name], + outputs=[transposed], + perm=[0, 2, 1], + name="transpose_to_cxn", + ) + model.graph.node.append(node) + model.graph.output.clear() + model.graph.output.extend( + [ + helper.make_tensor_value_info( + transposed, TensorProto.FLOAT, [shape[0], shape[2], shape[1]] + ) + ] + ) + onnx.checker.check_model(model) + onnx.save(model, str(out_onnx)) + print("[OK] output layout fixed by transpose to CxN") + return + + raise RuntimeError( + "unsupported output shape for YOLOv8 in this project; " + "need rank-3 tensor and channels=4+num_classes" + ) + + +def build_rknn( + onnx_path: Path, + rknn_path: Path, + target: str, + imgsz: int, + quant: bool, + dataset: str, + mean: float, + std: float, +): + try: + from rknn.api import RKNN + except Exception as e: + raise RuntimeError(f"missing dependency 'rknn-toolkit2': {e}") + + rknn = RKNN(verbose=True) + rknn.config( + target_platform=target, + mean_values=[[mean, mean, mean]], + std_values=[[std, std, std]], + ) + + ret = rknn.load_onnx(model=str(onnx_path), input_size_list=[[imgsz, imgsz, 3]]) + if ret != 0: + raise RuntimeError(f"rknn.load_onnx failed: {ret}") + + if quant: + if not dataset: + raise RuntimeError("--quant requires --dataset") + ret = rknn.build(do_quantization=True, dataset=dataset) + else: + ret = rknn.build(do_quantization=False) + if ret != 0: + raise RuntimeError(f"rknn.build failed: {ret}") + + ret = rknn.export_rknn(str(rknn_path)) + if ret != 0: + raise RuntimeError(f"rknn.export_rknn failed: {ret}") + rknn.release() + print(f"[OK] RKNN: {rknn_path}") + + +def main(): + ap = argparse.ArgumentParser(description="YOLOv8 .pt -> .rknn for rk3588 media-server") + ap.add_argument("--pt", required=True, help="path to yolov8 .pt") + ap.add_argument("--onnx", default="", help="output onnx path") + ap.add_argument("--onnx_fixed", default="", help="onnx path after layout fix") + ap.add_argument("--out", required=True, help="output .rknn path") + ap.add_argument("--imgsz", type=int, default=640) + ap.add_argument("--opset", type=int, default=12) + ap.add_argument("--target", default="rk3588") + ap.add_argument("--quant", action="store_true", help="enable int8 quantization") + ap.add_argument("--dataset", default="", help="dataset txt for quantization") + ap.add_argument("--mean", type=float, default=0.0, help="default 0.0") + ap.add_argument("--std", type=float, default=255.0, help="default 255.0") + args = ap.parse_args() + + pt_path = Path(args.pt).resolve() + if not pt_path.exists(): + raise SystemExit(f"pt file not found: {pt_path}") + + onnx_path = Path(args.onnx).resolve() if args.onnx else pt_path.with_suffix(".onnx").resolve() + onnx_fixed = ( + Path(args.onnx_fixed).resolve() + if args.onnx_fixed + else onnx_path.with_name(onnx_path.stem + "_cxn.onnx").resolve() + ) + out_path = Path(args.out).resolve() + out_path.parent.mkdir(parents=True, exist_ok=True) + + export_pt_to_onnx(pt_path, onnx_path, args.imgsz, args.opset) + ensure_v8_output_layout_cxn(onnx_path, onnx_fixed) + build_rknn( + onnx_path=onnx_fixed, + rknn_path=out_path, + target=args.target, + imgsz=args.imgsz, + quant=args.quant, + dataset=args.dataset, + mean=args.mean, + std=args.std, + ) + + print( + "[NEXT] config ai_yolo with model_version='v8', and preprocess dst_w/dst_h == imgsz, keep_ratio=false" + ) + + +if __name__ == "__main__": + try: + main() + except Exception as e: + print("[ERROR]", e) + sys.exit(1) diff --git a/yolov8n.onnx b/yolov8n.onnx new file mode 100644 index 0000000..57d71d7 Binary files /dev/null and b/yolov8n.onnx differ diff --git a/yolov8n_cxn.onnx b/yolov8n_cxn.onnx new file mode 100644 index 0000000..8c802a1 Binary files /dev/null and b/yolov8n_cxn.onnx differ