user3048747
user3048747

Reputation: 309

How to save ndarray in Rust as image?

I have ndarray

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

and I use image library

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

How could I save my ndarray as image ? If there is better image library then image for this I could use it.

Upvotes: 7

Views: 3774

Answers (1)

E_net4
E_net4

Reputation: 29972

The easiest way is to ensure that the array follows is in standard layout (C-contiguous) with the image dimensions in the order (height, width, channel) order (HWC), or in an equivalent memory layout. This is necessary because image expects rows to be contiguous in memory.

Then, build a RgbImage using the type's from_raw function.

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

Example of use:

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")?;

The output image:

the output image: a gradient over the red and green channels


Below are a few related helper functions, in case they are necessary.

Ndarrays can be converted to standard layout by calling the method as_standard_layout, available since version 0.13.0. Before this version, you would need to collect each array element into a vector and rebuild the array, like so:

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

Moreover, converting an ndarray in the layout (width, height, channel) to (height, width, channel) is also possible by swapping the first two axes and making the array C-contiguous afterwards:

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

    arr.as_standard_layout().to_owned()
}

Upvotes: 9

Related Questions