"""看帖识字 Kantie Shizi v1 — 书法单字识别（推理代码）。
输入：一张切好的单字图（任意尺寸，彩色或灰度）；输出：最可能的字（共 4,167 个字种，繁体按馆方釋文写法）及概率。
预处理与训练一致：四边取纸色补成正方形 → 224 双三次 → 灰度 → 64 双线性 → 每图减均值除标准差。
用法：python3 kantie.py 字图.png [--topk 5]"""
import json, sys, os, numpy as np
from PIL import Image
D = os.path.dirname(os.path.abspath(__file__))

def preprocess(img):
    a = np.asarray(img.convert('RGB')); h, w, _ = a.shape
    bg = np.median(np.concatenate([a[0], a[-1], a[:, 0], a[:, -1]]), 0).astype(np.uint8)
    n = max(h, w); c = np.empty((n, n, 3), np.uint8); c[:] = bg
    c[(n - h) // 2:(n - h) // 2 + h, (n - w) // 2:(n - w) // 2 + w] = a
    g = np.asarray(Image.fromarray(c).resize((224, 224), Image.BICUBIC).convert('L').resize((64, 64), Image.BILINEAR)).astype(np.float32) / 255
    return ((g - g.mean()) / (g.std() + 1e-3))[None, None]

def build_net(n_cls):
    import torch.nn as nn
    def blk(i, o): return nn.Sequential(nn.Conv2d(i, o, 3, padding=1), nn.BatchNorm2d(o), nn.ReLU(), nn.Conv2d(o, o, 3, padding=1), nn.BatchNorm2d(o), nn.ReLU(), nn.MaxPool2d(2))
    return nn.Sequential(blk(1, 64), blk(64, 128), blk(128, 256), blk(256, 384), nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Dropout(0.3), nn.Linear(384, n_cls))

class Kantie:
    def __init__(self, d=D, backend='onnx'):
        self.cls = json.load(open(os.path.join(d, 'classes.json'), encoding='utf-8'))
        if backend == 'onnx':
            import onnxruntime as ort
            self.s = ort.InferenceSession(os.path.join(d, 'kantie_v1.onnx'), providers=['CPUExecutionProvider']); self.torch = None
        else:
            import torch
            from safetensors.torch import load_file
            self.torch = build_net(len(self.cls)); self.torch.load_state_dict(load_file(os.path.join(d, 'kantie_v1.safetensors'))); self.torch.eval()
    def predict(self, img, topk=5):
        x = preprocess(img if isinstance(img, Image.Image) else Image.open(img))
        if self.torch is None: z = self.s.run(None, {'image': x})[0][0]
        else:
            import torch
            with torch.no_grad(): z = self.torch(torch.from_numpy(x))[0].numpy()
        p = np.exp(z - z.max()); p /= p.sum(); i = np.argsort(-p)[:topk]
        return [(self.cls[k], float(p[k])) for k in i]

if __name__ == '__main__':
    k = int(sys.argv[sys.argv.index('--topk') + 1]) if '--topk' in sys.argv else 5
    for ch, pr in Kantie().predict(sys.argv[1], k): print(ch, round(pr, 4))