pagesift_ocr/layout/
mod.rs

1//! 版面检测：PP-DocLayoutV3（HGNetV2-L 骨干 + RT-DETR 式混合编码器 + 可变形注意力解码器），
2//! 输入整页图，输出带类别的区域框（正文、标题、公式、公式编号、表格、图……），按预测的阅读顺序排。
3//! 结构照 transformers 的 `modeling_pp_doclayout_v3.py`、`modeling_hgnet_v2.py`，预处理与后处理照
4//! `image_processing_pp_doclayout_v3.py`（均 Apache-2.0）。掩码转多边形不做。
5//!
6//! 流程：预处理（[`preprocess`]，拉伸到 800×800）→ 骨干（[`hgnet`]）→ 混合编码器（[`encoder`]，
7//! 含掩码特征）→ 选查询（`query`，掩码外接框作初始参考框）→ 解码器（[`decoder`]，可变形注意力在
8//! `deform`）→ 后处理（`postprocess`）。
9
10pub mod config;
11pub mod decoder;
12pub(crate) mod deform;
13pub mod encoder;
14pub mod hgnet;
15pub(crate) mod layers;
16pub(crate) mod postprocess;
17pub mod preprocess;
18pub(crate) mod query;
19
20use std::path::Path;
21
22use crate::Image;
23
24/// 一个版面区域，坐标是原图像素。
25#[derive(Debug, Clone, PartialEq)]
26pub struct Region {
27    pub x0: f32,
28    pub y0: f32,
29    pub x1: f32,
30    pub y1: f32,
31    /// 类别名（config.json 的 `id2label`，例如 `formula`、`formula_number`、`text`）。
32    pub label: String,
33    /// 细分类名（inference.yml 的 `label_list`，例如 `display_formula`、`inline_formula`）；
34    /// 没有那个文件时同 `label`。
35    pub kind: String,
36    /// 分数，0–1。
37    pub score: f32,
38}
39
40/// [协议] `post_process_object_detection` 的 threshold 缺省值，也是 inference.yml 的 draw_threshold
41const SCORE_THRESHOLD: f32 = 0.5;
42
43/// 载入好的版面模型。
44#[derive(Debug, Clone)]
45pub struct Layout {
46    pre: preprocess::Preprocess,
47    backbone: hgnet::Backbone,
48    encoder: encoder::Encoder,
49    decoder: decoder::Decoder,
50    labels: Vec<String>,
51    kinds: Vec<String>,
52    /// 分数阈值；缺省同参考实现，可按需调整。
53    pub threshold: f32,
54}
55
56impl Layout {
57    /// `dir` 是模型目录：model.safetensors、config.json、preprocessor_config.json，可选 inference.yml。
58    pub fn load(dir: &Path) -> Result<Layout, String> {
59        let cfg = config::Config::from_json(&crate::read_json(dir, "config.json")?)?;
60        let pre =
61            preprocess::Preprocess::from_json(&crate::read_json(dir, "preprocessor_config.json")?)?;
62        let w = crate::read_weights(dir)?;
63        let kinds = std::fs::read_to_string(dir.join("inference.yml"))
64            .map(|s| config::parse_label_list(&s))
65            .unwrap_or_default();
66        Ok(Layout {
67            pre,
68            // [协议] 权重名前缀
69            backbone: hgnet::Backbone::load(&w, "model.backbone.model")?,
70            encoder: encoder::Encoder::load(&w, &cfg)?,
71            decoder: decoder::Decoder::load(&w, &cfg)?,
72            kinds: if kinds.len() == cfg.labels.len() {
73                kinds
74            } else {
75                cfg.labels.clone()
76            },
77            labels: cfg.labels,
78            threshold: SCORE_THRESHOLD,
79        })
80    }
81
82    /// 检测整页图的版面区域，按预测的阅读顺序返回。
83    pub fn detect(&self, img: &Image) -> Result<Vec<Region>, String> {
84        let raw = self.raw(img)?;
85        let (q, classes) = raw.logits.tc();
86        if classes != self.labels.len() {
87            return Err(format!(
88                "the model outputs {classes} classes but the config has {}",
89                self.labels.len()
90            ));
91        }
92        let ranks = postprocess::order_ranks(&raw.order.data, q);
93        let dets = postprocess::select(
94            &raw.logits.data,
95            classes,
96            &raw.boxes,
97            &ranks,
98            self.threshold,
99            img.width as f32,
100            img.height as f32,
101        );
102        Ok(dets
103            .into_iter()
104            .map(|d| Region {
105                x0: d.xyxy[0],
106                y0: d.xyxy[1],
107                x1: d.xyxy[2],
108                y1: d.xyxy[3],
109                label: self.labels[d.class].clone(),
110                kind: self.kinds[d.class].clone(),
111                score: d.score,
112            })
113            .collect())
114    }
115
116    /// 网络的原始输出（后处理之前），用来和参考实现逐数对照。
117    pub fn raw(&self, img: &Image) -> Result<decoder::Raw, String> {
118        if img.width == 0 || img.height == 0 {
119            return Err("empty image".into());
120        }
121        // [协议] `Image` 只有灰度（1）、RGB（3）两种
122        if img.channels != 1 && img.channels != 3 {
123            return Err(format!(
124                "only grayscale or RGB images are supported, got {} channels",
125                img.channels
126            ));
127        }
128        if img.data.len() != img.width * img.height * img.channels {
129            return Err("image data length does not match its width and height".into());
130        }
131        let x = self.pre.run(img);
132        let feats = self.backbone.forward(&x);
133        let enc = self.encoder.forward(&feats);
134        drop(feats);
135        Ok(self.decoder.forward(&enc.feats, &enc.mask_feat))
136    }
137}
pagesift_ocr/layout/mod.rs

pagesift_ocr/layout/
mod.rs
