pagesift_ocr/det/
mod.rs

1//! 文字检测（PP-OCRv6 small 和 medium det：LCNetV4 骨干 + 颈部（small 是 RepLKFPN，medium 是
2//! LK-PAN + 类内关系块）+ DB 头 + DB 后处理）。
3//!
4//! 流程：预处理（[`preprocess`]）→ 骨干（[`crate::lcnet`]）→ 颈部与头（[`neck`]）得到与
5//! 网络输入同尺寸的文字概率图 → DB 后处理（[`postprocess`]，几何在 [`geometry`]、轮廓在
6//! [`contour`]）→ 映射回原图坐标、排成阅读顺序。
7
8pub mod contour;
9pub mod geometry;
10pub mod neck;
11pub mod neck_medium;
12pub mod postprocess;
13pub mod preprocess;
14
15use std::path::Path;
16
17use pagesift_infer::Tensor;
18use pagesift_infer::exec::{Cpu, Exec};
19
20use crate::{Image, Quad, lcnet};
21use postprocess::{DbParams, ProbMap};
22use preprocess::Preprocess;
23
24/// [策略] 阅读顺序里“同一行”的容差（原图像素），同 PaddleOCR `sorted_boxes` 的 10
25const LINE_TOL: f32 = 10.0;
26
27/// 两档模型的颈部。
28enum AnyNeck {
29    Small(neck::Neck),
30    Medium(neck_medium::MediumNeck),
31}
32
33/// 载入好的检测模型。
34pub struct Det {
35    pre: Preprocess,
36    backbone: lcnet::Backbone,
37    neck: AnyNeck,
38    head: neck::Head,
39    /// DB 后处理参数；缺省同参考实现，可按需调整。
40    pub params: DbParams,
41}
42
43impl Det {
44    /// `dir` 是模型目录：model.safetensors、config.json、preprocessor_config.json。
45    pub fn load(dir: &Path) -> Result<Det, String> {
46        let cfg = crate::read_json(dir, "config.json")?;
47        let pre = Preprocess::from_json(&crate::read_json(dir, "preprocessor_config.json")?)?;
48        let w = crate::read_weights(dir)?;
49        let bcfg = lcnet::Config::from_json(&cfg)?;
50        // [协议] 权重名前缀：model.backbone.encoder、model.neck、head；
51        // model_type 区分两档结构（config.json）
52        let neck = match cfg["model_type"].as_str().unwrap_or("") {
53            "pp_ocrv6_medium_det" => {
54                AnyNeck::Medium(neck_medium::MediumNeck::load(&w, "model.neck")?)
55            }
56            _ => AnyNeck::Small(neck::Neck::load(
57                &w,
58                "model.neck",
59                &neck::NeckConfig::from_json(&cfg)?,
60            )?),
61        };
62        Ok(Det {
63            pre,
64            backbone: lcnet::Backbone::load(&w, "model.backbone.encoder", &bcfg)?,
65            neck,
66            head: neck::Head::load(&w, "head", neck::head_kernels(&cfg)?)?,
67            params: DbParams::default(),
68        })
69    }
70
71    /// 整页图 → 文字概率图（网络输入尺寸，按行存放）及其宽、高。
72    pub fn prob_map(&self, page: &Image) -> Result<(Vec<f32>, usize, usize), String> {
73        let x = self.preprocess(page)?;
74        let prob = self.prob_on(&Cpu, &x)?;
75        let (_, h, w) = prob.chw();
76        Ok((prob.data, w, h))
77    }
78
79    /// 预处理整页图成网络输入 `[3, H, W]`。
80    pub fn preprocess(&self, page: &Image) -> Result<Tensor, String> {
81        if page.width == 0 || page.height == 0 {
82            return Err("empty image".into());
83        }
84        if page.channels != 1 && page.channels != 3 {
85            return Err(format!(
86                "only grayscale or RGB images are supported, got {} channels",
87                page.channels
88            ));
89        }
90        Ok(self.pre.run(page))
91    }
92
93    /// 网络部分：`[3, H, W]` → `[1, H, W]` 的概率图，在执行器上算。
94    pub fn prob_on<E: Exec>(&self, e: &E, x: &E::T) -> Result<E::T, String> {
95        let feats = self.backbone.forward_on(e, x);
96        let fused = match &self.neck {
97            AnyNeck::Small(n) => n.forward_on(e, &feats)?,
98            AnyNeck::Medium(n) => n.forward_on(e, &feats)?,
99        };
100        drop(feats);
101        let prob = self.head.forward_on(e, &fused);
102        let (got, want) = (e.shape(&prob), e.shape(x));
103        if got[1..] != want[1..] {
104            return Err(format!(
105                "probability map size {got:?} does not match the input {want:?}"
106            ));
107        }
108        Ok(prob)
109    }
110
111    /// 检测整页图里的文字行，返回各行的四边形（原图像素坐标），从上到下、从左到右排好。
112    pub fn detect(&self, page: &Image) -> Result<Vec<Quad>, String> {
113        Ok(self
114            .detect_scored(page)?
115            .into_iter()
116            .map(|(q, _)| q)
117            .collect())
118    }
119
120    /// 同 [`Det::detect`]，并给出每个框的平均概率。
121    pub fn detect_scored(&self, page: &Image) -> Result<Vec<(Quad, f32)>, String> {
122        let (prob, w, h) = self.prob_map(page)?;
123        Ok(self.boxes_from_prob(&prob, w, h, page.width, page.height))
124    }
125
126    /// 只做后处理：`w × h` 的概率图 → 原图（`page_w × page_h`）坐标下排好序的框和分数。
127    pub fn boxes_from_prob(
128        &self,
129        prob: &[f32],
130        w: usize,
131        h: usize,
132        page_w: usize,
133        page_h: usize,
134    ) -> Vec<(Quad, f32)> {
135        let pm = ProbMap {
136            data: prob,
137            width: w,
138            height: h,
139        };
140        let mut boxes = postprocess::boxes_from_prob(&pm, page_w, page_h, &self.params);
141        postprocess::sort_reading_order(&mut boxes, LINE_TOL);
142        boxes
143    }
144}
pagesift_ocr/det/mod.rs

pagesift_ocr/det/
mod.rs
