fix: correct Math decode - output is [num_positions, batch, num_chars]
This commit is contained in:
+17
-15
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user