File size: 2,545 Bytes
b3fcb68 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | #!/usr/bin/env python3
"""Convert PaddlePaddle {lang}_PP-OCRv5_mobile_rec_onnx to NCNN for PaddleOCR-ncnn-CPP."""
from __future__ import annotations
import argparse, re, subprocess, sys
from pathlib import Path
import yaml
def fix_param(text: str) -> str:
lines = []
for line in text.splitlines():
if line.startswith("Input"):
line = line.replace(
"Input in0 0 1 in0",
"Input input 0 1 input",
)
lines.append(line)
continue
parts = line.split()
if len(parts) >= 4:
for i in range(len(parts) - 1, 3, -1):
if parts[i] == "in0": parts[i] = "input"
elif parts[i] == "out0": parts[i] = "output"
line = " ".join(parts)
lines.append(line)
return "\n".join(lines) + "\n"
def num_classes(param_text: str) -> int:
num = 0
for line in param_text.splitlines():
if line.startswith("Gemm"):
m = re.search(r" 8=(\d+)", line)
if m: num = int(m.group(1))
return num
def main():
ap = argparse.ArgumentParser()
ap.add_argument("lang", help="e.g. arabic, ta, en")
ap.add_argument("--onnx-dir", type=Path, default=Path("onnx"))
ap.add_argument("--out-dir", type=Path, default=Path("."))
ap.add_argument("--pnnx", default="pnnx")
args = ap.parse_args()
lang = args.lang
onnx = args.onnx_dir / "inference.onnx"
yml = args.onnx_dir / "inference.yml"
if not onnx.exists():
print("missing", onnx, file=sys.stderr); sys.exit(1)
subprocess.run([args.pnnx, str(onnx), "inputshape=[1,3,48,320]", "inputshape2=[1,3,48,640]", "fp16=0"], check=True, cwd=args.onnx_dir)
param_text = fix_param((args.onnx_dir / "inference.ncnn.param").read_text())
ncls = num_classes(param_text)
chars = yaml.safe_load(yml.read_text(encoding="utf-8"))["PostProcess"]["character_dict"]
keys = [""] + [str(c) for c in chars]
while len(keys) < ncls: keys.append(" ")
keys = keys[:ncls]
args.out_dir.mkdir(parents=True, exist_ok=True)
(args.out_dir / f"{lang}_PP-OCRv5_mobile_rec.param").write_text(param_text)
(args.out_dir / f"{lang}_PP-OCRv5_mobile_rec.bin").write_bytes((args.onnx_dir / "inference.ncnn.bin").read_bytes())
(args.out_dir / f"ppocr_keys_{lang}.txt").write_text("\n".join(keys) + "\n", encoding="utf-8")
print(f"wrote {lang}: classes={ncls}")
if __name__ == "__main__":
main()
|