layout_check/
layout_check.rs1use std::fmt::Write;
5use std::time::Instant;
6
7fn 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 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}