The most effective way to random flip every image in a numpy array

Question

For example I have a numpy array a with shape 1000*3*256*256.
In other words, a is an array of 1000 images, the size of each image is 3*256*256.
I want to random flip every image, so my question is how to effectively do this? Thanks!

Asked By: kli_nlpr

||

Source

Answer 1

Basics: array[slice(a,b,c)] is equivalent to array[a:b:c], and to reverse (“flip”) an array use slice(None, None, -1), which is the same as array[::-1].

So let’s build the random flips for each image:

>>> import random
>> flips = [(slice(None, None, None),
...          slice(None, None, random.choice([-1, None])),
...          slice(None, None, random.choice([-1, None])))
...          for _ in xrange(a.shape[0])]

The first slice is for the channel, the second is for the Y axis and the third is for the X axis. Let’s build some test data:

>>> import numpy as np
>>> a = np.array(range(3*2*5*5)).reshape(3,2,5,5)

We can apply each random flip individually to each image:

>>> flips[0]
(slice(None, None, None), slice(None, None, -1), slice(None, None, None))
>>> a[0]
array([[[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19],
        [20, 21, 22, 23, 24]],

       [[25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34],
        [35, 36, 37, 38, 39],
        [40, 41, 42, 43, 44],
        [45, 46, 47, 48, 49]]])
>>> a[0][flips[0]]
array([[[20, 21, 22, 23, 24],
        [15, 16, 17, 18, 19],
        [10, 11, 12, 13, 14],
        [ 5,  6,  7,  8,  9],
        [ 0,  1,  2,  3,  4]],

       [[45, 46, 47, 48, 49],
        [40, 41, 42, 43, 44],
        [35, 36, 37, 38, 39],
        [30, 31, 32, 33, 34],
        [25, 26, 27, 28, 29]]])

As you can see, flips[0] flips the image vertically. Now it’s simple to do it for each image:

>>> random_flipped = np.array([img[flip] for img, flip in zip(a, flips)])

Answered By: BlackBear

Answer 2

Thanks to the awesome answer of @BlackBear, I was able to get starting. I noticed, however, such functionality would probably run very often and thus might benefit from some performance tweaks.

In thinking of how to improve the performance I tackled two things:

use a "pure numpy" implementation
(building upon 1.) use just-in-time compilation through numba

This is what I came up with:

import numpy as np
from numba import njit

def np_flip(arr: np.ndarray, axis=None) -> np.ndarray:
    """Flip image np arrays with a (batch, row, col, band) configuration."""
    forw = np.int64(1)
    rev = np.int64(-1)
    flip_ax0 = np.random.choice(np.array([forw, rev]))
    flip_ax1 = np.random.choice(np.array([forw, rev]))

    flips = tuple(
        (
            slice(None, None, None),        # TF batch
            slice(None, None, flip_ax0),    # image rows
            slice(None, None, flip_ax1),    # image cols
            slice(None, None, None),        # image bands
        )
    )
    return arr[flips]

# njit the function, also possible via @njit decorator
njit_np_flip = njit(np_flip)

Benchmarks

import scipy.ndimage

arr = np.random.randint(0, 255, size=(2, 4, 4, 3))
print(arr.shape)
# (2, 4, 4, 3)
arr = scipy.ndimage.zoom(input=arr, zoom=(1, 128, 128, 1), order=0)
print(arr.shape)
# (2, 512, 512, 3)

# @BlackBear's answer
def py_np_flip(arr: np.ndarray, axis=None):
    """Python-Heavy version of np_flip."""
    # see https://stackoverflow.com/a/36716579/3250126
    flips = [
        (
            slice(None, None, np.random.choice([-1, None])),
            slice(None, None, np.random.choice([-1, None])),
            slice(None, None, None),
        )
        for _ in range(arr.shape[0])
    ]
    return np.array([img[flip] for img, flip in zip(arr, flips)])

# @BlackBear's answer
%timeit arr_flipped_np = np.apply_over_axes(py_np_flip, arr, axes=0)
# 2.18 ms ± 34.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# "pure numpy" implementation
%timeit arr_flipped_np = np.apply_over_axes(np_flip, arr, axes=0)
# 19.2 µs ± 2.34 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

# njit(ed) "pure numpy" implementeation
%timeit arr_flipped_np = np.apply_over_axes(njit_np_flip, arr, axes=0)
# 1.7 µs ± 28.9 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

1600 x speed increase

> 2.8 ms / 1.7 µs

  (2.8 × millisecond) / (1.7 × microsecond) ≈ 1647.0588

Answered By: loki

The most effective way to random flip every image in a numpy array

Question:

Answers:

Benchmarks

1600 x speed increase