Skip to main content

det_check/
det_check.rs

1//! 检测检查:载入检测模型,逐张读裸像素图跑检测,把框写成文本(每行 8 个数:四角 x y 和分数)。
2//!     cargo run --profile release-fast -p pagesift-ocr --example det_check -- \
3//!         <模型目录> <输出目录> <裸图1> <宽1> <高1> [<裸图2> <宽2> <高2> ...]
4//! 裸图是按行存放、无行间填充的 8 位像素,灰度或 RGB(由文件大小判断)。
5//! 输出 `<输出目录>/<裸图文件名>.boxes.txt`,耗时打到标准错误。
6use std::io::Write;
7use std::time::Instant;
8
9fn main() {
10    let a: Vec<String> = std::env::args().skip(1).collect();
11    if a.len() < 5 || !(a.len() - 2).is_multiple_of(3) {
12        eprintln!("用法:det_check <模型目录> <输出目录> <裸图> <宽> <高> ...");
13        std::process::exit(2);
14    }
15    let t = Instant::now();
16    let det = pagesift_ocr::det::Det::load(std::path::Path::new(&a[0])).unwrap();
17    eprintln!("载入 {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
18    let out_dir = std::path::Path::new(&a[1]);
19    for c in a[2..].chunks(3) {
20        let data = std::fs::read(&c[0]).unwrap();
21        let (w, h): (usize, usize) = (c[1].parse().unwrap(), c[2].parse().unwrap());
22        let channels = data.len() / (w * h);
23        assert!(channels == 1 || channels == 3, "文件大小与宽高不符");
24        let img = pagesift_ocr::Image {
25            width: w,
26            height: h,
27            channels,
28            data,
29        };
30        let name = std::path::Path::new(&c[0])
31            .file_name()
32            .unwrap()
33            .to_string_lossy();
34        let t = Instant::now();
35        let (prob, pw, ph) = det.prob_map(&img).unwrap();
36        let net_ms = t.elapsed().as_secs_f64() * 1e3;
37        let boxes = det.boxes_from_prob(&prob, pw, ph, w, h);
38        let ms = t.elapsed().as_secs_f64() * 1e3;
39        eprintln!("{name}: 网络 {net_ms:.0} ms,后处理 {:.0} ms", ms - net_ms);
40        // 设了 DET_DUMP_PROB 就把概率图另存成 8 位 PGM
41        if std::env::var_os("DET_DUMP_PROB").is_some() {
42            let mut pgm = format!("P5\n{pw} {ph}\n255\n").into_bytes();
43            pgm.extend(prob.iter().map(|v| (v * 255.0).round() as u8));
44            std::fs::write(out_dir.join(format!("{name}.prob.pgm")), pgm).unwrap();
45        }
46        let mut f = std::fs::File::create(out_dir.join(format!("{name}.boxes.txt"))).unwrap();
47        for (q, s) in &boxes {
48            for (x, y) in q {
49                write!(f, "{x} {y} ").unwrap();
50            }
51            writeln!(f, "{s:.4}").unwrap();
52        }
53        eprintln!("{name}: {w}×{h},{} 个框,{ms:.0} ms", boxes.len());
54    }
55}