df_ocr_switcher/engine/
ppocr.rs

1//! Engine PaddleOCR PP-OCRv6 via ppocr-rs.
2//!
3//! Pipeline:
4//!   DocOrientationClassifier → rotate_upright
5//!   LayoutAnalyzer (PP-DocLayoutV3)
6//!   detect_with_layout (OCR testo)
7//!   [opt] TableStructureRecognizer (SLANeXt wired)
8//!   [opt] FormulaRecognizer (PP-FormulaNet-plus-L)
9
10use std::path::PathBuf;
11
12use image::{imageops, RgbImage};
13use ppocr_rs::{
14    DocOrientation, DocOrientationClassifier, FormulaRecognizer,
15    LayoutAnalyzer, LayoutBox, ModelHub, OcrLite, OcrOptions, Point,
16    PpOcrVersion, PpStructureModel, SemanticClass,
17    TableCellBox, TableStructureRecognizer, TextBlockWithLayout,
18};
19
20use crate::engine::{OcrBlock, OcrFormula, OcrLayoutBox, OcrPageResult, OcrTable, OcrWord};
21use crate::error::{Error, Result};
22
23use super::OcrEngine;
24
25// ─── Config ──────────────────────────────────────────────────────────────────
26
27/// Path espliciti per i modelli OCR (det/rec/dict).
28pub struct OcrModelPaths {
29    pub det:  PathBuf,
30    pub rec:  PathBuf,
31    pub dict: PathBuf,
32}
33
34/// Path espliciti per il modello table structure (SLANeXt wired/wireless o SLANet_plus).
35pub struct TableModelPaths {
36    /// Path al file `.onnx` (SLANeXt wired di default, oppure SLANet_plus).
37    pub structure_onnx: PathBuf,
38    /// Path al dizionario `table_structure_dict.txt`.
39    pub structure_dict: PathBuf,
40    /// Input size in pixel. `None` = default 512 (SLANeXt).
41    /// Usare `Some(488)` per SLANet_plus.
42    pub input_size: Option<u32>,
43}
44
45/// Configurazione per `PpOcrEngine`.
46pub struct PpOcrEngineConfig {
47    /// Versione PP-OCRv6: Tiny (default), Small, Medium.
48    pub tier: PpOcrVersion,
49    /// Path esplicito `inference.onnx` orientamento. Se `None` usa ModelHub.
50    pub ori_model: Option<PathBuf>,
51    /// Path esplicito per det/rec/dict OCR. Se `None` usa ModelHub.
52    pub ocr_models: Option<OcrModelPaths>,
53    /// Thread inferenza ONNX (default: 4).
54    pub num_threads: usize,
55    /// Path esplicito per SLANeXt table structure. Se `None` tenta ModelHub.
56    pub table_models: Option<TableModelPaths>,
57    /// Abilita table structure recognition. Default `true`.
58    /// Se i modelli non sono in cache, la feature viene disabilitata silenziosamente.
59    pub enable_tables: bool,
60    /// Abilita il decoder formula (PP-FormulaNet → LaTeX). Default `false`:
61    /// il decoder autoregressive è costoso (~3 s/formula su CPU).
62    /// Attivare solo quando si vuole output LaTeX effettivo.
63    pub enable_formula_decoder: bool,
64}
65
66impl Default for PpOcrEngineConfig {
67    fn default() -> Self {
68        Self {
69            tier:                  PpOcrVersion::V6Tiny,
70            ori_model:             None,
71            ocr_models:            None,
72            num_threads:           4,
73            table_models:          None,
74            enable_tables:         true,
75            enable_formula_decoder: false,
76        }
77    }
78}
79
80// ─── Engine ──────────────────────────────────────────────────────────────────
81
82pub struct PpOcrEngine {
83    ocr:         OcrLite,
84    ori:         Option<DocOrientationClassifier>,
85    table_rec:   Option<TableStructureRecognizer>,
86    formula_rec: Option<FormulaRecognizer>,
87}
88
89impl PpOcrEngine {
90    /// Costruisce l'engine caricando i modelli.
91    ///
92    /// Table e formula recognizer sono opzionali: se i modelli non sono in
93    /// cache e `fetch-models` non è attivo, la feature viene disabilitata
94    /// con un warning su stderr ma l'engine rimane utilizzabile.
95    pub fn new(cfg: PpOcrEngineConfig) -> Result<Self> {
96        let hub = ModelHub::with_default_cache()?;
97
98        // ── Modelli OCR ───────────────────────────────────────────────────
99        let (det, rec, dict) = if let Some(p) = cfg.ocr_models {
100            (p.det, p.rec, p.dict)
101        } else {
102            let paths = hub.ensure(cfg.tier)?;
103            (paths.det_onnx, paths.rec_onnx, paths.dict_txt)
104        };
105
106        let mut ocr = OcrLite::new();
107        ocr.init_models_no_angle(
108            det.to_str().ok_or_else(|| Error::Other("det path non UTF-8".into()))?,
109            rec.to_str().ok_or_else(|| Error::Other("rec path non UTF-8".into()))?,
110            dict.to_str().ok_or_else(|| Error::Other("dict path non UTF-8".into()))?,
111            cfg.num_threads,
112        )?;
113
114        // ── Modello orientamento (opzionale — fallisce silenziosamente) ──────
115        let ori = if let Some(ori_path) = cfg.ori_model {
116            match DocOrientationClassifier::from_path(&ori_path) {
117                Ok(clf) => Some(clf),
118                Err(e) => {
119                    eprintln!("[df-ocr-switcher] orientamento disabilitato: {e}");
120                    None
121                }
122            }
123        } else {
124            match hub.ensure_single(PpStructureModel::DocOrientation) {
125                Ok(sp) => match DocOrientationClassifier::from_path(&sp.onnx) {
126                    Ok(clf) => Some(clf),
127                    Err(e)  => {
128                        eprintln!("[df-ocr-switcher] orientamento disabilitato: {e}");
129                        None
130                    }
131                },
132                Err(_) => None, // modello non in cache → skip silenzioso
133            }
134        };
135
136        // ── Table structure recognizer (opzionale) ────────────────────────
137        let table_rec = if cfg.enable_tables {
138            match load_table_recognizer(&hub, cfg.table_models) {
139                Ok(rec) => Some(rec),
140                Err(e)  => {
141                    eprintln!("[df-ocr-switcher] table recognition disabilitato: {e}");
142                    None
143                }
144            }
145        } else {
146            None
147        };
148
149        // ── Formula recognizer (opzionale, decoder disabilitato di default) ─
150        let formula_rec = match load_formula_recognizer(&hub, cfg.enable_formula_decoder) {
151            Ok(mut rec) => {
152                rec.decoder_enabled = cfg.enable_formula_decoder;
153                Some(rec)
154            }
155            Err(e) => {
156                if cfg.enable_formula_decoder {
157                    eprintln!("[df-ocr-switcher] formula recognition disabilitato: {e}");
158                }
159                None
160            }
161        };
162
163        Ok(Self { ocr, ori, table_rec, formula_rec })
164    }
165}
166
167impl OcrEngine for PpOcrEngine {
168    fn process_page(&mut self, img: &RgbImage, layout: &mut LayoutAnalyzer) -> Result<OcrPageResult> {
169        // ── Step 1: orientamento ──────────────────────────────────────────
170        let (page_angle, upright) = if let Some(clf) = &self.ori {
171            let (orient, _conf) = clf.classify(img)?;
172            let degrees = orient.degrees();
173            let rotated = rotate_upright(img, orient);
174            (degrees, rotated)
175        } else {
176            (0u32, img.clone())
177        };
178        let (w, h) = upright.dimensions();
179
180        // ── Step 2: layout-aware OCR ──────────────────────────────────────
181        let opts = OcrOptions {
182            return_word_box:     true,
183            use_doc_orientation: false,
184            ..OcrOptions::default()
185        };
186        let result = self.ocr.detect_with_layout(
187            &upright, layout,
188            10, 960, 0.6, 0.3, 1.6,
189            false, false,
190            opts,
191        )?;
192
193        // ── Step 3: table pipeline ────────────────────────────────────────
194        let tables = if let Some(rec) = &self.table_rec {
195            extract_tables(&upright, &result.layout_boxes, &result.blocks, rec)
196        } else {
197            vec![]
198        };
199
200        // ── Step 4: formula pipeline ──────────────────────────────────────
201        let formulas = if let Some(rec) = &self.formula_rec {
202            extract_formulas(&upright, &result.layout_boxes, rec)
203        } else {
204            vec![]
205        };
206
207        // ── Step 5: converti in OcrPageResult ────────────────────────────
208        let layout_boxes = result.layout_boxes.iter().map(lb_to_ocr).collect();
209        let (blocks, words) = blocks_and_words(&result.blocks);
210
211        Ok(OcrPageResult {
212            page_angle, page_width: w, page_height: h,
213            layout_boxes, blocks, words, tables, formulas,
214        })
215    }
216}
217
218// ─── Table pipeline ───────────────────────────────────────────────────────────
219
220fn extract_tables(
221    img:    &RgbImage,
222    lbs:    &[LayoutBox],
223    blocks: &[TextBlockWithLayout],
224    rec:    &TableStructureRecognizer,
225) -> Vec<OcrTable> {
226    lbs.iter()
227        .enumerate()
228        .filter(|(_, lb)| lb.class.semantic() == SemanticClass::Table)
229        .filter_map(|(i, lb)| {
230            let crop = crop_lb(img, lb);
231            match rec.recognize(&crop) {
232                Ok(structure) => {
233                    let gfm = table_to_gfm(&structure.cell_boxes, blocks, lb);
234                    Some(OcrTable { layout_idx: i as i32, gfm })
235                }
236                Err(e) => {
237                    eprintln!("[df-ocr-switcher] SLANeXt errore su tabella {i}: {e}");
238                    None
239                }
240            }
241        })
242        .collect()
243}
244
245/// Converte le cell_boxes + testo OCR full-page in GFM Markdown table.
246///
247/// Le coordinate delle cell_boxes sono relative al crop tabella
248/// → vanno traslate in coordinate di pagina aggiungendo (lb.x, lb.y).
249fn table_to_gfm(
250    cells:  &[TableCellBox],
251    blocks: &[TextBlockWithLayout],
252    lb:     &LayoutBox,
253) -> String {
254    if cells.is_empty() {
255        return String::new();
256    }
257
258    // Ordina per centroide Y poi X
259    let mut indexed: Vec<(usize, f32, f32)> = cells.iter()
260        .enumerate()
261        .map(|(i, c)| (i, (c.x1 + c.x2) * 0.5, (c.y1 + c.y2) * 0.5))
262        .collect();
263    indexed.sort_by(|a, b| a.2.partial_cmp(&b.2)
264        .unwrap_or(std::cmp::Ordering::Equal)
265        .then(a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)));
266
267    // Stima numero colonne: massima densità in Y (prima riga)
268    let first_y = indexed.first().map(|e| e.2).unwrap_or(0.0);
269    let n_cols  = indexed.iter().filter(|e| (e.2 - first_y).abs() < 20.0).count().max(1);
270
271    let mut md = String::new();
272    for (row_i, chunk) in indexed.chunks(n_cols).enumerate() {
273        // Ordina celle della riga per X
274        let mut row = chunk.to_vec();
275        row.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
276
277        md.push('|');
278        for &(ci, _, _) in &row {
279            let text = collect_cell_text(blocks, &cells[ci], lb.x, lb.y);
280            md.push_str(&format!(" {} |", text.replace('|', "\\|")));
281        }
282        md.push('\n');
283
284        if row_i == 0 {
285            md.push('|');
286            for _ in 0..row.len() { md.push_str("---|"); }
287            md.push('\n');
288        }
289    }
290    md
291}
292
293/// Trova il testo dei blocchi OCR il cui centroide cade dentro la cella.
294/// `lb_x / lb_y` = offset del crop tabella in coordinate pagina.
295fn collect_cell_text(
296    blocks: &[TextBlockWithLayout],
297    cell:   &TableCellBox,
298    lb_x:   u32,
299    lb_y:   u32,
300) -> String {
301    let cx1 = lb_x as f32 + cell.x1;
302    let cy1 = lb_y as f32 + cell.y1;
303    let cx2 = lb_x as f32 + cell.x2;
304    let cy2 = lb_y as f32 + cell.y2;
305
306    let mut parts: Vec<(u32, &str)> = blocks.iter()
307        .filter(|b| {
308            let bx = b.centroid_x as f32;
309            let by = b.centroid_y as f32;
310            bx >= cx1 && bx <= cx2 && by >= cy1 && by <= cy2
311        })
312        .map(|b| (b.centroid_y, b.block.text.trim()))
313        .collect();
314    parts.sort_by_key(|(y, _)| *y);
315    parts.iter()
316        .filter(|(_, t)| !t.is_empty())
317        .map(|(_, t)| *t)
318        .collect::<Vec<_>>()
319        .join(" ")
320}
321
322// ─── Formula pipeline ─────────────────────────────────────────────────────────
323
324fn extract_formulas(
325    img:  &RgbImage,
326    lbs:  &[LayoutBox],
327    rec:  &FormulaRecognizer,
328) -> Vec<OcrFormula> {
329    lbs.iter()
330        .enumerate()
331        .filter(|(_, lb)| lb.class.semantic() == SemanticClass::Equation)
332        .filter_map(|(i, lb)| {
333            let crop = crop_lb(img, lb);
334            match rec.recognize(&crop) {
335                Ok(fr) => Some(OcrFormula { layout_idx: i as i32, latex: fr.latex }),
336                Err(e) => {
337                    eprintln!("[df-ocr-switcher] formula rec errore su regione {i}: {e}");
338                    None
339                }
340            }
341        })
342        .collect()
343}
344
345// ─── Model loaders ───────────────────────────────────────────────────────────
346
347fn load_table_recognizer(
348    hub:   &ModelHub,
349    paths: Option<TableModelPaths>,
350) -> std::result::Result<TableStructureRecognizer, Box<dyn std::error::Error>> {
351    let (onnx, dict, input_size) = if let Some(p) = paths {
352        (p.structure_onnx, p.structure_dict, p.input_size)
353    } else {
354        let sp = hub.ensure_single(PpStructureModel::TableStructureWired)?;
355        let onnx = sp.onnx;
356        let dict = sp.dict_txt.ok_or("SLANeXt: dict_txt mancante")?;
357        (onnx, dict, None)
358    };
359    let rec = TableStructureRecognizer::from_path_with_dict(&onnx, Some(&dict))?;
360    Ok(if let Some(sz) = input_size { rec.with_input_size(sz) } else { rec })
361}
362
363fn load_formula_recognizer(
364    hub:     &ModelHub,
365    enabled: bool,
366) -> std::result::Result<FormulaRecognizer, Box<dyn std::error::Error>> {
367    if !enabled {
368        // Carica in modalità stub (decoder_enabled = false): nessuna inferenza
369        // ma struttura inizializzata. Se i modelli non sono in cache → skip.
370        let sp = hub.ensure_single(PpStructureModel::FormulaRec)?;
371        let tok = sp.tokenizer_json.as_deref();
372        Ok(FormulaRecognizer::from_paths(&sp.onnx, tok)?)
373    } else {
374        let sp = hub.ensure_single(PpStructureModel::FormulaRec)?;
375        let tok = sp.tokenizer_json.as_deref();
376        Ok(FormulaRecognizer::from_paths(&sp.onnx, tok)?)
377    }
378}
379
380// ─── Helper ──────────────────────────────────────────────────────────────────
381
382/// Crop di un LayoutBox dall'immagine.
383fn crop_lb(img: &RgbImage, lb: &LayoutBox) -> RgbImage {
384    let x = lb.x.min(img.width().saturating_sub(1));
385    let y = lb.y.min(img.height().saturating_sub(1));
386    let w = lb.w.min(img.width().saturating_sub(x));
387    let h = lb.h.min(img.height().saturating_sub(y));
388    imageops::crop_imm(img, x, y, w, h).to_image()
389}
390
391fn rotate_upright(img: &RgbImage, orient: DocOrientation) -> RgbImage {
392    match orient {
393        DocOrientation::Deg0   => img.clone(),
394        DocOrientation::Deg90  => imageops::rotate270(img),
395        DocOrientation::Deg180 => imageops::rotate180(img),
396        DocOrientation::Deg270 => imageops::rotate90(img),
397    }
398}
399
400fn aabb(pts: &[Point]) -> (u32, u32, u32, u32) {
401    let x1 = pts.iter().map(|p| p.x).min().unwrap_or(0);
402    let y1 = pts.iter().map(|p| p.y).min().unwrap_or(0);
403    let x2 = pts.iter().map(|p| p.x).max().unwrap_or(0);
404    let y2 = pts.iter().map(|p| p.y).max().unwrap_or(0);
405    (x1, y1, x2, y2)
406}
407
408fn lb_to_ocr(lb: &LayoutBox) -> OcrLayoutBox {
409    OcrLayoutBox {
410        class_name:    format!("{:?}", lb.class),
411        semantic:      lb.class.semantic(),
412        x1: lb.xmin(), y1: lb.ymin(),
413        x2: lb.xmax(), y2: lb.ymax(),
414        reading_order: lb.reading_order,
415    }
416}
417
418fn blocks_and_words(
419    src: &[TextBlockWithLayout],
420) -> (Vec<OcrBlock>, Vec<OcrWord>) {
421    let mut blocks = Vec::with_capacity(src.len());
422    let mut words  = Vec::new();
423
424    for blk in src {
425        let layout_idx = blk.layout_index.map(|i| i as i32).unwrap_or(-1);
426        let (x1, y1, x2, y2) = aabb(&blk.block.box_points);
427
428        blocks.push(OcrBlock {
429            text:       blk.block.text.clone(),
430            x1, y1, x2, y2,
431            confidence: blk.block.text_score,
432            layout_idx,
433        });
434
435        if blk.block.words.is_empty() {
436            words.push(OcrWord {
437                text: blk.block.text.clone(),
438                x1, y1, x2, y2,
439                confidence: blk.block.text_score,
440                layout_idx,
441            });
442        } else {
443            for w in &blk.block.words {
444                let (wx1, wy1, wx2, wy2) = aabb(&w.box_points);
445                words.push(OcrWord {
446                    text: w.text.clone(),
447                    x1: wx1, y1: wy1, x2: wx2, y2: wy2,
448                    confidence: w.score,
449                    layout_idx,
450                });
451            }
452        }
453    }
454
455    (blocks, words)
456}
df_ocr_switcher/engine/ppocr.rs

df_ocr_switcher/engine/
ppocr.rs
