Skip to main content

df_ocr_switcher/
pipeline.rs

1//! `DocPipeline` — orchestratore principale.
2//!
3//! Gestisce:
4//! - Scelta engine (ppocr-rs o tesseract-engine)
5//! - Caricamento LayoutAnalyzer (PP-DocLayoutV3, condiviso tra pagine)
6//! - Dispatch per tipo di input: singola immagine, TIFF multi-pagina
7//! - Generazione output: Markdown o JSON raw
8
9use 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/// Output format richiesto.
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
21pub enum OutputFormat {
22    #[default]
23    Markdown,
24    Json,
25    // SearchablePdf,  // TODO feature = searchable-pdf
26}
27
28/// Tipo di input rilevato dall'estensione file.
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30enum InputKind { Image, Tiff }
31
32/// Orchestratore principale. Tiene lo stato condiviso tra pagine:
33/// `LayoutAnalyzer` (sessione ONNX già caricata) e l'engine scelto.
34pub struct DocPipeline {
35    engine: Box<dyn OcrEngine>,
36    layout: LayoutAnalyzer,
37}
38
39impl DocPipeline {
40    /// Costruisce la pipeline.
41    ///
42    /// - `engine`: istanza già inizializzata di `PpOcrEngine` o `TesseractEngine`.
43    /// - `layout_model`: path al file `PP-DocLayoutV3.onnx`.
44    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    /// Processa un singolo file (immagine o TIFF multi-pagina).
57    /// Ritorna il testo nel formato richiesto.
58    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    /// Processa un'immagine già in memoria.
71    pub fn process_image(&mut self, img: &RgbImage) -> Result<OcrPageResult> {
72        self.engine.process_page(img, &mut self.layout)
73    }
74
75    // ── Interni ──────────────────────────────────────────────────────────────
76
77    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
91// ─── Helper ──────────────────────────────────────────────────────────────────
92
93fn 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}