use std::collections::VecDeque;
use std::error::Error;
use std::path::PathBuf;
use ocrs::{ImageSource, OcrEngine, OcrEngineParams};
use rten::Model;
#[allow(unused)]
use rten_tensor::prelude::*;
struct Args {
image: String,
}
fn parse_args() -> Result<Args, lexopt::Error> {
use lexopt::prelude::*;
let mut values = VecDeque::new();
let mut parser = lexopt::Parser::from_env();
while let Some(arg) = parser.next()? {
match arg {
Value(val) => values.push_back(val.string()?),
Long("help") => {
println!(
"Usage: {bin_name} <image>",
bin_name = parser.bin_name().unwrap_or("hello_ocrs")
);
std::process::exit(0);
}
_ => return Err(arg.unexpected()),
}
}
let image = values.pop_front().ok_or("missing `image` arg")?;
Ok(Args { image })
}
fn file_path(path: &str) -> PathBuf {
let mut abs_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
abs_path.push(path);
abs_path
}
fn main() -> Result<(), Box<dyn Error>> {
let args = parse_args()?;
let detection_model_path = file_path("examples/text-detection.rten");
let rec_model_path = file_path("examples/text-recognition.rten");
let detection_model = Model::load_file(detection_model_path)?;
let recognition_model = Model::load_file(rec_model_path)?;
let engine = OcrEngine::new(OcrEngineParams {
detection_model: Some(detection_model),
recognition_model: Some(recognition_model),
..Default::default()
})?;
let img = image::open(&args.image).map(|image| image.into_rgb8())?;
let img_source = ImageSource::from_bytes(img.as_raw(), img.dimensions())?;
let ocr_input = engine.prepare_input(img_source)?;
let word_rects = engine.detect_words(&ocr_input)?;
let line_rects = engine.find_text_lines(&ocr_input, &word_rects);
let line_texts = engine.recognize_text(&ocr_input, &line_rects)?;
for line in line_texts
.iter()
.flatten()
.filter(|l| l.to_string().len() > 1)
{
println!("{}", line);
}
Ok(())
}