CNWei/ddddocr-rs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(ocr): 优化 preprocess_image 逻辑实现，修复部分BUG

Browse Source 
- 优化 `OcrBuilder` 重名名为 `OcrPredictor`
 - 优化 `OcrPredictor` 的 `preprocess_image` 支持多种图像管道。
 - 修复 `OcrPredictor` 引发的并发BUG。
This commit is contained in:
CNWei
2026-07-01 20:22:17 +08:00
parent 48c2cbedb0
commit b352fc344f
4 changed files with 159 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
src/lib.rs
View File

@@ -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<String> {
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<Vec<ColorRange>>,
color_filter_custom_ranges: Option<Vec<ColorRange>>,
}
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<ColorRange>) -> Self {
self.color_filter_colors = Some(value);
self
}
pub fn color_filter_custom_ranges(mut self, value: Vec<ColorRange>) -> 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<Vec<ColorRange>>,
// color_filter_custom_ranges: Option<Vec<ColorRange>>,
// }
// 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<ColorRange>) -> Self {
// self.color_filter_colors = Some(value);
// self
// }
// pub fn color_filter_custom_ranges(mut self, value: Vec<ColorRange>) -> Self {
// self.color_filter_custom_ranges = Some(value);
// self
// }
// pub fn build(self) -> Classification {
// Classification {}
// }
// }
#[cfg(test)]
mod tests {

4
src/model_metadata.rs
View File

@@ -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)
}
/// 通用的静态切片转换构造器

151
src/models/ocr.rs
View File

@@ -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<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
@@ -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::<i64>()?;
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<f32>) -> anyhow::Result<String> {
// 直接在原始分值（Logits）上进行 Argmax，数学结果与 Softmax 后完全一致
let indices: Vec<i64> = 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<Option<Vec<HsvRange>>, String>,
/// 颜色过滤：自定义RGB范围
color_filter_custom_ranges: Option<Vec<ColorRange>>,
/// 字符集范围
charset_restrict: Option<Vec<usize>>,
}
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<ColorRange>) -> 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<String> {
impl<'a> OcrPredictor<'a> {
pub fn predict(self, image: &DynamicImage) -> anyhow::Result<String> {
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<Tensor> {
// 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;
// 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;
let gray_img = convert_to_grayscale(&current_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();
(w, h)
}
Resize::Square(size) => {
// 单字识别模型，直接缩放为正方形
(size, size)
}
};
// 执行缩放
let resized_img = resize_image(&current_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(&current_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<usize> {
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<usize> {
// 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<usize>) {
// 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 {

19
src/utils/image_processor.rs
View File

@@ -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)
}