Files
ddddocr-rs/ddddocr-core/src/ocr/builder.rs

80 lines
2.3 KiB
Rust

//! OCR 构建器。
use crate::ocr::executor::Ocr;
// use ddddocr_tract::session::OcrSession;
use crate::traits::OcrEngine;
use crate::ocr::color_filter::ColorFilter;
use crate::ocr::token_filter::TokenFilter;
/// OCR 构建器:配置识别选项后绑定引擎会话构建 [`crate::Ocr`]。
#[derive(Default)]
pub struct OcrBuilder {
/// 是否修复PNG格式问题
png_fix: bool,
/// 是否返回概率信息
probability: bool,
/// 颜色过滤:保留的颜色列表
color_filter: Option<Box<dyn ColorFilter + Send + Sync>>,
/// 字符集范围
charset_restrict: Option<Box<dyn TokenFilter + Send + Sync>>,
}
impl OcrBuilder {
// 初始化任务,设置默认参数
pub fn new() -> Self {
Self {
png_fix: false, // 默认值
probability: false,
color_filter: None,
charset_restrict: None,
}
}
pub fn png_fix(mut self, value: bool) -> Self {
self.png_fix = value;
self
}
pub fn probability(mut self, value: bool) -> Self {
self.probability = value;
self
}
pub fn color_filter<T>(mut self, filter: T) -> Self
where
T: ColorFilter + Send + Sync + 'static,
{
self.color_filter = Some(Box::new(filter));
self
}
pub fn charset_restrict<T>(mut self, restrict: T) -> Self
where
T: TokenFilter + Send + Sync + 'static,
{
self.charset_restrict = Some(Box::new(restrict));
self
}
/// 绑定引擎会话并构建 OCR 识别器。
pub fn build_with<E: OcrEngine>(self, runtime: &E) -> Ocr<'_> {
// 1. 原地解析颜色过滤器
let final_color_ranges = match &self.color_filter {
Some(filter) => filter.collect_to_vec(),
None => Ok(None),
};
// 2. 原地解析字符集过滤
let tokens = &runtime.metadata().charset.tokens;
let final_charset_indices = match &self.charset_restrict {
Some(restrict) => restrict.apply_to_charset(tokens),
None => None,
};
// Ocr::new(session, self)
Ocr {
runtime,
png_fix: self.png_fix, // 原地解构出来
probability: self.probability,
final_color_ranges,
final_charset_indices,
}
}
}