df_ocr_switcher/
pipeline.rs1use std::path::Path;
10
11use image::RgbImage;
12use ppocr_rs::LayoutAnalyzer;
13
14use crate::engine::{OcrEngine, OcrPageResult};
15use crate::error::{Error, Result};
16use crate::input::{load_image, load_tiff_pages};
17use crate::output::to_markdown;
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
21pub enum OutputFormat {
22 #[default]
23 Markdown,
24 Json,
25 }
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30enum InputKind { Image, Tiff }
31
32pub struct DocPipeline {
35 engine: Box<dyn OcrEngine>,
36 layout: LayoutAnalyzer,
37}
38
39impl DocPipeline {
40 pub fn new(engine: Box<dyn OcrEngine>, layout_model: impl AsRef<Path>) -> Result<Self> {
45 let layout_model = layout_model.as_ref();
46 if !layout_model.exists() {
47 return Err(Error::ModelNotFound(format!(
48 "PP-DocLayoutV3.onnx non trovato: {}",
49 layout_model.display()
50 )));
51 }
52 let layout = LayoutAnalyzer::from_path(layout_model)?;
53 Ok(Self { engine, layout })
54 }
55
56 pub fn process_file(
59 &mut self,
60 path: impl AsRef<Path>,
61 format: OutputFormat,
62 ) -> Result<String> {
63 let path = path.as_ref();
64 let kind = detect_kind(path);
65 let pages = self.load_pages(path, kind)?;
66 let results = self.process_pages(pages)?;
67 Ok(render(&results, format))
68 }
69
70 pub fn process_image(&mut self, img: &RgbImage) -> Result<OcrPageResult> {
72 self.engine.process_page(img, &mut self.layout)
73 }
74
75 fn load_pages(&self, path: &Path, kind: InputKind) -> Result<Vec<RgbImage>> {
78 match kind {
79 InputKind::Tiff => load_tiff_pages(path),
80 InputKind::Image => load_image(path).map(|img| vec![img]),
81 }
82 }
83
84 fn process_pages(&mut self, pages: Vec<RgbImage>) -> Result<Vec<OcrPageResult>> {
85 pages.iter()
86 .map(|img| self.engine.process_page(img, &mut self.layout))
87 .collect()
88 }
89}
90
91fn detect_kind(path: &Path) -> InputKind {
94 match path.extension().and_then(|e| e.to_str()).map(|s| s.to_ascii_lowercase()).as_deref() {
95 Some("tif" | "tiff") => InputKind::Tiff,
96 _ => InputKind::Image,
97 }
98}
99
100fn render(results: &[OcrPageResult], format: OutputFormat) -> String {
101 match format {
102 OutputFormat::Markdown => to_markdown(results),
103 OutputFormat::Json => serde_json::to_string_pretty(results)
104 .unwrap_or_else(|e| format!("{{\"error\":\"{e}\"}}") ),
105 }
106}