如何在 Rust 中将 ndarray 保存为图像？

Question

我有 ndarray

let mut a = Array3::<u8>::from_elem((50, 40, 3), 3);

我使用图像库

 let mut imgbuf = image::ImageBuffer::new(50, 40);

如何将 ndarray 保存为图像？如果有更好的图像库，那么我可以使用它。

Answer 1

最简单的方法是确保数组遵循标准布局（C 连续），图像尺寸按顺序(height, width, channel)(HWC) 排列，或者采用等效的内存布局。这是必要的，因为image期望行在内存中是连续的。

然后，构建一个RgbImage使用该类型的from_raw函数。

use image::RgbImage;
use ndarray::Array3;

fn array_to_image(arr: Array3<u8>) -> RgbImage {
    assert!(arr.is_standard_layout());

    let (height, width, _) = arr.dim();
    let raw = arr.into_raw_vec();

    RgbImage::from_raw(width as u32, height as u32, raw)
        .expect("container should have the right size for the image dimensions")
}

使用示例：

let mut array: Array3<u8> = Array3::zeros((200, 250, 3)); // 250x200 RGB

for ((x, y, z), v) in array.indexed_iter_mut() {
    *v = match z {
        0 => y as u8,
        1 => x as u8,
        2 => 0,
        _ => unreachable!(),
    };
}

let image = array_to_image(array);
image.save("out.png")?;

输出图像：

下面是一些相关的辅助函数，以备必要时使用。

Ndarrays 可以通过调用 方法转换为标准布局as_standard_layout，该方法从 0.13.0 版本开始可用。在此版本之前，您需要将每个数组元素收集到一个向量中并重建数组，如下所示：

fn to_standard_layout<A, D>(arr: Array<A, D>) -> Array<A, D>
where
    A: Clone,
    D: Dimension,
{
    let v: Vec<_> = arr.iter().cloned().collect();
    let dim = arr.dim();
    Array::from_shape_vec(dim, v).unwrap()
}

此外，通过交换前两个轴并随后使数组 C 连续，也可以将布局中的 ndarray 转换(width, height, channel)为：(height, width, channel)

fn wh_to_hw(mut arr: Array3<u8>) -> Array3<u8> {
    arr.swap_axes(0, 1);

    arr.as_standard_layout().to_owned()
}