refactor: 优化 slide_model.rs
- 新增 cv2.rs 模拟 opencv
This commit is contained in:
169
src/image_io.rs
169
src/image_io.rs
@@ -1,9 +1,15 @@
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use anyhow::{Context, Result, anyhow, bail};
|
||||
use base64::{Engine as _, engine::general_purpose};
|
||||
use image::{DynamicImage, GenericImageView, ImageBuffer, ImageFormat, Luma, Rgb, RgbImage, Rgba};
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use tract_onnx::prelude::tract_ndarray::{Array3, ArrayD, ArrayViewD};
|
||||
#[derive(Debug)]
|
||||
pub enum ColorMode {
|
||||
RGB,
|
||||
RGBA,
|
||||
L,
|
||||
}
|
||||
/// 定义支持的输入类型枚举
|
||||
pub enum ImageInput {
|
||||
Bytes(Vec<u8>),
|
||||
@@ -60,41 +66,7 @@ pub fn get_img_base64<P: AsRef<Path>>(image_path: P) -> Result<String> {
|
||||
|
||||
Ok(b64_string)
|
||||
}
|
||||
/// 处理 PNG 图像的 RGBA 透明背景，将透明部分设置为白色
|
||||
///
|
||||
/// 对应 Python 版 png_rgba_black_preprocess
|
||||
pub fn png_rgba_black_preprocess(img: &DynamicImage) -> Result<DynamicImage> {
|
||||
// 1. 获取原图尺寸
|
||||
let (width, height) = (img.width(), img.height());
|
||||
|
||||
// 2. 创建一个等尺寸的纯白色 RGB 图像作为底色
|
||||
// ImageBuffer::<Rgb<u8>, Vec<u8>>
|
||||
let mut white_bg = ImageBuffer::from_fn(width, height, |_, _| {
|
||||
Rgb([255, 255, 255])
|
||||
});
|
||||
|
||||
// 3. 将原图复合到底色上
|
||||
// 我们需要处理原图，将其转为 RGBA 确保有 alpha 通道可以参考
|
||||
let rgba_img = img.to_rgba8();
|
||||
|
||||
// 遍历每一个像素进行复合（模拟 Python 的 paste 逻辑）
|
||||
for (x, y, pixel) in rgba_img.enumerate_pixels() {
|
||||
let alpha = pixel[3] as f32 / 255.0;
|
||||
if alpha > 0.0 {
|
||||
// 获取底色像素（白色）
|
||||
let bg_pixel = white_bg.get_pixel_mut(x, y);
|
||||
|
||||
// 简单的 Alpha 复合公式：输出 = 源 * alpha + 背景 * (1 - alpha)
|
||||
for i in 0..3 {
|
||||
let fg = pixel[i] as f32;
|
||||
let bg = bg_pixel[i] as f32;
|
||||
bg_pixel[i] = (fg * alpha + bg * (1.0 - alpha)) as u8;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(DynamicImage::ImageRgb8(white_bg))
|
||||
}
|
||||
/// 封装数组转图像的逻辑，对齐 Python 版 _numpy_to_pil_image
|
||||
fn numpy_to_pil_image(array: ArrayViewD<u8>) -> Result<DynamicImage> {
|
||||
let shape = array.shape();
|
||||
@@ -143,11 +115,11 @@ fn numpy_to_pil_image(array: ArrayViewD<u8>) -> Result<DynamicImage> {
|
||||
|
||||
/// 对应 Python 的 png_rgba_black_preprocess
|
||||
/// 将带有透明通道的图片转换为白色背景的 RGB 图片
|
||||
#[allow(dead_code)]
|
||||
|
||||
pub fn png_rgba_white_preprocess(img: &DynamicImage) -> DynamicImage {
|
||||
// 1. 检查是否包含透明通道，如果没有，直接克隆并返回
|
||||
if !img.color().has_alpha() {
|
||||
return img.clone();
|
||||
return DynamicImage::ImageRgb8(img.to_rgb8());
|
||||
}
|
||||
|
||||
let (width, height) = img.dimensions();
|
||||
@@ -160,83 +132,87 @@ pub fn png_rgba_white_preprocess(img: &DynamicImage) -> DynamicImage {
|
||||
|
||||
// 4. 遍历像素并手动进行 Alpha 混合
|
||||
// 对应 Python 的 image.paste(img, ..., mask=img)
|
||||
for (x, y, pixel) in rgba_img.enumerate_pixels() {
|
||||
let alpha = pixel[3] as f32 / 255.0;
|
||||
// 使用 enumerate_pixels_mut 同时获取坐标和背景像素的可变引用，减少查找开销
|
||||
for (x, y, bg_pixel) in background.enumerate_pixels_mut() {
|
||||
// 安全性说明：x, y 源自 background 尺寸，与 rgba_img 一致，get_pixel 是安全的
|
||||
let src_pixel = rgba_img.get_pixel(x, y);
|
||||
let alpha_u8 = src_pixel[3];
|
||||
|
||||
if alpha >= 1.0 {
|
||||
// 完全不透明，直接覆盖
|
||||
background.put_pixel(x, y, Rgb([pixel[0], pixel[1], pixel[2]]));
|
||||
} else if alpha > 0.0 {
|
||||
// 半透明，执行 Alpha 混合公式： (src * alpha) + (dst * (1 - alpha))
|
||||
let bg_pixel = background.get_pixel(x, y);
|
||||
let r = (pixel[0] as f32 * alpha + bg_pixel[0] as f32 * (1.0 - alpha)) as u8;
|
||||
let g = (pixel[1] as f32 * alpha + bg_pixel[1] as f32 * (1.0 - alpha)) as u8;
|
||||
let b = (pixel[2] as f32 * alpha + bg_pixel[2] as f32 * (1.0 - alpha)) as u8;
|
||||
background.put_pixel(x, y, Rgb([r, g, b]));
|
||||
match alpha_u8 {
|
||||
// 情况 A：完全不透明，直接覆盖背景色
|
||||
255 => {
|
||||
bg_pixel.0 = [src_pixel[0], src_pixel[1], src_pixel[2]];
|
||||
}
|
||||
// 情况 B：完全透明，保持背景色（白色），无需操作
|
||||
0 => {
|
||||
continue;
|
||||
}
|
||||
// 情况 C：半透明，进行 Alpha 混合计算
|
||||
_ => {
|
||||
let alpha = alpha_u8 as f32 / 255.0;
|
||||
let inv_alpha = 1.0 - alpha;
|
||||
|
||||
bg_pixel[0] = (src_pixel[0] as f32 * alpha + 255.0 * inv_alpha).round() as u8;
|
||||
bg_pixel[1] = (src_pixel[1] as f32 * alpha + 255.0 * inv_alpha).round() as u8;
|
||||
bg_pixel[2] = (src_pixel[2] as f32 * alpha + 255.0 * inv_alpha).round() as u8;
|
||||
}
|
||||
}
|
||||
// alpha == 0 的情况不需要处理，因为背景已经是白色了
|
||||
}
|
||||
|
||||
DynamicImage::ImageRgb8(background)
|
||||
}
|
||||
pub fn image_to_numpy(image: &DynamicImage, target_mode: &str) -> Result<Array3<u8>> {
|
||||
// 1. 模式转换 (对应 image.convert(target_mode))
|
||||
pub fn image_to_numpy(image: &DynamicImage, mode: ColorMode) -> Result<Array3<u8>> {
|
||||
// 1. 模式转换 (对应 image.convert(target_mode)),此函数在时保留看后续优化是否需要替代image_to_ndarray
|
||||
// Rust image 库通过 to_rgb8, to_luma8 等方法实现转换
|
||||
let (width, height) = image.dimensions();
|
||||
|
||||
match target_mode {
|
||||
"RGB" => {
|
||||
let rgb_img = image.to_rgb8();
|
||||
let raw = rgb_img.into_raw();
|
||||
// shape 为 [Height, Width, Channels] -> [H, W, 3]
|
||||
Array3::from_shape_vec((height as usize, width as usize, 3), raw)
|
||||
.map_err(|e| anyhow!("Failed to build ndarray: {}", e))
|
||||
},
|
||||
"L" | "GRAY" => {
|
||||
let gray_img = image.to_luma8();
|
||||
let raw = gray_img.into_raw();
|
||||
// shape 为 [H, W, 1]
|
||||
Array3::from_shape_vec((height as usize, width as usize, 1), raw)
|
||||
.map_err(|e| anyhow!("Failed to build ndarray: {}", e))
|
||||
},
|
||||
"RGBA" => {
|
||||
let rgba_img = image.to_rgba8();
|
||||
let raw = rgba_img.into_raw();
|
||||
// shape 为 [H, W, 4]
|
||||
Array3::from_shape_vec((height as usize, width as usize, 4), raw)
|
||||
.map_err(|e| anyhow!("Failed to build ndarray: {}", e))
|
||||
},
|
||||
_ => Err(anyhow!("Unsupported target_mode: {}", target_mode)),
|
||||
}
|
||||
let (channels, raw) = match mode {
|
||||
ColorMode::RGB => (3, image.to_rgb8().into_raw()),
|
||||
ColorMode::L => (1, image.to_luma8().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))
|
||||
}
|
||||
|
||||
pub fn numpy_to_image(array: ArrayViewD<u8>, mode: &str) -> Result<DynamicImage> {
|
||||
pub fn numpy_to_image(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());
|
||||
}
|
||||
|
||||
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)
|
||||
let standard = array.as_standard_layout();
|
||||
let (raw_data, _) = standard.to_owned().into_raw_vec_and_offset();
|
||||
|
||||
let height = shape[0] as u32;
|
||||
let width = shape[1] as u32;
|
||||
|
||||
match mode {
|
||||
"L" => {
|
||||
ImageBuffer::<Luma<u8>, _>::from_raw(width, height, raw_data)
|
||||
.map(DynamicImage::ImageLuma8)
|
||||
.ok_or_else(|| anyhow!("Failed to create Luma image"))
|
||||
},
|
||||
"RGB" => {
|
||||
ImageBuffer::<Rgb<u8>, _>::from_raw(width, height, raw_data)
|
||||
.map(DynamicImage::ImageRgb8)
|
||||
.ok_or_else(|| anyhow!("Failed to create RGB image"))
|
||||
},
|
||||
"RGBA" => {
|
||||
ImageBuffer::<Rgba<u8>, _>::from_raw(width, height, raw_data)
|
||||
.map(DynamicImage::ImageRgba8)
|
||||
.ok_or_else(|| anyhow!("Failed to create RGBA image"))
|
||||
},
|
||||
_ => Err(anyhow!("Unsupported mode: {}", mode)),
|
||||
ColorMode::L => ImageBuffer::<Luma<u8>, _>::from_raw(width, height, raw_data)
|
||||
.map(DynamicImage::ImageLuma8),
|
||||
ColorMode::RGB => ImageBuffer::<Rgb<u8>, _>::from_raw(width, height, raw_data)
|
||||
.map(DynamicImage::ImageRgb8),
|
||||
ColorMode::RGBA => ImageBuffer::<Rgba<u8>, _>::from_raw(width, height, raw_data)
|
||||
.map(DynamicImage::ImageRgba8),
|
||||
}
|
||||
.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();
|
||||
@@ -251,6 +227,7 @@ pub fn image_to_ndarray(img: &DynamicImage) -> Array3<u8> {
|
||||
Array3::from_shape_vec((height as usize, width as usize, 3), raw_data)
|
||||
.expect("Failed to construct ndarray from image") // 建议显式报错，而不是返回全黑图
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
fn save_rust_result(result: &ImageBuffer<Luma<f32>, Vec<f32>>, filename: &str) {
|
||||
let (width, height) = result.dimensions();
|
||||
|
||||
Reference in New Issue
Block a user