Rust/Python/NumPy: Iterate over ndarray and collect into Array3 (for image)

help
1

Hello!

I'm trying to create an iterator over an image (coming from Python) that is stored in a NumPy array. I need to go through every pixel and do some operations (basically transform every pixel into something else).

I think someone somewhere should have done something similar.

My attempt looks something like this:

    fn transform_image<'py>(
        py: Python<'py>,
        x: PyReadonlyArray3<'py, f64>,
    ) -> PyResult<&'py PyArray3<f64>> {
        let x = x.as_array();
        let a: ndarray::Array3<f64> = x
            .axis_iter(Axis(0))
            .enumerate()
            .map(|(x, row)| {
                let c: ArrayVec<_, 2> = row
                    .axis_iter(Axis(0))
                    .enumerate()
                    .map(|(y, channel)| {
                       // channel is rgb as [r,g,b].
                        [x as f64, y as f64]
                    })
                    .collect();
                c.into_inner().unwrap()
            })
            .collect_vec()
            .into();
        Ok(a.into_pyarray(py))
    }

This complies but gets runtime error: pyo3_runtime.PanicException: ArrayVec: capacity exceeded in extend/from_iter

Does anyone who has a clue how to do this effectively?

Solution (non-iterator)

Maybe the problem is not suitable using iterators as the arrays are not dynamically allocated. There is probably a good way of doing it with iterators (and found in my search map_inplace() that may or may not be useful).

Just a dirty solution, that first allocates the array with zeros and then mutates all of them:

    #[pyfn(m)]
    fn transform_image<'py>(
        py: Python<'py>,
        x: PyReadonlyArray3<'py, f64>,
    ) -> PyResult<&'py PyArray3<f64>> {
        let x = x.as_array();
        let shape = x.shape();
        let h = shape[0];
        let w = shape[1];
        let new_channel_size = 2;
        let mut transformed_image = ndarray::Array3::<f64>::zeros((h, w, new_channel_size));
        for (x, row) in x.axis_iter(Axis(0)).enumerate() {
            for (y, channel) in row.axis_iter(Axis(0)).enumerate() {
                let mut c = transformed_image.slice_mut(s![x, y, ..]);
                c[0] = x as f64;
                c[1] = y as f64;
            }
        }
        Ok(transformed_image.into_pyarray(py))
    }

In terms of readability, I find this to be more understandable.

2

It means that ArrayVec panicked, and pyo3 raised the panic as a Python exception. In this case the capacity of the arrayvec was too small, and you need to increase it. Or use a normal Vec.

1 Like
3

Hmm yes. My bad, I interpreted size 2 of the ArrayVec<_,2> wrong. But if using vec, how do I turn it into a Array3 efficiently? It doesn't compile if changing into:

        let a: ndarray::Array3<f64> = x
            .axis_iter(Axis(0))
            .enumerate()
            .map(|(x, row)| {
                row.axis_iter(Axis(0))
                    .enumerate()
                    .map(|(y, channel)| [x as f64, y as f64])
                    .collect_vec()
            })
            .collect_vec()
            .into();

With compile error: the trait bound Vec<[f64; 2]>: FixedInitializer is not satisfied. Is implementing FixedInitializer the way to go? Or is there some other way? Or a more correct way of iterating over the array? And collecting into 2 vectors (inner and outer) seems like a bad solution overall. The documentation is rather sparse and I failed to find a solution to it.

4

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.