debug: print argmax results for Math CTC decode
This commit is contained in:
+40
-29
@@ -406,52 +406,63 @@ func preprocessMath(img image.Image) ([]float32, error) {
|
||||
}
|
||||
|
||||
func decodeMathFromOutput(output []float32) string {
|
||||
// Math 模型输出: [num_positions, batch, num_chars] = [51, 1, 19]
|
||||
// 有 51 个位置,每个位置有 19 个字符概率(blank + 18 chars)
|
||||
// 对每个位置取 argmax,得到 51 个字符索引
|
||||
numPositions := 51
|
||||
batch := 1
|
||||
numChars := 19
|
||||
// Math 模型输出: [T, B, C] = [51, 1, 19]
|
||||
// 51 个时间步,batch=1,19 个字符类别(blank + 18 chars)
|
||||
// 原版 Python: transpose(1, 0, 2) -> [B, T, C],然后 argmax(axis=2)
|
||||
// 但实际上对于 [T, B, C] 行优先存储,argmax 在 C 维度不需要改变读取顺序
|
||||
const T = 51
|
||||
const B = 1
|
||||
const C = 19
|
||||
|
||||
if len(output) != numPositions*batch*numChars {
|
||||
// 尝试推断
|
||||
if len(output) != T*B*C {
|
||||
total := len(output)
|
||||
if total%numChars != 0 {
|
||||
if total%C != 0 {
|
||||
return ""
|
||||
}
|
||||
numPositions = total / numChars
|
||||
}
|
||||
|
||||
CHARS := "0123456789+-*/÷×=?"
|
||||
|
||||
// 对每个位置取 argmax
|
||||
// output[pos*batch*numChars + b*numChars + c] 存储的是位置 pos, batch b 的字符 c 的概率
|
||||
result := ""
|
||||
lastIdx := -1
|
||||
|
||||
for pos := 0; pos < numPositions; pos++ {
|
||||
// 先收集所有时间步的 argmax 结果,用于调试
|
||||
allChars := make([]int, T)
|
||||
for t := 0; t < T; t++ {
|
||||
maxIdx := 0
|
||||
maxProb := float32(-math.MaxFloat32)
|
||||
|
||||
// 对这个位置的 19 个字符取最大值
|
||||
for c := 0; c < numChars; c++ {
|
||||
idx := pos*numChars + c // batch=1 简化
|
||||
for c := 0; c < C; c++ {
|
||||
// [T, B, C] 行优先存储: output[t*B*C + b*C + c]
|
||||
idx := t*B*C + c
|
||||
if idx < len(output) && output[idx] > maxProb {
|
||||
maxProb = output[idx]
|
||||
maxIdx = c
|
||||
}
|
||||
}
|
||||
|
||||
// CTC: 跳过 blank (index 0) 和重复
|
||||
if maxIdx != 0 && maxIdx != lastIdx {
|
||||
// 字符索引从 1 开始,对应 CHARS[0:]
|
||||
if maxIdx-1 < len(CHARS) {
|
||||
result += string(CHARS[maxIdx-1])
|
||||
}
|
||||
}
|
||||
lastIdx = maxIdx
|
||||
allChars[t] = maxIdx
|
||||
}
|
||||
|
||||
// 打印所有位置的识别结果
|
||||
debugStr := "DEBUG argmax: "
|
||||
for _, idx := range allChars {
|
||||
if idx > 0 && idx <= len(CHARS) {
|
||||
debugStr += string(CHARS[idx-1])
|
||||
} else {
|
||||
debugStr += "_"
|
||||
}
|
||||
}
|
||||
fmt.Println(debugStr)
|
||||
|
||||
// CTC 解码: 跳过 blank (index 0) 和连续重复
|
||||
result := ""
|
||||
lastIdx := -1
|
||||
for _, idx := range allChars {
|
||||
if idx != 0 && idx != lastIdx {
|
||||
if idx-1 < len(CHARS) {
|
||||
result += string(CHARS[idx-1])
|
||||
}
|
||||
}
|
||||
lastIdx = idx
|
||||
}
|
||||
|
||||
fmt.Printf("DEBUG CTC decoded: '%s'\n", result)
|
||||
return result
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user