refactor(ocr): 优化 color_filter.rs
- 重构 `OcrBuilder` 移除is_pixel_matched,filter_image。 - 优化 `OcrBuilder` 的color_filter方法(部分逻辑转移给merge_to_vec) 。 - 新增 `ColorFilter` 特征增加merge_to_vec方法。
This commit is contained in:
@@ -1,18 +1,21 @@
|
||||
use std::borrow::Cow;
|
||||
use crate::charset::{TokenFilter, ValidationCtx};
|
||||
use crate::model_metadata::ModelMetadata;
|
||||
use crate::models::base::ModelArgs;
|
||||
use crate::models::loader::{ModelLoader, ModelSession, ModelType};
|
||||
use crate::utils::color_filter::{ColorFilter, HsvRange};
|
||||
use crate::utils::color_filter::{filter_image, ColorFilter, HsvRange};
|
||||
use crate::utils::image_io::png_rgba_white_preprocess;
|
||||
use crate::utils::image_processor::{convert_to_grayscale, resize_image};
|
||||
use anyhow::Context;
|
||||
use image::DynamicImage;
|
||||
use anyhow::{anyhow, Result};
|
||||
use image::{DynamicImage, ImageBuffer, Rgb};
|
||||
use std::collections::HashSet;
|
||||
use tract_onnx::prelude::tract_ndarray::s;
|
||||
use tract_onnx::prelude::{
|
||||
DatumType, Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp, tract_ndarray, tvec,
|
||||
};
|
||||
|
||||
// 引入 cv_ops 模块中的 OpenCV HSV 转换算子
|
||||
use crate::utils::cv_ops::rgb_to_opencv_hsv;
|
||||
// 颜色过滤的自定义范围:(低值RGB, 高值RGB)
|
||||
pub type ColorRange = ((u8, u8, u8), (u8, u8, u8));
|
||||
|
||||
@@ -185,19 +188,16 @@ impl<'a> OcrBuilder<'a> {
|
||||
existing_vec.reserve(total_capacity);
|
||||
|
||||
// 尝试追加倒入
|
||||
filter.append_ranges(existing_vec) ;
|
||||
|
||||
filter.append_ranges(existing_vec);
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
if let Some(v) =&mut ranges {
|
||||
if let Some(v) = &mut ranges {
|
||||
v.sort_unstable();
|
||||
v.dedup();
|
||||
}
|
||||
// 5. 更新状态
|
||||
self.color_filter = Ok(ranges);
|
||||
},
|
||||
}
|
||||
Err(_) => return self,
|
||||
};
|
||||
self
|
||||
@@ -252,7 +252,27 @@ impl<'a> OcrBuilder<'a> {
|
||||
}
|
||||
impl<'a> OcrBuilder<'a> {
|
||||
pub fn predict(&self, image: &DynamicImage) -> anyhow::Result<String> {
|
||||
let tensor = self.preprocess_image(image)?;
|
||||
println!("当前颜色过滤器状态: {:?}", self.color_filter);
|
||||
// =====================================================================
|
||||
// 管道节点 1: 颜色过滤流水线
|
||||
// 使用 Cow (Copy-On-Write) 智能指针。
|
||||
// 如果未开启过滤,img_cow 内部只是持有原图的【只读借用】,发生【零内存分配】!
|
||||
// =====================================================================
|
||||
let img_cow = match &self.color_filter {
|
||||
Err(err_msg) => {
|
||||
return Err(anyhow::anyhow!("颜色过滤器初始化失败,全链路短路: {}", err_msg));
|
||||
}
|
||||
Ok(None) => {
|
||||
// 核心优化点:直接借用原图,不发生任何克隆
|
||||
Cow::Borrowed(image)
|
||||
}
|
||||
Ok(Some(ranges)) => {
|
||||
// 只有真正需要过滤时,才在内部提取像素并生成清洗后的 Owned 新图
|
||||
let filtered_img = filter_image(image, ranges)?;
|
||||
Cow::Owned(filtered_img)
|
||||
}
|
||||
};
|
||||
let tensor = self.preprocess_image(&img_cow)?;
|
||||
|
||||
let raw_tensor = self.ocr.inference(tensor)?;
|
||||
let raw_indices = self.ocr.extract_indices_from_tensor(&raw_tensor)?;
|
||||
@@ -266,15 +286,17 @@ impl<'a> OcrBuilder<'a> {
|
||||
/// 负责:透明背景修复 -> 灰度化 -> 按比例 Resize -> 归一化 -> 4维张量转换
|
||||
fn preprocess_image(&self, img: &DynamicImage) -> anyhow::Result<Tensor> {
|
||||
// A. 修复 PNG 透明背景 (内部逻辑你之前已实现)
|
||||
let _ = if self.png_fix && img.color().has_alpha() {
|
||||
png_rgba_white_preprocess(img)
|
||||
let current_img = if self.png_fix && img.color().has_alpha() {
|
||||
// 只有满足条件才去触发分配,生成新图
|
||||
Cow::Owned(png_rgba_white_preprocess(img))
|
||||
} else {
|
||||
img.clone()
|
||||
// 正常情况下,仅仅是再次安全借用,无开销
|
||||
Cow::Borrowed(img)
|
||||
};
|
||||
|
||||
let h = 64u32;
|
||||
let w = (img.width() as f32 * (h as f32 / img.height() as f32)) as u32;
|
||||
let gray_img = convert_to_grayscale(img);
|
||||
let w = (current_img.width() as f32 * (h as f32 / current_img.height() as f32)) as u32;
|
||||
let gray_img = convert_to_grayscale(¤t_img);
|
||||
let resized = resize_image(&gray_img, w, h);
|
||||
// resized.save("debug_preprocessed.png").unwrap();
|
||||
// 1. 预处理:转灰度 -> Resize -> 归一化
|
||||
@@ -291,6 +313,10 @@ impl<'a> OcrBuilder<'a> {
|
||||
|
||||
Ok(tensor)
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
}
|
||||
impl<'a> OcrBuilder<'a> {
|
||||
pub fn get_valid_indices(&self) -> HashSet<usize> {
|
||||
|
||||
Reference in New Issue
Block a user