Как сохранить ndarray в Rust как изображение?

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 ).

Кроме того, преобразование ndarray в компоновке (width, height, channel) в (height, width, channel) также возможно, если поменять местами первые две оси и сделать массив C-смежным:

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

    to_standard_layout(arr)
}