如何将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转换为标准布局，如下所示：

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()
}

但是，在以后的版本中，借助as_standard_layout方法（#616），这将变得更加容易。

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

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

    to_standard_layout(arr)
}