refactor: 优化图像输入源设计,重构为零污染的 TryFromImage 特征
- 移除原有的 ImageInput 枚举,避免运行时匹配与所有权限制 - 引入自定义 TryFromImage 特征,专用于将不同来源安全转换为 ImageSource - 优化错误处理新增 InvalidBase64Header 错误信息 - 迁移 load_image_from_input 到 image_helper.rs 为后续剥离到业务层做准备
This commit is contained in:
@@ -21,8 +21,12 @@ use thiserror::Error;
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
pub enum DdddError {
|
||||
// 【新增】专门处理文件读取、路径不存在等原生 I/O 错误
|
||||
#[error("系统网络或文件 I/O 异常: {0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
|
||||
#[error("图像预处理失败: {0}")]
|
||||
PreprocessError(String),
|
||||
PreprocessError(#[from] ImagePreprocessReason),
|
||||
|
||||
#[error("模型推理引擎内部发生异常: {0}")]
|
||||
EngineError(#[from] anyhow::Error),
|
||||
@@ -46,5 +50,27 @@ pub enum DdddError {
|
||||
InternalError(String),
|
||||
}
|
||||
|
||||
/// 专门服务于预处理的子错误枚举,保留全部底层上下文
|
||||
#[derive(Error, Debug)]
|
||||
pub enum ImagePreprocessReason {
|
||||
#[error("图片加载或文件 I/O 失败: {0}")]
|
||||
ImageIo(#[from] image::ImageError),
|
||||
|
||||
#[error("图片转矩阵矩阵(ndarray)失败: {0}")]
|
||||
NdarrayError(#[from] ndarray::ShapeError),
|
||||
|
||||
#[error("Base64 解码失败: {0}")]
|
||||
Base64(#[from] base64::DecodeError),
|
||||
|
||||
#[error("Base64 头部格式不正确,缺少 ';base64,' 分隔符")]
|
||||
InvalidBase64Header,
|
||||
|
||||
#[error("不支持的通道数: {0}")]
|
||||
UnsupportedChannels(usize),
|
||||
|
||||
#[error("其他预处理错误: {0}")]
|
||||
Custom(String),
|
||||
}
|
||||
|
||||
/// 统一用我们自己的 DdddError 包装 Result
|
||||
pub type Result<T> = std::result::Result<T, DdddError>;
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use crate::utils::image_io::image_to_ndarray;
|
||||
use crate::utils::image_io::{image_to_ndarray,ColorMode};
|
||||
use crate::utils::image_processor;
|
||||
use crate::utils::image_processor::{abs_diff, min_max_loc, ndarray_to_luma8, rgb_to_gray};
|
||||
use anyhow::{Result, anyhow};
|
||||
@@ -42,8 +42,8 @@ impl Slider {
|
||||
background_image: &DynamicImage,
|
||||
simple_target: bool,
|
||||
) -> Result<SlideResult> {
|
||||
let target_array = image_to_ndarray(target_image);
|
||||
let background_array = image_to_ndarray(background_image);
|
||||
let target_array = image_to_ndarray(target_image,ColorMode::RGB)?;
|
||||
let background_array = image_to_ndarray(background_image,ColorMode::RGB)?;
|
||||
|
||||
self.perform_slide_match(target_array.view(), background_array.view(), simple_target)
|
||||
}
|
||||
@@ -55,8 +55,8 @@ impl Slider {
|
||||
background_image: &DynamicImage,
|
||||
) -> Result<SlideResult> {
|
||||
// 1. 转换为 ndarray (HWC RGB)
|
||||
let target_array = image_to_ndarray(target_image);
|
||||
let background_array = image_to_ndarray(background_image);
|
||||
let target_array = image_to_ndarray(target_image,ColorMode::RGB)?;
|
||||
let background_array = image_to_ndarray(background_image,ColorMode::RGB)?;
|
||||
|
||||
// 2. 执行比较逻辑 (对应 _perform_slide_comparison)
|
||||
self.perform_slide_comparison(target_array.view(), background_array.view())
|
||||
|
||||
@@ -1,5 +1,8 @@
|
||||
pub mod image_io;
|
||||
pub mod image_processor;
|
||||
mod tensor_transform;
|
||||
mod image_helper;
|
||||
|
||||
// 对外统一暴露干净的 API 语义层
|
||||
pub use tensor_transform::normalize_ocr_logits;
|
||||
pub use image_io::{ColorMode};
|
||||
|
||||
102
ddddocr-core/src/utils/image_helper.rs
Normal file
102
ddddocr-core/src/utils/image_helper.rs
Normal file
@@ -0,0 +1,102 @@
|
||||
use crate::error::{DdddError, ImagePreprocessReason, Result};
|
||||
use crate::utils::image_io::{base64_to_image, ndarray_to_pil_image};
|
||||
use image::DynamicImage;
|
||||
use ndarray::ArrayViewD;
|
||||
use std::path::Path;
|
||||
use std::path::PathBuf;
|
||||
|
||||
pub struct Base64<'a>(pub &'a str);
|
||||
/// 专属图像输入源转换器
|
||||
pub struct ImageSource {
|
||||
inner: DynamicImage,
|
||||
}
|
||||
// 1. 将 into_inner 优化为 into_image,符合 Rust 官方命名规范
|
||||
impl ImageSource {
|
||||
/// 消耗当前包装器,获取最终的 DynamicImage
|
||||
pub fn into_image(self) -> DynamicImage {
|
||||
self.inner
|
||||
}
|
||||
}
|
||||
|
||||
pub trait TryFromImage<T>: Sized {
|
||||
// 唯一的转换入口,通过目标类型来调用
|
||||
fn try_from_image(value: T) -> Result<Self>;
|
||||
}
|
||||
|
||||
// 1. 本身是 DynamicImage
|
||||
impl TryFromImage<DynamicImage> for ImageSource {
|
||||
fn try_from_image(img: DynamicImage) -> Result<Self> {
|
||||
Ok(Self { inner: img })
|
||||
}
|
||||
}
|
||||
// 2.2 路径类型 A: &str (最常用)
|
||||
impl<'a> TryFromImage<&'a str> for ImageSource {
|
||||
fn try_from_image(path_or_b64: &'a str) -> Result<Self> {
|
||||
// 1. 嗅探:如果包含 Base64 特征
|
||||
if path_or_b64.starts_with("data:image/") && path_or_b64.contains(";base64,") {
|
||||
// 提取出真正的 base64 数据部分
|
||||
let (_, clean_b64) = path_or_b64.split_once(";base64,").ok_or_else(|| {
|
||||
// 返回一个明确的、可读性极佳的格式错误
|
||||
DdddError::PreprocessError(ImagePreprocessReason::InvalidBase64Header)
|
||||
})?;
|
||||
// 转换为 Base64 包装器,并复用其 TryFromImage 实现
|
||||
Self::try_from_image(Base64(clean_b64))
|
||||
} else {
|
||||
// 2. 否则,老老实实当作本地路径打开
|
||||
let img = image::open(path_or_b64).map_err(ImagePreprocessReason::from)?;
|
||||
Ok(Self { inner: img })
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 2.3 路径类型 B: &Path (标准借用)
|
||||
impl<'a> TryFromImage<&'a Path> for ImageSource {
|
||||
fn try_from_image(path: &'a Path) -> Result<Self> {
|
||||
let img = image::open(path).map_err(ImagePreprocessReason::from)?;
|
||||
Ok(Self { inner: img })
|
||||
}
|
||||
}
|
||||
// 2.4 路径类型 C: PathBuf / String (拥有所有权,透传给借用)
|
||||
impl TryFromImage<PathBuf> for ImageSource {
|
||||
fn try_from_image(path: PathBuf) -> Result<Self> {
|
||||
Self::try_from_image(path.as_path())
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFromImage<String> for ImageSource {
|
||||
fn try_from_image(path_or_b64: String) -> Result<Self> {
|
||||
Self::try_from_image(path_or_b64.as_str())
|
||||
}
|
||||
}
|
||||
// 2. 支持带有生命周期的借用:直接支持 &[u8](不强绑生命周期到 ImageSource 结构体上!)
|
||||
impl<'a> TryFromImage<&'a [u8]> for ImageSource {
|
||||
fn try_from_image(bytes: &'a [u8]) -> Result<Self> {
|
||||
let img = image::load_from_memory(bytes).map_err(ImagePreprocessReason::from)?;
|
||||
Ok(Self { inner: img })
|
||||
}
|
||||
}
|
||||
|
||||
// 4. 完美支持 ndarray 的借用 ArrayViewD
|
||||
impl<'a> TryFromImage<ArrayViewD<'a, u8>> for ImageSource {
|
||||
fn try_from_image(array: ArrayViewD<'a, u8>) -> Result<Self> {
|
||||
let img = ndarray_to_pil_image(array)?;
|
||||
Ok(Self { inner: img })
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> TryFromImage<Base64<'a>> for ImageSource {
|
||||
fn try_from_image(b64_str: Base64<'a>) -> Result<Self> {
|
||||
let img = base64_to_image(b64_str.0)?;
|
||||
Ok(Self { inner: img })
|
||||
}
|
||||
}
|
||||
|
||||
/// 模拟 Python 的 load_image_from_input
|
||||
#[allow(dead_code)]
|
||||
pub fn load_image_from_input<I>(input: I) -> Result<DynamicImage>
|
||||
where
|
||||
ImageSource: TryFromImage<I>,
|
||||
{
|
||||
let img = ImageSource::try_from_image(input)?.into_image();
|
||||
Ok(img)
|
||||
}
|
||||
@@ -1,43 +1,18 @@
|
||||
use anyhow::{Context, Result, anyhow, bail};
|
||||
use crate::error::{DdddError, ImagePreprocessReason, Result};
|
||||
use base64::{Engine as _, engine::general_purpose};
|
||||
use image::{DynamicImage, GenericImageView, ImageBuffer, ImageFormat, Luma, Rgb, RgbImage, Rgba};
|
||||
use image::{DynamicImage, GenericImageView, ImageBuffer, Luma, Rgb, Rgba};
|
||||
use ndarray::{Array3, ArrayD, ArrayViewD};
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use ndarray::{Array3, ArrayD, ArrayViewD};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum ColorMode {
|
||||
RGB,
|
||||
RGBA,
|
||||
L,
|
||||
}
|
||||
/// 定义支持的输入类型枚举
|
||||
pub enum ImageInput {
|
||||
Bytes(Vec<u8>),
|
||||
Array(ArrayD<u8>), // 对应 numpy 数组
|
||||
Path(PathBuf),
|
||||
Base64(String),
|
||||
DynamicImage(DynamicImage),
|
||||
}
|
||||
/// 模拟 Python 的 load_image_from_input
|
||||
#[allow(dead_code)]
|
||||
pub fn load_image_from_input(img_input: ImageInput) -> Result<DynamicImage> {
|
||||
match img_input {
|
||||
// 2. 处理字节流 (Bytes)
|
||||
ImageInput::Bytes(bytes) => {
|
||||
image::load_from_memory(&bytes).context("Failed to load utils from bytes")
|
||||
}
|
||||
// 1. 已经是 DynamicImage
|
||||
ImageInput::DynamicImage(i) => Ok(i),
|
||||
// 5. 处理 ndarray (Numpy-like)
|
||||
// 假设输入是 HWC 格式的 Array3<u8>
|
||||
ImageInput::Array(a) => numpy_to_pil_image(a.view()),
|
||||
// 4. 处理 Base64 字符串
|
||||
ImageInput::Base64(b) => base64_to_image(&b),
|
||||
// 3. 处理文件路径 (Path)
|
||||
ImageInput::Path(p) => image::open(p).context("Failed to open utils from path"),
|
||||
}
|
||||
}
|
||||
fn base64_to_image(b64_str: &str) -> Result<DynamicImage> {
|
||||
/// 将base64编码的图片转换为 DynamicImage
|
||||
pub fn base64_to_image(b64_str: &str) -> Result<DynamicImage> {
|
||||
// 过滤掉可能存在的 base64 前缀,例如 "data:utils/png;base64,"
|
||||
let clean_b64 = if let Some(pos) = b64_str.find(",") {
|
||||
&b64_str[pos + 1..]
|
||||
@@ -47,18 +22,18 @@ fn base64_to_image(b64_str: &str) -> Result<DynamicImage> {
|
||||
|
||||
let bytes = general_purpose::STANDARD
|
||||
.decode(clean_b64.trim())
|
||||
.map_err(|e| anyhow!("Base64 decode error: {}", e))?;
|
||||
.map_err(ImagePreprocessReason::from)?;
|
||||
|
||||
image::load_from_memory(&bytes).context("Failed to load utils from decoded base64")
|
||||
let img = image::load_from_memory(&bytes).map_err(ImagePreprocessReason::from)?;
|
||||
Ok(img)
|
||||
}
|
||||
|
||||
/// 读取图片文件并转换为 base64 编码字符串
|
||||
/// 对应 Python 版 get_img_base64
|
||||
pub fn get_img_base64<P: AsRef<Path>>(image_path: P) -> Result<String> {
|
||||
// 对应 Python 版 get_img_base64
|
||||
pub fn img_base64<P: AsRef<Path>>(image_path: P) -> Result<String> {
|
||||
// 1. 读取文件原始字节流
|
||||
// 使用 AsRef<Path> 泛型可以让函数同时支持 String, &str, PathBuf 等类型
|
||||
let image_data = fs::read(&image_path)
|
||||
.with_context(|| format!("Failed to read utils file: {:?}", image_path.as_ref()))?;
|
||||
let image_data = fs::read(&image_path)?;
|
||||
|
||||
// 2. 进行 Base64 编码
|
||||
// 使用 STANDARD 引擎对齐 Python 的 base64.b64encode
|
||||
@@ -67,55 +42,50 @@ pub fn get_img_base64<P: AsRef<Path>>(image_path: P) -> Result<String> {
|
||||
Ok(b64_string)
|
||||
}
|
||||
|
||||
/// 封装数组转图像的逻辑,对齐 Python 版 _numpy_to_pil_image
|
||||
fn numpy_to_pil_image(array: ArrayViewD<u8>) -> Result<DynamicImage> {
|
||||
/// 封装数组转图像的逻辑,
|
||||
// 对应 Python 版 _numpy_to_pil_image
|
||||
pub(crate) fn ndarray_to_pil_image(array: ArrayViewD<u8>) -> Result<DynamicImage> {
|
||||
let shape = array.shape();
|
||||
let dim = shape.len();
|
||||
|
||||
// 1. 确保数据在内存中是连续的 (C order / Standard Layout)
|
||||
// 如果 arr 是经过切片或转置的,这一步会进行必要的内存拷贝
|
||||
let standard = array.as_standard_layout();
|
||||
let (raw_data, _offset) = standard.to_owned().into_raw_vec_and_offset();
|
||||
// let standard = array.as_standard_layout();
|
||||
// let (raw_data, _offset) = standard.to_owned().into_raw_vec_and_offset();
|
||||
|
||||
match dim {
|
||||
let color_mode = match dim {
|
||||
// 对应 Python: len(array.shape) == 2 (灰度图 H, W)
|
||||
2 => {
|
||||
let (h, w) = (shape[0], shape[1]);
|
||||
ImageBuffer::<Luma<u8>, _>::from_raw(w as u32, h as u32, raw_data)
|
||||
.map(DynamicImage::ImageLuma8)
|
||||
.ok_or_else(|| anyhow!("Failed to create Luma utils from 2D array"))
|
||||
}
|
||||
2 => ColorMode::L,
|
||||
|
||||
// 对应 Python: len(array.shape) == 3 (H, W, C)
|
||||
3 => {
|
||||
let (h, w, c) = (shape[0], shape[1], shape[2]);
|
||||
let (_h, _w, c) = (shape[0], shape[1], shape[2]);
|
||||
match c {
|
||||
// 对应 Python: array.shape[2] == 1 (单通道 H, W, 1)
|
||||
1 => ImageBuffer::<Luma<u8>, _>::from_raw(w as u32, h as u32, raw_data)
|
||||
.map(DynamicImage::ImageLuma8),
|
||||
|
||||
1 => ColorMode::L,
|
||||
// 对应 Python: array.shape[2] == 3 (RGB H, W, 3)
|
||||
3 => ImageBuffer::<Rgb<u8>, _>::from_raw(w as u32, h as u32, raw_data)
|
||||
.map(DynamicImage::ImageRgb8),
|
||||
|
||||
3 => ColorMode::RGB,
|
||||
// 对应 Python: array.shape[2] == 4 (RGBA H, W, 4)
|
||||
4 => ImageBuffer::<Rgba<u8>, _>::from_raw(w as u32, h as u32, raw_data)
|
||||
.map(DynamicImage::ImageRgba8),
|
||||
|
||||
4 => ColorMode::RGBA,
|
||||
_ => {
|
||||
return Err(anyhow!("不支持的通道数: {}", c));
|
||||
return Err(DdddError::PreprocessError(
|
||||
ImagePreprocessReason::UnsupportedChannels(c),
|
||||
));
|
||||
}
|
||||
}
|
||||
.ok_or_else(|| anyhow!("转换彩色图失败"))
|
||||
}
|
||||
|
||||
_ => Err(anyhow!("不支持的数组维度: {},仅支持 2D 或 3D", dim)),
|
||||
}
|
||||
_ => {
|
||||
return Err(DdddError::DimensionMismatch {
|
||||
expected: "2D (H,W) 或 3D (H,W,C)".to_string(),
|
||||
actual: shape.to_vec(),
|
||||
});
|
||||
}
|
||||
};
|
||||
from_ndarray(array, color_mode)
|
||||
}
|
||||
|
||||
/// 对应 Python 的 png_rgba_black_preprocess
|
||||
/// 将带有透明通道的图片转换为白色背景的 RGB 图片
|
||||
|
||||
/// 处理PNG图片的RGBA透明背景,将透明部分设置为白色背景
|
||||
// 对应 Python 的 png_rgba_black_preprocess
|
||||
pub fn png_rgba_white_preprocess(img: &DynamicImage) -> DynamicImage {
|
||||
// 1. 检查是否包含透明通道,如果没有,直接克隆并返回
|
||||
if !img.color().has_alpha() {
|
||||
@@ -161,104 +131,71 @@ pub fn png_rgba_white_preprocess(img: &DynamicImage) -> DynamicImage {
|
||||
|
||||
DynamicImage::ImageRgb8(background)
|
||||
}
|
||||
pub fn image_to_numpy(image: &DynamicImage, mode: ColorMode) -> Result<Array3<u8>> {
|
||||
/// 将 DynamicImage 转换为 array 数组
|
||||
pub fn image_to_ndarray(image: &DynamicImage, mode: ColorMode) -> Result<Array3<u8>> {
|
||||
// 1. 模式转换 (对应 utils.convert(target_mode)),此函数在时保留看后续优化是否需要替代image_to_ndarray
|
||||
// Rust utils 库通过 to_rgb8, to_luma8 等方法实现转换
|
||||
let (width, height) = image.dimensions();
|
||||
|
||||
let (channels, raw) = match mode {
|
||||
ColorMode::RGB => (3, image.to_rgb8().into_raw()),
|
||||
ColorMode::L => (1, image.to_luma8().into_raw()),
|
||||
ColorMode::RGB => (3, image.to_rgb8().into_raw()),
|
||||
ColorMode::RGBA => (4, image.to_rgba8().into_raw()),
|
||||
};
|
||||
|
||||
Array3::from_shape_vec((height as usize, width as usize, channels), raw)
|
||||
.map_err(|e| anyhow!("Failed to build ndarray: {}", e))
|
||||
let array = Array3::from_shape_vec((height as usize, width as usize, channels), raw)
|
||||
.map_err(ImagePreprocessReason::from)?;
|
||||
Ok(array)
|
||||
}
|
||||
/// 将 array 数组转换为 DynamicImage
|
||||
pub fn ndarray_to_image(array: ArrayViewD<u8>, mode: ColorMode) -> Result<DynamicImage> {
|
||||
let shape = array.shape();
|
||||
|
||||
// 基础边界检查:至少要有 H 和 W 两个维度
|
||||
if shape.len() < 2 {
|
||||
return Err(DdddError::DimensionMismatch {
|
||||
expected: "At least 2D array [H, W]".to_string(),
|
||||
actual: shape.to_vec(),
|
||||
});
|
||||
}
|
||||
from_ndarray(array, mode)
|
||||
}
|
||||
|
||||
pub fn numpy_to_image(array: ArrayViewD<u8>, mode: ColorMode) -> Result<DynamicImage> {
|
||||
fn from_ndarray(array: ArrayViewD<u8>, mode: ColorMode) -> Result<DynamicImage> {
|
||||
let shape = array.shape();
|
||||
// 1. 基础维度检查 (必须是 H, W, C 三维数组)
|
||||
if shape.len() != 3 {
|
||||
bail!("Expected a 3D array (H, W, C), but got {}D", shape.len());
|
||||
}
|
||||
|
||||
// 映射:ndarray 的 shape 默认是 [Height, Width, (Channels)]
|
||||
// image 库的 from_raw 接收 (width, height)
|
||||
let height = shape[0] as u32;
|
||||
let width = shape[1] as u32;
|
||||
let channels = shape[2];
|
||||
// 2. 检查通道数是否与模式匹配
|
||||
let expected_channels = match mode {
|
||||
ColorMode::L => 1,
|
||||
ColorMode::RGB => 3,
|
||||
ColorMode::RGBA => 4,
|
||||
};
|
||||
if channels != expected_channels {
|
||||
bail!(
|
||||
"Mode {:?} expects {} channels, but array has {}",
|
||||
mode,
|
||||
expected_channels,
|
||||
channels
|
||||
);
|
||||
}
|
||||
// 确保数据连续性 (C-order)
|
||||
|
||||
// 1. 确保数据在内存中是连续的 (C order)
|
||||
let standard = array.as_standard_layout();
|
||||
let (raw_data, _) = standard.to_owned().into_raw_vec_and_offset();
|
||||
|
||||
let raw_len = raw_data.len();
|
||||
// 2. 重新解释内存并构建 ImageBuffer
|
||||
match mode {
|
||||
ColorMode::L => ImageBuffer::<Luma<u8>, _>::from_raw(width, height, raw_data)
|
||||
.map(DynamicImage::ImageLuma8),
|
||||
.map(DynamicImage::ImageLuma8).ok_or_else(|| {
|
||||
DdddError::PreprocessError(ImagePreprocessReason::Custom(format!(
|
||||
"Failed to construct Luma image: buffer size {} does not match expected {} ({}x{}x1)",
|
||||
raw_len, width * height * 1, width, height
|
||||
)))
|
||||
}),
|
||||
ColorMode::RGB => ImageBuffer::<Rgb<u8>, _>::from_raw(width, height, raw_data)
|
||||
.map(DynamicImage::ImageRgb8),
|
||||
.map(DynamicImage::ImageRgb8).ok_or_else(|| {
|
||||
DdddError::PreprocessError(ImagePreprocessReason::Custom(format!(
|
||||
"Failed to construct RGB image: buffer size {} does not match expected {} ({}x{}x3)",
|
||||
raw_len, width * height * 3, width, height
|
||||
)))
|
||||
}),
|
||||
ColorMode::RGBA => ImageBuffer::<Rgba<u8>, _>::from_raw(width, height, raw_data)
|
||||
.map(DynamicImage::ImageRgba8),
|
||||
.map(DynamicImage::ImageRgba8).ok_or_else(|| {
|
||||
DdddError::PreprocessError(ImagePreprocessReason::Custom(format!(
|
||||
"Failed to construct RGBA image: buffer size {} does not match expected {} ({}x{}x4)",
|
||||
raw_len, width * height * 4, width, height
|
||||
)))
|
||||
}),
|
||||
}
|
||||
.ok_or_else(|| anyhow!("Failed to construct ImageBuffer. Buffer size might be incorrect."))
|
||||
}
|
||||
pub fn image_to_ndarray(img: &DynamicImage) -> Array3<u8> {
|
||||
let (width, height) = img.dimensions();
|
||||
|
||||
// 1. 强制转为 RGB8 (丢弃 Alpha 通道,与 Python 的 target_mode='RGB' 对齐)
|
||||
let rgb_img = img.to_rgb8();
|
||||
|
||||
// 2. 获取原始像素数据
|
||||
let raw_data = rgb_img.into_raw();
|
||||
|
||||
// 3. 构造数组 (通道数改为 3)
|
||||
Array3::from_shape_vec((height as usize, width as usize, 3), raw_data)
|
||||
.expect("Failed to construct ndarray from utils") // 建议显式报错,而不是返回全黑图
|
||||
}
|
||||
|
||||
#[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);
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use image::{DynamicImage, GrayImage, ImageBuffer, Luma, imageops::FilterType};
|
||||
use image::{imageops::FilterType, DynamicImage, GrayImage, ImageBuffer, Luma};
|
||||
|
||||
use ndarray::{Array2, Array3, ArrayView2, ArrayView3, azip};
|
||||
use ndarray::{azip, Array2, Array3, ArrayView2, ArrayView3};
|
||||
use std::cmp::{max, min};
|
||||
|
||||
// 模拟openCV
|
||||
@@ -174,8 +174,6 @@ pub fn rgb_to_opencv_hsv(r: u8, g: u8, b: u8) -> (u8, u8, u8) {
|
||||
(h_opencv, s_opencv, v_opencv)
|
||||
}
|
||||
|
||||
|
||||
|
||||
/// 对应 Python 的 convert_to_grayscale
|
||||
/// 将图像转换为灰度图 (L模式)
|
||||
pub fn convert_to_grayscale(image: &DynamicImage) -> GrayImage {
|
||||
|
||||
@@ -6,16 +6,22 @@ pub fn normalize_ocr_logits(array: ndarray::ArrayD<f32>, shape: &[usize]) -> Res
|
||||
let (steps, classes, data_dyn_view) = match shape.len() {
|
||||
3 => {
|
||||
if shape[1] == 1 {
|
||||
// 形状: [Steps, 1, Classes]
|
||||
(shape[0], shape[2], array)
|
||||
} else if shape[0] == 1 {
|
||||
// 形状: [1, Steps, Classes]
|
||||
(shape[1], shape[2], array)
|
||||
} else {
|
||||
// 默认取第一个 batch: [Batch, Steps, Classes]
|
||||
// 使用 ndarray 的 s! 宏,对应 Python 的 output[0, :, :]
|
||||
let sliced = array.slice_move(s![0, .., ..]);
|
||||
(shape[1], shape[2], sliced.into_dyn())
|
||||
}
|
||||
}
|
||||
// 形状: [Steps, Classes]
|
||||
2 => (shape[0], shape[1], array),
|
||||
// 形状: [Classes] -> 单字符输出(对应 Python 的 ndim == 0 保护逻辑)
|
||||
// 我们把它虚构成一个 [1, Classes] 的 2D 矩阵来复用后面的 argmax 逻辑
|
||||
1 => (1, shape[0], array),
|
||||
_ => {
|
||||
return Err(DdddError::DimensionMismatch {
|
||||
|
||||
Reference in New Issue
Block a user