pagesift_ocr/
math.rs

1//! 行间公式：版面模型找出公式框，公式模型把框里的图认成 LaTeX。框里的普通文字行由流水线去掉。
2
3use std::path::Path;
4
5use pagesift_infer::Weights;
6
7use crate::formula::{Decoder, Tokenizer, to_latex};
8use crate::formula_enc::{Config, Encoder, preprocess};
9use crate::layout::{Layout, Region};
10use crate::{Image, Line, read_json};
11
12/// [协议] PP-DocLayoutV3 的 `inference.yml` 里行间公式的细分类名
13const DISPLAY_FORMULA: &str = "display_formula";
14
15/// 版面检测加公式识别。
16pub struct Math {
17    layout: Layout,
18    encoder: Encoder,
19    decoder: Decoder,
20    tokenizer: Tokenizer,
21}
22
23impl Math {
24    pub fn load(layout_dir: &Path, formula_dir: &Path) -> Result<Math, String> {
25        let cfg = read_json(formula_dir, "config.json")?;
26        let generation = read_json(formula_dir, "generation_config.json")?;
27        // 编码器和解码器的权重在同一个文件里，读一次，建好就放掉
28        let w = crate::read_weights(formula_dir)?;
29        let decoder = Decoder::from_weights(&w, &cfg["text_config"], &generation)?;
30        let encoder = Encoder::from_weights(&w, Config::from_json(&cfg)?)?;
31        drop::<Weights>(w);
32        Ok(Math {
33            layout: Layout::load(layout_dir)?,
34            encoder,
35            decoder,
36            tokenizer: Tokenizer::load(formula_dir)?,
37        })
38    }
39
40    /// 这页的行间公式框。
41    pub(crate) fn formula_boxes(&self, page: &Image) -> Result<Vec<Region>, String> {
42        let mut regions = self.layout.detect(page)?;
43        regions.retain(|r| r.kind == DISPLAY_FORMULA);
44        Ok(regions)
45    }
46
47    /// 认出一个框，得到一行 `$$…$$`；没认出内容时返回 `None`。
48    pub(crate) fn recognize(&self, page: &Image, r: &Region) -> Result<Option<Line>, String> {
49        let crop = crop_rect(page, r);
50        if crop.width == 0 || crop.height == 0 {
51            return Ok(None);
52        }
53        let memory = self.encoder.encode(&preprocess(&crop))?;
54        let ids = self.decoder.generate(&memory, self.decoder.max_length())?;
55        let latex = to_latex(&self.tokenizer.decode(&ids));
56        let latex = latex.trim();
57        Ok((!latex.is_empty()).then(|| Line {
58            quad: [(r.x0, r.y0), (r.x1, r.y0), (r.x1, r.y1), (r.x0, r.y1)],
59            text: format!("$${latex}$$"),
60            score: r.score,
61            formula: true,
62        }))
63    }
64}
65
66/// 按框裁图（不像文字行那样按高宽比转向：公式常比宽还高）。
67fn crop_rect(page: &Image, r: &Region) -> Image {
68    let (x0, y0) = (r.x0.max(0.0) as usize, r.y0.max(0.0) as usize);
69    let x1 = (r.x1.ceil().max(0.0) as usize).min(page.width);
70    let y1 = (r.y1.ceil().max(0.0) as usize).min(page.height);
71    let (w, h) = (x1.saturating_sub(x0), y1.saturating_sub(y0));
72    let c = page.channels;
73    let mut data = Vec::with_capacity(w * h * c);
74    // 框整个在页外时宽或高是 0，没有可读的行
75    for y in (y0..y0 + h).filter(|_| w > 0) {
76        let at = (y * page.width + x0) * c;
77        data.extend_from_slice(&page.data[at..at + w * c]);
78    }
79    Image {
80        width: w,
81        height: h,
82        channels: c,
83        data,
84    }
85}
86
87/// 中心落在某个框里的行算在公式里。
88pub(crate) fn inside_any(quad: &crate::Quad, boxes: &[Region]) -> bool {
89    let n = quad.len() as f32;
90    let cx = quad.iter().map(|p| p.0).sum::<f32>() / n;
91    let cy = quad.iter().map(|p| p.1).sum::<f32>() / n;
92    boxes
93        .iter()
94        .any(|r| cx >= r.x0 && cx <= r.x1 && cy >= r.y0 && cy <= r.y1)
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    fn region(x0: f32, y0: f32, x1: f32, y1: f32) -> Region {
102        Region {
103            x0,
104            y0,
105            x1,
106            y1,
107            label: "formula".into(),
108            kind: DISPLAY_FORMULA.into(),
109            score: 1.0,
110        }
111    }
112
113    #[test]
114    fn crop_stays_in_the_page() {
115        let page = Image {
116            width: 4,
117            height: 3,
118            channels: 1,
119            data: (0..12).collect(),
120        };
121        let c = crop_rect(&page, &region(1.0, 1.0, 3.0, 9.0));
122        assert_eq!((c.width, c.height), (2, 2));
123        assert_eq!(c.data, [5, 6, 9, 10]);
124        let c = crop_rect(&page, &region(9.0, 0.0, 12.0, 2.0));
125        assert_eq!(c.width, 0);
126    }
127
128    #[test]
129    fn lines_inside_a_formula_box_are_found_by_center() {
130        let boxes = [region(10.0, 10.0, 20.0, 20.0)];
131        let q = |x: f32, y: f32| {
132            [
133                (x - 1.0, y - 1.0),
134                (x + 1.0, y - 1.0),
135                (x + 1.0, y + 1.0),
136                (x - 1.0, y + 1.0),
137            ]
138        };
139        assert!(inside_any(&q(15.0, 15.0), &boxes));
140        assert!(!inside_any(&q(25.0, 15.0), &boxes));
141    }
142}
pagesift_ocr/math.rs

pagesift_ocr/
math.rs
