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()