from rknn.api import RKNN

def convert(onnx_path, rknn_path):
    print(f"\n=== Converting {onnx_path} → {rknn_path} ===")

    rknn = RKNN()

    # Ottimizzazioni consigliate da Rockchip
    rknn.config(
        mean_values=[[127.5, 127.5, 127.5]],
        std_values=[[128.0, 128.0, 128.0]],
        target_platform="rk3588"
    )

    print("[1] Loading ONNX model...")
    ret = rknn.load_onnx(model=onnx_path)
    if ret != 0:
        print("Error loading ONNX")
        return

    print("[2] Building RKNN model...")
    ret = rknn.build(do_quantization=False)
    if ret != 0:
        print("Error building RKNN")
        return

    print("[3] Exporting RKNN model...")
    ret = rknn.export_rknn(rknn_path)
    if ret != 0:
        print("Error exporting RKNN")
        return

    print("[OK] Conversion completed!")

# Convert SCRFD
convert(
    "models/onnx/scrfd_2.5g.onnx",
    "models/rknn/scrfd.rknn"
)

# Convert ArcFace
convert(
    "models/onnx/arcface_r100.onnx",
    "models/rknn/arcface.rknn"
)