face-rknn/tools/conversion/convert_scrfd_rknn.py

45 lines
833 B
Python

from rknn.api import RKNN
ONNX_MODEL = "SCRFD_500M_KPS_640.onnx"
RKNN_MODEL = "SCRFD_500M_KPS_640.rknn"
print("=" * 60)
print("SCRFD ONNX -> RKNN")
print("=" * 60)
rknn = RKNN(verbose=True)
print("\n[1] Configuring RKNN...")
ret = rknn.config(
target_platform="rk3588"
)
if ret != 0:
raise RuntimeError("rknn.config() failed")
print("\n[2] Loading ONNX...")
ret = rknn.load_onnx(
model=ONNX_MODEL
)
if ret != 0:
raise RuntimeError("rknn.load_onnx() failed")
print("\n[3] Building RKNN (FP16, no quantization)...")
ret = rknn.build(
do_quantization=False
)
if ret != 0:
raise RuntimeError("rknn.build() failed")
print("\n[4] Exporting RKNN...")
ret = rknn.export_rknn(RKNN_MODEL)
if ret != 0:
raise RuntimeError("rknn.export_rknn() failed")
print("\nSUCCESS")
print(RKNN_MODEL)
rknn.release()