from pathlib import Path
from rknn.api import RKNN


# プロジェクトルートから見た ONNX モデルのパス
ONNX_MODEL = "./models/onnx/PP-OCRv6_small_rec/inference.onnx"

# 出力する RKNN モデルのパス
RKNN_MODEL = "./models/rknn/ppocrv6_small_rec_rv1126b_fp.rknn"

# PP-OCRv6 認識モデルの固定入力サイズ
# 一般的な PP-OCR 認識モデルでは、高さ 48、幅 320 を使用します。
REC_INPUT_SIZE = [1, 3, 48, 320]


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

    # 認識モデルでは、入力画像をおおむね [-1, 1] 範囲へ正規化する想定です。
    # C/C++ 側では、RGB 変換と resize / padding のみを行い、
    # mean/std は 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")

    # ONNX の入力は動的形状になっている可能性があるため、
    # RKNN 変換時に固定入力サイズ [1, 3, 48, 320] を指定します。
    ret = rknn.load_onnx(
        model=ONNX_MODEL,
        inputs=["x"],
        input_size_list=[REC_INPUT_SIZE],
    )
    if ret != 0:
        raise RuntimeError("load_onnx failed")

    print("--> Build RKNN without quantization")

    # 認識モデルは文字精度への影響を避けるため、v1.0 では非量子化で変換します。
    # そのため、dataset は不要です。
    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()
