57 lines
1.3 KiB
Rust
57 lines
1.3 KiB
Rust
use std::fs;
|
|
use std::error::Error;
|
|
use std::time::Instant;
|
|
|
|
use rten_tensor::{NdTensor, AsView};
|
|
use rten::Model;
|
|
use ocrs::{OcrEngine, OcrEngineParams};
|
|
|
|
pub fn print_words(image: NdTensor<f32, 3>) -> Result<(), Box<dyn Error>> {
|
|
|
|
let begin = Instant::now();
|
|
|
|
println!("{:#?}", begin.elapsed());
|
|
let begin = Instant::now();
|
|
ocr(image).unwrap();
|
|
println!("{:#?}", begin.elapsed());
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn ocr(data: NdTensor<f32, 3>) -> Result<(), Box<dyn Error>> {
|
|
|
|
let detection_model_data = fs::read("text-detection.rten")?;
|
|
let rec_model_data = fs::read("text-recognition.rten")?;
|
|
|
|
let detection_model = Model::load(&detection_model_data)?;
|
|
let rec_model = Model::load(&rec_model_data)?;
|
|
|
|
let ocr_engine = OcrEngine::new(OcrEngineParams {
|
|
detection_model: Some(detection_model),
|
|
recognition_model: Some(rec_model),
|
|
..Default::default()
|
|
})?;
|
|
|
|
let input = ocr_engine.prepare_input(data.view())?;
|
|
|
|
let word_rects = ocr_engine.detect_words(&input)?;
|
|
|
|
let line_rects = ocr_engine.find_text_lines(&input, &word_rects);
|
|
|
|
let line_texts = ocr_engine.recognize_text(&input, &line_rects)?;
|
|
|
|
for line in line_texts
|
|
.iter()
|
|
.flatten()
|
|
.filter(|l| l.to_string().len() > 1)
|
|
{
|
|
println!("{}", line);
|
|
}
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
|
|
}
|