use std::{
fs,
path::Path,
sync::{OnceLock, RwLock},
};
use regex::RegexBuilder;
use serde::{Deserialize, Serialize};
use super::{
core::{
default_ocr_case_sensitive, default_ocr_thread_count, default_ocr_threshold, MatchPoint,
MatchResult, Roi,
},
error::{McvError, Result},
frame::{RgbaFrame, VisionTemplate},
};
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct OcrModelFiles {
pub model_dir: String,
pub det_model: String,
pub rec_model: String,
pub charset: String,
}
impl OcrModelFiles {
pub fn new(
model_dir: impl AsRef<Path>,
det_model: impl Into<String>,
rec_model: impl Into<String>,
charset: impl Into<String>,
) -> Self {
Self {
model_dir: model_dir.as_ref().to_string_lossy().into_owned(),
det_model: det_model.into(),
rec_model: rec_model.into(),
charset: charset.into(),
}
}
pub fn resolve(&self) -> Result<ResolvedOcrModelFiles> {
let model_dir = std::path::PathBuf::from(&self.model_dir);
let det_model = model_dir.join(&self.det_model);
let rec_model = model_dir.join(&self.rec_model);
let charset = model_dir.join(&self.charset);
for path in [&det_model, &rec_model, &charset] {
if !path.is_file() {
return Err(McvError::MissingOcrModel(path.display().to_string()));
}
}
Ok(ResolvedOcrModelFiles {
model_dir,
det_model,
rec_model,
charset,
})
}
}
static DEFAULT_OCR_MODELS: OnceLock<RwLock<Option<OcrModelFiles>>> = OnceLock::new();
fn default_ocr_models_store() -> &'static RwLock<Option<OcrModelFiles>> {
DEFAULT_OCR_MODELS.get_or_init(|| RwLock::new(None))
}
pub fn set_default_ocr_models(models: OcrModelFiles) -> Result<()> {
models.resolve()?;
let mut current = default_ocr_models_store()
.write()
.unwrap_or_else(|poisoned| poisoned.into_inner());
*current = Some(models);
Ok(())
}
pub fn default_ocr_models() -> Option<OcrModelFiles> {
default_ocr_models_store()
.read()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.clone()
}
#[derive(Debug, Clone)]
pub struct ResolvedOcrModelFiles {
pub model_dir: std::path::PathBuf,
pub det_model: std::path::PathBuf,
pub rec_model: std::path::PathBuf,
pub charset: std::path::PathBuf,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct OcrEngineOptions {
pub model: OcrModelFiles,
pub threshold: Option<f32>,
pub thread_count: Option<i32>,
pub roi: Option<Roi>,
}
pub type OcrPoint = MatchPoint;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct OcrText {
pub text: String,
pub confidence: f32,
pub bbox: Roi,
pub points: Option<[OcrPoint; 4]>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct OcrReport {
pub text: String,
pub results: Vec<OcrText>,
}
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum OcrMatchMode {
#[default]
Contains,
Exact,
Regex,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct OcrTemplateOptions {
pub pattern: Option<String>,
pub match_mode: Option<OcrMatchMode>,
pub case_sensitive: Option<bool>,
pub max_count: Option<usize>,
}
impl Default for OcrTemplateOptions {
fn default() -> Self {
Self {
pattern: None,
match_mode: Some(OcrMatchMode::Contains),
case_sensitive: None,
max_count: Some(1),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct OcrFindOptions {
#[serde(flatten)]
pub ocr: OcrEngineOptions,
pub pattern: Option<String>,
pub match_mode: Option<OcrMatchMode>,
pub case_sensitive: Option<bool>,
pub max_count: Option<usize>,
}
impl OcrFindOptions {
pub fn matcher_options(&self) -> OcrTemplateOptions {
OcrTemplateOptions {
pattern: self.pattern.clone(),
match_mode: self.match_mode,
case_sensitive: self.case_sensitive,
max_count: self.max_count,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct OcrMatchResult {
pub top_left: (i32, i32),
pub bottom_right: (i32, i32),
pub center: (i32, i32),
pub score: f64,
pub text: String,
pub confidence: f32,
pub bbox: Roi,
pub points: Option<[OcrPoint; 4]>,
}
impl OcrMatchResult {
pub fn as_match_result(&self) -> MatchResult {
MatchResult::from_box(self.bbox.x, self.bbox.y, self.bbox.width, self.bbox.height)
}
}
impl From<OcrText> for OcrMatchResult {
fn from(value: OcrText) -> Self {
let result = MatchResult::from_box(
value.bbox.x,
value.bbox.y,
value.bbox.width,
value.bbox.height,
);
let top_left = result.top_left();
let bottom_right = result.bottom_right();
let center = result.center();
Self {
top_left,
bottom_right,
center,
score: f64::from(value.confidence).clamp(0.0, 1.0),
text: value.text,
confidence: value.confidence,
bbox: value.bbox,
points: value.points,
}
}
}
pub struct OcrTextMatcher {
pattern: Option<String>,
mode: OcrMatchMode,
case_sensitive: bool,
regex: Option<regex::Regex>,
}
impl OcrTextMatcher {
pub fn new(options: OcrTemplateOptions) -> Result<Self> {
let mode = options.match_mode.unwrap_or_default();
let case_sensitive = options
.case_sensitive
.unwrap_or_else(default_ocr_case_sensitive);
let regex = match (&options.pattern, mode) {
(Some(pattern), OcrMatchMode::Regex) => Some(
RegexBuilder::new(pattern)
.case_insensitive(!case_sensitive)
.build()
.map_err(|err| McvError::Ocr(format!("invalid OCR regex: {err}")))?,
),
_ => None,
};
Ok(Self {
pattern: options.pattern,
mode,
case_sensitive,
regex,
})
}
pub fn is_match(&self, text: &str) -> bool {
let Some(pattern) = self.pattern.as_deref() else {
return true;
};
match self.mode {
OcrMatchMode::Contains => {
if self.case_sensitive {
text.contains(pattern)
} else {
text.to_lowercase().contains(&pattern.to_lowercase())
}
}
OcrMatchMode::Exact => {
if self.case_sensitive {
text == pattern
} else {
text.eq_ignore_ascii_case(pattern)
|| text.to_lowercase() == pattern.to_lowercase()
}
}
OcrMatchMode::Regex => self
.regex
.as_ref()
.is_some_and(|regex| regex.is_match(text)),
}
}
pub fn filter(&self, results: Vec<OcrText>, max_count: Option<usize>) -> Vec<OcrMatchResult> {
if max_count == Some(0) {
return Vec::new();
}
let mut matches = Vec::new();
for item in results {
if self.is_match(&item.text) {
matches.push(OcrMatchResult::from(item));
if max_count.is_some_and(|max_count| matches.len() >= max_count) {
break;
}
}
}
matches
}
}
pub struct OcrTemplate {
options: OcrFindOptions,
engine: Option<ocr_rs::OcrEngine>,
}
#[derive(Debug, Clone)]
pub struct OcrTemplateBuilder {
options: OcrFindOptions,
}
impl OcrTemplate {
pub fn new(pattern: impl Into<String>) -> Result<Self> {
let pattern = pattern.into();
if pattern.trim().is_empty() {
return Err(McvError::Ocr(
"OCR template text must not be empty".to_string(),
));
}
let models = default_ocr_models().ok_or(McvError::MissingDefaultOcrModels)?;
Self::builder(models).contains(pattern).open()
}
pub fn with_options(options: OcrFindOptions) -> Self {
Self {
options,
engine: None,
}
}
pub fn builder(model: OcrModelFiles) -> OcrTemplateBuilder {
let matcher = OcrTemplateOptions::default();
OcrTemplateBuilder {
options: OcrFindOptions {
ocr: OcrEngineOptions {
model,
threshold: None,
thread_count: None,
roi: None,
},
pattern: matcher.pattern,
match_mode: matcher.match_mode,
case_sensitive: matcher.case_sensitive,
max_count: matcher.max_count,
},
}
}
pub fn recognize_path(&mut self, image_path: impl AsRef<Path>) -> Result<OcrReport> {
let image_bytes = fs::read(image_path.as_ref())?;
let image = image::load_from_memory(&image_bytes)?;
self.recognize_image(&image)
}
pub fn recognize_image(&mut self, image: &::image::DynamicImage) -> Result<OcrReport> {
let options = self.options.ocr.clone();
let engine = self.engine()?;
recognize_ocr_image_with_engine(engine, image, &options)
}
pub fn find_path(&mut self, image_path: impl AsRef<Path>) -> Result<Option<OcrMatchResult>> {
let results = self.find_all_path(image_path)?;
Ok(results.into_iter().next())
}
pub fn find_image(&mut self, image: &::image::DynamicImage) -> Result<Option<OcrMatchResult>> {
let results = self.find_all_image(image)?;
Ok(results.into_iter().next())
}
pub fn find(&mut self, image: &mut RgbaFrame<'_>) -> Result<Option<MatchResult>> {
self.find_with_roi(image, None)
}
pub fn find_with_roi(
&mut self,
image: &mut RgbaFrame<'_>,
roi: Option<Roi>,
) -> Result<Option<MatchResult>> {
let mut options = self.options.clone();
options.ocr.roi = roi.or(options.ocr.roi);
let engine = self.engine()?;
let report = recognize_ocr_image_with_engine(engine, image.dynamic_image()?, &options.ocr)?;
let matcher = OcrTextMatcher::new(options.matcher_options())?;
Ok(matcher
.filter(report.results, options.max_count)
.into_iter()
.next()
.map(|result| result.as_match_result()))
}
pub fn find_all_path(&mut self, image_path: impl AsRef<Path>) -> Result<Vec<OcrMatchResult>> {
let report = self.recognize_path(image_path)?;
let matcher = OcrTextMatcher::new(self.options.matcher_options())?;
Ok(matcher.filter(report.results, self.options.max_count))
}
pub fn find_all_image(&mut self, image: &::image::DynamicImage) -> Result<Vec<OcrMatchResult>> {
let report = self.recognize_image(image)?;
let matcher = OcrTextMatcher::new(self.options.matcher_options())?;
Ok(matcher.filter(report.results, self.options.max_count))
}
fn engine(&mut self) -> Result<&ocr_rs::OcrEngine> {
if self.engine.is_none() {
self.engine = Some(create_ocr_engine(&self.options.ocr)?);
}
Ok(self.engine.as_ref().expect("OCR engine initialized"))
}
}
impl OcrTemplateBuilder {
pub fn model(mut self, model: OcrModelFiles) -> Self {
self.options.ocr.model = model;
self
}
pub fn threshold(mut self, threshold: f32) -> Result<Self> {
resolve_ocr_threshold(Some(threshold))?;
self.options.ocr.threshold = Some(threshold);
Ok(self)
}
pub fn threads(mut self, thread_count: i32) -> Result<Self> {
if thread_count <= 0 {
return Err(McvError::InvalidThreadCount);
}
self.options.ocr.thread_count = Some(thread_count);
Ok(self)
}
pub fn roi(mut self, roi: Roi) -> Self {
self.options.ocr.roi = Some(roi);
self
}
pub fn optional_roi(mut self, roi: Option<Roi>) -> Self {
self.options.ocr.roi = roi;
self
}
pub fn any_text(mut self) -> Self {
self.options.pattern = None;
self
}
pub fn pattern(mut self, pattern: impl Into<String>) -> Self {
self.options.pattern = Some(pattern.into());
self
}
pub fn match_mode(mut self, match_mode: OcrMatchMode) -> Self {
self.options.match_mode = Some(match_mode);
self
}
pub fn contains(self, pattern: impl Into<String>) -> Self {
self.pattern(pattern).match_mode(OcrMatchMode::Contains)
}
pub fn exact(self, pattern: impl Into<String>) -> Self {
self.pattern(pattern).match_mode(OcrMatchMode::Exact)
}
pub fn regex(self, pattern: impl Into<String>) -> Result<Self> {
let builder = self.pattern(pattern).match_mode(OcrMatchMode::Regex);
OcrTextMatcher::new(builder.options.matcher_options())?;
Ok(builder)
}
pub fn case_sensitive(mut self, case_sensitive: bool) -> Self {
self.options.case_sensitive = Some(case_sensitive);
self
}
pub fn case_insensitive(self) -> Self {
self.case_sensitive(false)
}
pub fn max_count(mut self, max_count: usize) -> Self {
self.options.max_count = Some(max_count);
self
}
pub fn all_matches(mut self) -> Self {
self.options.max_count = None;
self
}
pub fn matcher_options(mut self, options: OcrTemplateOptions) -> Self {
self.options.pattern = options.pattern;
self.options.match_mode = options.match_mode;
self.options.case_sensitive = options.case_sensitive;
self.options.max_count = options.max_count;
self
}
pub fn open(self) -> Result<OcrTemplate> {
self.options.ocr.model.resolve()?;
resolve_ocr_threshold(self.options.ocr.threshold)?;
resolve_ocr_thread_count(self.options.ocr.thread_count)?;
if self
.options
.ocr
.roi
.is_some_and(|roi| roi.width <= 0 || roi.height <= 0)
{
return Err(McvError::InvalidRoi(
"OCR ROI width and height must be positive".to_string(),
));
}
OcrTextMatcher::new(self.options.matcher_options())?;
Ok(OcrTemplate::with_options(self.options))
}
}
impl VisionTemplate for OcrTemplate {
fn find_with_roi(
&mut self,
image: &mut RgbaFrame<'_>,
roi: Option<Roi>,
) -> Result<Option<MatchResult>> {
OcrTemplate::find_with_roi(self, image, roi)
}
}
pub fn find_ocr_text_path(
image_path: impl AsRef<Path>,
options: OcrFindOptions,
) -> Result<Vec<OcrMatchResult>> {
OcrTemplate::with_options(options).find_all_path(image_path)
}
pub fn find_ocr_text_image(
image: &::image::DynamicImage,
options: OcrFindOptions,
) -> Result<Vec<OcrMatchResult>> {
OcrTemplate::with_options(options).find_all_image(image)
}
pub fn recognize_ocr_path(
image_path: impl AsRef<Path>,
options: OcrEngineOptions,
) -> Result<OcrReport> {
let image_bytes = fs::read(image_path.as_ref())?;
let image = image::load_from_memory(&image_bytes)?;
recognize_ocr_image(&image, options)
}
pub fn recognize_ocr_image(
image: &::image::DynamicImage,
options: OcrEngineOptions,
) -> Result<OcrReport> {
let engine = create_ocr_engine(&options)?;
recognize_ocr_image_with_engine(&engine, image, &options)
}
fn create_ocr_engine(options: &OcrEngineOptions) -> Result<ocr_rs::OcrEngine> {
let model = options.model.resolve()?;
preload_mnn_runtime()?;
let threshold = resolve_ocr_threshold(options.threshold)?;
let thread_count = resolve_ocr_thread_count(options.thread_count)?;
let config = ocr_rs::OcrEngineConfig::new()
.with_threads(thread_count)
.with_min_result_confidence(threshold);
ocr_rs::OcrEngine::new(
model.det_model,
model.rec_model,
model.charset,
Some(config),
)
.map_err(|err| McvError::Ocr(err.to_string()))
}
fn resolve_ocr_threshold(threshold: Option<f32>) -> Result<f32> {
let threshold = threshold.unwrap_or_else(default_ocr_threshold);
super::core::validate_threshold_value(f64::from(threshold))?;
Ok(threshold)
}
fn resolve_ocr_thread_count(thread_count: Option<i32>) -> Result<i32> {
let thread_count = thread_count.unwrap_or_else(default_ocr_thread_count);
if thread_count <= 0 {
return Err(McvError::InvalidThreadCount);
}
Ok(thread_count)
}
fn recognize_ocr_image_with_engine(
engine: &ocr_rs::OcrEngine,
image: &::image::DynamicImage,
options: &OcrEngineOptions,
) -> Result<OcrReport> {
let cropped;
let (offset_x, offset_y, image) = if let Some(roi) = options.roi {
let image_width = image.width() as i32;
let image_height = image.height() as i32;
let roi = roi.clamp(image_width, image_height).ok_or_else(|| {
McvError::InvalidRoi(
"OCR ROI is outside the image or has non-positive size".to_string(),
)
})?;
cropped = image.crop_imm(
roi.x as u32,
roi.y as u32,
roi.width as u32,
roi.height as u32,
);
(roi.x, roi.y, &cropped)
} else {
(0, 0, image)
};
let raw_results = engine
.recognize(image)
.map_err(|err| McvError::Ocr(err.to_string()))?;
let mut results = Vec::with_capacity(raw_results.len());
for item in raw_results {
let rect = item.bbox.rect;
let x = rect.left() + offset_x;
let y = rect.top() + offset_y;
let points = item.bbox.points.map(|points| {
points.map(|point| OcrPoint {
x: point.x + offset_x as f32,
y: point.y + offset_y as f32,
})
});
results.push(OcrText {
text: item.text,
confidence: item.confidence,
bbox: Roi {
x,
y,
width: rect.width() as i32,
height: rect.height() as i32,
},
points,
});
}
let text = results
.iter()
.map(|result| result.text.as_str())
.collect::<Vec<_>>()
.join("\n");
Ok(OcrReport { text, results })
}
#[cfg(windows)]
fn preload_mnn_runtime() -> Result<()> {
use std::sync::OnceLock;
use std::{collections::HashSet, sync::Mutex};
static PRELOADED_DLLS: OnceLock<Mutex<HashSet<std::path::PathBuf>>> = OnceLock::new();
let Some(dll_path) = std::env::var_os("MNN_DLL_PATH")
.map(std::path::PathBuf::from)
.filter(|path| path.is_file())
else {
return Ok(());
};
let dll_path = dll_path.canonicalize()?;
let cache = PRELOADED_DLLS.get_or_init(|| Mutex::new(HashSet::new()));
{
let loaded = cache
.lock()
.map_err(|_| McvError::Ocr("MNN DLL preload cache mutex poisoned".to_string()))?;
if loaded.contains(&dll_path) {
return Ok(());
}
}
let library = unsafe { libloading::Library::new(&dll_path) }
.map_err(|err| McvError::Ocr(format!("failed to load {}: {err}", dll_path.display())))?;
std::mem::forget(library);
let mut loaded = cache
.lock()
.map_err(|_| McvError::Ocr("MNN DLL preload cache mutex poisoned".to_string()))?;
loaded.insert(dll_path);
Ok(())
}
#[cfg(not(windows))]
fn preload_mnn_runtime() -> Result<()> {
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::{SystemTime, UNIX_EPOCH};
fn ocr_text(text: &str) -> OcrText {
OcrText {
text: text.to_string(),
confidence: 0.9,
bbox: Roi::full(10, 10),
points: None,
}
}
#[test]
fn zero_max_count_returns_no_matches() -> Result<()> {
let matcher = OcrTextMatcher::new(OcrTemplateOptions::default())?;
let matches = matcher.filter(vec![ocr_text("first"), ocr_text("second")], Some(0));
assert!(matches.is_empty());
Ok(())
}
#[test]
fn explicit_threshold_must_be_finite_and_in_range() {
assert!(resolve_ocr_threshold(Some(f32::NAN)).is_err());
assert!(resolve_ocr_threshold(Some(-0.01)).is_err());
assert!(resolve_ocr_threshold(Some(1.01)).is_err());
assert_eq!(resolve_ocr_threshold(Some(0.75)).unwrap(), 0.75);
}
#[test]
fn global_case_and_thread_defaults_are_used_when_unspecified() -> Result<()> {
crate::mcv::set_default_ocr_case_sensitive(false);
let matcher = OcrTextMatcher::new(OcrTemplateOptions::default())?;
assert!(!matcher.case_sensitive);
crate::mcv::set_default_ocr_thread_count(7)?;
assert_eq!(resolve_ocr_thread_count(None)?, 7);
assert!(crate::mcv::set_default_ocr_thread_count(0).is_err());
assert_eq!(resolve_ocr_thread_count(Some(2))?, 2);
crate::mcv::set_default_ocr_case_sensitive(true);
crate::mcv::set_default_ocr_thread_count(4)?;
Ok(())
}
#[test]
fn builder_creates_template_from_custom_models() -> Result<()> {
let unique = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("system time")
.as_nanos();
let model_dir = std::env::temp_dir().join(format!(
"mcv-rs-ocr-builder-{}-{unique}",
std::process::id()
));
fs::create_dir_all(&model_dir)?;
for file_name in ["detector.mnn", "recognizer.mnn", "characters.txt"] {
fs::write(model_dir.join(file_name), [])?;
}
let models = OcrModelFiles::new(
&model_dir,
"detector.mnn",
"recognizer.mnn",
"characters.txt",
);
crate::mcv::set_default_ocr_models(models.clone())?;
let simple_template = OcrTemplate::new("Start")?;
assert_eq!(simple_template.options.pattern.as_deref(), Some("Start"));
let template = OcrTemplate::builder(models)
.contains("Start")
.case_insensitive()
.threshold(0.8)?
.threads(2)?
.max_count(3)
.open()?;
assert_eq!(template.options.pattern.as_deref(), Some("Start"));
assert!(matches!(
template.options.match_mode,
Some(OcrMatchMode::Contains)
));
assert_eq!(template.options.case_sensitive, Some(false));
assert_eq!(template.options.max_count, Some(3));
assert_eq!(template.options.ocr.threshold, Some(0.8));
assert_eq!(template.options.ocr.thread_count, Some(2));
fs::remove_dir_all(model_dir)?;
Ok(())
}
#[test]
fn simple_constructor_rejects_empty_text() {
assert!(OcrTemplate::new(" ").is_err());
}
#[test]
fn inherent_frame_find_api_is_available() {
let _find = OcrTemplate::find;
let _find_with_roi = OcrTemplate::find_with_roi;
}
#[test]
fn builder_rejects_invalid_threads_and_regex() {
let models =
OcrModelFiles::new("models", "detector.mnn", "recognizer.mnn", "characters.txt");
assert!(matches!(
OcrTemplate::builder(models.clone()).threads(0),
Err(McvError::InvalidThreadCount)
));
assert!(OcrTemplate::builder(models).regex("[").is_err());
}
}