File size: 2,138 Bytes
99fe130 | 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 63 64 65 | """随 HF 模型发布的自包含 NCNN 推理入口。"""
from __future__ import annotations
import json
from pathlib import Path
import ncnn
import numpy as np
from PIL import Image
from normalize import normalize
class HCCRRecognizer:
def __init__(
self,
param_path: str | Path,
bin_path: str | Path,
charset_path: str | Path,
num_threads: int = 4,
) -> None:
self.net = ncnn.Net()
self.net.opt.use_vulkan_compute = False
self.net.opt.num_threads = num_threads
self.net.load_param(str(param_path))
self.net.load_model(str(bin_path))
data = json.loads(Path(charset_path).read_text(encoding="utf-8"))
self.idx_to_char = {int(i): ch for ch, i in data["char_to_idx"].items()}
@classmethod
def from_pretrained(cls, model_dir: str | Path = ".", num_threads: int = 4):
root = Path(model_dir)
return cls(
root / "ncnn/model.ncnn.param",
root / "ncnn/model.ncnn.bin",
root / "ncnn/charset.json",
num_threads=num_threads,
)
@staticmethod
def preprocess(image: Image.Image) -> np.ndarray:
gray = np.asarray(image.convert("L"), dtype=np.uint8)
normalized = normalize(gray)
return ((255 - normalized).astype(np.float32) / 255.0)[None, :, :]
def predict(self, image: Image.Image, k: int = 10) -> list[tuple[str, float]]:
x = self.preprocess(image)
if not x.any():
return []
extractor = self.net.create_extractor()
extractor.input("in0", ncnn.Mat(x))
ret, out = extractor.extract("out0")
if ret != 0:
raise RuntimeError(f"ncnn extract failed: {ret}")
logits = np.asarray(out)
probs = np.exp(logits - logits.max())
probs /= probs.sum()
indices = np.argsort(-probs)[:k]
return [(self.idx_to_char[int(i)], float(probs[i])) for i in indices]
def predict_file(self, path: str | Path, k: int = 10) -> list[tuple[str, float]]:
with Image.open(path) as image:
return self.predict(image, k=k)
|