fix: correct Math model output transpose from [C,B,T] to [B,T,C]
This commit is contained in:
+6
-133
@@ -392,54 +392,22 @@ func preprocessMath(img image.Image) ([]float32, error) {
|
||||
}
|
||||
|
||||
func decodeMathFromOutput(output []float32) string {
|
||||
// Math 模型输出固定: numChars=51, batch=1, timesteps=19, total=969
|
||||
// 可能的格式: [51, 1, 19] 或 [19, 1, 51]
|
||||
// Math 模型输出: [C, B, T] = [51, 1, 19]
|
||||
// 需要转置为 [B, T, C] = [1, 19, 51]
|
||||
numChars := 51
|
||||
batch := 1
|
||||
timesteps := 19
|
||||
totalElems := numChars * batch * timesteps // 969
|
||||
|
||||
if len(output) != totalElems {
|
||||
// 尝试推断 timesteps
|
||||
if len(output) != numChars*batch*timesteps {
|
||||
timesteps = len(output) / numChars
|
||||
if timesteps <= 0 {
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// 尝试两种格式
|
||||
// 格式1: [T, B, C] = [19, 1, 51]
|
||||
result1 := decodeMathFormat(output, timesteps, batch, numChars, true)
|
||||
// 格式2: [C, B, T] = [51, 1, 19]
|
||||
result2 := decodeMathFormat(output, timesteps, batch, numChars, false)
|
||||
|
||||
// 返回更长的结果(更可能是正确的)
|
||||
if len(result1) >= len(result2) {
|
||||
return result1
|
||||
}
|
||||
return result2
|
||||
}
|
||||
|
||||
func decodeMathFormat(output []float32, timesteps, batch, numChars int, isTBC bool) string {
|
||||
var data []float32
|
||||
|
||||
if isTBC {
|
||||
// 格式: [T, B, C] -> [B, T, C]
|
||||
data = make([]float32, batch*timesteps*numChars)
|
||||
for t := 0; t < timesteps; t++ {
|
||||
for b := 0; b < batch; b++ {
|
||||
for c := 0; c < numChars; c++ {
|
||||
srcIdx := t*batch*numChars + b*numChars + c
|
||||
dstIdx := b*timesteps*numChars + t*numChars + c
|
||||
if srcIdx < len(output) && dstIdx < len(data) {
|
||||
data[dstIdx] = output[srcIdx]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// 格式: [C, B, T] -> [B, T, C]
|
||||
data = make([]float32, batch*timesteps*numChars)
|
||||
// 转置: [C, B, T] -> [B, T, C]
|
||||
// output[c*batch*timesteps + b*timesteps + t] -> data[b*timesteps*numChars + t*numChars + c]
|
||||
data := make([]float32, batch*timesteps*numChars)
|
||||
for c := 0; c < numChars; c++ {
|
||||
for b := 0; b < batch; b++ {
|
||||
for t := 0; t < timesteps; t++ {
|
||||
@@ -451,101 +419,6 @@ func decodeMathFormat(output []float32, timesteps, batch, numChars int, isTBC bo
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// CTC 解码
|
||||
result := ""
|
||||
lastIdx := -1
|
||||
|
||||
for t := 0; t < timesteps; t++ {
|
||||
maxIdx := 0
|
||||
maxProb := float32(-math.MaxFloat32)
|
||||
|
||||
for c := 0; c < numChars; c++ {
|
||||
idx := t*numChars + c
|
||||
if idx < len(data) && data[idx] > maxProb {
|
||||
maxProb = data[idx]
|
||||
maxIdx = c
|
||||
}
|
||||
}
|
||||
|
||||
if maxIdx != 0 && maxIdx != lastIdx {
|
||||
if maxIdx-1 < len(mathChars) {
|
||||
result += string(mathChars[maxIdx-1])
|
||||
}
|
||||
}
|
||||
lastIdx = maxIdx
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
func decodeMath(output []float32, outputShape []int64) string {
|
||||
// 原版逻辑:
|
||||
// if preds.shape[1] == 1: # Batch size is 1 at dim 1 implies [T, B, C]
|
||||
// preds = np.transpose(preds, (1, 0, 2)) # [B, T, C]
|
||||
// preds_idx = np.argmax(preds, axis=2) # [B, T]
|
||||
|
||||
if len(outputShape) < 3 || len(output) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
// 输出可能是 [T, B, C] 或 [C, B, T]
|
||||
// 判断依据:如果 shape[1] == 1,说明 batch 在中间,格式是 [T, B, C]
|
||||
// 否则格式是 [C, B, T] 或 [num_chars, batch, timesteps]
|
||||
|
||||
dim0 := int(outputShape[0])
|
||||
dim1 := int(outputShape[1])
|
||||
dim2 := int(outputShape[2])
|
||||
|
||||
// 处理动态维度
|
||||
if dim0 <= 0 { dim0 = 51 }
|
||||
if dim1 <= 0 { dim1 = 1 }
|
||||
if dim2 <= 0 { dim2 = len(output) / (dim0 * dim1) }
|
||||
|
||||
var timesteps, batch, numChars int
|
||||
var data []float32
|
||||
|
||||
// 判断输出格式
|
||||
if dim1 == 1 {
|
||||
// 格式: [T, B, C] = [19, 1, 51]
|
||||
// 需要转置为 [B, T, C] = [1, 19, 51]
|
||||
timesteps = dim0
|
||||
batch = dim1
|
||||
numChars = dim2
|
||||
|
||||
data = make([]float32, batch*timesteps*numChars)
|
||||
for t := 0; t < timesteps; t++ {
|
||||
for b := 0; b < batch; b++ {
|
||||
for c := 0; c < numChars; c++ {
|
||||
srcIdx := t*batch*numChars + b*numChars + c
|
||||
dstIdx := b*timesteps*numChars + t*numChars + c
|
||||
if srcIdx < len(output) && dstIdx < len(data) {
|
||||
data[dstIdx] = output[srcIdx]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// 格式: [C, B, T] = [51, 1, 19]
|
||||
// 需要转置为 [B, T, C] = [1, 19, 51]
|
||||
numChars = dim0
|
||||
batch = dim1
|
||||
timesteps = dim2
|
||||
|
||||
data = make([]float32, batch*timesteps*numChars)
|
||||
for c := 0; c < numChars; c++ {
|
||||
for b := 0; b < batch; b++ {
|
||||
for t := 0; t < timesteps; t++ {
|
||||
srcIdx := c*batch*timesteps + b*timesteps + t
|
||||
dstIdx := b*timesteps*numChars + t*numChars + c
|
||||
if srcIdx < len(output) && dstIdx < len(data) {
|
||||
data[dstIdx] = output[srcIdx]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// CTC 解码: 对每个 timestep 取 argmax
|
||||
result := ""
|
||||
|
||||
Reference in New Issue
Block a user