refactor: 重构 OcrBuilder 过滤机制,实现零克隆、高性能的局部白名单
- 将有状态的有效索引缓存 `valid_indices` 从全局 `Charset` 剥离至局部 `OcrBuilder` - 解码函数 `ctc_decode_to_string` 内部复用策略计算,通过 `has_any_match` 布尔开关实现全量免检短路加速 - 优化内存分配,根据 `CharsetRestrict` 策略动态精准计算 `HashSet` 初始容量,规避大词表空置浪费 - 增强鲁棒性,在策略与字库完全无交集时,自动触发智能降级,一键恢复全量识别并保持接口自洽
This commit is contained in:
@@ -600,11 +600,11 @@ impl Add for CharsetRestrict {
|
|||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub struct Charset {
|
pub struct Charset {
|
||||||
// 使用 Cow 统一静态切片和动态读取的 Vec<String>,内部实现真正的零拷贝
|
// 使用 Cow 统一静态切片和动态读取的 Vec<String>,内部实现真正的零拷贝
|
||||||
tokens: Vec<Cow<'static, str>>,
|
pub tokens: Vec<Cow<'static, str>>,
|
||||||
// 反向查找表,保证字符转索引为 O(1)
|
// 反向查找表,保证字符转索引为 O(1)
|
||||||
char_to_idx: HashMap<Cow<'static, str>, usize>,
|
pub char_to_idx: HashMap<Cow<'static, str>, usize>,
|
||||||
// 当前处于激活状态的有效索引缓存 (用于 CTC 解码前的过滤加速)
|
// 当前处于激活状态的有效索引缓存 (用于 CTC 解码前的过滤加速)
|
||||||
valid_indices: HashSet<usize>,
|
// pub valid_indices: HashSet<usize>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Charset {
|
impl Charset {
|
||||||
@@ -617,48 +617,14 @@ impl Charset {
|
|||||||
// char_to_idx.entry(token.to_string()).or_insert(idx);
|
// char_to_idx.entry(token.to_string()).or_insert(idx);
|
||||||
}
|
}
|
||||||
|
|
||||||
// 默认初始化时,所有索引均为有效状态
|
|
||||||
let valid_indices = (0..tokens.len()).collect();
|
|
||||||
|
|
||||||
Self {
|
Self {
|
||||||
tokens,
|
tokens,
|
||||||
char_to_idx,
|
char_to_idx,
|
||||||
valid_indices,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// --- 业务策略方法 ---
|
// --- 业务策略方法 ---
|
||||||
|
|
||||||
/// 根据传入的 CharsetRange 枚举策略,动态更新有效索引
|
|
||||||
pub fn apply_range_policy(&mut self, policy: &CharsetRestrict) -> bool {
|
|
||||||
let mut has_any_match = false;
|
|
||||||
// 3. 清空原有的索引
|
|
||||||
self.valid_indices.clear();
|
|
||||||
// 4. 执行 O(1) 级别的求交集过滤
|
|
||||||
for (idx, token) in self.tokens.iter().enumerate() {
|
|
||||||
let token_str = token.as_ref();
|
|
||||||
// CTC Blank 空字符串无条件放行,其余交给超高性能的 matches
|
|
||||||
if token_str.is_empty() {
|
|
||||||
self.valid_indices.insert(idx);
|
|
||||||
} else if policy.matches(token_str) {
|
|
||||||
self.valid_indices.insert(idx);
|
|
||||||
has_any_match = true;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
// 🛡️ 终极防御:如果除了 Blank 之外,没有任何一个字符被匹配到
|
|
||||||
if !has_any_match {
|
|
||||||
// 策略 C:智能降级,一键恢复全量字符集,防止模型“交白卷”
|
|
||||||
self.reset_range_policy();
|
|
||||||
println!("警告:当前限制策略与模型字符集完全没有交集!已自动恢复全量识别。");
|
|
||||||
return false; // 返回 false 提示外部:策略未实际生效,已降级
|
|
||||||
}
|
|
||||||
true
|
|
||||||
}
|
|
||||||
/// 清除范围限制,恢复完整字符集
|
|
||||||
pub fn reset_range_policy(&mut self) {
|
|
||||||
self.valid_indices = (0..self.tokens.len()).collect();
|
|
||||||
}
|
|
||||||
|
|
||||||
/// 将字符转为索引,不存在返回 -1 (保持与原 Python 库行为一致)
|
/// 将字符转为索引,不存在返回 -1 (保持与原 Python 库行为一致)
|
||||||
pub fn char_to_index(&self, char_str: &str) -> i32 {
|
pub fn char_to_index(&self, char_str: &str) -> i32 {
|
||||||
if let Some(&idx) = self.char_to_idx.get(char_str) {
|
if let Some(&idx) = self.char_to_idx.get(char_str) {
|
||||||
@@ -673,20 +639,6 @@ impl Charset {
|
|||||||
self.tokens.get(index).map(|cow| cow.as_ref())
|
self.tokens.get(index).map(|cow| cow.as_ref())
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 获取有效字符索引列表 (用于外部验证或过滤)
|
|
||||||
pub fn get_valid_indices(&self) -> Vec<usize> {
|
|
||||||
let mut indices: Vec<usize> = self.valid_indices.iter().copied().collect();
|
|
||||||
indices.sort_unstable();
|
|
||||||
indices
|
|
||||||
}
|
|
||||||
/// 🌟 【按需延迟打印】:当用户真的需要“知道当前有哪些限制字符”时,一秒反查并打印
|
|
||||||
/// 这里的 &str 完美借用了自 tokens,依然是彻底的零拷贝!
|
|
||||||
pub fn get_valid_tokens(&self) -> Vec<&str> {
|
|
||||||
self.get_valid_indices()
|
|
||||||
.iter()
|
|
||||||
.map(|&idx| self.tokens[idx].as_ref())
|
|
||||||
.collect()
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn is_valid_char(&self, char_str: &str) -> bool {
|
pub fn is_valid_char(&self, char_str: &str) -> bool {
|
||||||
self.char_to_idx.get(char_str).is_some()
|
self.char_to_idx.get(char_str).is_some()
|
||||||
@@ -695,9 +647,6 @@ impl Charset {
|
|||||||
self.tokens.len()
|
self.tokens.len()
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn valid_size(&self) -> usize {
|
|
||||||
self.valid_indices.len()
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// ==========================================
|
// ==========================================
|
||||||
@@ -707,9 +656,8 @@ impl std::fmt::Display for Charset {
|
|||||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
write!(
|
write!(
|
||||||
f,
|
f,
|
||||||
"Charset [Total Size: {}, Active Range Size: {}]",
|
"Charset [Total Size: {}",
|
||||||
self.size(),
|
self.size(),
|
||||||
self.valid_size()
|
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
11
src/lib.rs
11
src/lib.rs
@@ -14,6 +14,7 @@ use crate::models::ocr::ColorRange;
|
|||||||
use models::det::Det;
|
use models::det::Det;
|
||||||
use models::loader::ModelSession;
|
use models::loader::ModelSession;
|
||||||
use models::ocr::Ocr;
|
use models::ocr::Ocr;
|
||||||
|
use crate::model_metadata::ModelMetadata;
|
||||||
|
|
||||||
pub enum ModelSpec {
|
pub enum ModelSpec {
|
||||||
/// 默认 OCR (使用内置路径)
|
/// 默认 OCR (使用内置路径)
|
||||||
@@ -22,7 +23,7 @@ pub enum ModelSpec {
|
|||||||
/// 自定义 OCR (路径由用户提供)
|
/// 自定义 OCR (路径由用户提供)
|
||||||
CustomOcrModel {
|
CustomOcrModel {
|
||||||
path: String,
|
path: String,
|
||||||
charset: Vec<String>,
|
model_metadata: ModelMetadata,
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
impl ModelSpec {
|
impl ModelSpec {
|
||||||
@@ -61,9 +62,9 @@ impl DdddOcrBuilder {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// 设置自定义 OCR 路径
|
/// 设置自定义 OCR 路径
|
||||||
pub fn custom_ocr(mut self, path: String, charset: Vec<String>) -> Self {
|
pub fn custom_ocr(mut self, path: String, model_metadata: ModelMetadata) -> Self {
|
||||||
// 直接重写枚举,替换掉之前的 Ocr 或 Det
|
// 直接重写枚举,替换掉之前的 Ocr 或 Det
|
||||||
self.mode = ModelSpec::CustomOcrModel { path, charset };
|
self.mode = ModelSpec::CustomOcrModel { path, model_metadata };
|
||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -72,10 +73,10 @@ impl DdddOcrBuilder {
|
|||||||
let runtime = match self.mode {
|
let runtime = match self.mode {
|
||||||
ModelSpec::OcrModel => Runtime::Ocr(Ocr::new(
|
ModelSpec::OcrModel => Runtime::Ocr(Ocr::new(
|
||||||
ModelSpec::DEFAULT_OCR_PATH.into(),
|
ModelSpec::DEFAULT_OCR_PATH.into(),
|
||||||
get_default_charset(),
|
ModelMetadata::from_builtin_beta(),
|
||||||
)?),
|
)?),
|
||||||
ModelSpec::DetModel => Runtime::Det(Det::new(ModelSpec::DEFAULT_DET_PATH.into())?),
|
ModelSpec::DetModel => Runtime::Det(Det::new(ModelSpec::DEFAULT_DET_PATH.into())?),
|
||||||
ModelSpec::CustomOcrModel { path, charset } => Runtime::Ocr(Ocr::new(path, charset)?),
|
ModelSpec::CustomOcrModel { path, model_metadata } => Runtime::Ocr(Ocr::new(path, model_metadata)?),
|
||||||
};
|
};
|
||||||
|
|
||||||
Ok(DdddOcr { runtime })
|
Ok(DdddOcr { runtime })
|
||||||
|
|||||||
@@ -34,7 +34,7 @@ struct ModelMetadataDto {
|
|||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub struct ModelMetadata {
|
pub struct ModelMetadata {
|
||||||
/// 字符集管理器
|
/// 字符集管理器
|
||||||
pub charset: Charset,
|
pub charset: Charset,
|
||||||
/// 是否为单字识别模型
|
/// 是否为单字识别模型
|
||||||
pub word: bool,
|
pub word: bool,
|
||||||
/// 预处理的缩放策略
|
/// 预处理的缩放策略
|
||||||
|
|||||||
@@ -1,23 +1,23 @@
|
|||||||
|
use crate::charset::CharsetRestrict;
|
||||||
|
use crate::model_metadata::ModelMetadata;
|
||||||
use crate::models::base::ModelArgs;
|
use crate::models::base::ModelArgs;
|
||||||
use crate::models::loader::{ModelLoader, ModelSession, ModelType};
|
use crate::models::loader::{ModelLoader, ModelSession, ModelType};
|
||||||
use crate::utils::image_io::png_rgba_white_preprocess;
|
use crate::utils::image_io::png_rgba_white_preprocess;
|
||||||
use crate::utils::image_processor::{convert_to_grayscale, resize_image};
|
use crate::utils::image_processor::{convert_to_grayscale, resize_image};
|
||||||
use anyhow::Context;
|
use anyhow::Context;
|
||||||
use image::DynamicImage;
|
use image::DynamicImage;
|
||||||
|
use std::collections::HashSet;
|
||||||
use tract_onnx::prelude::tract_ndarray::s;
|
use tract_onnx::prelude::tract_ndarray::s;
|
||||||
use tract_onnx::prelude::{
|
use tract_onnx::prelude::{
|
||||||
DatumType, Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp, tract_ndarray, tvec,
|
DatumType, Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp, tract_ndarray, tvec,
|
||||||
};
|
};
|
||||||
use crate::charset::CharsetRange;
|
|
||||||
|
|
||||||
// 颜色过滤的自定义范围:(低值RGB, 高值RGB)
|
// 颜色过滤的自定义范围:(低值RGB, 高值RGB)
|
||||||
pub type ColorRange = ((u8, u8, u8), (u8, u8, u8));
|
pub type ColorRange = ((u8, u8, u8), (u8, u8, u8));
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
pub struct Ocr {
|
pub struct Ocr {
|
||||||
pub session: RunnableModel<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
|
pub session: RunnableModel<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
|
||||||
pub charset: Vec<String>,
|
pub model_metadata: ModelMetadata,
|
||||||
}
|
}
|
||||||
impl ModelSession for Ocr {
|
impl ModelSession for Ocr {
|
||||||
fn get_model_type(&self) -> ModelType {
|
fn get_model_type(&self) -> ModelType {
|
||||||
@@ -28,9 +28,12 @@ impl ModelSession for Ocr {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
impl Ocr {
|
impl Ocr {
|
||||||
pub fn new(model_path: String, charset: Vec<String>) -> Result<Self, anyhow::Error> {
|
pub fn new(model_path: String, model_metadata: ModelMetadata) -> Result<Self, anyhow::Error> {
|
||||||
let session = ModelLoader::load_model(&model_path)?.session;
|
let session = ModelLoader::load_model(&model_path)?.session;
|
||||||
Ok(Self { session, charset })
|
Ok(Self {
|
||||||
|
session,
|
||||||
|
model_metadata,
|
||||||
|
})
|
||||||
}
|
}
|
||||||
pub fn predict<'a>(&'a self, image: &'a DynamicImage) -> OcrBuilder<'a> {
|
pub fn predict<'a>(&'a self, image: &'a DynamicImage) -> OcrBuilder<'a> {
|
||||||
OcrBuilder::new(self, image)
|
OcrBuilder::new(self, image)
|
||||||
@@ -50,7 +53,7 @@ pub struct OcrBuilder<'a> {
|
|||||||
/// 颜色过滤:自定义RGB范围
|
/// 颜色过滤:自定义RGB范围
|
||||||
color_filter_custom_ranges: Option<Vec<ColorRange>>,
|
color_filter_custom_ranges: Option<Vec<ColorRange>>,
|
||||||
/// 字符集范围
|
/// 字符集范围
|
||||||
charset_range: Option<CharsetRange>,
|
charset_restrict: Option<CharsetRestrict>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<'a> OcrBuilder<'a> {
|
impl<'a> OcrBuilder<'a> {
|
||||||
@@ -63,7 +66,7 @@ impl<'a> OcrBuilder<'a> {
|
|||||||
probability: false,
|
probability: false,
|
||||||
color_filter_colors: None,
|
color_filter_colors: None,
|
||||||
color_filter_custom_ranges: None,
|
color_filter_custom_ranges: None,
|
||||||
charset_range: None
|
charset_restrict: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
pub fn png_fix(mut self, value: bool) -> Self {
|
pub fn png_fix(mut self, value: bool) -> Self {
|
||||||
@@ -78,21 +81,22 @@ impl<'a> OcrBuilder<'a> {
|
|||||||
self.color_filter_custom_ranges = Some(value);
|
self.color_filter_custom_ranges = Some(value);
|
||||||
self
|
self
|
||||||
}
|
}
|
||||||
pub fn charset_range(mut self, range: CharsetRange) -> Self {
|
pub fn charset_restrict(mut self, restrict: CharsetRestrict) -> Self {
|
||||||
self.charset_range = Some(range);
|
self.charset_restrict = Some(restrict);
|
||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn run(&self) -> Result<String, anyhow::Error> {
|
pub fn run(&self) -> anyhow::Result<String> {
|
||||||
let tensor = self.preprocess_image(self.image, self.png_fix)?;
|
let tensor = self.preprocess_image(self.image, self.png_fix)?;
|
||||||
//
|
|
||||||
// let result = self.session.run(tvec!(tensor.into()))?;
|
let raw_tensor = self.inference(tensor)?;
|
||||||
// // 3. 解析结果
|
let raw_indices = self.extract_indices_from_tensor(&raw_tensor)?;
|
||||||
// // let output = result[0].to_array_view::<i64>()?;
|
|
||||||
let output = self.inference(tensor)?;
|
// 步骤 2: 将索引切片 `&[i64]` 传给解码器进行 CTC 去重和字符映射
|
||||||
let output2 = self.process_text_output(&output)?;
|
let final_text = self.ctc_decode_to_string(&raw_indices);
|
||||||
Ok(self.ctc_decode_indices(&output2))
|
|
||||||
// Ok("ocr result".to_string())
|
println!("最终识别出的验证码是: {}", final_text);
|
||||||
|
Ok(final_text)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 对应 Python 的 _preprocess_image
|
/// 对应 Python 的 _preprocess_image
|
||||||
@@ -136,8 +140,9 @@ impl<'a> OcrBuilder<'a> {
|
|||||||
println!("模型输出原始数据: {:?}", result);
|
println!("模型输出原始数据: {:?}", result);
|
||||||
Ok(result.remove(0).into_tensor())
|
Ok(result.remove(0).into_tensor())
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 核心解析逻辑:将模型输出的各种维度/类型的 Tensor 转为字符索引序列
|
/// 核心解析逻辑:将模型输出的各种维度/类型的 Tensor 转为字符索引序列
|
||||||
fn process_text_output(&self, raw_tensor: &Tensor) -> anyhow::Result<Vec<i64>> {
|
fn extract_indices_from_tensor(&self, raw_tensor: &Tensor) -> anyhow::Result<Vec<i64>> {
|
||||||
let shape = raw_tensor.shape();
|
let shape = raw_tensor.shape();
|
||||||
println!("模型输出shape数据: {:?}", shape);
|
println!("模型输出shape数据: {:?}", shape);
|
||||||
let datum_type = raw_tensor.datum_type();
|
let datum_type = raw_tensor.datum_type();
|
||||||
@@ -172,6 +177,9 @@ impl<'a> OcrBuilder<'a> {
|
|||||||
// 形状: [Steps, Classes] -> 已经剥离了 Batch 维度
|
// 形状: [Steps, Classes] -> 已经剥离了 Batch 维度
|
||||||
(shape[0], shape[1], view.into_dyn())
|
(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)),
|
_ => return Err(anyhow::anyhow!("不支持的输出维度: {:?}", shape)),
|
||||||
};
|
};
|
||||||
let array_2d = data_view.to_shape((steps, classes))?;
|
let array_2d = data_view.to_shape((steps, classes))?;
|
||||||
@@ -197,29 +205,107 @@ impl<'a> OcrBuilder<'a> {
|
|||||||
)),
|
)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
fn ctc_decode_indices(&self, predicted_indices: &[i64]) -> String {
|
/// 获取有效字符索引列表 (用于外部验证或过滤)
|
||||||
|
pub fn get_valid_indices(&self) -> HashSet<usize> {
|
||||||
|
let (_, valid_indices) = self.valid_indices();
|
||||||
|
valid_indices
|
||||||
|
}
|
||||||
|
|
||||||
|
fn valid_indices(&self) -> (bool, HashSet<usize>) {
|
||||||
|
let charset = &self.ocr.model_metadata.charset;
|
||||||
|
let tokens = &charset.tokens;
|
||||||
|
/// 根据传入的 CharsetRestrict 枚举策略,动态更新有效索引
|
||||||
|
// 1. 🧠 零克隆防御战:在局部判断并动态构建专属的白名单判定表
|
||||||
|
let mut valid_indices = HashSet::new();
|
||||||
|
let mut has_any_match = false;
|
||||||
|
if let Some(ref policy) = self.charset_restrict {
|
||||||
|
// 🧠 性能精算:根据限制策略的类型,智能分配合适的初始容量,1个字节都不浪费!
|
||||||
|
let estimated_capacity = match policy {
|
||||||
|
CharsetRestrict::Digit => 16,
|
||||||
|
CharsetRestrict::Lowercase | CharsetRestrict::Uppercase => 32,
|
||||||
|
CharsetRestrict::CustomList(vec) => vec.len() + 1, // 动态匹配列表大小
|
||||||
|
_ => 128, // 组合子(Or)等复杂情况,给个 128 黄金保底值
|
||||||
|
};
|
||||||
|
// 🚀 精准开辟内存,完美避开 8120 个槽位的巨大空置浪费
|
||||||
|
valid_indices = HashSet::with_capacity(estimated_capacity);
|
||||||
|
|
||||||
|
for (idx, token) in tokens.iter().enumerate() {
|
||||||
|
let token_str = token.as_ref();
|
||||||
|
// CTC Blank 空字符串无条件放行,其余交给超高性能的 matches
|
||||||
|
if token_str.is_empty() {
|
||||||
|
valid_indices.insert(idx);
|
||||||
|
} else if policy.matches(token_str) {
|
||||||
|
valid_indices.insert(idx);
|
||||||
|
has_any_match = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 终极防御:如果除了 Blank 之外,没有任何一个字符被匹配到
|
||||||
|
if !has_any_match {
|
||||||
|
valid_indices = (0..tokens.len()).collect();
|
||||||
|
println!("警告:当前限制策略与模型字符集完全没有交集!已自动恢复全量识别。");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
(has_any_match, valid_indices)
|
||||||
|
}
|
||||||
|
/// 🌟 【按需延迟打印】:当用户真的需要“知道当前有哪些限制字符”时,一秒反查并打印
|
||||||
|
/// 这里的 &str 完美借用了自 tokens,依然是彻底的零拷贝!
|
||||||
|
pub fn get_valid_tokens(&self) -> Vec<&str> {
|
||||||
|
let charset = &self.ocr.model_metadata.charset;
|
||||||
|
let tokens = &charset.tokens;
|
||||||
|
self.get_valid_indices()
|
||||||
|
.iter()
|
||||||
|
.map(|&idx| tokens[idx].as_ref())
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
pub fn valid_size(&self) -> usize {
|
||||||
|
self.get_valid_indices().len()
|
||||||
|
}
|
||||||
|
fn ctc_decode_to_string(&self, predicted_indices: &[i64]) -> String {
|
||||||
println!("indices模型输出原始数据: {:?}", predicted_indices);
|
println!("indices模型输出原始数据: {:?}", predicted_indices);
|
||||||
|
let charset = &self.ocr.model_metadata.charset;
|
||||||
|
let tokens = &charset.tokens;
|
||||||
|
// let valid_indices = &charset.valid_indices;
|
||||||
|
|
||||||
|
let (has_any_match, valid_indices) = self.valid_indices();
|
||||||
|
|
||||||
// 对应 _ctc_decode_indices 的逻辑:去重、去 blank (0)
|
// 对应 _ctc_decode_indices 的逻辑:去重、去 blank (0)
|
||||||
let mut res = String::new();
|
let mut res = String::new();
|
||||||
let mut prev_idx: i64 = -1;
|
let mut prev_idx: i64 = -1;
|
||||||
|
|
||||||
for &idx in predicted_indices {
|
for &idx in predicted_indices {
|
||||||
// 1. 跳过连续重复的索引
|
// 1. CTC 去重:如果是连续重复的,直接跳过
|
||||||
// 2. 跳过 blank 字符 (假设索引 0 是 blank)
|
if idx == prev_idx {
|
||||||
if idx != prev_idx && idx != 0 {
|
continue;
|
||||||
if let Ok(u_idx) = usize::try_from(idx) {
|
}
|
||||||
if let Some(char_str) = self.ocr.charset.get(u_idx) {
|
// 【关键核心】只要不是连续重复,立刻更新 prev_idx 状态,绝对不能被后续的过滤短路!
|
||||||
res.push_str(char_str);
|
prev_idx = idx;
|
||||||
} else {
|
|
||||||
// 保护逻辑:如果模型预测的索引超出了字符集范围
|
// 2. CTC 过滤 Blank (0)
|
||||||
eprintln!("警告: 预测索引 {} 超出字符集范围", u_idx);
|
if idx == 0 {
|
||||||
}
|
continue;
|
||||||
|
}
|
||||||
|
// 3. 类型安全转换
|
||||||
|
let u_idx = match usize::try_from(idx) {
|
||||||
|
Ok(u) => u,
|
||||||
|
Err(_) => continue,
|
||||||
|
};
|
||||||
|
|
||||||
|
// 4. 终极性能:既然你的有效索引库必然有全量或部分数据,这里直接进行 O(1) 包含校验
|
||||||
|
// 注意:去掉原本错误的 `!`
|
||||||
|
if has_any_match {
|
||||||
|
if !valid_indices.contains(&u_idx) {
|
||||||
|
continue; // 不在有效字符集内,安全跳过
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
prev_idx = idx;
|
|
||||||
|
// 5. 字符映射
|
||||||
|
if let Some(char_str) = tokens.get(u_idx) {
|
||||||
|
res.push_str(char_str);
|
||||||
|
} else {
|
||||||
|
eprintln!("警告: 预测索引 {} 超出字符集范围", u_idx);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
println!("最终识别出的验证码是: {}", res);
|
|
||||||
res
|
res
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -66,7 +66,7 @@ fn test_full_classification() {
|
|||||||
let ocr = DdddOcrBuilder::new().build().expect("模型加载失败");
|
let ocr = DdddOcrBuilder::new().build().expect("模型加载失败");
|
||||||
|
|
||||||
// 2. 加载测试图片
|
// 2. 加载测试图片
|
||||||
let img = image::open("samples/code2.png").expect("测试图片不存在");
|
let img = image::open("samples/code3.png").expect("测试图片不存在");
|
||||||
|
|
||||||
// 3. 执行识别
|
// 3. 执行识别
|
||||||
let result = ocr.classification(&img).expect("识别过程出错");
|
let result = ocr.classification(&img).expect("识别过程出错");
|
||||||
|
|||||||
Reference in New Issue
Block a user