FlanChanXwO/phpwind-captcha-ocr

1#!/usr/bin/env python32"""用原始图训练的 ONNX 模型识别验证码。预处理: RGB 160x64 /255 -> (1,3,64,160) -> 4位数字3"""4import json5import os6import sys7 8import numpy as np9import onnxruntime as ort10from PIL import Image11 12MODEL = os.path.expanduser("~/sp_captcha_assets/models/captcha_1000_raw.onnx")13_sess = None14 15 16def _get_sess():17    global _sess18    if _sess is None:19        _sess = ort.InferenceSession(MODEL, providers=["CPUExecutionProvider"])20    return _sess21 22 23def predict(path):24    im = Image.open(path).convert("RGB").resize((160, 64), Image.BILINEAR)25    a = np.asarray(im, dtype=np.float32) / 255.026    x = a.transpose(2, 0, 1)[None]27    out = _get_sess().run(None, {"input": x})[0]28    return "".join(str(int(out[0, p].argmax())) for p in range(4))29 30 31def main():32    out = {}33    for f in sys.argv[1:]:34        try:35            out[f] = {"onnx": predict(f)}36        except Exception as e:37            out[f] = {"onnx": "", "err": str(e)}38    print(json.dumps(out, ensure_ascii=False))39 40 41if __name__ == "__main__":42    main()43