From ae4bf8dd3699005dd2fd34026fd6037e9cf6fda9 Mon Sep 17 00:00:00 2001
From: Admin <admin@viaeon.com>
Date: Fri, 17 Jul 2026 12:19:55 +0000
Subject: [PATCH] debug: print argmax results for Math CTC decode

---
 internal/captcha/handler.go | 69 +++++++++++++++++++++----------------
 1 file changed, 40 insertions(+), 29 deletions(-)

diff --git a/internal/captcha/handler.go b/internal/captcha/handler.go
index 7add83e..134a7f8 100644
--- a/internal/captcha/handler.go
+++ b/internal/captcha/handler.go
@@ -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
 }