| """Learned stroke cost for HanziFlow. |
| |
| Needs: numpy, onnxruntime, and the three files from the training notebook: |
| stroke_encoder.onnx, template_embeddings.npz, calibration.json |
| Strokes must be in the Make Me a Hanzi frame (1024 box, y up) and resampled to 32 points, |
| exactly as in the matcher notebook (prep()). |
| """ |
| import json |
| import numpy as np |
| import onnxruntime as ort |
|
|
|
|
| class LearnedCost: |
| def __init__(self, onnx_path="stroke_encoder.onnx", npz_path="template_embeddings.npz", calib_path="calibration.json"): |
| self.sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"]) |
| z = np.load(npz_path) |
| self.chars = [str(c) for c in z["chars"]] |
| self.templates = {c: z[f"emb_{i}"] for i, c in enumerate(self.chars)} |
| cal = json.load(open(calib_path)) |
| self.scale, self.tau, self.good = cal["scale"], cal["tau"], cal["good"] |
| self.alpha = cal.get("alpha_recommended", 0.5) |
|
|
| @staticmethod |
| def features(S): |
| |
| S = np.asarray(S, np.float32) |
| x = (S[..., 0] - 512) / 512 |
| y = (S[..., 1] - 388) / 512 |
| g = np.gradient(S, axis=-2) |
| nrm = np.linalg.norm(g, axis=-1) + 1e-6 |
| rx = (S[..., 0] - S[..., :1, 0]) / 512 |
| ry = (S[..., 1] - S[..., :1, 1]) / 512 |
| return np.stack([x, y, g[..., 0] / nrm, g[..., 1] / nrm, rx, ry], axis=-1).astype(np.float32) |
|
|
| def embed(self, S): |
| return self.sess.run(None, {"stroke_features": self.features(S)})[0] |
|
|
| def costs(self, char, U): |
| """Learned cost of every drawn stroke (rows) against every template stroke of `char` (columns). |
| Direction-agnostic: a backwards stroke is flagged by the geometric part of the matcher, not here.""" |
| te = self.templates[char] |
| fwd = np.linalg.norm(self.embed(U)[:, None, :] - te[None], axis=-1) |
| bwd = np.linalg.norm(self.embed(np.ascontiguousarray(U[:, ::-1]))[:, None, :] - te[None], axis=-1) |
| return np.minimum(fwd, bwd) / self.scale |
|
|
| def cost_fn(self, char, geometric_cost_matrix, alpha=None): |
| """Drop-in `cost_fn` for align_full(): blends the geometric cost with the learned one. |
| alpha = 0 -> geometric only, 1 -> learned only, 0.5 -> half and half (recommended).""" |
| a = self.alpha if alpha is None else alpha |
| if a == 0: |
| return None |
|
|
| def fn(U, T): |
| C_geo, rev = geometric_cost_matrix(U, T) |
| return (1 - a) * C_geo + a * self.costs(char, U), rev |
| return fn |
|
|
