pagesift_ocr/det/
neck.rs

1//! 颈部 RepLKFPN 与 DB 头。结构照 transformers 的 `modeling_pp_ocrv6_small_det.py`
2//! （Apache-2.0，PaddlePaddle 与 HuggingFace）：
3//! - 颈部：每段 1×1 卷积 + SE 残差统一到 96 通道 → 自顶向下最近邻上采样相加 →
4//!   每段大核深度卷积 + 逐点卷积（到 24 通道）+ SE 残差 → 都上采样到步长 4 → 按 p5、p4、p3、p2 拼接；
5//! - 头：3×3 卷积 + BN + ReLU → 转置卷积 ×2 + BN + ReLU → 转置卷积 ×2 → sigmoid，得到原尺寸概率图。
6
7use pagesift_infer::exec::{Cpu, Exec};
8use pagesift_infer::{Activation, BN_EPS, BatchNorm, Conv2d, ConvTranspose2d, Tensor, Weights};
9
10/// 颈部和头的配置（config.json 顶层）。
11#[derive(Debug, Clone)]
12pub struct NeckConfig {
13    pub in_channels: Vec<usize>,
14    pub out_channels: usize,
15    pub kernel_size: usize,
16}
17
18/// 读配置里的 `kernel_list`（头的三个核大小）。
19// [协议] kernel_list 固定三项：conv_down、conv_up、conv_final
20pub fn head_kernels(v: &serde_json::Value) -> Result<[usize; 3], String> {
21    let a = v["kernel_list"].as_array().ok_or("missing kernel_list")?;
22    let k: Vec<usize> = a
23        .iter()
24        .map(|x| {
25            x.as_u64()
26                .map(|n| n as usize)
27                .ok_or("kernel_list must be integers")
28        })
29        .collect::<Result<_, _>>()?;
30    k.try_into()
31        .map_err(|_| "kernel_list must have 3 numbers".into())
32}
33
34impl NeckConfig {
35    pub fn from_json(v: &serde_json::Value) -> Result<NeckConfig, String> {
36        let ints = |key: &str| -> Result<Vec<usize>, String> {
37            v[key]
38                .as_array()
39                .ok_or(format!("missing {key}"))?
40                .iter()
41                .map(|x| {
42                    x.as_u64()
43                        .map(|n| n as usize)
44                        .ok_or(format!("{key} must be an integer"))
45                })
46                .collect()
47        };
48        if v["interpolate_mode"].as_str().unwrap_or("nearest") != "nearest" {
49            return Err("the neck only supports nearest-neighbor upsampling".into());
50        }
51        let in_channels = ints("layer_list_out_channels")?;
52        if in_channels.len() != 4 {
53            return Err("layer_list_out_channels must have 4 stages".into());
54        }
55        Ok(NeckConfig {
56            in_channels,
57            // [协议] neck_out_channels 缺省 96、dilated_kernel_size 缺省 7（PPOCRV6SmallDetConfig）
58            out_channels: v["neck_out_channels"].as_u64().unwrap_or(96) as usize,
59            kernel_size: v["dilated_kernel_size"].as_u64().unwrap_or(7) as usize,
60        })
61    }
62}
63
64/// 颈部的 SE：全局平均 → 1×1 → ReLU → 1×1 → clamp(0.2x + 0.5, 0, 1)。
65/// 两处都以 `h + se(h)` 的形式出现，等价于逐通道乘 `1 + s`，这里直接原地乘。
66#[derive(Debug, Clone)]
67struct SeResidual {
68    c1: Conv2d,
69    c2: Conv2d,
70}
71
72impl SeResidual {
73    fn load(w: &Weights, p: &str) -> Result<SeResidual, String> {
74        let c = |n: &str| Conv2d::load(w, &format!("{p}.{n}"), (1, 1), (0, 0), (1, 1), 1);
75        Ok(SeResidual {
76            c1: c("conv1")?,
77            c2: c("conv2")?,
78        })
79    }
80
81    fn apply_on<E: Exec>(&self, e: &E, x: &mut E::T) {
82        let pooled = e.global_avg(x);
83        let mut s = e.conv(&self.c1, &pooled);
84        e.act(Activation::Relu, &mut s);
85        let mut s = e.conv(&self.c2, &s);
86        e.se_gate(&mut s);
87        e.scale_channels(x, &s);
88    }
89}
90
91/// `insert_conv`：1×1 卷积（无偏置）+ SE 残差。
92#[derive(Debug, Clone)]
93struct InsertConv {
94    conv: Conv2d,
95    se: SeResidual,
96}
97
98impl InsertConv {
99    fn forward_on<E: Exec>(&self, e: &E, x: &E::T) -> E::T {
100        let mut h = e.conv(&self.conv, x);
101        self.se.apply_on(e, &mut h);
102        h
103    }
104}
105
106/// `input_conv`：大核深度卷积（带偏置）→ 逐点卷积到 1/4 通道 → SE 残差。
107#[derive(Debug, Clone)]
108struct InputConv {
109    depthwise: Conv2d,
110    pointwise: Conv2d,
111    se: SeResidual,
112}
113
114impl InputConv {
115    fn forward_on<E: Exec>(&self, e: &E, x: &E::T) -> E::T {
116        let mut h = e.conv(&self.pointwise, &e.conv(&self.depthwise, x));
117        self.se.apply_on(e, &mut h);
118        h
119    }
120}
121
122/// 颈部。
123#[derive(Debug, Clone)]
124pub struct Neck {
125    insert: Vec<InsertConv>,
126    input: Vec<InputConv>,
127}
128
129/// 最近邻放大整数倍（`F.interpolate(scale_factor=s, mode="nearest")`），`s > 1`。
130pub(super) fn upsample_on<E: Exec>(e: &E, x: &E::T, s: usize) -> E::T {
131    let shape = e.shape(x);
132    e.upsample(x, shape[1] * s, shape[2] * s)
133}
134
135impl Neck {
136    /// `p` 是权重名前缀（PP-OCRv6 里是 `model.neck`）。
137    pub fn load(w: &Weights, p: &str, cfg: &NeckConfig) -> Result<Neck, String> {
138        let k = cfg.kernel_size;
139        let c = cfg.out_channels;
140        let mut insert = Vec::new();
141        let mut input = Vec::new();
142        for i in 0..cfg.in_channels.len() {
143            let q = format!("{p}.insert_conv.{i}");
144            insert.push(InsertConv {
145                conv: Conv2d::load(w, &format!("{q}.in_conv"), (1, 1), (0, 0), (1, 1), 1)?,
146                se: SeResidual::load(w, &format!("{q}.squeeze_excitation_block"))?,
147            });
148            let q = format!("{p}.input_conv.{i}");
149            input.push(InputConv {
150                // [协议] 深度卷积补边 kernel_size // 2、groups = 输入通道、不带空洞
151                depthwise: Conv2d::load(
152                    w,
153                    &format!("{q}.depthwise_convolution"),
154                    (1, 1),
155                    (k / 2, k / 2),
156                    (1, 1),
157                    c,
158                )?,
159                pointwise: Conv2d::load(
160                    w,
161                    &format!("{q}.pointwise_convolution"),
162                    (1, 1),
163                    (0, 0),
164                    (1, 1),
165                    1,
166                )?,
167                se: SeResidual::load(w, &format!("{q}.squeeze_excitation_module"))?,
168            });
169        }
170        Ok(Neck { insert, input })
171    }
172
173    /// 输入骨干的 4 段特征（步长 4、8、16、32），输出步长 4 的融合特征。
174    pub fn forward(&self, feats: &[Tensor]) -> Result<Tensor, String> {
175        self.forward_on(&Cpu, feats)
176    }
177
178    /// 同 [`Neck::forward`]，在执行器上算。
179    pub fn forward_on<E: Exec>(&self, e: &E, feats: &[E::T]) -> Result<E::T, String> {
180        if feats.len() != self.insert.len() {
181            return Err(format!(
182                "the backbone gave {} feature stages, the neck needs {}",
183                feats.len(),
184                self.insert.len()
185            ));
186        }
187        let mut fused: Vec<E::T> = self
188            .insert
189            .iter()
190            .zip(feats)
191            .map(|(m, f)| m.forward_on(e, f))
192            .collect();
193        // 自顶向下：p4 += up(p5)，p3 += up(p4)，p2 += up(p3)
194        for i in (0..fused.len() - 1).rev() {
195            let up = upsample_on(e, &fused[i + 1], 2);
196            if e.shape(&up) != e.shape(&fused[i]) {
197                return Err(format!(
198                    "neck upsampling size mismatch: {:?} and {:?}",
199                    e.shape(&up),
200                    e.shape(&fused[i])
201                ));
202            }
203            e.add_assign(&mut fused[i], &up);
204        }
205        let processed: Vec<E::T> = self
206            .input
207            .iter()
208            .zip(&fused)
209            .enumerate()
210            .map(|(i, (m, f))| {
211                let h = m.forward_on(e, f);
212                if i == 0 {
213                    h
214                } else {
215                    upsample_on(e, &h, 1 << i)
216                }
217            })
218            .collect();
219        // [协议] 按 [p5, p4, p3, p2] 的顺序拼通道
220        let refs: Vec<&E::T> = processed.iter().rev().collect();
221        Ok(e.concat(&refs))
222    }
223}
224
225/// DB 头：输出与网络输入同尺寸的文字概率图。
226#[derive(Debug, Clone)]
227pub struct Head {
228    down: Conv2d,
229    up: ConvTranspose2d,
230    fin: ConvTranspose2d,
231}
232
233impl Head {
234    /// `p` 是权重名前缀（`head`），`k` 是配置里的 `kernel_list`。
235    // [协议] kernel_list 固定三项，见 head_kernels
236    pub fn load(w: &Weights, p: &str, k: [usize; 3]) -> Result<Head, String> {
237        // [协议] conv_down 补边 kernel_list[0] // 2、无偏置；两个转置卷积步长 2、无补边
238        let mut down = Conv2d::load(
239            w,
240            &format!("{p}.conv_down.convolution"),
241            (1, 1),
242            (k[0] / 2, k[0] / 2),
243            (1, 1),
244            1,
245        )?;
246        down.fold(&BatchNorm::load(w, &format!("{p}.conv_down.norm"), BN_EPS)?);
247        let mut up = ConvTranspose2d::load(w, &format!("{p}.conv_up.convolution"), (2, 2), (0, 0))?;
248        up.fold(&BatchNorm::load(w, &format!("{p}.conv_up.norm"), BN_EPS)?);
249        let fin = ConvTranspose2d::load(w, &format!("{p}.conv_final"), (2, 2), (0, 0))?;
250        for (name, got, want) in [
251            ("conv_up", up.weight.shape[2], k[1]),
252            ("conv_final", fin.weight.shape[2], k[2]),
253        ] {
254            if got != want {
255                return Err(format!(
256                    "{name} has kernel size {got}, the config says {want}"
257                ));
258            }
259        }
260        Ok(Head { down, up, fin })
261    }
262
263    /// 返回 `[1, H, W]` 的概率图（已过 sigmoid）。
264    pub fn forward(&self, x: &Tensor) -> Tensor {
265        self.forward_on(&Cpu, x)
266    }
267
268    /// 同 [`Head::forward`]，在执行器上算。
269    pub fn forward_on<E: Exec>(&self, e: &E, x: &E::T) -> E::T {
270        let mut h = e.conv(&self.down, x);
271        e.act(Activation::Relu, &mut h);
272        let mut h = e.conv_transpose(&self.up, &h);
273        e.act(Activation::Relu, &mut h);
274        let mut h = e.conv_transpose(&self.fin, &h);
275        e.act(Activation::Sigmoid, &mut h);
276        h
277    }
278}
pagesift_ocr/det/neck.rs

pagesift_ocr/det/
neck.rs
