diff --git a/internal/captcha/handler.go b/internal/captcha/handler.go index f955418..b24606d 100644 --- a/internal/captcha/handler.go +++ b/internal/captcha/handler.go @@ -392,35 +392,24 @@ func preprocessMath(img image.Image) ([]float32, error) { } func decodeMathFromOutput(output []float32) string { - // Math 模型输出: [C, B, T] = [51, 1, 19] - // 需要转置为 [B, T, C] = [1, 19, 51] - numChars := 51 + // Math 模型输出: [T, B, C] = [51, 1, 19] + // T = 51 个时间步, B = 1 batch, C = 19 个字符类别 (blank + 18 chars) + // 对每个 timestep 取 argmax,得到字符索引序列 + timesteps := 51 batch := 1 - timesteps := 19 - - if len(output) != numChars*batch*timesteps { - timesteps = len(output) / numChars - if timesteps <= 0 { + numChars := 19 // blank(0) + 18 chars + + if len(output) != timesteps*batch*numChars { + // 尝试推断 + total := len(output) + if total%numChars != 0 { return "" } + timesteps = total / numChars } - - // 转置: [C, B, T] -> [B, T, C] - // output[c*batch*timesteps + b*timesteps + t] -> data[b*timesteps*numChars + t*numChars + c] - data := make([]float32, batch*timesteps*numChars) - for c := 0; c < numChars; c++ { - for b := 0; b < batch; b++ { - for t := 0; t < timesteps; t++ { - srcIdx := c*batch*timesteps + b*timesteps + t - dstIdx := b*timesteps*numChars + t*numChars + c - if srcIdx < len(output) && dstIdx < len(data) { - data[dstIdx] = output[srcIdx] - } - } - } - } - - // CTC 解码: 对每个 timestep 取 argmax + + // 对每个 timestep 取 argmax + // output[t*batch*numChars + b*numChars + c] 存储的是 timestep t, batch b 的字符 c 的概率 result := "" lastIdx := -1 @@ -428,16 +417,18 @@ func decodeMathFromOutput(output []float32) string { maxIdx := 0 maxProb := float32(-math.MaxFloat32) + // 对这个 timestep 的所有字符取最大值 for c := 0; c < numChars; c++ { - idx := t*numChars + c - if idx < len(data) && data[idx] > maxProb { - maxProb = data[idx] + idx := t*numChars + c // 因为 batch=1,所以简化了 + if idx < len(output) && output[idx] > maxProb { + maxProb = output[idx] maxIdx = c } } // CTC: 跳过 blank (index 0) 和重复 if maxIdx != 0 && maxIdx != lastIdx { + // 字符索引从 1 开始,对应 mathChars[0:] if maxIdx-1 < len(mathChars) { result += string(mathChars[maxIdx-1]) }