ddddocr-rs/ddddocr-core/src/utils/image_convert.rs

use crate::error::{DdddError, ImagePreprocessReason, Result};
use image::{DynamicImage, GenericImageView, ImageBuffer, Luma, Rgb, Rgba};
use ndarray::{Array3, ArrayViewD};

#[derive(Debug)]
pub enum ColorMode {
    RGB,
    RGBA,
    L,
}

/// 封装数组转图像的逻辑，
// 对应 Python 版 _numpy_to_pil_image
pub fn ndarray_to_hwc_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 color_mode = match dim {
        // 对应 Python: len(array.shape) == 2 (灰度图 H, W)
        2 => ColorMode::L,

        // 对应 Python: len(array.shape) == 3 (H, W, C)
        3 => {
            let (_h, _w, c) = (shape[0], shape[1], shape[2]);
            match c {
                // 对应 Python: array.shape[2] == 1 (单通道 H, W, 1)
                1 => ColorMode::L,
                // 对应 Python: array.shape[2] == 3 (RGB H, W, 3)
                3 => ColorMode::RGB,
                // 对应 Python: array.shape[2] == 4 (RGBA H, W, 4)
                4 => ColorMode::RGBA,
                _ => {
                    return Err(DdddError::Preprocess(
                        ImagePreprocessReason::UnsupportedChannels(c),
                    ));
                }
            }
        }
        _ => {
            return Err(DdddError::Preprocess(
                ImagePreprocessReason::InvalidImageDimensions {
                    expected: "2D (H,W) 或 3D (H,W,C)".to_string(),
                    actual: shape.to_vec(),
                },
            ));
        }
    };
    from_ndarray(array, color_mode)
}

/// 处理PNG图片的RGBA透明背景，将透明部分设置为白色背景
// 对应 Python 的 png_rgba_black_preprocess
pub fn png_rgba_white_preprocess(img: &DynamicImage) -> DynamicImage {
    // 1. 检查是否包含透明通道，如果没有，直接克隆并返回
    if !img.color().has_alpha() {
        return DynamicImage::ImageRgb8(img.to_rgb8());
    }

    let (width, height) = img.dimensions();

    // 2. 创建一个新的 RGB 图像缓冲，默认填充为白色 (255, 255, 255)
    let mut background = ImageBuffer::from_pixel(width, height, Rgb([255u8, 255u8, 255u8]));

    // 3. 获取原图的 RGBA 视图
    let rgba_img = img.to_rgba8();

    // 4. 遍历像素并手动进行 Alpha 混合
    // 对应 Python 的 utils.paste(img, ..., mask=img)
    // 使用 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];

        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;
            }
        }
    }

    DynamicImage::ImageRgb8(background)
}
/// 将 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::L => (1, image.to_luma8().into_raw()),
        ColorMode::RGB => (3, image.to_rgb8().into_raw()),
        ColorMode::RGBA => (4, image.to_rgba8().into_raw()),
    };

    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::Preprocess(
            ImagePreprocessReason::InvalidImageDimensions {
                expected: "至少为 2D array [H, W]".to_string(),
                actual: shape.to_vec(),
            },
        ));
    }
    from_ndarray(array, mode)
}

fn from_ndarray(array: ArrayViewD<u8>, mode: ColorMode) -> Result<DynamicImage> {
    let shape = array.shape();

    // 映射：ndarray 的 shape 默认是 [Height, Width, (Channels)]
    // image 库的 from_raw 接收 (width, height)
    let height = shape[0] as u32;
    let width = shape[1] as u32;

    // 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();

    // 获取当前模式对应的通道数
    let channels = match mode {
        ColorMode::L => 1,
        ColorMode::RGB => 3,
        ColorMode::RGBA => 4,
    };

    let expected_len = (width * height) as usize * channels;

    // 构造通用错误闭包，避免 match 分支中重复编写冗长的错误对象
    let make_err = || {
        DdddError::Preprocess(ImagePreprocessReason::BufferLengthMismatch {
            expected: expected_len,
            actual: raw_len,
            width,
            height,
            channels,
        })
    };

    // 2. 重新解释内存并构建 ImageBuffer
    match mode {
        ColorMode::L => ImageBuffer::<Luma<u8>, _>::from_raw(width, height, raw_data)
            .map(DynamicImage::ImageLuma8)
            .ok_or_else(make_err),
        ColorMode::RGB => ImageBuffer::<Rgb<u8>, _>::from_raw(width, height, raw_data)
            .map(DynamicImage::ImageRgb8)
            .ok_or_else(make_err),
        ColorMode::RGBA => ImageBuffer::<Rgba<u8>, _>::from_raw(width, height, raw_data)
            .map(DynamicImage::ImageRgba8)
            .ok_or_else(make_err),
    }
}