use image::{GenericImageView, Rgb};
use imageproc::drawing::draw_hollow_rect_mut;
use imageproc::rect::Rect;
use std::env;
use ocr_rs::{DetOptions, OcrEngine, OcrEngineConfig};
fn main() {
let args: Vec<String> = env::args().collect();
if args.len() < 5 {
eprintln!("用法: visualize_det <det_model> <rec_model> <keys> <image> [output]");
return;
}
let det_model = &args[1];
let rec_model = &args[2];
let keys_path = &args[3];
let image_path = &args[4];
let output_path = args.get(5).map(|s| s.as_str()).unwrap_or("det_result.png");
println!("加载模型...");
let config = OcrEngineConfig::new().with_det_options(DetOptions::fast());
let engine =
OcrEngine::new(det_model, rec_model, keys_path, Some(config))
.expect("创建引擎失败");
println!("加载图像: {}", image_path);
let image = image::open(image_path).expect("加载图像失败");
let (width, height) = image.dimensions();
println!("图像尺寸: {}x{}", width, height);
println!("执行检测...");
let boxes = engine.detect(&image).expect("检测失败");
println!("检测到 {} 个边界框", boxes.len());
let mut output_image = image.to_rgb8();
let colors = [
Rgb([255u8, 0, 0]), Rgb([0, 255, 0]), Rgb([0, 0, 255]), Rgb([255, 255, 0]), Rgb([255, 0, 255]), Rgb([0, 255, 255]), Rgb([255, 128, 0]), Rgb([128, 0, 255]), ];
for (i, text_box) in boxes.iter().enumerate() {
let color = colors[i % colors.len()];
let rect = Rect::at(text_box.rect.left(), text_box.rect.top())
.of_size(text_box.rect.width(), text_box.rect.height());
draw_hollow_rect_mut(&mut output_image, rect, color);
if text_box.rect.left() > 0 && text_box.rect.top() > 0 {
let rect2 = Rect::at(text_box.rect.left() - 1, text_box.rect.top() - 1)
.of_size(text_box.rect.width() + 2, text_box.rect.height() + 2);
draw_hollow_rect_mut(&mut output_image, rect2, color);
}
println!(
"[{:2}] ({:4}, {:4}) {:4}x{:3}",
i,
text_box.rect.left(),
text_box.rect.top(),
text_box.rect.width(),
text_box.rect.height()
);
}
output_image.save(output_path).expect("保存图像失败");
println!("\n结果已保存到: {}", output_path);
}