ppocr_rs/
formula_rec.rs

1//! Formula recognition via PP-FormulaNet_plus-L ONNX (PP-StructureV3 v2).
2//!
3//! Converte il crop di una formula matematica (identificata da
4//! `LayoutAnalyzer` come `DisplayFormula` o `InlineFormula`) in una stringa
5//! LaTeX.
6//!
7//! ## Architettura del modello
8//!
9//! PP-FormulaNet_plus-L è un modello **encoder-decoder**:
10//! - **Encoder**: `Vary_VIT_B_Formula` (ViT 768-dim, patch 16×16)
11//! - **Decoder**: `CustomMBartForCausalLM` (8 layer, 16 head, vocab 50 000)
12//! - **Input**: immagine `[1, 3, 768, 768]`, ImageNet normalize
13//! - **Output**: sequenza token IDs LaTeX (BOS=0, PAD=1, EOS=2)
14//!
15//! L'export ONNX può essere in due modalità:
16//! - **Singolo forward** → output `[1, T, 50 000]` logit o `[1, T]` int64
17//! - **Autoregresivo step-by-step** → input aggiuntivo `decoder_input_ids`
18//!
19//! [`FormulaRecognizer`] rileva automaticamente la modalità all'init.
20//!
21//! ## Tokenizer
22//!
23//! BPE da `tokenizer.json` (HuggingFace format). Scaricato insieme all'ONNX
24//! da `ModelHub::ensure_single(PpStructureModel::FormulaRec)`.
25//! Campo `model.vocab`: `{ token_string → id }` → invertito a `Vec<String>`.
26//!
27//! ## Uso
28//!
29//! ```no_run
30//! use ppocr_rs::{ModelHub, PpStructureModel, FormulaRecognizer};
31//! let hub   = ModelHub::with_default_cache().unwrap();
32//! let paths = hub.ensure_single(PpStructureModel::FormulaRec).unwrap();
33//! let rec   = FormulaRecognizer::from_paths(
34//!     &paths.onnx,
35//!     paths.tokenizer_json.as_deref(),
36//! ).unwrap();
37//! // image: crop della regione formula
38//! // let result = rec.recognize(&image).unwrap();
39//! // println!("{}", result.latex);
40//! ```
41
42use crate::ocr_error::OcrError;
43use ndarray::{Array, Array4};
44use ort::{
45    inputs,
46    session::{builder::GraphOptimizationLevel, Session},
47    value::Tensor,
48};
49use std::path::Path;
50
51/// Dimensione dell'input del modello (768×768, fissa dall'architettura ViT).
52pub const FORMULA_INPUT_SIZE: u32 = 768;
53
54/// Token ID speciali (MBart convention).
55const BOS_ID: i64 = 0;
56const PAD_ID: i64 = 1;
57const EOS_ID: i64 = 2;
58
59/// Lunghezza massima sequenza generata (config PP-FormulaNet_plus-L).
60const MAX_NEW_TOKENS: usize = 2560;
61
62// ─── Result ───────────────────────────────────────────────────────────────────
63
64/// Risultato del riconoscimento formula.
65#[derive(Debug, Clone)]
66pub struct FormulaResult {
67    /// Stringa LaTeX decodificata (senza `$` o `\[` delimitatori — solo il
68    /// contenuto della formula).
69    pub latex: String,
70    /// Confidence media sul token sequence (score medio dei logit argmax).
71    /// `0.0` se il modello emette token IDs direttamente (senza logit).
72    pub score: f32,
73}
74
75// ─── Tokenizer ────────────────────────────────────────────────────────────────
76
77struct Tokenizer {
78    /// id → token string (indice diretto).
79    id_to_token: Vec<String>,
80}
81
82impl Tokenizer {
83    /// Carica da `tokenizer.json` (formato HuggingFace BPE).
84    ///
85    /// Struttura attesa: `{ "model": { "vocab": { token: id, ... } } }`.
86    fn from_json(path: &Path) -> Result<Self, OcrError> {
87        let content = std::fs::read_to_string(path)
88            .map_err(|e| OcrError::ModelHubError(format!("tokenizer.json: {e}")))?;
89        let json: serde_json::Value = serde_json::from_str(&content)
90            .map_err(|e| OcrError::ModelHubError(format!("parse tokenizer.json: {e}")))?;
91
92        // Supporta due path di vocab comuni nei tokenizer.json HuggingFace:
93        // 1. model.vocab (BPE)
94        // 2. vocab (SentencePiece / Unigram)
95        let vocab = json.get("model")
96            .and_then(|m| m.get("vocab"))
97            .or_else(|| json.get("vocab"))
98            .and_then(|v| v.as_object())
99            .ok_or_else(|| OcrError::ModelHubError(
100                "tokenizer.json: campo 'model.vocab' non trovato".into()
101            ))?;
102
103        // Dimensione vocab: usa max_id + 1 per evitare OOB su id non contigui
104        let max_id = vocab.values()
105            .filter_map(|v| v.as_i64())
106            .max()
107            .unwrap_or(0) as usize;
108
109        let mut id_to_token = vec![String::new(); max_id + 1];
110        for (token, id_val) in vocab {
111            if let Some(id) = id_val.as_i64() {
112                let id = id as usize;
113                if id < id_to_token.len() {
114                    id_to_token[id] = token.clone();
115                }
116            }
117        }
118
119        eprintln!("[FormulaRec] tokenizer caricato: {} token", id_to_token.len());
120        Ok(Self { id_to_token })
121    }
122
123    /// Vocabolario embedded minimo (solo token speciali).
124    /// Usato quando `tokenizer.json` non è disponibile — produce output
125    /// con token IDs testuali invece di LaTeX.
126    fn fallback() -> Self {
127        let mut id_to_token = vec![String::new(); 50_000];
128        id_to_token[BOS_ID as usize] = "<s>".to_string();
129        id_to_token[PAD_ID as usize] = "<pad>".to_string();
130        id_to_token[EOS_ID as usize] = "</s>".to_string();
131        eprintln!(
132            "[FormulaRec] WARN: tokenizer.json non disponibile — \
133             output sarà token IDs testuale. Scarica il modello con \
134             ModelHub::ensure_single(PpStructureModel::FormulaRec)."
135        );
136        Self { id_to_token }
137    }
138
139    fn decode(&self, ids: &[i64]) -> String {
140        let mut out = String::new();
141        let mut prev_needs_space = false;
142
143        for &id in ids {
144            if id == BOS_ID || id == PAD_ID { continue; }
145            if id == EOS_ID { break; }
146
147            let token = match self.id_to_token.get(id as usize) {
148                Some(t) if !t.is_empty() => t.as_str(),
149                _ => continue,
150            };
151
152            // SentencePiece/BPE: '▁' (U+2581) indica inizio parola → spazio
153            if let Some(stripped) = token.strip_prefix('▁') {
154                if prev_needs_space && !out.is_empty() {
155                    out.push(' ');
156                }
157                out.push_str(stripped);
158                prev_needs_space = true;
159            } else if let Some(stripped) = token.strip_prefix('Ġ') {
160                // GPT-2 style space marker
161                if !out.is_empty() { out.push(' '); }
162                out.push_str(stripped);
163                prev_needs_space = false;
164            } else {
165                out.push_str(token);
166                prev_needs_space = false;
167            }
168        }
169        out.trim().to_string()
170    }
171
172    fn vocab_size(&self) -> usize { self.id_to_token.len() }
173}
174
175// ─── InferenceMode ────────────────────────────────────────────────────────────
176
177/// Modalità di inference rilevata all'init.
178#[derive(Debug, Clone, Copy)]
179enum InferenceMode {
180    /// Il modello produce l'intera sequenza in un unico forward pass.
181    /// Output: logit `[1, T, V]` float32 o IDs `[1, T]` int64.
182    SinglePass,
183    /// Il modello richiede `decoder_input_ids` — genera un token alla volta.
184    Autoregressive,
185}
186
187// ─── FormulaRecognizer ────────────────────────────────────────────────────────
188
189/// Riconosce formule matematiche in immagini crop → stringa LaTeX.
190pub struct FormulaRecognizer {
191    session:   Session,
192    tokenizer: Tokenizer,
193    mode:      InferenceMode,
194    /// Nome del tensore immagine input (rilevato dai metadati ONNX).
195    img_input: String,
196    /// Abilita il decoder autoregressive. **Default `false`**: quando
197    /// disabilitato [`recognize`] restituisce `latex = ""` senza eseguire
198    /// alcuna inferenza. Impostare a `true` solo quando si vuole produrre
199    /// output LaTeX effettivo (il decoder è costoso: fino a 2560 step).
200    pub decoder_enabled: bool,
201}
202
203impl FormulaRecognizer {
204    /// Carica il modello ONNX e il tokenizer da path espliciti.
205    ///
206    /// `tokenizer_path`: path a `tokenizer.json`. Se `None` usa un tokenizer
207    /// fallback che emette token IDs anziché LaTeX — scarica il modello
208    /// completo via `ModelHub::ensure_single(PpStructureModel::FormulaRec)`.
209    pub fn from_paths(
210        model_path:     impl AsRef<Path>,
211        tokenizer_path: Option<&Path>,
212    ) -> Result<Self, OcrError> {
213        let session = Session::builder()?
214            .with_optimization_level(GraphOptimizationLevel::Level3)?
215            .commit_from_file(model_path)?;
216
217        let tokenizer = match tokenizer_path {
218            Some(p) => Tokenizer::from_json(p)?,
219            None    => Tokenizer::fallback(),
220        };
221
222        let mode = detect_inference_mode(&session);
223        let img_input = find_image_input(&session);
224
225        eprintln!(
226            "[FormulaRec] modalità: {mode:?}, input immagine: '{img_input}', vocab: {}",
227            tokenizer.vocab_size()
228        );
229
230        Ok(Self { session, tokenizer, mode, img_input, decoder_enabled: false })
231    }
232
233    /// Riconosce la formula nell'immagine crop → stringa LaTeX.
234    ///
235    /// Se `self.decoder_enabled == false` (default) restituisce
236    /// `FormulaResult { latex: String::new(), score: 0.0 }` senza alcuna
237    /// inferenza. Il decoder autoregressive può impiegare fino a 2560 step
238    /// (≈1.5 s/GPU, ≈3 s/CPU per formula complessa).
239    ///
240    /// `image` deve essere il crop della regione formula (da `LayoutAnalyzer`
241    /// con `LayoutClass::DisplayFormula` o `InlineFormula`).
242    pub fn recognize(&self, image: &image::RgbImage) -> Result<FormulaResult, OcrError> {
243        if !self.decoder_enabled {
244            return Ok(FormulaResult { latex: String::new(), score: 0.0 });
245        }
246
247        let blob = preprocess(image);
248
249        match self.mode {
250            InferenceMode::SinglePass  => self.recognize_single_pass(blob),
251            InferenceMode::Autoregressive => self.recognize_autoregressive(blob),
252        }
253    }
254
255    // ── Single-pass ─────────────────────────────────────────────────────────
256
257    fn recognize_single_pass(&self, blob: Array4<f32>) -> Result<FormulaResult, OcrError> {
258        let outputs = self.session.run(
259            inputs![self.img_input.clone() => Tensor::from_array(blob)?]?
260        )?;
261
262        let (_, first) = outputs.iter().next()
263            .ok_or_else(|| OcrError::ModelOutput("FormulaNet: nessun output".into()))?;
264        let (shape, data_f32) = crate::compat::tensor_extract_with_shape_f32(&first)?;
265
266        eprintln!("[FormulaNet] single-pass output shape: {shape:?}");
267
268        match shape.as_slice() {
269            // Output: logit [1, T, V]
270            [1, t_len, v_len] => {
271                let t = *t_len as usize;
272                let v = *v_len as usize;
273                let mut ids: Vec<i64>  = Vec::with_capacity(t);
274                let mut scores: Vec<f32> = Vec::with_capacity(t);
275                for step in 0..t {
276                    let base = step * v;
277                    let (idx, score) = argmax_slice(&data_f32[base..base + v]);
278                    let id = idx as i64;
279                    if id == EOS_ID { break; }
280                    if id == PAD_ID || id == BOS_ID { continue; }
281                    ids.push(id);
282                    scores.push(score);
283                }
284                let latex = self.tokenizer.decode(&ids);
285                let score = mean_score(&scores);
286                Ok(FormulaResult { latex, score })
287            }
288
289            // Output: IDs [1, T] già decodificati (generate baked-in)
290            [1, _t] => {
291                // Reinterpreta data_f32 come int via round
292                let ids: Vec<i64> = data_f32.iter().map(|&f| f.round() as i64).collect();
293                let latex = self.tokenizer.decode(&ids);
294                Ok(FormulaResult { latex, score: 0.0 })
295            }
296
297            other => Err(OcrError::ModelOutput(format!(
298                "FormulaNet single-pass: shape output non riconosciuta: {other:?}"
299            ))),
300        }
301    }
302
303    // ── Autoregressive ───────────────────────────────────────────────────────
304
305    fn recognize_autoregressive(&self, blob: Array4<f32>) -> Result<FormulaResult, OcrError> {
306        // Trova il nome del tensore decoder_input_ids
307        let dec_input_name = self.session.inputs.iter()
308            .find(|i| {
309                let n = i.name.to_lowercase();
310                n.contains("decoder") || n.contains("input_ids") || n.contains("tgt")
311            })
312            .map(|i| i.name.clone())
313            .unwrap_or_else(|| "decoder_input_ids".to_string());
314
315        let mut ids: Vec<i64>   = vec![BOS_ID];
316        let mut scores: Vec<f32> = Vec::new();
317
318        for _ in 0..MAX_NEW_TOKENS {
319            // decoder_input_ids: [1, current_len] int64
320            let dec_len = ids.len();
321            let dec_arr: ndarray::Array2<i64> = ndarray::Array2::from_shape_vec(
322                (1, dec_len),
323                ids.clone(),
324            ).map_err(|e| OcrError::ModelInput(format!("dec_arr shape: {e}")))?;
325
326            let outputs = self.session.run(inputs![
327                self.img_input.clone() => Tensor::from_array(blob.clone())?,
328                dec_input_name.clone() => Tensor::from_array(dec_arr)?,
329            ]?)?;
330
331            let (_, last_out) = outputs.iter().next()
332                .ok_or_else(|| OcrError::ModelOutput("FormulaNet auto: nessun output".into()))?;
333            let (shape, data) = crate::compat::tensor_extract_with_shape_f32(&last_out)?;
334
335            // Logit dell'ultimo token: shape [1, current_len, V] → slice dell'ultimo step
336            let next_id = match shape.as_slice() {
337                [1, _t, v] => {
338                    let v = *v as usize;
339                    let last_step_base = (dec_len - 1) * v;
340                    let (idx, score) = argmax_slice(&data[last_step_base..last_step_base + v]);
341                    scores.push(score);
342                    idx as i64
343                }
344                [1, v] => {
345                    let (idx, score) = argmax_slice(&data[..(*v as usize)]);
346                    scores.push(score);
347                    idx as i64
348                }
349                other => return Err(OcrError::ModelOutput(format!(
350                    "FormulaNet auto step: shape {other:?}"
351                ))),
352            };
353
354            if next_id == EOS_ID || next_id == PAD_ID { break; }
355            ids.push(next_id);
356        }
357
358        // Rimuovi BOS dalla decodifica
359        let decode_ids: Vec<i64> = ids.into_iter().skip(1).collect();
360        let latex = self.tokenizer.decode(&decode_ids);
361        let score = mean_score(&scores);
362        Ok(FormulaResult { latex, score })
363    }
364}
365
366// ─── Preprocessing ────────────────────────────────────────────────────────────
367
368fn preprocess(image: &image::RgbImage) -> Array4<f32> {
369    let s = FORMULA_INPUT_SIZE;
370    let resized = image::imageops::resize(image, s, s, image::imageops::FilterType::Triangle);
371    let mean = [0.485f32, 0.456, 0.406];
372    let std  = [0.229f32, 0.224, 0.225];
373    let mut blob: Array4<f32> = Array::zeros((1, 3, s as usize, s as usize));
374    for y in 0..s as usize {
375        for x in 0..s as usize {
376            let p = resized.get_pixel(x as u32, y as u32);
377            for c in 0..3usize {
378                blob[[0, c, y, x]] = (p[c] as f32 / 255.0 - mean[c]) / std[c];
379            }
380        }
381    }
382    blob
383}
384
385// ─── Helpers ─────────────────────────────────────────────────────────────────
386
387fn detect_inference_mode(session: &Session) -> InferenceMode {
388    let has_dec_input = session.inputs.iter().any(|i| {
389        let n = i.name.to_lowercase();
390        n.contains("decoder") || n.contains("input_ids") || n.contains("tgt")
391    });
392    if has_dec_input {
393        InferenceMode::Autoregressive
394    } else {
395        InferenceMode::SinglePass
396    }
397}
398
399fn find_image_input(session: &Session) -> String {
400    session.inputs.iter()
401        .find(|i| {
402            let n = i.name.to_lowercase();
403            n == "x" || n == "image" || n == "img" || n.contains("pixel")
404        })
405        .or_else(|| session.inputs.first())
406        .map(|i| i.name.clone())
407        .unwrap_or_else(|| "x".to_string())
408}
409
410fn argmax_slice(s: &[f32]) -> (usize, f32) {
411    s.iter().copied().enumerate()
412        .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
413        .unwrap_or((0, 0.0))
414}
415
416fn mean_score(v: &[f32]) -> f32 {
417    if v.is_empty() { 0.0 } else { v.iter().sum::<f32>() / v.len() as f32 }
418}
ppocr_rs/formula_rec.rs

ppocr_rs/
formula_rec.rs
