FlanChanXwO/phpwind-captcha-ocr

1#!/usr/bin/env python32"""训练 PHPWind 纯数字验证码识别小模型（CTC），导出 ONNX。3输入: labels.json {filename: {label, ...}} + 图片目录4模型: 小 CNN -> 高度聚合 -> (T=20, C=11) 时序 -> CTC 解码(10数字+blank)5输出: onnx 模型 + 训练日志6 7预处理(Go 侧需完全一致):8  灰度 -> resize 到 (160,64) BILINEAR -> float32 /255 -> (1,1,64,160)9"""10import json11import os12import random13import sys14 15import numpy as np16import torch17import torch.nn as nn18import torch.nn.functional as F19from PIL import Image20 21W, H = 160, 6422NUM_CLASSES = 11  # 0-9 + blank(10)23DIGITS = "0123456789"24 25 26class Dataset:27    def __init__(self, items, aug=False, w=W, h=H):28        self.items = items29        self.aug, self.w, self.h = aug, w, h30 31    def __len__(self):32        return len(self.items)33 34    def __getitem__(self, i):35        path, lbl = self.items[i]36        im = Image.open(path).convert("L").resize((self.w, self.h), Image.BILINEAR)37        if self.aug:38            im = augment(im)39        a = np.asarray(im, dtype=np.float32) / 255.040        x = torch.from_numpy(a).unsqueeze(0)  # (1,h,w)41        target = torch.tensor([DIGITS.index(c) for c in lbl], dtype=torch.long)42        return x, target, len(lbl)43 44 45def augment(im):46    """纯 PIL 数据增强：旋转/缩放/平移 + 高斯噪声。背景保持白色。"""47    if random.random() < 0.8:48        im = im.rotate(random.uniform(-6, 6), resample=Image.BILINEAR, fillcolor=255)49    scale = random.uniform(0.92, 1.08)50    tx, ty = random.uniform(-3, 3), random.uniform(-2, 2)51    w, h = im.size52    # 以中心为锚点缩放 + 平移（输出坐标 -> 输入坐标 的仿射矩阵）53    a, b, c, d, e, f = (scale, 0, -w * scale / 2 + w / 2 + tx,54                        0, scale, -h * scale / 2 + h / 2 + ty)55    im = im.transform((w, h), Image.AFFINE, (a, b, c, d, e, f),56                      resample=Image.BILINEAR, fillcolor=255)57    arr = np.asarray(im, dtype=np.float32)58    if random.random() < 0.6:59        arr += np.random.normal(0, 10, arr.shape)60    return Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8), "L")61 62 63class CaptchaNet(nn.Module):64    def __init__(self, num_classes=NUM_CLASSES):65        super().__init__()66        self.features = nn.Sequential(67            nn.Conv2d(1, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2),68            nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2),69            nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(), nn.MaxPool2d(2),70            nn.Conv2d(128, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(),71            nn.Conv2d(256, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(),72        )73        self.head = nn.Conv2d(256, num_classes, 1)74 75    def forward(self, x):76        x = self.features(x)          # (B,256,8,20)77        x = self.head(x)              # (B,11,8,20)78        x = x.mean(dim=2)             # 高度聚合 (B,11,20)79        x = x.permute(0, 2, 1)        # (B,20,11)80        return x81 82 83def ctc_decode(logits):84    """logits: (B,T,C) -> 字符串列表"""85    out = []86    idxs = logits.argmax(dim=-1)  # (B,T)87    for row in idxs:88        prev = -189        s = []90        for t in row.tolist():91            if t != prev and t != 10:  # 去重+跳过blank92                s.append(DIGITS[t])93            prev = t94        out.append("".join(s))95    return out96 97 98def main():99    labels_json = sys.argv[1]100    imgdir = sys.argv[2]101    out_model = sys.argv[3] if len(sys.argv) > 3 else "/tmp/sp_captest/captcha.onnx"102    labels = {k: v.get("label") for k, v in json.load(open(labels_json)).items()}103 104    all_items = [(os.path.join(imgdir, fn), lbl) for fn, lbl in labels.items() if lbl]105    random.Random(42).shuffle(all_items)106    val_frac = float(sys.argv[4]) if len(sys.argv) > 4 else 0.15107    n_val = int(len(all_items) * val_frac)108    val_items, train_items = all_items[:n_val], all_items[n_val:]109    print(f"train={len(train_items)} val={len(val_items)} total={len(all_items)}")110 111    train_ds, val_ds = Dataset(train_items, aug=True), Dataset(val_items, aug=False)112    tr = torch.utils.data.DataLoader(train_ds, batch_size=16, shuffle=True,113                                     collate_fn=lambda b: collate(b))114    va = torch.utils.data.DataLoader(val_ds, batch_size=16, shuffle=False,115                                     collate_fn=lambda b: collate(b))116 117    device = "cpu"118    model = CaptchaNet().to(device)119    opt = torch.optim.Adam(model.parameters(), lr=1e-3)120    sched = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, factor=0.5, patience=8)121    crit = nn.CTCLoss(blank=10, zero_infinity=True)122 123    best = 0.0124    epochs = int(sys.argv[5]) if len(sys.argv) > 5 else 200125    for epoch in range(epochs):126        model.train()127        tot = 0.0128        for x, target, tl in tr:129            x, target = x.to(device), target.to(device)130            tl = torch.tensor(tl)131            logits = model(x)                     # (B,T,C)132            lp = F.log_softmax(logits, dim=2)133            input_lengths = torch.full((x.size(0),), logits.size(1), dtype=torch.long)134            loss = crit(lp.permute(1, 0, 2), target, input_lengths, tl)135            opt.zero_grad(); loss.backward(); opt.step()136            tot += loss.item() * x.size(0)137        # eval138        model.eval()139        acc = 0.0; n = 0140        if len(va) > 0:141            with torch.no_grad():142                for x, target, tl in va:143                    preds = ctc_decode(model(x.to(device)))144                    for p, (_, lbl) in zip(preds, [(None, lbl) for _, lbl in val_ds.items[n:n + x.size(0)]]):145                        acc += (p == lbl)146                        n += 1147            acc /= max(1, n)148            sched.step(acc)149            if acc > best:150                best = acc151                torch.save(model.state_dict(), out_model + ".pt")152        if epoch % 10 == 0 or epoch == 199:153            print(f"epoch {epoch} loss={tot/len(train_ds):.3f} val_acc={acc:.2%} best={best:.2%}", flush=True)154    if len(va) == 0:  # 无验证集：保存最后权重155        torch.save(model.state_dict(), out_model + ".pt")156    print(f"DONE best_val_acc={best:.2%}")157 158    # 加载最优权重导出 ONNX159    model.load_state_dict(torch.load(out_model + ".pt"))160    model.eval()161    dummy = torch.randn(1, 1, H, W)162    torch.onnx.export(model, dummy, out_model,163                      input_names=["input"], output_names=["logits"],164                      dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}},165                      opset_version=13, external_data=False)166    print(f"ONNX exported -> {out_model}")167    # 导出标签 JSON（校验用）168    with open(out_model + ".labels.json", "w") as f:169        json.dump({"charset": list(DIGITS), "blank": 10, "w": W, "h": H}, f)170 171 172def collate(batch):173    xs = torch.stack([b[0] for b in batch])174    ts = torch.cat([b[1] for b in batch])175    tl = [b[2] for b in batch]176    return xs, ts, tl177 178 179if __name__ == "__main__":180    main()181