use std::path::Path;
use crate::error::Result;
use crate::preprocess;
use crate::recognize::Recognizer;
use crate::segment;
use crate::types::{BoundingBox, OcrResult, OcrWord};
#[derive(Debug, Clone)]
pub struct OcrEngine {
pub language: String,
pub preprocessing: bool,
recognizer: Recognizer,
}
impl Default for OcrEngine {
fn default() -> Self {
Self {
language: "eng".to_string(),
preprocessing: false,
recognizer: Recognizer::new(),
}
}
}
impl OcrEngine {
pub fn new() -> Self {
Self::default()
}
pub fn language(mut self, lang: impl Into<String>) -> Self {
self.language = lang.into();
self
}
pub fn preprocessing(mut self, enable: bool) -> Self {
self.preprocessing = enable;
self
}
pub fn recognize_file(&self, path: &Path) -> Result<OcrResult> {
let img = image::open(path)?;
self.run(&img)
}
pub fn recognize_bytes(&self, data: &[u8]) -> Result<OcrResult> {
let img = image::load_from_memory(data)?;
self.run(&img)
}
pub fn recognize_image(&self, img: &image::DynamicImage) -> Result<OcrResult> {
self.run(img)
}
fn run(&self, img: &image::DynamicImage) -> Result<OcrResult> {
let processed = if self.preprocessing {
preprocess::preprocess(img)
} else {
img.clone()
};
let mut binary = preprocess::binarize(&processed);
preprocess::clean_noise(&mut binary, 3);
let height = binary.len();
let width = binary.first().map(|r| r.len()).unwrap_or(0);
if width == 0 || height == 0 {
return Ok(OcrResult {
text: String::new(),
words: Vec::new(),
confidence: 0.0,
});
}
let lines = segment::find_lines(&binary, width, height);
if lines.is_empty() {
return Ok(OcrResult {
text: String::new(),
words: Vec::new(),
confidence: 0.0,
});
}
let mut text_parts: Vec<String> = Vec::new();
let mut all_words: Vec<OcrWord> = Vec::new();
let mut total_conf = 0.0_f32;
let mut char_count = 0u32;
for line in &lines {
let chars = segment::find_chars(&binary, line, width);
if chars.is_empty() {
continue;
}
let gaps = segment::compute_gaps(&chars);
let space_threshold = space_gap_threshold(&gaps);
let mut line_text = String::new();
let mut cur_word_chars: Vec<(char, f32, BoundingBox)> = Vec::new();
for (i, seg) in chars.iter().enumerate() {
if i > 0
&& !gaps.is_empty()
&& i - 1 < gaps.len()
&& gaps[i - 1] as f64 > space_threshold
&& !line_text.is_empty()
{
line_text.push(' ');
if !cur_word_chars.is_empty() {
all_words.push(build_word(&cur_word_chars));
cur_word_chars.clear();
}
}
let char_pixels = extract_char(&binary, seg, line);
let (ch, conf) = self.recognizer.recognize(&char_pixels);
line_text.push(ch);
total_conf += conf;
char_count += 1;
cur_word_chars.push((ch, conf, BoundingBox {
x: seg.x as u32,
y: line.y_start as u32,
width: seg.width as u32,
height: (line.y_end - line.y_start) as u32,
}));
}
if !cur_word_chars.is_empty() {
all_words.push(build_word(&cur_word_chars));
}
text_parts.push(line_text);
}
let text = text_parts.join("\n");
let confidence = if char_count > 0 {
total_conf / char_count as f32
} else {
0.0
};
Ok(OcrResult {
text,
words: all_words,
confidence,
})
}
}
fn space_gap_threshold(gaps: &[usize]) -> f64 {
if gaps.is_empty() {
return 0.0;
}
let mut sorted: Vec<usize> = gaps.to_vec();
sorted.sort();
let mid = sorted.len() / 2;
let median = if sorted.len().is_multiple_of(2) {
(sorted[mid - 1] + sorted[mid]) as f64 / 2.0
} else {
sorted[mid] as f64
};
median * 3.0
}
fn extract_char(
grid: &[Vec<bool>],
seg: &segment::CharSeg,
line: &segment::LineInfo,
) -> Vec<Vec<bool>> {
grid[line.y_start..line.y_end]
.iter()
.map(|row| row[seg.x..seg.x + seg.width].to_vec())
.collect()
}
fn build_word(chars: &[(char, f32, BoundingBox)]) -> OcrWord {
let text: String = chars.iter().map(|(c, _, _)| *c).collect();
let confidence: f32 = chars.iter().map(|(_, c, _)| c).sum::<f32>() / chars.len() as f32;
let min_x = chars.iter().map(|(_, _, bb)| bb.x).min().unwrap_or(0);
let min_y = chars.iter().map(|(_, _, bb)| bb.y).min().unwrap_or(0);
let max_x = chars.iter().map(|(_, _, bb)| bb.x + bb.width).max().unwrap_or(0);
let max_y = chars.iter().map(|(_, _, bb)| bb.y + bb.height).max().unwrap_or(0);
OcrWord {
text,
confidence,
bounding_box: BoundingBox {
x: min_x,
y: min_y,
width: max_x - min_x,
height: max_y - min_y,
},
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::render;
#[test]
fn recognize_rendered_hello() {
let img = render::render_text("HELLO", 6);
let engine = OcrEngine::new();
let result = engine.run(&img).unwrap();
assert_eq!(result.text, "HELLO");
}
#[test]
fn recognize_rendered_digits() {
let img = render::render_text("0123456789", 6);
let engine = OcrEngine::new();
let result = engine.run(&img).unwrap();
assert_eq!(result.text, "0123456789");
}
#[test]
fn recognize_rendered_multiline() {
let img = render::render_text("ABC\nDEF", 6);
let engine = OcrEngine::new();
let result = engine.run(&img).unwrap();
assert_eq!(result.text, "ABC\nDEF");
}
#[test]
fn recognize_rendered_with_spaces() {
let img = render::render_text("HI OK", 6);
let engine = OcrEngine::new();
let result = engine.run(&img).unwrap();
assert_eq!(result.text, "HI OK");
}
#[test]
fn recognize_lowercase() {
let img = render::render_text("abc", 6);
let engine = OcrEngine::new();
let result = engine.run(&img).unwrap();
assert_eq!(result.text, "abc");
}
#[test]
fn empty_image() {
let img = image::DynamicImage::ImageLuma8(
image::GrayImage::from_pixel(100, 100, image::Luma([255u8])),
);
let engine = OcrEngine::new();
let result = engine.run(&img).unwrap();
assert!(result.text.is_empty());
assert!(result.words.is_empty());
}
}