fix: correct Math CRNN output transpose and CTC decoding
This commit is contained in:
@@ -303,7 +303,13 @@ func (h *Handler) Math(imageBase64 string) (string, error) {
|
||||
return "", err
|
||||
}
|
||||
|
||||
expr := decodeMath(output)
|
||||
// 获取输出形状
|
||||
outputShape, err := sess.GetOutputShape(0)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("获取输出形状失败: %v", err)
|
||||
}
|
||||
|
||||
expr := decodeMath(output, outputShape)
|
||||
if expr == "" {
|
||||
return "", fmt.Errorf("无法识别表达式")
|
||||
}
|
||||
@@ -334,10 +340,28 @@ func preprocessMath(img image.Image) ([]float32, error) {
|
||||
return pixels, nil
|
||||
}
|
||||
|
||||
func decodeMath(output []float32) string {
|
||||
numChars := len(mathChars) + 1
|
||||
timesteps := len(output) / numChars
|
||||
func decodeMath(output []float32, outputShape []int64) string {
|
||||
// 输出格式: [num_chars, batch, timesteps] = [51, 1, 19]
|
||||
// 需要转置为 [batch, timesteps, num_chars] = [1, 19, 51]
|
||||
numChars := int(outputShape[0]) // 51
|
||||
batch := int(outputShape[1]) // 1
|
||||
timesteps := int(outputShape[2]) // 19
|
||||
|
||||
// 转置: output[c*batch*t + b*t + t] -> transposed[b*timesteps*numChars + t*numChars + c]
|
||||
transposed := make([]float32, batch*timesteps*numChars)
|
||||
for b := 0; b < batch; b++ {
|
||||
for t := 0; t < timesteps; t++ {
|
||||
for c := 0; c < numChars; c++ {
|
||||
srcIdx := c*batch*timesteps + b*timesteps + t
|
||||
dstIdx := b*timesteps*numChars + t*numChars + c
|
||||
if srcIdx < len(output) && dstIdx < len(transposed) {
|
||||
transposed[dstIdx] = output[srcIdx]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// CTC 解码: 对每个 timestep 取 argmax
|
||||
result := ""
|
||||
lastIdx := 0
|
||||
|
||||
@@ -347,12 +371,13 @@ func decodeMath(output []float32) string {
|
||||
|
||||
for c := 0; c < numChars; c++ {
|
||||
idx := t*numChars + c
|
||||
if idx < len(output) && output[idx] > maxProb {
|
||||
maxProb = output[idx]
|
||||
if idx < len(transposed) && transposed[idx] > maxProb {
|
||||
maxProb = transposed[idx]
|
||||
maxIdx = c
|
||||
}
|
||||
}
|
||||
|
||||
// CTC: 跳过 blank (index 0) 和重复
|
||||
if maxIdx != 0 && maxIdx != lastIdx {
|
||||
if maxIdx-1 < len(mathChars) {
|
||||
result += string(mathChars[maxIdx-1])
|
||||
|
||||
Reference in New Issue
Block a user