ppocr_rs/
table_classifier.rs

1//! PP-LCNet classifiers: table type (wired/wireless) e document orientation.
2//!
3//! Entrambi i modelli condividono la stessa architettura PP-LCNet_x1_0 ma
4//! dimensioni di input diverse:
5//! - `TableTypeClassifier` — `"x"` shape `[1, 3, 48, 192]`
6//! - `DocOrientationClassifier` — `"x"` shape `[1, 3, 224, 224]`
7//! - Normalizzazione ImageNet: mean=[0.485,0.456,0.406] std=[0.229,0.224,0.225]
8//! - Output: logits → argmax
9//!
10//! ## Modelli HuggingFace
11//!
12//! | Modello                            | Classi | HF repo                                      |
13//! |------------------------------------|--------|----------------------------------------------|
14//! | `PP-LCNet_x1_0_table_cls_onnx`     | 2      | `PaddlePaddle/PP-LCNet_x1_0_table_cls_onnx` |
15//! | `PP-LCNet_x1_0_doc_ori_onnx`       | 4      | `PaddlePaddle/PP-LCNet_x1_0_doc_ori_onnx`   |
16//!
17//! ## Pipeline consigliata per tabelle
18//!
19//! 1. Identifica la regione tabella via `LayoutAnalyzer`.
20//! 2. Usa `TableTypeClassifier` per discriminare wired vs wireless.
21//! 3. Carica `CellDetector` (wired: RT-DETR-L wired, wireless: RT-DETR-L wireless).
22//! 4. Usa `TableStructureRecognizer` con il modello SLANeXt corrispondente.
23
24use crate::ocr_error::OcrError;
25use ndarray::{Array, Array4};
26use ort::{
27    inputs,
28    session::{builder::GraphOptimizationLevel, Session},
29    value::Tensor,
30};
31use std::path::Path;
32
33// ─── Table type ───────────────────────────────────────────────────────────────
34
35/// Tipo di tabella classificato da `PP-LCNet_x1_0_table_cls_onnx`.
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub enum TableType {
38    /// Tabella con bordi/griglia visibili.
39    Wired,
40    /// Tabella senza bordi espliciti (struttura implicita nel layout).
41    Wireless,
42}
43
44impl std::fmt::Display for TableType {
45    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46        match self {
47            TableType::Wired    => write!(f, "wired"),
48            TableType::Wireless => write!(f, "wireless"),
49        }
50    }
51}
52
53/// Classificatore tabella wired/wireless.
54///
55/// Usa `PP-LCNet_x1_0_table_cls_onnx` — modello ultra-leggero (7 MB),
56/// input `[1, 3, 48, 192]`, output 2 logits.
57#[derive(Debug)]
58pub struct TableTypeClassifier {
59    session: Session,
60}
61
62impl TableTypeClassifier {
63    pub fn from_path(model_path: impl AsRef<Path>) -> Result<Self, OcrError> {
64        let session = Session::builder()?
65            .with_optimization_level(GraphOptimizationLevel::Level3)?
66            .commit_from_file(model_path)?;
67        Ok(Self { session })
68    }
69
70    /// Classifica l'immagine come `Wired` o `Wireless`.
71    /// Ritorna `(tipo, confidence)`.
72    pub fn classify(&self, image: &image::RgbImage) -> Result<(TableType, f32), OcrError> {
73        let blob = preprocess_lcnet(image);
74        let name = self.session.inputs[0].name.clone();
75        let outputs = self.session.run(
76            inputs![name => Tensor::from_array(blob)?]?
77        )?;
78        let (_, first) = outputs.iter().next()
79            .ok_or_else(|| OcrError::ModelOutput("TableTypeCls: nessun output".into()))?;
80        let (_, raw) = crate::compat::tensor_extract_with_shape_f32(&first)?;
81        let (idx, score) = argmax_f32(&raw);
82        Ok((if idx == 0 { TableType::Wired } else { TableType::Wireless }, score))
83    }
84}
85
86// ─── Document orientation ─────────────────────────────────────────────────────
87
88/// Orientamento di una pagina documento (multipli di 90°).
89///
90/// Indica la rotazione CORRENTE della pagina. Per raddrizzarla applicare
91/// una rotazione di `(360 - degrees())°`.
92#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub enum DocOrientation {
94    /// Pagina già diritta (angolo 0°).
95    Deg0,
96    /// Pagina ruotata di 90° orario.
97    Deg90,
98    /// Pagina capovolta (180°).
99    Deg180,
100    /// Pagina ruotata di 90° antiorario (270° orario).
101    Deg270,
102}
103
104impl DocOrientation {
105    /// Gradi di rotazione corrente (0, 90, 180, 270).
106    pub fn degrees(self) -> u32 {
107        match self {
108            Self::Deg0   =>   0,
109            Self::Deg90  =>  90,
110            Self::Deg180 => 180,
111            Self::Deg270 => 270,
112        }
113    }
114}
115
116impl std::fmt::Display for DocOrientation {
117    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
118        write!(f, "{}°", self.degrees())
119    }
120}
121
122/// Classificatore orientamento documento.
123///
124/// Usa `PP-LCNet_x1_0_doc_ori_onnx` — modello 7 MB, input `[1, 3, 224, 224]`,
125/// output 4 logits (classe 0=0°, 1=90°, 2=180°, 3=270°).
126///
127/// Utile per scansioni ruotate dove la classificazione per-riga `do_angle`
128/// (`AngleNet`) non è sufficiente: un documento capovolto avrà tutte le righe
129/// con orientamento "coerente" ma globalmente sbagliato.
130#[derive(Debug)]
131pub struct DocOrientationClassifier {
132    session: Session,
133}
134
135impl DocOrientationClassifier {
136    pub fn from_path(model_path: impl AsRef<Path>) -> Result<Self, OcrError> {
137        let session = Session::builder()?
138            .with_optimization_level(GraphOptimizationLevel::Level3)?
139            .commit_from_file(model_path)?;
140        Ok(Self { session })
141    }
142
143    /// Classifica l'orientamento della pagina.
144    /// Ritorna `(orientamento, confidence)`.
145    pub fn classify(&self, image: &image::RgbImage) -> Result<(DocOrientation, f32), OcrError> {
146        // PP-LCNet_x1_0_doc_ori_onnx usa 224×224, non 48×192 come il table_cls.
147        let blob = preprocess_lcnet_sq224(image);
148        let name = self.session.inputs[0].name.clone();
149        let outputs = self.session.run(
150            inputs![name => Tensor::from_array(blob)?]?
151        )?;
152        let (_, first) = outputs.iter().next()
153            .ok_or_else(|| OcrError::ModelOutput("DocOriCls: nessun output".into()))?;
154        let (_, raw) = crate::compat::tensor_extract_with_shape_f32(&first)?;
155        let (idx, score) = argmax_f32(&raw);
156        let orient = match idx {
157            0 => DocOrientation::Deg0,
158            1 => DocOrientation::Deg90,
159            2 => DocOrientation::Deg180,
160            _ => DocOrientation::Deg270,
161        };
162        Ok((orient, score))
163    }
164}
165
166// ─── Shared preprocessing ─────────────────────────────────────────────────────
167
168/// Preprocessing PP-LCNet per `PP-LCNet_x1_0_table_cls_onnx`: resize 192×48.
169/// `ClsResizeImg(image_shape=[3, 48, 192])` + normalizzazione ImageNet.
170fn preprocess_lcnet(image: &image::RgbImage) -> Array4<f32> {
171    preprocess_lcnet_wh(image, 192, 48)
172}
173
174/// Preprocessing PP-LCNet per `PP-LCNet_x1_0_doc_ori_onnx`: resize 224×224.
175/// `ClsResizeImg(image_shape=[3, 224, 224])` + normalizzazione ImageNet.
176fn preprocess_lcnet_sq224(image: &image::RgbImage) -> Array4<f32> {
177    preprocess_lcnet_wh(image, 224, 224)
178}
179
180fn preprocess_lcnet_wh(image: &image::RgbImage, w: u32, h: u32) -> Array4<f32> {
181    let resized = image::imageops::resize(image, w, h, image::imageops::FilterType::Triangle);
182    let mean = [0.485f32, 0.456, 0.406];
183    let std  = [0.229f32, 0.224, 0.225];
184    let mut blob: Array4<f32> = Array::zeros((1, 3, h as usize, w as usize));
185    for y in 0..h as usize {
186        for x in 0..w as usize {
187            let p = resized.get_pixel(x as u32, y as u32);
188            for c in 0..3usize {
189                blob[[0, c, y, x]] = (p[c] as f32 / 255.0 - mean[c]) / std[c];
190            }
191        }
192    }
193    blob
194}
195
196/// Ritorna `(indice_max, valore_max)` su un vettore di logit.
197fn argmax_f32(logits: &[f32]) -> (usize, f32) {
198    logits.iter().copied().enumerate()
199        .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
200        .unwrap_or((0, 0.0))
201}
ppocr_rs/table_classifier.rs

ppocr_rs/
table_classifier.rs
