pagesift_ocr/layout/
config.rs

1//! 模型配置（config.json）里推理要用的字段，缺省值取 transformers `PPDocLayoutV3Config` 的类属性。
2//! 只支持 PP-DocLayoutV3 发布权重用到的那种结构，遇到别的取值直接报错，不默默算错。
3
4use pagesift_infer::Activation;
5
6/// 推理要用的配置。
7#[derive(Debug, Clone)]
8pub struct Config {
9    pub d_model: usize,
10    pub encoder_heads: usize,
11    pub decoder_heads: usize,
12    pub decoder_points: usize,
13    pub num_queries: usize,
14    /// 编码器、解码器输入各层特征的步长（相对网络输入）。
15    pub feat_strides: Vec<usize>,
16    /// 做 AIFI（层内自注意力）的那一层在 `feat_strides` 里的下标。
17    pub aifi_level: usize,
18    /// CSP 块、下采样卷积等的激活（`activation_function`）。
19    pub act: Activation,
20    pub encoder_act: Activation,
21    pub decoder_act: Activation,
22    pub layer_norm_eps: f32,
23    pub pos_temperature: f64,
24    pub pointer_dim: usize,
25    /// 类别号 → 名字（`id2label`）。
26    pub labels: Vec<String>,
27}
28
29/// [协议] `PPDocLayoutV3Config.feat_strides` 的缺省值
30const DEFAULT_STRIDES: [usize; 3] = [8, 16, 32];
31/// [协议] `encoder_attention_heads`、`decoder_attention_heads` 的缺省值
32const DEFAULT_HEADS: usize = 8;
33/// [协议] `decoder_n_points` 的缺省值
34const DEFAULT_POINTS: usize = 4;
35/// [协议] `num_queries` 的缺省值
36const DEFAULT_QUERIES: usize = 300;
37
38fn activation(name: &str) -> Result<Activation, String> {
39    match name {
40        "relu" => Ok(Activation::Relu),
41        "silu" | "swish" => Ok(Activation::Silu),
42        "gelu" => Ok(Activation::Gelu),
43        _ => Err(format!("unsupported activation {name}")),
44    }
45}
46
47impl Config {
48    pub fn from_json(v: &serde_json::Value) -> Result<Config, String> {
49        let int = |key: &str, def: usize| -> Result<usize, String> {
50            match &v[key] {
51                serde_json::Value::Null => Ok(def),
52                x => x
53                    .as_u64()
54                    .map(|n| n as usize)
55                    .ok_or(format!("{key} must be an integer")),
56            }
57        };
58        let list = |key: &str, def: &[usize]| -> Result<Vec<usize>, String> {
59            match v[key].as_array() {
60                None => Ok(def.to_vec()),
61                Some(a) => a
62                    .iter()
63                    .map(|x| {
64                        x.as_u64()
65                            .map(|n| n as usize)
66                            .ok_or(format!("{key} must be an integer"))
67                    })
68                    .collect(),
69            }
70        };
71        let act = |key: &str, def: &str| activation(v[key].as_str().unwrap_or(def));
72        let flag = |key: &str, def: bool| v[key].as_bool().unwrap_or(def);
73
74        // 只实现发布权重的结构：后归一化、查询取自编码器输出、掩码框作初始参考框、CSP 不扩通道
75        if flag("normalize_before", false) {
76            return Err("normalize_before is not supported".into());
77        }
78        if flag("learn_initial_query", false) {
79            return Err("learn_initial_query is not supported".into());
80        }
81        if !flag("mask_enhanced", true) {
82            return Err(
83                "only mask_enhanced is supported (mask boxes as the initial reference boxes)"
84                    .into(),
85            );
86        }
87        if v["hidden_expansion"].as_f64().unwrap_or(1.0) != 1.0 {
88            return Err("only hidden_expansion = 1 is supported".into());
89        }
90
91        // [协议] 缺省值取 PPDocLayoutV3Config：d_model 256、注意力头 8、采样点 4、查询 300、
92        // 步长 (8, 16, 32)、AIFI 在第 2 层、指针头宽 64、位置编码温度 10000
93        let d_model = int("d_model", 256)?;
94        let feat_strides = match v["feat_strides"].as_array() {
95            Some(_) => list("feat_strides", &DEFAULT_STRIDES)?,
96            None => list("feature_strides", &DEFAULT_STRIDES)?,
97        };
98        let proj = list("encode_proj_layers", &[2])?;
99        if proj.len() != 1 || int("encoder_layers", 1)? != 1 {
100            return Err("only one AIFI layer is supported".into());
101        }
102        let levels = int("num_feature_levels", feat_strides.len())?;
103        if levels != feat_strides.len() || int("encoder_hidden_dim", d_model)? != d_model {
104            return Err("only decoder layers that take encoder outputs directly are supported (same layer count and channels)".into());
105        }
106        let labels = parse_labels(&v["id2label"])?;
107        Ok(Config {
108            d_model,
109            encoder_heads: int("encoder_attention_heads", DEFAULT_HEADS)?,
110            decoder_heads: int("decoder_attention_heads", DEFAULT_HEADS)?,
111            decoder_points: int("decoder_n_points", DEFAULT_POINTS)?,
112            num_queries: int("num_queries", DEFAULT_QUERIES)?,
113            aifi_level: proj[0],
114            feat_strides,
115            act: act("activation_function", "silu")?,
116            encoder_act: act("encoder_activation_function", "gelu")?,
117            decoder_act: act("decoder_activation_function", "relu")?,
118            // [协议] PPDocLayoutV3Config 的 layer_norm_eps 缺省 1e-5
119            layer_norm_eps: v["layer_norm_eps"].as_f64().unwrap_or(1e-5) as f32,
120            pos_temperature: v["positional_encoding_temperature"]
121                .as_f64()
122                .unwrap_or(10000.0),
123            pointer_dim: int("global_pointer_head_size", 64)?,
124            labels,
125        })
126    }
127}
128
129/// `id2label`：键是类别号的字符串。
130fn parse_labels(v: &serde_json::Value) -> Result<Vec<String>, String> {
131    let obj = v.as_object().ok_or("config.json has no id2label")?;
132    let mut labels = vec![String::new(); obj.len()];
133    for (k, name) in obj {
134        let i: usize = k
135            .parse()
136            .map_err(|_| format!("id2label key {k} is not an integer"))?;
137        let slot = labels
138            .get_mut(i)
139            .ok_or(format!("id2label class ids are not contiguous at {i}"))?;
140        *slot = name
141            .as_str()
142            .ok_or("id2label values must be strings")?
143            .to_string();
144    }
145    Ok(labels)
146}
147
148/// inference.yml 里的 `label_list`（PaddleX 的细分类名，例如把 `formula` 分成 `display_formula`、
149/// `inline_formula`）。只认 `label_list:` 下一行一个 `- 名字` 的写法；没有或格式不对时返回空。
150pub fn parse_label_list(yml: &str) -> Vec<String> {
151    let mut out = Vec::new();
152    let mut inside = false;
153    for line in yml.lines() {
154        if line.trim_end() == "label_list:" {
155            inside = true;
156            continue;
157        }
158        if !inside {
159            continue;
160        }
161        match line.trim_start().strip_prefix("- ") {
162            // 列表项与键同级或更深缩进；遇到下一个键就结束
163            Some(name) => out.push(name.trim().to_string()),
164            None => break,
165        }
166    }
167    out
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    #[test]
175    fn 读细分类名() {
176        let yml = "arch: DETR\nlabel_list:\n- abstract\n- display_formula\nHpi:\n  x: 1\n";
177        assert_eq!(parse_label_list(yml), ["abstract", "display_formula"]);
178        assert!(parse_label_list("arch: DETR\n").is_empty());
179    }
180
181    #[test]
182    fn 读配置() {
183        let v: serde_json::Value = serde_json::from_str(
184            r#"{"d_model": 256, "feature_strides": [8, 16, 32], "id2label": {"1": "b", "0": "a"},
185                "encoder_activation_function": "gelu"}"#,
186        )
187        .unwrap();
188        let c = Config::from_json(&v).unwrap();
189        assert_eq!(c.labels, ["a", "b"]);
190        assert_eq!(c.feat_strides, [8, 16, 32]);
191        assert_eq!(c.aifi_level, 2);
192        assert_eq!(c.encoder_act, Activation::Gelu);
193    }
194}
pagesift_ocr/layout/config.rs

pagesift_ocr/layout/
config.rs
