CNWei
7d159c5702
- 在 ddddocr-core 中定义 ModelBuilder Trait 及其错误类型 - ddddocr-ort 支持 use_gpu、device_id 及 num_threads 链式配置与 CUDA 硬件加速 - ddddocr-tract 基于 multithread-mm 特性支持 CPU 线程数控制 - 支持基于 tract-linalg 配置推理线程数，显式引入 tract-linalg 的 multithread-mm 特性，控制 GEMM 算子并发 - 优化线程池加载策略，适配 Tokio 异步及 CLI 等多场景
249 lines
8.7 KiB
Rust
249 lines
8.7 KiB
Rust
use anyhow::Context;
|
||
use ddddocr_core::{DetectionResult, ModelBuilder};
|
||
use ddddocr_core::{Detector, ModelMetadata, Normalization, Slider};
|
||
// 假设你的包名是这个
|
||
use ddddocr_ort::{DetSession, OcrBuilder, OcrSession};
|
||
use image::{DynamicImage, ImageBuffer, Luma, Rgb};
|
||
use std::fs;
|
||
use std::path::Path;
|
||
|
||
mod char_slice;
|
||
use char_slice::CHARSET_BETA;
|
||
use ddddocr_core::Resize;
|
||
|
||
use ddddocr_ort::loader::OrtModelLoader;
|
||
|
||
fn load_image<P: AsRef<Path>>(path: P) -> anyhow::Result<image::DynamicImage> {
|
||
// 1. 先将泛型转为具体的 &Path 引用
|
||
let path_ref = path.as_ref();
|
||
|
||
// 2. 调用 open 时传入引用（utils::open 支持 AsRef<Path>）
|
||
image::open(path_ref).map_err(|e| {
|
||
// 3. 此时 path_ref 依然有效，可以安全地在闭包中使用
|
||
anyhow::anyhow!("无法加载图片 {:?}: {}", path_ref, e)
|
||
})
|
||
}
|
||
/// 将检测结果绘制在图像上并保存
|
||
fn save_debug_image(
|
||
dynamic_img: &DynamicImage, // 【优化点 1】直接传入解码好的引用，拒绝重复解码
|
||
bboxes: &[DetectionResult], // 【修改点 1】类型改为自定义结构体切片
|
||
output_path: &str,
|
||
) -> anyhow::Result<()> {
|
||
// 删除了原本的 let dynamic_img = image::load_from_memory(image_bytes)?;
|
||
let mut img = dynamic_img.to_rgb8();
|
||
let (width, height) = img.dimensions();
|
||
let red = Rgb([255u8, 0, 0]);
|
||
|
||
for bbox in bboxes {
|
||
// 【修改点 2】将原来的索引 bbox[0].. 改为结构体字段访问 .x1, .y1 ..
|
||
let x1 = bbox.x1.max(0).min(width as i32 - 1) as u32;
|
||
let y1 = bbox.y1.max(0).min(height as i32 - 1) as u32;
|
||
let x2 = bbox.x2.max(0).min(width as i32 - 1) as u32;
|
||
let y2 = bbox.y2.max(0).min(height as i32 - 1) as u32;
|
||
|
||
// 绘制横向线条
|
||
for x in x1..=x2 {
|
||
img.put_pixel(x, y1, red);
|
||
img.put_pixel(x, y2, red);
|
||
if y1 + 1 < height {
|
||
img.put_pixel(x, y1 + 1, red);
|
||
}
|
||
if y2.saturating_sub(1) > 0 {
|
||
img.put_pixel(x, y2 - 1, red);
|
||
}
|
||
}
|
||
// 绘制纵向线条
|
||
for y in y1..=y2 {
|
||
img.put_pixel(x1, y, red);
|
||
img.put_pixel(x2, y, red);
|
||
if x1 + 1 < width {
|
||
img.put_pixel(x1 + 1, y, red);
|
||
}
|
||
if x2.saturating_sub(1) > 0 {
|
||
img.put_pixel(x2 - 1, y, red);
|
||
}
|
||
}
|
||
}
|
||
|
||
img.save(output_path)?;
|
||
Ok(())
|
||
}
|
||
#[allow(dead_code)]
|
||
fn save_rust_result(result: &ImageBuffer<Luma<f32>, Vec<f32>>, filename: &str) {
|
||
let (width, height) = result.dimensions();
|
||
|
||
// 1. 寻找最值进行归一化
|
||
let mut max_val = f32::MIN;
|
||
let mut min_val = f32::MAX;
|
||
for p in result.pixels() {
|
||
if p.0[0] > max_val {
|
||
max_val = p.0[0];
|
||
}
|
||
if p.0[0] < min_val {
|
||
min_val = p.0[0];
|
||
}
|
||
}
|
||
|
||
// 2. 创建 8 位灰度图
|
||
let mut out_buf = ImageBuffer::new(width, height);
|
||
for y in 0..height {
|
||
for x in 0..width {
|
||
let val = result.get_pixel(x, y).0[0];
|
||
let normalized = if max_val > min_val {
|
||
((val - min_val) / (max_val - min_val) * 255.0) as u8
|
||
} else {
|
||
0u8
|
||
};
|
||
out_buf.put_pixel(x, y, Luma([normalized]));
|
||
}
|
||
}
|
||
|
||
// 3. 保存
|
||
DynamicImage::ImageLuma8(out_buf).save(filename).unwrap();
|
||
println!("Rust 结果热力图已保存至: {}", filename);
|
||
}
|
||
#[test]
|
||
fn test_full_classification() {
|
||
let model = OrtModelLoader::builder().use_gpu(true)
|
||
.model_for_path("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_sml2h3_f32.onnx")
|
||
.expect("模型加载失败");
|
||
let metadata = ModelMetadata::from_static_slice(
|
||
CHARSET_BETA,
|
||
false,
|
||
Resize::DynamicWidth(64),
|
||
1,
|
||
Normalization::MinusOneToOne,
|
||
);
|
||
// 1. 初始化模型
|
||
let ocr = OcrSession::new(model, metadata);
|
||
// 2. 加载测试图片
|
||
let img =
|
||
image::open("D:/CNWei/CNW/Rust/ddddocr-rs/samples/code2.png").expect("测试图片不存在");
|
||
|
||
// 3. 执行识别
|
||
// let result = Ocr::new(&ocr)
|
||
// .predict(&img)
|
||
// .expect("识别过程出错")
|
||
// .into_text();
|
||
let result = OcrBuilder::new()
|
||
.build(&ocr)
|
||
.predict(&img)
|
||
.expect("识别过程出错")
|
||
.into_text();
|
||
|
||
println!("识别结果: {}", result);
|
||
assert!(!result.is_empty());
|
||
}
|
||
#[test]
|
||
fn test_det_load() -> anyhow::Result<()> {
|
||
let det_model = OrtModelLoader::builder()
|
||
.model_for_path("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_det.onnx")
|
||
.expect("模型加载失败");
|
||
let det = DetSession::new(det_model);
|
||
let image_path = "D:/CNWei/CNW/Rust/ddddocr-rs/samples/det1.png";
|
||
let image_bytes =
|
||
fs::read(image_path).map_err(|e| anyhow::anyhow!("无法读取图片 {}: {}", image_path, e))?;
|
||
|
||
println!("图片读取成功，字节大小: {}", image_bytes.len());
|
||
|
||
// 【修改点 1】将字节流解码为统一的 DynamicImage
|
||
let img = image::load_from_memory(&image_bytes)
|
||
.map_err(|e| anyhow::anyhow!("图片解码失败: {}", e))?;
|
||
|
||
// 【修改点 2】传入统一的 &DynamicImage 引用
|
||
let bboxes = Detector::new(&det).predict(&img)?;
|
||
// println!("{:?}", det);
|
||
println!("检测到的目标数量: {}", bboxes.len());
|
||
|
||
if bboxes.is_empty() {
|
||
println!("未检测到任何目标。");
|
||
} else {
|
||
// 如果 save_debug_image 报错，记得去把它的入参类型和内部访问也改为 DetectionResult
|
||
save_debug_image(
|
||
&img,
|
||
&bboxes,
|
||
"D:/CNWei/CNW/Rust/ddddocr-rs/samples/result.jpg",
|
||
)?;
|
||
|
||
for (i, bbox) in bboxes.iter().enumerate() {
|
||
// 【修改点 3】将原来的 bbox[0].. 索引访问改为结构体字段访问
|
||
println!("目标 [{}]: {}", i, bbox);
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn test_real_slide_match() {
|
||
let engine = Slider::new();
|
||
|
||
// 1. 加载你准备好的测试图
|
||
// 假设图片放在项目根目录下的 assets 文件夹
|
||
let target_img = load_image("D:/CNWei/CNW/Rust/ddddocr-rs/samples/hua.png")
|
||
.expect("请确保 samples/hua.png 存在");
|
||
let bg_img = load_image("D:/CNWei/CNW/Rust/ddddocr-rs/samples/huatu.png")
|
||
.expect("请确保 samples/huatu.png 存在");
|
||
|
||
// 2. 执行匹配
|
||
// 如果是那种带有明显阴影边缘的复杂滑块，建议 simple_target 传 false
|
||
let start = std::time::Instant::now();
|
||
let result = engine
|
||
.slide_match(&target_img, &bg_img, false)
|
||
.expect("Slide match 执行失败");
|
||
let duration = start.elapsed();
|
||
|
||
// 3. 打印结果
|
||
println!("-------------------------------------------");
|
||
println!("{}", result);
|
||
println!("耗时: {:?}", duration);
|
||
println!("-------------------------------------------");
|
||
|
||
// 验证基本逻辑：坐标不应为 0 (除非匹配失败)
|
||
assert_eq!(result.target_x, 237);
|
||
assert_eq!(result.target_y, 77);
|
||
assert!(result.confidence > 0.0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_real_slide_comparison() {
|
||
let engine = Slider::new();
|
||
|
||
// 1. 加载你准备好的测试图
|
||
// 假设图片放在项目根目录下的 assets 文件夹
|
||
let target_img = load_image("D:/CNWei/CNW/Rust/ddddocr-rs/samples/ken.jpg")
|
||
.expect("请确保 samples/ken.jpg 存在");
|
||
let bg_img = load_image("D:/CNWei/CNW/Rust/ddddocr-rs/samples/kenyuan.jpg")
|
||
.expect("请确保 samples/kenyuan.jpg 存在");
|
||
|
||
// 2. 执行匹配
|
||
// 如果是那种带有明显阴影边缘的复杂滑块，建议 simple_target 传 false
|
||
let start = std::time::Instant::now();
|
||
let result = engine
|
||
.slide_comparison(&target_img, &bg_img)
|
||
.expect("Slide match 执行失败");
|
||
let duration = start.elapsed();
|
||
|
||
// 3. 打印结果
|
||
println!("-------------------------------------------");
|
||
println!("滑块匹配测试结果:");
|
||
println!("检测坐标: [x: {}, y: {}]", result.target_x, result.target_y);
|
||
println!("置信度: {:.4}", result.confidence);
|
||
println!("耗时: {:?}", duration);
|
||
println!("-------------------------------------------");
|
||
|
||
// 验证基本逻辑：坐标不应为 0 (除非匹配失败)
|
||
assert_eq!(result.target_x, 171);
|
||
assert_eq!(result.target_y, 90);
|
||
assert!(result.confidence > 0.0);
|
||
}
|
||
#[test]
|
||
fn test_resolve_shape_logic_direct() {
|
||
// 创建一个哑 ModelLoader 实例（session 用不上，因为我们直接测私有方法）
|
||
let loader = OrtModelLoader::builder()
|
||
.model_for_path(
|
||
// "D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_sml2h3_f32.onnx",
|
||
"D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_huashi666_i64.onnx",
|
||
)
|
||
.expect("建立测试模型图失败");
|
||
}
|