use std::path::PathBuf;
use crate::types::{ClsConfig, DetConfig, RecConfig, OrientConfig, UnwarpConfig, GlobalConfig};
#[derive(Clone, Debug)]
pub struct RustOConfig {
pub det: DetConfig,
pub rec: RecConfig,
pub global: GlobalConfig,
pub orient: Option<OrientConfig>,
pub unwarp: Option<UnwarpConfig>,
pub cls: Option<ClsConfig>,
}
impl RustOConfig {
pub fn new_ppv5<P: Into<PathBuf>>(
det_model_path: P,
rec_model_path: P,
dict_path: P,
) -> Self {
let det_path = det_model_path.into();
let rec_path = rec_model_path.into();
let dict = dict_path.into();
let det_config = DetConfig::ppv5(det_path);
let mut rec_config = RecConfig::ppv5(rec_path);
rec_config.rec_keys_path = Some(dict);
Self {
det: det_config,
rec: rec_config,
global: GlobalConfig::default(),
orient: None,
unwarp: None,
cls: None,
}
}
pub fn with_orientation<P: Into<PathBuf>>(mut self, model_path: P) -> Self {
self.orient = Some(OrientConfig::default(model_path.into()));
self.global.use_orient = true;
self
}
pub fn with_orientation_threshold<P: Into<PathBuf>>(mut self, model_path: P, threshold: f32) -> Self {
let mut config = OrientConfig::default(model_path.into());
config.confidence_threshold = threshold;
self.orient = Some(config);
self.global.use_orient = true;
self
}
pub fn with_unwarp<P: Into<PathBuf>>(mut self, model_path: P) -> Self {
self.unwarp = Some(UnwarpConfig::default(model_path.into()));
self.global.use_unwarp = true;
self
}
pub fn with_cls<P: Into<PathBuf>>(mut self, model_path: P) -> Self {
self.cls = Some(ClsConfig::default(model_path.into()));
self.global.use_cls = true;
self
}
pub fn with_debug_images(mut self, enabled: bool) -> Self {
self.global.debug_images = enabled;
self
}
pub fn with_text_score(mut self, score: f32) -> Self {
self.global.text_score = score;
self
}
pub fn with_min_height(mut self, height: f32) -> Self {
self.global.min_height = height;
self
}
pub fn with_max_side_len(mut self, len: f32) -> Self {
self.global.max_side_len = len;
self
}
pub fn with_detection(mut self, enabled: bool) -> Self {
self.global.use_det = enabled;
self
}
pub fn with_recognition(mut self, enabled: bool) -> Self {
self.global.use_rec = enabled;
self
}
}