from pathlib import Path
from rknn.api import RKNN


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

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

# INT8 量子化用のキャリブレーション画像リスト
DATASET = "./datasets/ppocrv6_det_calib.txt"

# PP-OCRv6 検出モデルの固定入力サイズ
# PaddleOCR 側の text_det_limit_side_len=736 に合わせる
DET_INPUT_SIZE = [1, 3, 736, 736]


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 モデル")
    check_file(DATASET, "量子化用データセットリスト")

    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=[[123.675, 116.28, 103.53]],
        std_values=[[58.395, 57.12, 57.375]],
        optimization_level=3,
    )

    print("--> Load ONNX")

    # ONNX の入力名はエラーメッセージから 'x' と確認できる
    # 動的入力形状を固定入力形状 [1, 3, 736, 736] に変換する
    ret = rknn.load_onnx(
        model=ONNX_MODEL,
        inputs=["x"],
        input_size_list=[DET_INPUT_SIZE],
    )
    if ret != 0:
        raise RuntimeError("load_onnx failed")

    print("--> Build RKNN with INT8 quantization")
    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)


if __name__ == "__main__":
    main()

