feat: 重构 CTC 解码逻辑

- 重构 ctc_decode 为关联函数并优化内存分配。
- 增加 单元测试和集成测试

src/lib.rs

@@ -1,18 +1,16 @@
-mod model;
-mod utils;
-
 mod charset;
 mod image_io;
 mod image_processor;
+mod model;
+mod utils;
 
+use crate::image_io::png_rgba_white_preprocess;
+use crate::image_processor::{convert_to_grayscale, resize_image};
 use anyhow::{Context, Result};
 use image::{DynamicImage, imageops::FilterType};
 use tract_onnx::prelude::*;
 // 关键点：直接使用 tract 重导出的 ndarray
-use crate::image_io::png_rgba_white_preprocess;
-use crate::image_processor::{convert_to_grayscale, resize_image};
-use tract_onnx::prelude::tract_itertools::Itertools;
-
+use tract_onnx::prelude::tract_ndarray::s;
 pub struct DdddOcr {
     session: RunnableModel<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
 }
@@ -37,14 +35,14 @@ impl DdddOcr {
         // 3. 解析结果
         // let output = result[0].to_array_view::<i64>()?;
         let output = self.inference(tensor)?;
-        let output2 = self.extract_indices(&output)?;
-        Ok(self.decode_ctc(&output2))
+        let output2 = self.process_text_output(&output)?;
+        Ok(Self::ctc_decode_indices(&output2))
     }
     /// 对应 Python 的 _preprocess_image
     /// 负责：透明背景修复 -> 灰度化 -> 按比例 Resize -> 归一化 -> 4维张量转换
     fn preprocess_image(&self, img: &DynamicImage, png_fix: bool) -> Result<Tensor> {
         // A. 修复 PNG 透明背景 (内部逻辑你之前已实现)
-        let processed_img = if png_fix && img.color().has_alpha() {
+        let _ = if png_fix && img.color().has_alpha() {
             png_rgba_white_preprocess(img)
         } else {
             img.clone()
@@ -54,6 +52,7 @@ impl DdddOcr {
         let w = (img.width() as f32 * (h as f32 / img.height() as f32)) as u32;
         let gray_img = convert_to_grayscale(img);
         let resized = resize_image(&gray_img, w, h);
+        // resized.save("debug_preprocessed.png").unwrap();
         // 1. 预处理：转灰度 -> Resize -> 归一化
         // let resized = img.resize_exact(w, h, FilterType::Lanczos3).to_luma8();
 
@@ -76,12 +75,15 @@ impl DdddOcr {
             .session
             .run(tvec!(tensor.into()))
             .context("执行模型推理失败")?;
-
+        println!("模型输出原始数据: {:?}", result);
         Ok(result.remove(0).into_tensor())
     }
     /// 核心解析逻辑：将模型输出的各种维度/类型的 Tensor 转为字符索引序列
-    fn extract_indices(&self, raw_tensor: &Tensor) -> Result<Vec<i64>> {
+    fn process_text_output(&self, raw_tensor: &Tensor) -> Result<Vec<i64>> {
         let shape = raw_tensor.shape();
+        println!("模型输出shape数据: {:?}", shape);
+        let datum_type = raw_tensor.datum_type();
+        println!("模型输出datum_type数据: {:?}", datum_type);
 
         match raw_tensor.datum_type() {
             // 情况 1: huashi666 式模型，直接输出 i64 索引 (通常是模型内部做好了 Argmax)
@@ -93,32 +95,43 @@ impl DdddOcr {
             // 情况 2: sml2h3 原版模型，输出 F32 概率矩阵
             DatumType::F32 => {
                 let view = raw_tensor.to_array_view::<f32>()?;
-
-                // 处理典型的 CTC 输出形状 [TimeSteps, Batch:1, Classes]
-                if shape.len() == 3 {
-                    let steps = shape[0];
-                    let classes = shape[2];
-
-                    // 将一维视图重新整理为二维 [steps, classes]
-                    let array_2d = view.to_shape((steps, classes))?;
-
-                    // 对每一行执行 Argmax (寻找概率最大的字符索引)
-                    let indices = array_2d
-                        .outer_iter()
-                        .map(|row| {
-                            row.iter()
-                                .enumerate()
-                                .max_by(|(_, a), (_, b)| {
-                                    a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)
-                                })
-                                .map(|(idx, _)| idx as i64)
-                                .unwrap_or(0)
-                        })
-                        .collect();
-                    Ok(indices)
-                } else {
-                    Err(anyhow::anyhow!("不支持的 F32 输出形状: {:?}", shape))
-                }
+                let (steps, classes, data_view) = match shape.len() {
+                    3 => {
+                        if shape[1] == 1 {
+                            // 形状: [Steps, 1, Classes] -> 你的原有逻辑
+                            (shape[0], shape[2], view.into_dyn())
+                        } else if shape[0] == 1 {
+                            // 形状: [1, Steps, Classes] -> 另一种常见导出格式
+                            (shape[1], shape[2], view.into_dyn())
+                        } else {
+                            // 默认取第一个 batch: [Batch, Steps, Classes]
+                            // 使用 slice 对应 Python 的 output[0, :, :]
+                            let sliced = view.slice(s![0, .., ..]);
+                            (shape[1], shape[2], sliced.into_dyn())
+                        }
+                    }
+                    2 => {
+                        // 形状: [Steps, Classes] -> 已经剥离了 Batch 维度
+                        (shape[0], shape[1], view.into_dyn())
+                    }
+                    _ => return Err(anyhow::anyhow!("不支持的输出维度: {:?}", shape)),
+                };
+                let array_2d = data_view.to_shape((steps, classes))?;
+                //
+                // 对每一行执行 Argmax (寻找概率最大的字符索引)
+                let indices = array_2d
+                    .outer_iter()
+                    .map(|row| {
+                        row.iter()
+                            .enumerate()
+                            .max_by(|(_, a), (_, b)| {
+                                a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)
+                            })
+                            .map(|(idx, _)| idx as i64)
+                            .unwrap_or(0)
+                    })
+                    .collect();
+                Ok(indices)
             }
             _ => Err(anyhow::anyhow!(
                 "不支持的模型输出数据类型: {:?}",
@@ -126,20 +139,44 @@ impl DdddOcr {
             )),
         }
     }
-    fn decode_ctc(&self, indices: &[i64]) -> String {
-        use crate::charset::CHARSET_BETA;
-        let mut res = String::new();
-        let mut last_idx: i64 = -1;
+    fn ctc_decode_indices(predicted_indices: &[i64]) -> String {
+        println!("indices模型输出原始数据: {:?}", predicted_indices);
 
-        for &idx in indices {
-            // ddddocr 的 blank 通常是 0
-            if idx != 0 && idx != last_idx {
-                if let Some(&char_str) = CHARSET_BETA.get(idx as usize) {
-                    res.push_str(char_str);
+        use crate::charset::CHARSET_BETA;
+        // 对应 _ctc_decode_indices 的逻辑：去重、去 blank (0)
+        let mut res = String::new();
+        let mut prev_idx: i64 = -1;
+
+        for &idx in predicted_indices {
+            // 1. 跳过连续重复的索引
+            // 2. 跳过 blank 字符 (假设索引 0 是 blank)
+            if idx != prev_idx && idx != 0 {
+                if let Ok(u_idx) = usize::try_from(idx) {
+                    if let Some(&char_str) = CHARSET_BETA.get(u_idx) {
+                        res.push_str(char_str);
+                    }
                 }
             }
-            last_idx = idx;
+            prev_idx = idx;
         }
+        println!("最终识别出的验证码是: {}", res);
         res
     }
 }
+
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn test_ctc_decode_indices() {
+        // 模拟一个 DdddOcr 实例（如果 decode 不依赖 session，可以设为相关函数）
+        // 这里假设你的 decode_ctc 是公开或内部可访问的
+        let input = vec![1, 1, 0, 1, 2, 2, 0, 2];
+        // 逻辑：[1, 1] -> 1, [0] -> 跳过, [1] -> 1, [2, 2] -> 2, [0] -> 跳过, [2] -> 2
+        // 预期结果索引应该是 [1, 1, 2, 2] 对应的字符
+        // 具体的断言取决于你的 CHARSET_BETA
+        // let result = dddd.ctc_decode_indices(&input);
+        // assert_eq!(result, "AABB");
+    }
+}