from pathlib import Path
from rknn.api import RKNN


ONNX_MODEL = "./models/onnx/PP-OCRv6_small_rec/inference.onnx"
RKNN_MODEL = "./models/rknn/ppocrv6_small_rec_rv1126b_fp_no_softmax.rknn"

REC_INPUT_SIZE = [1, 3, 48, 320]

# Final Softmax.2 の入力テンソルを出力として指定します。
# CTC decode では softmax 後ではなく logits の argmax を使用しても結果は同じです。
REC_OUTPUT_NAME = "p2o.pd_op.add.79.0"


def check_file(path: str, desc: str):
    p = Path(path)
    if not p.exists():
        raise FileNotFoundError(f"{desc} が見つかりません: {path}")


def main():
    check_file(ONNX_MODEL, "ONNX モデル")

    Path("./models/rknn").mkdir(parents=True, exist_ok=True)
    Path("./logs").mkdir(parents=True, exist_ok=True)

    rknn = RKNN(verbose=True)

    print("--> Config RKNN")
    rknn.config(
        target_platform="rv1126b",
        mean_values=[[127.5, 127.5, 127.5]],
        std_values=[[127.5, 127.5, 127.5]],
        optimization_level=3,
    )

    print("--> Load ONNX without final Softmax")
    ret = rknn.load_onnx(
        model=ONNX_MODEL,
        inputs=["x"],
        input_size_list=[REC_INPUT_SIZE],
        outputs=[REC_OUTPUT_NAME],
    )
    if ret != 0:
        raise RuntimeError("load_onnx failed")

    print("--> Build RKNN without quantization")
    ret = rknn.build(do_quantization=False)
    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)


if __name__ == "__main__":
    main()
