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