Files
ddddocr-rs/ddddocr-core/src/ocr/executor.rs
CNWei 3499e89bf1 refactor(error): 规范化分层错误类型并优化异常捕捉
- 新增 tracing 记录异常,移除不必要的 Result
- 重构 错误处理架构
2026-07-20 20:18:38 +08:00

549 lines
23 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

use crate::ocr::metadata::Resize;
use crate::ocr::color_filter::{HsvRange, apply_to_image};
// use ddddocr_tract::session::{ModelOutput, OcrSession};
use crate::utils::image_convert::png_rgba_white_preprocess;
use crate::utils::image_processor::{convert_to_grayscale, resize_image};
use image::DynamicImage;
use serde::Serialize;
use std::borrow::Cow;
use std::fmt;
// use tract_onnx::prelude::tract_ndarray::{ Ix2, s};
// use tract_onnx::prelude::{DatumType, Tensor, tract_ndarray};
// !!!【核心纠正】:彻底弃用 tract_ndarray全线转用标准 ndarray
use ndarray::ArrayView2;
// pub enum ModelOutput {
// Indices(ndarray::Array1<i64>), // 拥有完整所有权的 1维数组可任意传递和返回
// Logits(ndarray::Array2<f32>), // 拥有完整所有权的 2维矩阵可任意传递和返回
// }
use crate::error::{ImagePreprocessError, Result, TensorError};
use crate::{OcrEngine, OcrOutput};
use tracing::{ warn};
#[derive(Debug, Clone, Serialize)]
pub enum OcrResult {
/// 纯文本分支(对应 probability = false
Text(String),
/// 包含全量概率的分支(对应 probability = true
Probability {
text: String,
/// 满额概率矩阵 [Steps, Classes]
probabilities: Vec<Vec<f32>>,
/// 全局平均置信度
confidence: f64,
},
/// 不支持的模型或未知输出
Unsupported { message: String },
}
impl OcrResult {
/// 消费自身,直接提取最终文本
pub fn into_text(self) -> String {
match self {
OcrResult::Text(text) => text,
OcrResult::Probability { text, .. } => text,
OcrResult::Unsupported { message } => {
// 作为库,这里可以返回空,或者直接携带错误信息,取决于你的设计
format!("Error: {}", message)
}
}
}
}
impl fmt::Display for OcrResult {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
OcrResult::Text(text) => {
// 纯文本分支,直接输出文本内容
write!(f, "{}", text)
}
OcrResult::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, "]")
}
OcrResult::Unsupported { message } => {
// 错误分支,直观输出异常原因
write!(f, "未识别成功: {}", message)
}
}
}
}
pub struct Ocr<'a> {
pub(crate) session: &'a dyn OcrEngine,
pub(crate) png_fix: bool,
pub(crate) probability: bool,
/// 颜色过滤:保留的颜色列表
pub(crate) final_color_ranges: Result<Option<Vec<HsvRange>>, ImagePreprocessError>,
/// 字符集范围
pub(crate) final_charset_indices: Option<Vec<usize>>,
}
impl<'a> Ocr<'a> {
// 初始化任务,设置默认参数
pub fn new(session: &'a dyn OcrEngine) -> Self {
Ocr {
session,
png_fix: false, // 默认值
probability: false,
final_color_ranges: Ok(None),
final_charset_indices: None,
}
}
}
impl<'a> Ocr<'a> {
pub fn predict(&self, image: &DynamicImage) -> Result<OcrResult> {
println!("当前颜色过滤器状态: {:?}", self.final_color_ranges);
// =====================================================================
// 管道节点 1: 颜色过滤流水线
// 使用 Cow (Copy-On-Write) 智能指针。
// 如果未开启过滤img_cow 内部只是持有原图的【只读借用】,发生【零内存分配】!
// =====================================================================
let img_cow = match &self.final_color_ranges {
Err(err_msg) => {
// return Err(anyhow::anyhow!(
// "颜色过滤器初始化失败,全链路短路: {}",
// err_msg
// ));
return Err(ImagePreprocessError::FilterConfigInvalid(
err_msg.to_string(),
))?;
}
Ok(None) => {
// 核心优化点:直接借用原图,不发生任何克隆
Cow::Borrowed(image)
}
Ok(Some(ranges)) => {
// 只有真正需要过滤时,才在内部提取像素并生成清洗后的 Owned 新图
let filtered_img = apply_to_image(image, ranges)?;
Cow::Owned(filtered_img)
}
};
let tensor = self.preprocess_image(&img_cow)?;
let raw_tensor = self.session.inference(tensor)?;
// 3. 后处理分流:直接返回 OcrResult
// let ocr_output = match raw_tensor.datum_type() {
// DatumType::I64 => self.process_i64_tensor(raw_tensor)?,
// DatumType::F32 => self.process_f32_tensor(raw_tensor)?,
// _ => OcrResult::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);
let ocr_output = self.process_model_output(raw_tensor)?;
Ok(ocr_output)
}
/// 对应 Python 的 _preprocess_image
/// 负责:透明背景修复 -> 灰度化 -> 按比例 Resize -> 归一化 -> 4维张量转换
fn preprocess_image(&self, img: &DynamicImage) -> Result<ndarray::Array4<f32>,ImagePreprocessError> {
// 1. 获取模型元数据配置
let meta = self.session.metadata();
let norm = &meta.normalization; // 获取归一化器
// A. 修复 PNG 透明背景 (内部逻辑你之前已实现)
let current_img = if self.png_fix && img.color().has_alpha() {
// 只有满足条件才去触发分配,生成新图
Cow::Owned(png_rgba_white_preprocess(img))
} else {
// 正常情况下,仅仅是再次安全借用,无开销
Cow::Borrowed(img)
};
// 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(&current_img, target_w, target_h);
// 4. 管道节点 3: 颜色通道转换(单通道灰度 vs 三通道 RGB与 4D 张量填充
let array4 = match meta.channel {
// --- 情况 A: 单通道(灰度图),对应 Python 的 len(shape) == 2 展开 ---
1 => {
let gray_img = convert_to_grayscale(&resized_img);
let array = 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]
// (pixel / 255.0 - 0.5) / 0.5
norm.normalize(pixel)
},
);
array
}
// --- 情况 B: 三通道RGB对应 Python 的 transpose(2, 0, 1) 的 CHW 布局 ---
3 => {
let rgb_img = resized_img.to_rgb8();
let array = 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]
// (pixel / 255.0 - 0.5) / 0.5
norm.normalize(pixel)
},
);
// Tensor::from(array)
array
}
// _ => return Err(anyhow::anyhow!("不支持的通道数配置: {}", meta.channel)),
_ => {
return Err(ImagePreprocessError::UnsupportedChannels(
meta.channel as usize,
));
}
};
Ok(array4)
// 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)
}
// 这段代码未来直接放入 ddddocr-core
fn process_model_output(&self, output: OcrOutput) -> Result<OcrResult,TensorError> {
match output {
OcrOutput::Indices(array1) => {
// 对应你原来的 process_i64_tensor
let slice = array1
.as_slice()
// .ok_or_else(|| anyhow::anyhow!("内存不连续,无法执行零拷贝解码"))?;
.ok_or_else(|| TensorError::NonContiguousMemory)?;
let final_text = self.ctc_decode_to_string(slice);
if self.probability {
Ok(OcrResult::Probability {
text: final_text,
probabilities: vec![],
confidence: 1.0,
})
} else {
Ok(OcrResult::Text(final_text))
}
}
OcrOutput::Logits(matrix_view) => {
// 对应你原来的 process_f32_tensor
// 注意:此时的 matrix_view 已经是干净的标准的 ndarray::Array2<f32>,且保证是 [Steps, Classes] 2D 形状
if self.probability {
let (probabilities_list, confidence, predicted_indices) =
self.compute_f32_full_probability(matrix_view.view());
let final_text = self.ctc_decode_to_string(&predicted_indices);
Ok(OcrResult::Probability {
text: final_text,
probabilities: probabilities_list,
confidence: confidence as f64,
})
} else {
let predicted_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();
let final_text = self.ctc_decode_to_string(&predicted_indices);
Ok(OcrResult::Text(final_text))
}
}
}
}
}
impl<'a> Ocr<'a> {
fn is_valid_indices(&self, idx: usize) -> bool {
if idx >= self.session.metadata().charset.size() {
return false;
}
match &self.final_charset_indices {
Some(v) => v.binary_search(&idx).is_ok(),
None => true,
}
}
/// 【按需延迟打印】:当用户真的需要“知道当前有哪些限制字符”时,一秒反查并打印
/// 这里的 &str 完美借用了自 tokens依然是彻底的零拷贝
pub fn valid_tokens(&self) -> Vec<&str> {
let charset = &self.session.metadata().charset;
let tokens = &charset.tokens;
match &self.final_charset_indices {
Some(indices) => indices
.iter()
.filter_map(|&idx| tokens.get(idx).map(|cow| cow.as_ref()))
.collect(),
// 如果是 None现场映射出全量 Token 视图给外部
None => tokens.iter().map(|cow| cow.as_ref()).collect(),
}
}
pub fn valid_size(&self) -> usize {
match &self.final_charset_indices {
Some(indices) => indices.len(),
None => self.session.metadata().charset.tokens.len(),
}
}
/// 变体 B 核心处理器:单次遍历 2D 视图,融合计算 Softmax、Argmax、置信度并输出概率大包
fn compute_f32_full_probability(
&self,
matrix_view: ArrayView2<f32>,
) -> (Vec<Vec<f32>>, f32, Vec<i64>) {
let steps = matrix_view.nrows();
let classes = matrix_view.ncols();
// 1. 预分配满额概率矩阵内存
let mut prob_matrix = ndarray::Array2::<f32>::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<Vec<f32>> =
prob_matrix.outer_iter().map(|row| row.to_vec()).collect();
(probabilities_list, confidence, predicted_indices)
}
/// 变体 A 专属提取器:直接从 I64 Tensor 零拷贝提取 CTC 文本与初始概率包
// fn process_i64_tensor(&self, raw_tensor: Tensor) -> anyhow::Result<OcrResult> {
// // 1. 拿到底层的动态维度只读视图
// let view = raw_tensor.to_array_view::<i64>()?;
//
// // 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(OcrResult::Probability {
// text: final_text,
// probabilities: vec![], // I64 模型物理上丢失了全量 Logits 分值网,降级处理
// confidence: 1.0, // 判定即百分之百置信
// })
// } else {
// Ok(OcrResult::Text(final_text))
// }
// }
// /// 变体二F32的总体管线负责降维并分流文本和概率
// fn process_f32_tensor(&self, raw_tensor: Tensor) -> anyhow::Result<OcrResult> {
// let shape = raw_tensor.shape();
// println!("模型输出shape数据: {:?}", shape);
// let view = raw_tensor.to_array_view::<f32>()?;
//
// // 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<f32> = 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(OcrResult::Probability {
// text: final_text,
// probabilities: probabilities_list,
// confidence: confidence as f64,
// })
// } else {
// // 走向 B2极速免 Softmax 提取纯文本(代码保持原地提取,简单短小不需要再拆)
// let predicted_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();
//
// let final_text = self.ctc_decode_to_string(&predicted_indices);
// Ok(OcrResult::Text(final_text))
// }
// }
/// 获取有效字符索引列表 (用于外部验证或过滤)
fn ctc_decode_to_string(&self, predicted_indices: &[i64]) -> String {
println!("indices模型输出原始数据: {:?}", predicted_indices);
let charset = &self.session.metadata().charset;
let tokens = &charset.tokens;
// let valid_indices = &charset.valid_indices;
// 对应 _ctc_decode_indices 的逻辑:去重、去 blank (0)
let mut res = String::new();
let mut prev_idx: i64 = -1;
for &idx in predicted_indices {
// 1. CTC 去重:如果是连续重复的,直接跳过
if idx == prev_idx {
continue;
}
// 【关键核心】只要不是连续重复,立刻更新 prev_idx 状态,绝对不能被后续的过滤短路!
prev_idx = idx;
// 2. CTC 过滤 Blank (0)
if idx == 0 {
continue;
}
// 3. 类型安全转换
let u_idx = match usize::try_from(idx) {
Ok(u) => u,
Err(_) => continue,
};
// 史诗级加速点:如果是 None说明没限制根本不进入分支直接放行
// 只有当有具体限制Some才去跑 4-5 次 CPU 寄存器级别的二分查找
if let Some(ref indices) = self.final_charset_indices {
if indices.binary_search(&u_idx).is_err() {
continue;
}
}
// 5. 字符映射
if let Some(char_str) = tokens.get(u_idx) {
res.push_str(char_str);
}
else {
warn!("警告: 预测索引 {} 超出字符集范围", u_idx);
}
}
res
}
}