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

refactor(predict): 重构预测流水线并优化模型元数据与输出架构

Browse Source 
- 优化 `predict` 核心方法：移除冗余日志与深层嵌套，将流程重塑为线性流水线。
- 重构 `compute_f32_full_probability`：解耦逻辑与外部状态，消除并发隐患与生命周期冲突。
- 增强 `ModelMetadata`：引入动态归一化配置并支持 Serde 序列化，解决特定模型漏字问题。
- 升级 `OcrOutput`：
  - 增加 `Unsupported` 变体以支持非致命异常的优雅降级。
  - 实现 `into_text(self)` 方法与 `Display` 特征（应用双重截断保护，防止日志刷屏）。

BREAKING CHANGE: `predict` 返回值由 `anyhow::Result<String>` 改为 `anyhow::Result<OcrOutput>`，将后处理和控制权移交上层。
This commit is contained in:
CNWei
2026-07-02 20:24:52 +08:00
parent b352fc344f
commit 22cc9709ad
4 changed files with 347 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
src/lib.rs
View File

@@ -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,22 +103,35 @@ impl DdddOcr {
pub fn classification(&self, img: &DynamicImage) -> Result<String> {
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")),
}
}
pub fn detection(&self, img: &[u8]) -> Result<Vec<Vec<i32>>> {
}
pub fn detection(&self, img: &[u8]) -> Result<Vec<Vec<i32>>> {
match &self.runtime {
Runtime::Det(s) => s.predict(img),
Runtime::Ocr(_) => Err(anyhow::anyhow!("当前模型是 OCR 模型，无法执行检测")),
}
}
}
}
// struct Classification {}

48
src/model_metadata.rs
View File

@@ -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,6 +49,13 @@ struct ModelMetadataDto {
#[serde(alias = "image")]
resize: Vec<i32>,
channel: u8,
/// 新增：允许在配置文件中指定归一化策略。
/// 使用 serde(default) 可以在不配置时提供一个默认值（比如默认 ZeroToOne）
#[serde(default = "default_normalization")]
normalization: Normalization,
}
fn default_normalization() -> Normalization {
Normalization::ZeroToOne
}
#[derive(Debug, Clone)]
@@ -41,18 +68,32 @@ pub struct ModelMetadata {
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<Cow<'static, str>> = 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,
})
}
}

353
src/models/ocr.rs
View File

@@ -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<Vec<f32>>,
/// 全局平均置信度
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<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
@@ -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<Vec<i64>> {
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::<i64>()?;
Ok(view.iter().cloned().collect())
}
// 情况 2: sml2h3 原版模型，输出 F32 概率矩阵
DatumType::F32 => {
let view = raw_tensor.to_array_view::<f32>()?;
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<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();
// 丢给现有的 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<String> {
pub fn predict(self, image: &DynamicImage) -> anyhow::Result<OcrOutput> {
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<Tensor> {
// 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(&current_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<f32>,
) -> (Vec<Vec<f32>>, f32, Vec<i64>) {
let steps = matrix_view.nrows();
let classes = matrix_view.ncols();
// 1. 预分配满额概率矩阵内存
let mut prob_matrix = tract_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 extract_from_i64_tensor(&self, raw_tensor: Tensor) -> anyhow::Result<OcrOutput> {
// 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(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<OcrOutput> {
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(OcrOutput::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(OcrOutput::Text(final_text))
}
}
/// 获取有效字符索引列表 (用于外部验证或过滤)
fn ctc_decode_to_string(&self, predicted_indices: &[i64]) -> String {
println!("indices模型输出原始数据: {:?}", predicted_indices);

15
src/utils/image_processor.rs
View File

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