1use 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 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}