fix: correct Math decode - output is [num_positions, batch, num_chars]
Build and Deploy / build (push) Successful in 3m3s
Build and Deploy / deploy (push) Successful in 11s

This commit is contained in:
2026-07-17 11:53:17 +00:00
parent f2bc0ecdb2
commit a3665ed63a
+17 -15
View File
@@ -392,34 +392,36 @@ func preprocessMath(img image.Image) ([]float32, error) {
} }
func decodeMathFromOutput(output []float32) string { func decodeMathFromOutput(output []float32) string {
// Math 模型输出: [T, B, C] = [51, 1, 19] // Math 模型输出: [num_positions, batch, num_chars] = [51, 1, 19]
// T = 51 个时间步, B = 1 batch, C = 19 个字符类别 (blank + 18 chars) // 51 个位置,每个位置有 19 个字符概率(blank + 18 chars
// 对每个 timestep 取 argmax,得到字符索引序列 // 对每个位置取 argmax,得到 51 个字符索引
timesteps := 51 numPositions := 51
batch := 1 batch := 1
numChars := 19 // blank(0) + 18 chars numChars := 19
if len(output) != timesteps*batch*numChars { if len(output) != numPositions*batch*numChars {
// 尝试推断 // 尝试推断
total := len(output) total := len(output)
if total%numChars != 0 { if total%numChars != 0 {
return "" return ""
} }
timesteps = total / numChars numPositions = total / numChars
} }
// 对每个 timestep 取 argmax CHARS := "0123456789+-*/÷×=?"
// output[t*batch*numChars + b*numChars + c] 存储的是 timestep t, batch b 的字符 c 的概率
// 对每个位置取 argmax
// output[pos*batch*numChars + b*numChars + c] 存储的是位置 pos, batch b 的字符 c 的概率
result := "" result := ""
lastIdx := -1 lastIdx := -1
for t := 0; t < timesteps; t++ { for pos := 0; pos < numPositions; pos++ {
maxIdx := 0 maxIdx := 0
maxProb := float32(-math.MaxFloat32) maxProb := float32(-math.MaxFloat32)
// 对这个 timestep 的所有字符取最大值 // 对这个位置的 19 个字符取最大值
for c := 0; c < numChars; c++ { for c := 0; c < numChars; c++ {
idx := t*numChars + c // 因为 batch=1,所以简化 idx := pos*numChars + c // batch=1 简化
if idx < len(output) && output[idx] > maxProb { if idx < len(output) && output[idx] > maxProb {
maxProb = output[idx] maxProb = output[idx]
maxIdx = c maxIdx = c
@@ -428,9 +430,9 @@ func decodeMathFromOutput(output []float32) string {
// CTC: 跳过 blank (index 0) 和重复 // CTC: 跳过 blank (index 0) 和重复
if maxIdx != 0 && maxIdx != lastIdx { if maxIdx != 0 && maxIdx != lastIdx {
// 字符索引从 1 开始,对应 mathChars[0:] // 字符索引从 1 开始,对应 CHARS[0:]
if maxIdx-1 < len(mathChars) { if maxIdx-1 < len(CHARS) {
result += string(mathChars[maxIdx-1]) result += string(CHARS[maxIdx-1])
} }
} }
lastIdx = maxIdx lastIdx = maxIdx