diff --git a/src/lib.rs b/src/lib.rs index 95c9702..da7124b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,20 +1,20 @@ mod charset; +mod model_metadata; pub mod models; pub mod utils; -mod model_metadata; -use anyhow::Result; +use anyhow::{Result, anyhow}; use image::DynamicImage; use std::fmt::{Display, Formatter}; // 关键点：直接使用 tract 重导出的 ndarray -use crate::charset::{ CharRestrict}; +use crate::charset::CharRestrict; +use crate::model_metadata::ModelMetadata; +use crate::utils::color_filter::{ColorPreset, HsvRange}; use models::det::Det; use models::loader::ModelSession; use models::ocr::Ocr; -use crate::model_metadata::ModelMetadata; -use crate::utils::color_filter::{ColorPreset, HsvRange}; pub enum ModelSpec { /// 默认 OCR (使用内置路径) @@ -64,7 +64,10 @@ impl DdddOcrBuilder { /// 设置自定义 OCR 路径 pub fn custom_ocr(mut self, path: String, model_metadata: ModelMetadata) -> Self { // 直接重写枚举，替换掉之前的 Ocr 或 Det - self.mode = ModelSpec::CustomOcrModel { path, model_metadata }; + self.mode = ModelSpec::CustomOcrModel { + path, + model_metadata, + }; self } @@ -76,7 +79,10 @@ impl DdddOcrBuilder { ModelMetadata::from_builtin_beta(), )?), ModelSpec::DetModel => Runtime::Det(Det::new(ModelSpec::DEFAULT_DET_PATH.into())?), - ModelSpec::CustomOcrModel { path, model_metadata } => Runtime::Ocr(Ocr::new(path, model_metadata)?), + ModelSpec::CustomOcrModel { + path, + model_metadata, + } => Runtime::Ocr(Ocr::new(path, model_metadata)?), }; Ok(DdddOcr { runtime }) @@ -97,23 +103,36 @@ impl DdddOcr { pub fn classification(&self, img: &DynamicImage) -> Result { match &self.runtime { // Runtime::Ocr(s) => s.predict(img).run(), - Runtime::Ocr(s) => s.predictor().predict(img), + // Runtime::Ocr(s) => s.predictor().probability(false).predict(img), + // Runtime::Ocr(s) => { + // let predictor = s.predictor(); + // let restricted = predictor.charset_restrict(&CharRestrict::Lowercase); + // let a = restricted.valid_tokens(); + // println!("{:?}", a); + // Ok("".to_string()) + // } + Runtime::Ocr(s) => { + let res = s.predictor().probability(true).predict(img)?; + println!("{}", res); + Ok("".to_string()) + } // Runtime::Ocr(s) => s.predictor().charset_restrict(&CharRestrict::Digit).predict(img), // Runtime::Ocr(s) => s.predictor().color_filter(&ColorPreset::Custom(vec![ // // 错误：下界 (82, 221, 14) 没问题 // // 但上界的 H 通道写成了 240，超过了 180 的法定上限！ // HsvRange::new((82, 221, 14), (240, 203, 82)), // ])).predict(img), - Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型，无法执行 OCR")), - } + + Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型，无法执行 OCR")), } - pub fn detection(&self, img: &[u8]) -> Result>> { - match &self.runtime { - Runtime::Det(s) => s.predict(img), - Runtime::Ocr(_) => Err(anyhow::anyhow!("当前模型是 OCR 模型，无法执行检测")), - } +} +pub fn detection(&self, img: &[u8]) -> Result>> { + match &self.runtime { + Runtime::Det(s) => s.predict(img), + Runtime::Ocr(_) => Err(anyhow::anyhow!("当前模型是 OCR 模型，无法执行检测")), } } +} // struct Classification {} // #[derive(Debug)] diff --git a/src/model_metadata.rs b/src/model_metadata.rs index 69a2df9..1fd5efa 100644 --- a/src/model_metadata.rs +++ b/src/model_metadata.rs @@ -10,6 +10,26 @@ use std::path::Path; // 1. 辅助定义的枚举与结构体 // ===================================================================== +#[derive(Debug, Clone, Copy, Deserialize)] +#[serde(rename_all = "snake_case")] // 支持 json 中写 "zero_to_one" 或 "minus_one_to_one" +pub enum Normalization { + /// 映射到 [0.0, 1.0] -> pixel / 255.0 + ZeroToOne, + /// 映射到 [-1.0, 1.0] -> (pixel / 255.0 - 0.5) / 0.5 + MinusOneToOne, +} + +impl Normalization { + /// 统一归一化计算逻辑 + #[inline(always)] + pub fn normalize(&self, pixel: f32) -> f32 { + match self { + Normalization::ZeroToOne => pixel / 255.0, + Normalization::MinusOneToOne => (pixel / 255.0 - 0.5) / 0.5, + } + } +} + /// 图像缩放策略枚举 #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Resize { @@ -29,30 +49,51 @@ struct ModelMetadataDto { #[serde(alias = "image")] resize: Vec, channel: u8, + /// 新增：允许在配置文件中指定归一化策略。 + /// 使用 serde(default) 可以在不配置时提供一个默认值（比如默认 ZeroToOne） + #[serde(default = "default_normalization")] + normalization: Normalization, +} +fn default_normalization() -> Normalization { + Normalization::ZeroToOne } #[derive(Debug, Clone)] pub struct ModelMetadata { /// 字符集管理器 - pub charset: Charset, + pub charset: Charset, /// 是否为单字识别模型 pub word: bool, /// 预处理的缩放策略 pub resize: Resize, /// 图像通道数 (1 或 3) pub channel: u8, + /// 新增：传递给核心业务使用的归一化配置 + pub normalization: Normalization, } impl ModelMetadata { // --- 优雅的工厂模式构造器 --- /// 从预设的旧版字符集创建 pub fn from_builtin_old() -> Self { - Self::from_static_slice(CHARSET_OLD, false, Resize::DynamicWidth(64), 1) + Self::from_static_slice( + CHARSET_OLD, + false, + Resize::DynamicWidth(64), + 1, + Normalization::ZeroToOne, + ) } /// 从预设的 Beta 版字符集创建 pub fn from_builtin_beta() -> Self { - Self::from_static_slice(CHARSET_BETA, false, Resize::DynamicWidth(64), 1) + Self::from_static_slice( + CHARSET_BETA, + false, + Resize::DynamicWidth(64), + 1, + Normalization::MinusOneToOne, + ) } /// 通用的静态切片转换构造器 @@ -61,6 +102,7 @@ impl ModelMetadata { word: bool, resize: Resize, channel: u8, + normalization: Normalization, ) -> Self { let tokens: Vec> = slice.iter().map(|&s| Cow::Borrowed(s)).collect(); Self { @@ -68,6 +110,7 @@ impl ModelMetadata { word, resize, channel, + normalization, } } @@ -117,6 +160,7 @@ impl ModelMetadata { word: dto.word, resize, channel: dto.channel, + normalization: dto.normalization, }) } } diff --git a/src/models/ocr.rs b/src/models/ocr.rs index 50a9bef..6b5da19 100644 --- a/src/models/ocr.rs +++ b/src/models/ocr.rs @@ -8,15 +8,97 @@ use crate::utils::image_processor::{convert_to_grayscale, resize_image}; use anyhow::Context; use anyhow::{Result, anyhow}; use image::{DynamicImage, ImageBuffer, Rgb}; +use serde::Serialize; use std::borrow::Cow; use std::collections::HashSet; -use tract_onnx::prelude::tract_ndarray::{s, ArrayView2}; +use std::fmt; +use tract_onnx::prelude::tract_ndarray::{ArrayView2, Ix2, s}; use tract_onnx::prelude::{ DatumType, Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp, tract_ndarray, tvec, }; // 引入 cv_ops 模块中的 OpenCV HSV 转换算子 use crate::utils::cv_ops::rgb_to_opencv_hsv; +/// 推理最终输出的强类型外壳（完全 Owned，无任何生命周期，可直接转 JSON） +#[derive(Debug, Clone, Serialize)] +pub enum OcrOutput { + /// 纯文本分支（对应 probability = false） + Text(String), + /// 包含全量概率的分支（对应 probability = true） + Probability { + text: String, + /// 满额概率矩阵 [Steps, Classes] + probabilities: Vec>, + /// 全局平均置信度 + confidence: f64, + }, + /// 不支持的模型或未知输出 + Unsupported { message: String }, +} +impl OcrOutput { + /// 消费自身，直接提取最终文本 + pub fn into_text(self) -> String { + match self { + OcrOutput::Text(text) => text, + OcrOutput::Probability { text, .. } => text, + OcrOutput::Unsupported { message } => { + // 作为库，这里可以返回空，或者直接携带错误信息，取决于你的设计 + format!("Error: {}", message) + } + } + } +} +impl fmt::Display for OcrOutput { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + OcrOutput::Text(text) => { + // 纯文本分支，直接输出文本内容 + write!(f, "{}", text) + } + OcrOutput::Probability { text,probabilities, confidence } => { + // 概率分支，友好地展示文本以及百分比形式的置信度 + // 1. 基本信息 + write!(f, "{} (置信度: {:.2}%)", text, confidence * 100.0)?; + + // 2. 概率矩阵流式安全打印 + write!(f, " [概率矩阵预览: ")?; + + let max_steps_to_show = 10; + let take_steps = probabilities.iter().take(max_steps_to_show); + + for (i, step_probs) in take_steps.enumerate() { + if i > 0 { + write!(f, ", ")?; + } + + // 为了防止单行内部数据过长，单行也做一下截断保护（比如每行最多显示前 3 个概率） + let max_classes_to_show = 3; + write!(f, "[")?; + for (j, prob) in step_probs.iter().take(max_classes_to_show).enumerate() { + if j > 0 { + write!(f, ", ")?; + } + write!(f, "{:.4}", prob)?; + } + if step_probs.len() > max_classes_to_show { + write!(f, ", ..")?; + } + write!(f, "]")?; + } + + // 如果总 Step 数量超过 10，末尾追加 .. 表示截断 + if probabilities.len() > max_steps_to_show { + write!(f, ", ..")?; + } + write!(f, "]") + } + OcrOutput::Unsupported { message } => { + // 错误分支，直观输出异常原因 + write!(f, "未识别成功: {}", message) + } + } + } +} pub struct Ocr { pub session: RunnableModel, Graph>>, @@ -50,89 +132,6 @@ impl Ocr { Ok(result.swap_remove(0).into_tensor()) } - /// 核心解析逻辑：将模型输出的各种维度/类型的 Tensor 转为字符索引序列 - fn extract_indices_from_tensor(&self, raw_tensor: &Tensor) -> anyhow::Result> { - let shape = raw_tensor.shape(); - println!("模型输出shape数据: {:?}", shape); - let datum_type = raw_tensor.datum_type(); - println!("模型输出datum_type数据: {:?}", datum_type); - - match datum_type { - // 情况 1: huashi666 式模型，直接输出 i64 索引 (通常是模型内部做好了 Argmax) - DatumType::I64 => { - let view = raw_tensor.to_array_view::()?; - - Ok(view.iter().cloned().collect()) - } - - // 情况 2: sml2h3 原版模型，输出 F32 概率矩阵 - DatumType::F32 => { - let view = raw_tensor.to_array_view::()?; - 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()) - } - // 形状: [Classes] -> 单字符输出（对应 Python 的 ndim == 0 保护逻辑） - // 我们把它虚构成一个 [1, Classes] 的 2D 矩阵来复用后面的 argmax 逻辑 - 1 => (1, shape[0], 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!( - "不支持的模型输出数据类型: {:?}", - datum_type - )), - } - } - /// 管道 2：纯文本解码流水线 (高性能版：免去 Softmax 计算) - fn process_text_pipeline(&self, matrix_view: ArrayView2) -> anyhow::Result { - // 直接在原始分值（Logits）上进行 Argmax，数学结果与 Softmax 后完全一致 - let indices: Vec = matrix_view - .outer_iter() - .map(|row| { - row.iter() - .enumerate() - .max_by(|(_, a), (_, b)| a.total_cmp(b)) - .map(|(idx, _)| idx as i64) - .unwrap_or(0) - }) - .collect(); - - // 丢给现有的 CTC 解码器去重并映射成字符串 - Ok(self.ctc_decode_to_string(&indices)) - } - pub fn predictor(&'_ self) -> OcrPredictor<'_> { OcrPredictor::new(self) } @@ -140,7 +139,6 @@ impl Ocr { pub struct OcrPredictor<'a> { ocr: &'a Ocr, - // image: &'a DynamicImage, /// 是否修复PNG格式问题 png_fix: bool, /// 是否返回概率信息 @@ -158,7 +156,6 @@ impl<'a> OcrPredictor<'a> { pub fn new(ocr: &'a Ocr) -> Self { Self { ocr, - // image, png_fix: false, // 默认值 probability: false, color_filter: Ok(None), @@ -169,11 +166,12 @@ impl<'a> OcrPredictor<'a> { self.png_fix = value; self } + pub fn probability(mut self, value: bool) -> Self { + self.probability = value; + self + } - // 反复调用color_filter怎么处理？ pub fn color_filter(mut self, filter: &dyn ColorFilter) -> Self { - // self.color_filter = Some(value); - // 一句话把活全包了！错误信息无缝传递，完美熔断 match filter.collect_to_vec() { Ok(new_ranges) => self.color_filter = Ok(new_ranges), @@ -186,13 +184,12 @@ impl<'a> OcrPredictor<'a> { pub fn charset_restrict(mut self, restrict: &dyn TokenFilter) -> Self { let charset = &self.ocr.model_metadata.charset; let tokens = &charset.tokens; - // let mut temp_indices = Vec::new(); self.charset_restrict = restrict.apply_to_charset(tokens); self } } impl<'a> OcrPredictor<'a> { - pub fn predict(self, image: &DynamicImage) -> anyhow::Result { + pub fn predict(self, image: &DynamicImage) -> anyhow::Result { println!("当前颜色过滤器状态: {:?}", self.color_filter); // ===================================================================== // 管道节点 1: 颜色过滤流水线 @@ -219,18 +216,28 @@ impl<'a> OcrPredictor<'a> { let tensor = self.preprocess_image(&img_cow)?; let raw_tensor = self.ocr.inference(tensor)?; - let raw_indices = self.ocr.extract_indices_from_tensor(&raw_tensor)?; - // 步骤 2: 将索引切片 `&[i64]` 传给解码器进行 CTC 去重和字符映射 - let final_text = self.ctc_decode_to_string(&raw_indices); - println!("最终识别出的验证码是: {}", final_text); - Ok(final_text) + // 3. 后处理分流：直接返回 OcrOutput + let ocr_output = match raw_tensor.datum_type() { + DatumType::I64 => self.extract_from_i64_tensor(raw_tensor)?, + DatumType::F32 => self.process_f32_pipeline(raw_tensor)?, + _ => OcrOutput::Unsupported { + message: format!("不支持的模型输出数据类型: {:?}", raw_tensor.datum_type()), + }, + }; + + // let raw_indices = self.ocr.extract_indices_from_tensor(&raw_tensor)?; + // // 步骤 2: 将索引切片 `&[i64]` 传给解码器进行 CTC 去重和字符映射 + // let final_text = self.ctc_decode_to_string(&raw_indices); + + Ok(ocr_output) } /// 对应 Python 的 _preprocess_image /// 负责：透明背景修复 -> 灰度化 -> 按比例 Resize -> 归一化 -> 4维张量转换 fn preprocess_image(&self, img: &DynamicImage) -> anyhow::Result { // 1. 获取模型元数据配置 let meta = &self.ocr.model_metadata; + let norm = &meta.normalization; // 获取归一化器 // A. 修复 PNG 透明背景 (内部逻辑你之前已实现) let current_img = if self.png_fix && img.color().has_alpha() { @@ -246,7 +253,8 @@ impl<'a> OcrPredictor<'a> { Resize::Fixed(w, h) => (w, h), Resize::DynamicWidth(h) => { // 高度固定，宽度根据原始比例动态计算：W_target = W_orig * (H_target / H_orig) - let w = (current_img.width() as f32 * (h as f32 / current_img.height() as f32)) as u32; + let w = + (current_img.width() as f32 * (h as f32 / current_img.height() as f32)) as u32; (w, h) } Resize::Square(size) => { @@ -267,7 +275,9 @@ impl<'a> OcrPredictor<'a> { (1, 1, target_h as usize, target_w as usize), |(_, _, y, x)| { let pixel = gray_img.get_pixel(x as u32, y as u32)[0] as f32; - pixel / 255.0 // 严格对齐 Python 归一化 [0.0, 1.0] + // pixel / 255.0 // 严格对齐 Python 归一化 [0.0, 1.0] + // (pixel / 255.0 - 0.5) / 0.5 + norm.normalize(pixel) }, ); Tensor::from(array) @@ -281,7 +291,9 @@ impl<'a> OcrPredictor<'a> { (1, 3, target_h as usize, target_w as usize), |(_, c, y, x)| { let pixel = rgb_img.get_pixel(x as u32, y as u32)[c] as f32; - pixel / 255.0 // 严格对齐 Python 归一化 [0.0, 1.0] + // pixel / 255.0 // 严格对齐 Python 归一化 [0.0, 1.0] + // (pixel / 255.0 - 0.5) / 0.5 + norm.normalize(pixel) }, ); Tensor::from(array) @@ -292,8 +304,6 @@ impl<'a> OcrPredictor<'a> { Ok(tensor) - - // let h = 64u32; // let w = (current_img.width() as f32 * (h as f32 / current_img.height() as f32)) as u32; // let gray_img = convert_to_grayscale(¤t_img); @@ -346,7 +356,148 @@ impl<'a> OcrPredictor<'a> { None => self.ocr.model_metadata.charset.tokens.len(), } } + /// 变体 B 核心处理器：单次遍历 2D 视图，融合计算 Softmax、Argmax、置信度并输出概率大包 + fn compute_f32_full_probability( + &self, + matrix_view: ArrayView2, + ) -> (Vec>, f32, Vec) { + let steps = matrix_view.nrows(); + let classes = matrix_view.ncols(); + // 1. 预分配满额概率矩阵内存 + let mut prob_matrix = tract_ndarray::Array2::::zeros((steps, classes)); + let mut predicted_indices = Vec::with_capacity(steps); + let mut confidence_sum = 0.0f32; + + // 2. 融合单次遍历 + for (step_idx, row) in matrix_view.outer_iter().enumerate() { + // 寻找当前 Step 的最大值和最大值索引 (Argmax) + let (row_max_idx, max_logit) = row + .iter() + .enumerate() + .max_by(|(_, a), (_, b)| a.total_cmp(b)) + .map(|(idx, &val)| (idx, val)) + .unwrap_or((0, 0.0)); + + predicted_indices.push(row_max_idx as i64); + + // 计算单行 exp 溢出防范和 + let mut exp_sum = 0.0f32; + for &val in row.iter() { + exp_sum += (val - max_logit).exp(); + } + + // 归一化 Softmax 顺序写入 + for (class_idx, &val) in row.iter().enumerate() { + prob_matrix[[step_idx, class_idx]] = (val - max_logit).exp() / exp_sum; + } + + // 当前 Step 最大概率在线累加 + confidence_sum += 1.0f32 / exp_sum; + } + + // 3. 统计全局平均置信度 + let confidence = if steps > 0 { + confidence_sum / steps as f32 + } else { + 1.0 + }; + + // 4. 将矩阵转化为标准安全序列化格式 [Steps, Classes] + let probabilities_list: Vec> = + prob_matrix.outer_iter().map(|row| row.to_vec()).collect(); + + (probabilities_list, confidence, predicted_indices) + } + /// 变体 A 专属提取器：直接从 I64 Tensor 零拷贝提取 CTC 文本与初始概率包 + fn extract_from_i64_tensor(&self, raw_tensor: Tensor) -> anyhow::Result { + // 1. 拿到底层的动态维度只读视图 + let view = raw_tensor.to_array_view::()?; + + // 2. 索要底层连续的只读切片引用 + let slice = view + .as_slice() + .ok_or_else(|| anyhow::anyhow!("I64 模型输出内存不连续，无法执行零拷贝解码"))?; + + // 3. 直接喂给 CTC 解码器（无任何物理克隆开销） + let final_text = self.ctc_decode_to_string(slice); + + // 4. 组装返回 + if self.probability { + Ok(OcrOutput::Probability { + text: final_text, + probabilities: vec![], // I64 模型物理上丢失了全量 Logits 分值网，降级处理 + confidence: 1.0, // 判定即百分之百置信 + }) + } else { + Ok(OcrOutput::Text(final_text)) + } + } + /// 变体二（F32）的总体管线：负责降维，并分流文本和概率 + fn process_f32_pipeline(&self, raw_tensor: Tensor) -> anyhow::Result { + let shape = raw_tensor.shape(); + println!("模型输出shape数据: {:?}", shape); + let view = raw_tensor.to_array_view::()?; + + // 1. 极其纯粹的、无拷贝的多维 Shape 压扁清洗 + let (steps, classes, data_dyn_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()) + } + } + // 形状: [Steps, Classes] -> 已经剥离了 Batch 维度 + 2 => (shape[0], shape[1], view.into_dyn()), + // 形状: [Classes] -> 单字符输出（对应 Python 的 ndim == 0 保护逻辑） + // 我们把它虚构成一个 [1, Classes] 的 2D 矩阵来复用后面的 argmax 逻辑 + 1 => (1, shape[0], view.into_dyn()), + _ => return Err(anyhow::anyhow!("不支持的输出维度: {:?}", shape)), + }; + let matrix_cow = data_dyn_view + .to_shape(Ix2(steps, classes)) + .map_err(|e| anyhow::anyhow!("转换为2D静态矩阵失败: {:?}", e))?; + + let matrix_view: ArrayView2 = matrix_cow.view(); + + // 2. 根据业务参数明确分流 + if self.probability { + // 走向 B1：调用刚刚拆分出来的“全量概率计算器” + let (probabilities_list, confidence, predicted_indices) = + self.compute_f32_full_probability(matrix_view); + // 5. 执行 CTC 解码 + let final_text = self.ctc_decode_to_string(&predicted_indices); + + Ok(OcrOutput::Probability { + text: final_text, + probabilities: probabilities_list, + confidence: confidence as f64, + }) + } else { + // 走向 B2：极速免 Softmax 提取纯文本（代码保持原地提取，简单短小不需要再拆） + let predicted_indices: Vec = matrix_view + .outer_iter() + .map(|row| { + row.iter() + .enumerate() + .max_by(|(_, a), (_, b)| a.total_cmp(b)) + .map(|(idx, _)| idx as i64) + .unwrap_or(0) + }) + .collect(); + + let final_text = self.ctc_decode_to_string(&predicted_indices); + Ok(OcrOutput::Text(final_text)) + } + } /// 获取有效字符索引列表 (用于外部验证或过滤) fn ctc_decode_to_string(&self, predicted_indices: &[i64]) -> String { println!("indices模型输出原始数据: {:?}", predicted_indices); diff --git a/src/utils/image_processor.rs b/src/utils/image_processor.rs index 7184f1a..2455dfe 100644 --- a/src/utils/image_processor.rs +++ b/src/utils/image_processor.rs @@ -27,4 +27,17 @@ pub fn resize_image( // FilterType::Lanczos3 与 Python Pillow 的 Image.LANCZOS 算法完全对齐，缩放质量最高 image.resize_exact(target_width, target_height, FilterType::Lanczos3) } - +pub fn resize_image1( + image: &GrayImage, + target_width: u32, + target_height: u32, + // resample 参数我们直接使用 FilterType，Lanczos3 是最接近 Python LANCZOS 的 +) -> GrayImage { + // 使用 resize 算法进行精确缩放 + image::imageops::resize( + image, + target_width, + target_height, + FilterType::Lanczos3 + ) +} \ No newline at end of file