75 lines
2.1 KiB
Rust
75 lines
2.1 KiB
Rust
use crate::ocr::executor::Ocr;
|
|
// use ddddocr_tract::session::OcrSession;
|
|
use crate::ocr::color_filter::ColorFilter;
|
|
use crate::ocr::token_filter::TokenFilter;
|
|
use crate::OcrEngine;
|
|
|
|
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
|
|
}
|
|
pub fn build(self, session: &dyn OcrEngine) -> Ocr<'_> {
|
|
// 1. 原地解析颜色过滤器
|
|
let final_color_ranges = match &self.color_filter {
|
|
Some(filter) => filter.collect_to_vec(),
|
|
None => Ok(None),
|
|
};
|
|
// 2. 原地解析字符集过滤
|
|
let tokens = &session.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 {
|
|
session,
|
|
png_fix: self.png_fix, // 原地解构出来
|
|
probability: self.probability,
|
|
final_color_ranges,
|
|
final_charset_indices,
|
|
}
|
|
}
|
|
}
|