from rknn.api import RKNN

ONNX_MODEL = "./models/onnx/PP-OCRv6_small_det/model.onnx"
RKNN_MODEL = "./models/rknn/ppocrv6_small_det_rv1126b_i8.rknn"
DATASET = "./dataset_det.txt"

rknn = RKNN(verbose=True)

rknn.config(
    target_platform="rv1126b",
    mean_values=[[123.675, 116.28, 103.53]],
    std_values=[[58.395, 57.12, 57.375]],
    optimization_level=3,
)

print("--> Load ONNX")
ret = rknn.load_onnx(model=ONNX_MODEL)
if ret != 0:
    raise RuntimeError("load_onnx failed")

print("--> Build RKNN")
ret = rknn.build(
    do_quantization=True,
    dataset=DATASET,
)
if ret != 0:
    raise RuntimeError("build failed")

print("--> Export RKNN")
ret = rknn.export_rknn(RKNN_MODEL)
if ret != 0:
    raise RuntimeError("export failed")

rknn.release()
print("done:", RKNN_MODEL)

