Upload scripts/train_captcha.py with huggingface_hub

scripts/train_captcha.py

@@ -1,3 +1,180 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:4d32be00cfa26e441c22668a065f8669182a3ff9c5c0498d6df75c3ea7cab9df
-size 6997
+#!/usr/bin/env python3
+"""训练南+纯数字验证码识别小模型（CTC），导出 ONNX。
+输入: labels.json {filename: {label, ...}} + 图片目录
+模型: 小 CNN -> 高度聚合 -> (T=20, C=11) 时序 -> CTC 解码(10数字+blank)
+输出: onnx 模型 + 训练日志
+
+预处理(Go 侧需完全一致):
+  灰度 -> resize 到 (160,64) BILINEAR -> float32 /255 -> (1,1,64,160)
+"""
+import json
+import os
+import random
+import sys
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+
+W, H = 160, 64
+NUM_CLASSES = 11  # 0-9 + blank(10)
+DIGITS = "0123456789"
+
+
+class Dataset:
+    def __init__(self, items, aug=False, w=W, h=H):
+        self.items = items
+        self.aug, self.w, self.h = aug, w, h
+
+    def __len__(self):
+        return len(self.items)
+
+    def __getitem__(self, i):
+        path, lbl = self.items[i]
+        im = Image.open(path).convert("L").resize((self.w, self.h), Image.BILINEAR)
+        if self.aug:
+            im = augment(im)
+        a = np.asarray(im, dtype=np.float32) / 255.0
+        x = torch.from_numpy(a).unsqueeze(0)  # (1,h,w)
+        target = torch.tensor([DIGITS.index(c) for c in lbl], dtype=torch.long)
+        return x, target, len(lbl)
+
+
+def augment(im):
+    """纯 PIL 数据增强：旋转/缩放/平移 + 高斯噪声。背景保持白色。"""
+    if random.random() < 0.8:
+        im = im.rotate(random.uniform(-6, 6), resample=Image.BILINEAR, fillcolor=255)
+    scale = random.uniform(0.92, 1.08)
+    tx, ty = random.uniform(-3, 3), random.uniform(-2, 2)
+    w, h = im.size
+    # 以中心为锚点缩放 + 平移（输出坐标 -> 输入坐标 的仿射矩阵）
+    a, b, c, d, e, f = (scale, 0, -w * scale / 2 + w / 2 + tx,
+                        0, scale, -h * scale / 2 + h / 2 + ty)
+    im = im.transform((w, h), Image.AFFINE, (a, b, c, d, e, f),
+                      resample=Image.BILINEAR, fillcolor=255)
+    arr = np.asarray(im, dtype=np.float32)
+    if random.random() < 0.6:
+        arr += np.random.normal(0, 10, arr.shape)
+    return Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8), "L")
+
+
+class CaptchaNet(nn.Module):
+    def __init__(self, num_classes=NUM_CLASSES):
+        super().__init__()
+        self.features = nn.Sequential(
+            nn.Conv2d(1, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2),
+            nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2),
+            nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(), nn.MaxPool2d(2),
+            nn.Conv2d(128, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(),
+            nn.Conv2d(256, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(),
+        )
+        self.head = nn.Conv2d(256, num_classes, 1)
+
+    def forward(self, x):
+        x = self.features(x)          # (B,256,8,20)
+        x = self.head(x)              # (B,11,8,20)
+        x = x.mean(dim=2)             # 高度聚合 (B,11,20)
+        x = x.permute(0, 2, 1)        # (B,20,11)
+        return x
+
+
+def ctc_decode(logits):
+    """logits: (B,T,C) -> 字符串列表"""
+    out = []
+    idxs = logits.argmax(dim=-1)  # (B,T)
+    for row in idxs:
+        prev = -1
+        s = []
+        for t in row.tolist():
+            if t != prev and t != 10:  # 去重+跳过blank
+                s.append(DIGITS[t])
+            prev = t
+        out.append("".join(s))
+    return out
+
+
+def main():
+    labels_json = sys.argv[1]
+    imgdir = sys.argv[2]
+    out_model = sys.argv[3] if len(sys.argv) > 3 else "/tmp/sp_captest/captcha.onnx"
+    labels = {k: v.get("label") for k, v in json.load(open(labels_json)).items()}
+
+    all_items = [(os.path.join(imgdir, fn), lbl) for fn, lbl in labels.items() if lbl]
+    random.Random(42).shuffle(all_items)
+    val_frac = float(sys.argv[4]) if len(sys.argv) > 4 else 0.15
+    n_val = int(len(all_items) * val_frac)
+    val_items, train_items = all_items[:n_val], all_items[n_val:]
+    print(f"train={len(train_items)} val={len(val_items)} total={len(all_items)}")
+
+    train_ds, val_ds = Dataset(train_items, aug=True), Dataset(val_items, aug=False)
+    tr = torch.utils.data.DataLoader(train_ds, batch_size=16, shuffle=True,
+                                     collate_fn=lambda b: collate(b))
+    va = torch.utils.data.DataLoader(val_ds, batch_size=16, shuffle=False,
+                                     collate_fn=lambda b: collate(b))
+
+    device = "cpu"
+    model = CaptchaNet().to(device)
+    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
+    sched = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, factor=0.5, patience=8)
+    crit = nn.CTCLoss(blank=10, zero_infinity=True)
+
+    best = 0.0
+    epochs = int(sys.argv[5]) if len(sys.argv) > 5 else 200
+    for epoch in range(epochs):
+        model.train()
+        tot = 0.0
+        for x, target, tl in tr:
+            x, target = x.to(device), target.to(device)
+            tl = torch.tensor(tl)
+            logits = model(x)                     # (B,T,C)
+            lp = F.log_softmax(logits, dim=2)
+            input_lengths = torch.full((x.size(0),), logits.size(1), dtype=torch.long)
+            loss = crit(lp.permute(1, 0, 2), target, input_lengths, tl)
+            opt.zero_grad(); loss.backward(); opt.step()
+            tot += loss.item() * x.size(0)
+        # eval
+        model.eval()
+        acc = 0.0; n = 0
+        if len(va) > 0:
+            with torch.no_grad():
+                for x, target, tl in va:
+                    preds = ctc_decode(model(x.to(device)))
+                    for p, (_, lbl) in zip(preds, [(None, lbl) for _, lbl in val_ds.items[n:n + x.size(0)]]):
+                        acc += (p == lbl)
+                        n += 1
+            acc /= max(1, n)
+            sched.step(acc)
+            if acc > best:
+                best = acc
+                torch.save(model.state_dict(), out_model + ".pt")
+        if epoch % 10 == 0 or epoch == 199:
+            print(f"epoch {epoch} loss={tot/len(train_ds):.3f} val_acc={acc:.2%} best={best:.2%}", flush=True)
+    if len(va) == 0:  # 无验证集：保存最后权重
+        torch.save(model.state_dict(), out_model + ".pt")
+    print(f"DONE best_val_acc={best:.2%}")
+
+    # 加载最优权重导出 ONNX
+    model.load_state_dict(torch.load(out_model + ".pt"))
+    model.eval()
+    dummy = torch.randn(1, 1, H, W)
+    torch.onnx.export(model, dummy, out_model,
+                      input_names=["input"], output_names=["logits"],
+                      dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}},
+                      opset_version=13, external_data=False)
+    print(f"ONNX exported -> {out_model}")
+    # 导出标签 JSON（校验用）
+    with open(out_model + ".labels.json", "w") as f:
+        json.dump({"charset": list(DIGITS), "blank": 10, "w": W, "h": H}, f)
+
+
+def collate(batch):
+    xs = torch.stack([b[0] for b in batch])
+    ts = torch.cat([b[1] for b in batch])
+    tl = [b[2] for b in batch]
+    return xs, ts, tl
+
+
+if __name__ == "__main__":
+    main()