diff --git a/src/lib.rs b/src/lib.rs index d5c68c8..95c9702 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,7 +10,6 @@ use std::fmt::{Display, Formatter}; // 关键点：直接使用 tract 重导出的 ndarray use crate::charset::{ CharRestrict}; -use crate::models::ocr::ColorRange; use models::det::Det; use models::loader::ModelSession; use models::ocr::Ocr; @@ -98,9 +97,9 @@ impl DdddOcr { pub fn classification(&self, img: &DynamicImage) -> Result { match &self.runtime { // Runtime::Ocr(s) => s.predict(img).run(), - Runtime::Ocr(s) => s.builder().predict(img), - // Runtime::Ocr(s) => s.builder().charset_restrict(&CharRestrict::Digit).predict(img), - // Runtime::Ocr(s) => s.builder().color_filter(&ColorPreset::Custom(vec![ + Runtime::Ocr(s) => s.predictor().predict(img), + // 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)), @@ -116,39 +115,39 @@ impl DdddOcr { } } -struct Classification {} -#[derive(Debug)] -struct ClassificationBuilder { - img: DynamicImage, - png_fix: bool, - color_filter_colors: Option>, - color_filter_custom_ranges: Option>, -} -impl ClassificationBuilder { - pub fn new(img: DynamicImage) -> Self { - ClassificationBuilder { - img, - png_fix: false, - color_filter_colors: None, - color_filter_custom_ranges: None, - } - } - pub fn png_fix(mut self, value: bool) -> Self { - self.png_fix = value; - self - } - pub fn color_filter_colors(mut self, value: Vec) -> Self { - self.color_filter_colors = Some(value); - self - } - pub fn color_filter_custom_ranges(mut self, value: Vec) -> Self { - self.color_filter_custom_ranges = Some(value); - self - } - pub fn build(self) -> Classification { - Classification {} - } -} +// struct Classification {} +// #[derive(Debug)] +// struct ClassificationBuilder { +// img: DynamicImage, +// png_fix: bool, +// color_filter_colors: Option>, +// color_filter_custom_ranges: Option>, +// } +// impl ClassificationBuilder { +// pub fn new(img: DynamicImage) -> Self { +// ClassificationBuilder { +// img, +// png_fix: false, +// color_filter_colors: None, +// color_filter_custom_ranges: None, +// } +// } +// pub fn png_fix(mut self, value: bool) -> Self { +// self.png_fix = value; +// self +// } +// pub fn color_filter_colors(mut self, value: Vec) -> Self { +// self.color_filter_colors = Some(value); +// self +// } +// pub fn color_filter_custom_ranges(mut self, value: Vec) -> Self { +// self.color_filter_custom_ranges = Some(value); +// self +// } +// pub fn build(self) -> Classification { +// Classification {} +// } +// } #[cfg(test)] mod tests { diff --git a/src/model_metadata.rs b/src/model_metadata.rs index dec971a..69a2df9 100644 --- a/src/model_metadata.rs +++ b/src/model_metadata.rs @@ -47,12 +47,12 @@ impl ModelMetadata { // --- 优雅的工厂模式构造器 --- /// 从预设的旧版字符集创建 pub fn from_builtin_old() -> Self { - Self::from_static_slice(CHARSET_OLD, false, Resize::Fixed(64, 64), 1) + Self::from_static_slice(CHARSET_OLD, false, Resize::DynamicWidth(64), 1) } /// 从预设的 Beta 版字符集创建 pub fn from_builtin_beta() -> Self { - Self::from_static_slice(CHARSET_BETA, false, Resize::Fixed(64, 64), 1) + Self::from_static_slice(CHARSET_BETA, false, Resize::DynamicWidth(64), 1) } /// 通用的静态切片转换构造器 diff --git a/src/models/ocr.rs b/src/models/ocr.rs index b2ac19d..50a9bef 100644 --- a/src/models/ocr.rs +++ b/src/models/ocr.rs @@ -1,5 +1,5 @@ use crate::charset::{TokenFilter, ValidationCtx}; -use crate::model_metadata::ModelMetadata; +use crate::model_metadata::{ModelMetadata, Resize}; use crate::models::base::ModelArgs; use crate::models::loader::{ModelLoader, ModelSession, ModelType}; use crate::utils::color_filter::{ColorFilter, HsvRange, filter_image}; @@ -10,14 +10,13 @@ use anyhow::{Result, anyhow}; use image::{DynamicImage, ImageBuffer, Rgb}; use std::borrow::Cow; use std::collections::HashSet; -use tract_onnx::prelude::tract_ndarray::s; +use tract_onnx::prelude::tract_ndarray::{s, ArrayView2}; 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; -// 颜色过滤的自定义范围：(低值RGB, 高值RGB) -pub type ColorRange = ((u8, u8, u8), (u8, u8, u8)); + pub struct Ocr { pub session: RunnableModel, Graph>>, @@ -48,7 +47,7 @@ impl Ocr { .run(tvec!(tensor.into())) .context("执行模型推理失败")?; println!("模型输出原始数据: {:?}", result); - Ok(result.remove(0).into_tensor()) + Ok(result.swap_remove(0).into_tensor()) } /// 核心解析逻辑：将模型输出的各种维度/类型的 Tensor 转为字符索引序列 @@ -58,10 +57,11 @@ impl Ocr { let datum_type = raw_tensor.datum_type(); println!("模型输出datum_type数据: {:?}", datum_type); - match raw_tensor.datum_type() { + match datum_type { // 情况 1: huashi666 式模型，直接输出 i64 索引 (通常是模型内部做好了 Argmax) DatumType::I64 => { let view = raw_tensor.to_array_view::()?; + Ok(view.iter().cloned().collect()) } @@ -111,17 +111,34 @@ impl Ocr { } _ => Err(anyhow::anyhow!( "不支持的模型输出数据类型: {:?}", - raw_tensor.datum_type() + 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(); - pub fn builder(&'_ self) -> OcrBuilder<'_> { - OcrBuilder::new(self) + // 丢给现有的 CTC 解码器去重并映射成字符串 + Ok(self.ctc_decode_to_string(&indices)) + } + + pub fn predictor(&'_ self) -> OcrPredictor<'_> { + OcrPredictor::new(self) } } -pub struct OcrBuilder<'a> { +pub struct OcrPredictor<'a> { ocr: &'a Ocr, // image: &'a DynamicImage, /// 是否修复PNG格式问题 @@ -131,13 +148,12 @@ pub struct OcrBuilder<'a> { probability: bool, /// 颜色过滤：保留的颜色列表 color_filter: Result>, String>, - /// 颜色过滤：自定义RGB范围 - color_filter_custom_ranges: Option>, + /// 字符集范围 charset_restrict: Option>, } -impl<'a> OcrBuilder<'a> { +impl<'a> OcrPredictor<'a> { // 初始化任务，设置默认参数 pub fn new(ocr: &'a Ocr) -> Self { Self { @@ -146,7 +162,6 @@ impl<'a> OcrBuilder<'a> { png_fix: false, // 默认值 probability: false, color_filter: Ok(None), - color_filter_custom_ranges: None, charset_restrict: None, } } @@ -168,10 +183,6 @@ impl<'a> OcrBuilder<'a> { self } - pub fn color_filter_custom_ranges(mut self, value: Vec) -> Self { - self.color_filter_custom_ranges = Some(value); - self - } pub fn charset_restrict(mut self, restrict: &dyn TokenFilter) -> Self { let charset = &self.ocr.model_metadata.charset; let tokens = &charset.tokens; @@ -180,8 +191,8 @@ impl<'a> OcrBuilder<'a> { self } } -impl<'a> OcrBuilder<'a> { - pub fn predict(&self, image: &DynamicImage) -> anyhow::Result { +impl<'a> OcrPredictor<'a> { + pub fn predict(self, image: &DynamicImage) -> anyhow::Result { println!("当前颜色过滤器状态: {:?}", self.color_filter); // ===================================================================== // 管道节点 1: 颜色过滤流水线 @@ -218,6 +229,9 @@ impl<'a> OcrBuilder<'a> { /// 对应 Python 的 _preprocess_image /// 负责：透明背景修复 -> 灰度化 -> 按比例 Resize -> 归一化 -> 4维张量转换 fn preprocess_image(&self, img: &DynamicImage) -> anyhow::Result { + // 1. 获取模型元数据配置 + let meta = &self.ocr.model_metadata; + // A. 修复 PNG 透明背景 (内部逻辑你之前已实现) let current_img = if self.png_fix && img.color().has_alpha() { // 只有满足条件才去触发分配，生成新图 @@ -227,41 +241,94 @@ impl<'a> OcrBuilder<'a> { Cow::Borrowed(img) }; - 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); - 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(); + // 3. 管道节点 2: 根据 Resize 策略计算目标宽高并进行缩放 + let (target_w, target_h) = match meta.resize { + 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; + (w, h) + } + Resize::Square(size) => { + // 单字识别模型，直接缩放为正方形 + (size, size) + } + }; + // 执行缩放 + let resized_img = resize_image(¤t_img, target_w, target_h); - // 使用 tract_ndarray 构造，避免版本冲突 - let array = - tract_ndarray::Array4::from_shape_fn((1, 1, h as usize, w as usize), |(_, _, y, x)| { - let pixel = resized.get_pixel(x as u32, y as u32)[0] as f32; - (pixel / 255.0 - 0.5) / 0.5 - }); + // 4. 管道节点 3: 颜色通道转换（单通道灰度 vs 三通道 RGB）与 4D 张量填充 + let tensor = match meta.channel { + // --- 情况 A: 单通道（灰度图），对应 Python 的 len(shape) == 2 展开 --- + 1 => { + let gray_img = convert_to_grayscale(&resized_img); - let tensor = Tensor::from(array); + let array = tract_ndarray::Array4::from_shape_fn( + (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] + }, + ); + Tensor::from(array) + } + + // --- 情况 B: 三通道（RGB），对应 Python 的 transpose(2, 0, 1) 的 CHW 布局 --- + 3 => { + let rgb_img = resized_img.to_rgb8(); + + let array = tract_ndarray::Array4::from_shape_fn( + (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] + }, + ); + Tensor::from(array) + } + + _ => return Err(anyhow::anyhow!("不支持的通道数配置: {}", meta.channel)), + }; 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); + // 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(); + // + // // 使用 tract_ndarray 构造，避免版本冲突 + // let array = + // tract_ndarray::Array4::from_shape_fn((1, 1, h as usize, w as usize), |(_, _, y, x)| { + // let pixel = resized.get_pixel(x as u32, y as u32)[0] as f32; + // (pixel / 255.0 - 0.5) / 0.5 + // }); + // + // let tensor = Tensor::from(array); + // + // Ok(tensor) } } -impl<'a> OcrBuilder<'a> { - pub fn get_valid_indices(&self) -> HashSet { - match &self.charset_restrict { - Some(indices) => indices.iter().cloned().collect(), - // 如果是 None，现场映射出全量索引集给外部 - None => (0..self.ocr.model_metadata.charset.tokens.len()).collect(), - } - } +impl<'a> OcrPredictor<'a> { + // pub fn get_valid_indices(&self) -> HashSet { + // match &self.charset_restrict { + // Some(indices) => indices.iter().cloned().collect(), + // // 如果是 None，现场映射出全量索引集给外部 + // None => (0..self.ocr.model_metadata.charset.tokens.len()).collect(), + // } + // } // compute_valid_indices // fn valid_indices(&self) -> (bool, HashSet) { // let charset = &self.ocr.model_metadata.charset; /// 【按需延迟打印】：当用户真的需要“知道当前有哪些限制字符”时，一秒反查并打印 /// 这里的 &str 完美借用了自 tokens，依然是彻底的零拷贝！ - pub fn get_valid_tokens(&self) -> Vec<&str> { + pub fn valid_tokens(&self) -> Vec<&str> { let charset = &self.ocr.model_metadata.charset; let tokens = &charset.tokens; match &self.charset_restrict { diff --git a/src/utils/image_processor.rs b/src/utils/image_processor.rs index cd39fae..7184f1a 100644 --- a/src/utils/image_processor.rs +++ b/src/utils/image_processor.rs @@ -11,17 +11,20 @@ pub fn convert_to_grayscale(image: &DynamicImage) -> GrayImage { /// 对应 Python 的 resize_image /// 调整图像尺寸。当前版本仅实现 keep_aspect_ratio=false pub fn resize_image( - image: &GrayImage, + image: &DynamicImage, target_width: u32, target_height: u32, // resample 参数我们直接使用 FilterType，Lanczos3 是最接近 Python LANCZOS 的 -) -> GrayImage { +) -> DynamicImage { // 使用 resize 算法进行精确缩放 - image::imageops::resize( - image, - target_width, - target_height, - FilterType::Lanczos3 - ) + // image::imageops::resize( + // image, + // target_width, + // target_height, + // FilterType::Lanczos3 + // ) + // image::imageops::resize 的最高层封装 + // FilterType::Lanczos3 与 Python Pillow 的 Image.LANCZOS 算法完全对齐，缩放质量最高 + image.resize_exact(target_width, target_height, FilterType::Lanczos3) }