Skip to main content

layout_check/
layout_check.rs

1//! 版面检测检查:载入模型,读一张 PGM/PPM(P5/P6)图,把区域写成文本(每行 x0 y0 x1 y1 分数 类别 细分类)。
2//!     cargo run --profile release-fast -p pagesift-ocr --example layout_check -- <模型目录> <图> <输出> [分数阈值]
3//! 耗时打到标准错误。
4use std::fmt::Write;
5use std::time::Instant;
6
7/// 读 P5/P6:跳过注释,头之后是一个空白再接像素。
8fn read_pnm(path: &str) -> pagesift_ocr::Image {
9    let b = std::fs::read(path).unwrap();
10    let mut i = 0;
11    let mut token = || {
12        loop {
13            while b[i].is_ascii_whitespace() {
14                i += 1;
15            }
16            if b[i] != b'#' {
17                break;
18            }
19            while b[i] != b'\n' {
20                i += 1;
21            }
22        }
23        let s = i;
24        while !b[i].is_ascii_whitespace() {
25            i += 1;
26        }
27        String::from_utf8_lossy(&b[s..i]).to_string()
28    };
29    let magic = token();
30    let channels = match magic.as_str() {
31        "P5" => 1,
32        "P6" => 3,
33        m => panic!("不支持的格式 {m}"),
34    };
35    let width: usize = token().parse().unwrap();
36    let height: usize = token().parse().unwrap();
37    let max: usize = token().parse().unwrap();
38    assert_eq!(max, 255, "只支持 8 位图");
39    let data = b[i + 1..i + 1 + width * height * channels].to_vec();
40    pagesift_ocr::Image {
41        width,
42        height,
43        channels,
44        data,
45    }
46}
47
48fn main() {
49    let a: Vec<String> = std::env::args().skip(1).collect();
50    assert!(
51        a.len() == 3 || a.len() == 4,
52        "用法:layout_check <模型目录> <图> <输出> [分数阈值]"
53    );
54    let t = Instant::now();
55    let mut m = pagesift_ocr::layout::Layout::load(std::path::Path::new(&a[0])).unwrap();
56    if let Some(th) = a.get(3) {
57        m.threshold = th.parse().unwrap();
58    }
59    eprintln!("载入 {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
60    let img = read_pnm(&a[1]);
61    // 设了 LAYOUT_RAW 就打印原始输出的左上角,和参考实现的集成测试对照
62    if std::env::var_os("LAYOUT_RAW").is_some() {
63        let r = m.raw(&img).unwrap();
64        let (_, k) = r.logits.tc();
65        let (q, _) = r.order.tc();
66        for i in 0..3 {
67            println!("logits {:?}", &r.logits.data[i * k..i * k + 3]);
68            println!("boxes {:?}", &r.boxes[i][..3]);
69            println!("order {:?}", &r.order.data[i * q..i * q + 3]);
70        }
71    }
72    let t = Instant::now();
73    let rs = m.detect(&img).unwrap();
74    let ms = t.elapsed().as_secs_f64() * 1e3;
75    eprintln!(
76        "{}: {}×{},{} 个区域,{ms:.0} ms",
77        a[1],
78        img.width,
79        img.height,
80        rs.len()
81    );
82    let mut s = String::new();
83    for r in &rs {
84        writeln!(
85            s,
86            "{:.4} {:.4} {:.4} {:.4} {:.4} {} {}",
87            r.x0, r.y0, r.x1, r.y1, r.score, r.label, r.kind
88        )
89        .unwrap();
90    }
91    std::fs::write(&a[2], s).unwrap();
92}