From 61f3223885e08b1f0eb343404e0ca796677ae887 Mon Sep 17 00:00:00 2001
From: Admin <admin@viaeon.com>
Date: Fri, 17 Jul 2026 10:03:40 +0000
Subject: [PATCH] fix: correct Math CRNN output transpose and CTC decoding

---
 internal/captcha/handler.go | 39 ++++++++++++++++++++++++++++++-------
 1 file changed, 32 insertions(+), 7 deletions(-)

diff --git a/internal/captcha/handler.go b/internal/captcha/handler.go
index a9ea1a7..390c72e 100644
--- a/internal/captcha/handler.go
+++ b/internal/captcha/handler.go
@@ -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
 
@@ -346,13 +370,14 @@ func decodeMath(output []float32) string {
 		maxProb := float32(-math.MaxFloat32)
 
 		for c := 0; c < numChars; c++ {
-			idx := t * numChars + c
-			if idx < len(output) && output[idx] > maxProb {
-				maxProb = output[idx]
+			idx := t*numChars + c
+			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])