Removing indexes from a 3D numpy array

Inspired by this post, we will extend it to 3D array case, while keeping in mind that row would mean the first axis for 3D image arrays and columns would be the second one. So, deleting one element per row, would give us one element less along the second axis.

The solution would look something like this -

def delete_per_row_3D(a, remove_idx):
    mask = np.ones(a.shape[:-1],dtype=bool)
    mask[np.arange(len(remove_idx)), remove_idx] = False
    return a[mask].reshape(np.array(a.shape)-[0,1,0])

Sample run -

In [43]: np.random.seed(0)
    ...: a = np.random.randint(10,100,(2,4,3))

In [44]: a
Out[44]: 
array([[[54, 57, 74],
        [77, 77, 19],
        [93, 31, 46],
        [97, 80, 98]],

       [[98, 22, 68],
        [75, 49, 97],
        [56, 98, 91],
        [47, 35, 87]]])

In [45]: remove_idx = np.array([2,0]) # length be a.shape[0],
         # as no. of rows in 3D image array would be that

In [46]: delete_per_row_3D(a, remove_idx)
Out[46]: 
array([[[54, 57, 74],
        [77, 77, 19],
        [97, 80, 98]],

       [[75, 49, 97],
        [56, 98, 91],
        [47, 35, 87]]])

Visualize with channels

To help the readers visualize it better, consider the three slices for RGB/BGR channels -

In [47]: a[...,0]
Out[47]: 
array([[54, 77, 93, 97],
       [98, 75, 56, 47]])

In [48]: a[...,1]
Out[48]: 
array([[57, 77, 31, 80],
       [22, 49, 98, 35]])

In [49]: a[...,2]
Out[49]: 
array([[74, 19, 46, 98],
       [68, 97, 91, 87]])

We want to delete the column-2 from the first row and column-0 from the second row. So, consider those being removed from all three channels, i.e. 93,31,46 and 98,22,68 removed. Thus, we will end up with the shown output earlier with one element less along the second axis.

Alternatives

Again from the earlier linked post, two more alternatives would be -

def delete_per_row_3D_v2(a, remove_idx):
    mask = remove_idx[:,None]!=np.arange(a.shape[1])
    return a[mask].reshape(np.array(a.shape)-[0,1,0])

def delete_per_row_3D_v3(a, remove_idx):
    m,n = a.shape[:-1]
    rm = remove_idx+n*np.arange(m)
    return np.delete(a.reshape(-1,a.shape[-1]),rm,axis=0).reshape(m,n-1,-1)

Push it to the next level with view-based approach

We will leverage view-based method to mask pixels in a group-based manner and hence achieve major benefits on memory and hence performance, as shown below -

def delete_per_row_3D_view(a, remove_idx):
    m,n,r = a.shape
    mask = np.ones((m,n),dtype=bool)
    mask[np.arange(len(remove_idx)), remove_idx] = False
    vd = np.dtype((np.void, a.dtype.itemsize * r))
    a_masked = a.view(vd).ravel()[mask.flat].view(a.dtype)
    return a_masked.reshape(m,-1,r)

Benchmarking

We are including all vectorized methods here (including all from this post and from @Paul Panzer's post).

1) OP's 640x340 3D array case :

In [180]: np.random.seed(0)
     ...: a = np.random.randint(0,256,(640,340,3)).astype(np.uint8)
     ...: remove_idx = np.random.randint(0,340,(640))

In [181]: %timeit delete_per_row_3D(a, remove_idx)
     ...: %timeit delete_per_row_3D_v2(a, remove_idx)
     ...: %timeit delete_per_row_3D_v3(a, remove_idx)
     ...: %timeit delete_per_row_3D_view(a, remove_idx)
     ...: %timeit pp(a, remove_idx)
     ...: %timeit pp_flat(a, remove_idx)
100 loops, best of 3: 5.09 ms per loop
100 loops, best of 3: 5.31 ms per loop
100 loops, best of 3: 3.58 ms per loop
1000 loops, best of 3: 211 µs per loop
100 loops, best of 3: 4.19 ms per loop
1000 loops, best of 3: 1.23 ms per loop

2) Small 64x34 3D array case :

In [182]: np.random.seed(0)
     ...: a = np.random.randint(0,256,(64,34,3)).astype(np.uint8)
     ...: remove_idx = np.random.randint(0,34,(64))

In [183]: %timeit delete_per_row_3D(a, remove_idx)
     ...: %timeit delete_per_row_3D_v2(a, remove_idx)
     ...: %timeit delete_per_row_3D_v3(a, remove_idx)
     ...: %timeit delete_per_row_3D_view(a, remove_idx)
     ...: %timeit pp(a, remove_idx)
     ...: %timeit pp_flat(a, remove_idx)
10000 loops, best of 3: 61.9 µs per loop
10000 loops, best of 3: 61.7 µs per loop
10000 loops, best of 3: 61.3 µs per loop
100000 loops, best of 3: 12.9 µs per loop
10000 loops, best of 3: 49.8 µs per loop
10000 loops, best of 3: 22.2 µs per loop

3) Large 6400x3400 3D array case :

In [184]: np.random.seed(0)
     ...: a = np.random.randint(0,256,(6400,3400,3)).astype(np.uint8)
     ...: remove_idx = np.random.randint(0,3400,(6400))

In [185]: %timeit delete_per_row_3D(a, remove_idx)
     ...: %timeit delete_per_row_3D_v2(a, remove_idx)
     ...: %timeit delete_per_row_3D_v3(a, remove_idx)
     ...: %timeit delete_per_row_3D_view(a, remove_idx)
     ...: %timeit pp(a, remove_idx)
     ...: %timeit pp_flat(a, remove_idx)
1 loop, best of 3: 649 ms per loop
1 loop, best of 3: 669 ms per loop
1 loop, best of 3: 434 ms per loop
10 loops, best of 3: 47.5 ms per loop
1 loop, best of 3: 415 ms per loop
10 loops, best of 3: 127 ms per loop

So, view-based method : delete_per_row_3D_view is blowing away all other approaches across all sizes.

Adapted from https://stackoverflow.com/a/58157198/7207392

import numpy as np
import operator as op

a = np.arange(60.0).reshape(4,5,3)
idxs = [(0,1), (1,3), (2, 1), (3,4)]

m,n,_ = a.shape

# extract column indices
# there are simpler ways but this is fast
columns = np.fromiter(map(op.itemgetter(1),idxs),int,m)

# build decimated array
result = np.where((columns[...,None]>np.arange(n-1))[...,None],a[:,:-1],a[:,1:])

If you want it as fast as possible it may be worthwhile temporarily merging the last two dimensions -(pp_flat vs pp in these timings):

Update: Never discount @Divakar (new entry delete_per_row_3D_view)...

If you are wondering about that zigzag just below 10^3 that is where I switch from image decimation to repeating pixels. The latter generates contiguous arrays. Interestingly a 4 times larger contiguous array is processed faster than a smaller noncontiguous one ... (More precisely, the method doesn't work on noncontiguous arrays, so we have to make a contiguous copy)

UPDATE: added a version that can handle noncontiguous arrays as long as the color triplets are contiguous.

Code:

from simple_benchmark import BenchmarkBuilder, MultiArgument
import numpy as np
import operator as op
from scipy.misc import face

B = BenchmarkBuilder()

@B.add_function()
def delete_per_row_3D(a, remove_idx):
    mask = np.ones(a.shape[:-1],dtype=bool)
    mask[np.arange(len(remove_idx)), remove_idx] = False
    return a[mask].reshape(np.array(a.shape)-[0,1,0])

@B.add_function()
def delete_per_row_3D_v2(a, remove_idx):
    mask = remove_idx[:,None]!=np.arange(a.shape[1])
    return a[mask].reshape(np.array(a.shape)-[0,1,0])

@B.add_function()
def delete_per_row_3D_v3(a, remove_idx):
    m,n = a.shape[:-1]
    rm = remove_idx+n*np.arange(m)
    return np.delete(a.reshape(-1,a.shape[-1]),rm,axis=0).reshape(m,n-1,-1)

@B.add_function()
def pp(a,columns):
    m,n,_ = a.shape
    return np.where((columns[...,None]>np.arange(n-1))[...,None],a[:,:-1],a[:,1:])

@B.add_function()
def pp_flat(a,columns):
    m,n,d = a.shape
    af = a.reshape(m,-1)
    return np.where(columns[...,None]>np.arange(n-1).repeat(d),af[...,:-d],af[...,d:]).reshape(m,n-1,d)


@B.add_function()
def delete_per_row_3D_view(a, remove_idx):
    m,n,r = a.shape
    mask = np.ones((m,n),dtype=bool)
    mask[np.arange(len(remove_idx)), remove_idx] = False
    vd = np.dtype((np.void, a.dtype.itemsize * r))
    a_masked = np.ascontiguousarray(a).view(vd).ravel()[mask.flat].view(a.dtype)
    return a_masked.reshape(m,-1,r)


from numpy.lib.stride_tricks import as_strided

@B.add_function()
def delete_per_row_3D_view_2(a, remove_idx):
    m,n,r = a.shape
    mask = np.ones((m,n),dtype=bool)
    mask[np.arange(len(remove_idx)), remove_idx] = False
    vd = np.dtype((np.void, a.dtype.itemsize * r))
    a_masked = as_strided(a[0,0,...].view(vd),a.shape[:-1],a.strides[:-1])[mask].view(a.dtype)
    return a_masked.reshape(m,-1,r)

@B.add_arguments('array size')
def argument_provider():
    orig = face()
    for dec in (128,64,32,16,8,4,2,1,-2,-4):
        if dec>0:
            a = orig[::dec,::dec]
        else:
            a = orig.repeat(-dec,axis=0).repeat(-dec,axis=1)
        dim_size,w,_ = a.shape
        idxs = np.random.randint(0,w,dim_size)
#        idxs = [*enumerate(np.random.randint(0,w,dim_size)),]
        yield dim_size, MultiArgument([a,idxs])

r = B.run()
r.plot()

import pylab
pylab.savefig('delunif.png')

you can delete the pixels by calculating 1D index of each pixel to be deleted from your index list like below and then delete elements using the index then reshaping back array to new size

ar = np.arange(60).reshape(5,4,3)
indx = [(0,3), (1,1), (2,0), (3,2), (4,1)]
indx_1d = [ar.shape[2]*ar.shape[1]*i+ar.shape[2]*j + k for item in indx for k,(i,j) in enumerate([item]*3)]

deleted = np.delete(ar, indx_1d).reshape(5,3,3)

print(ar)
print("\n\n******************************")
print(deleted)

Output

[[[ 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 50]
  [51 52 53]
  [54 55 56]
  [57 58 59]]]


******************************
[[[ 0  1  2]
  [ 3  4  5]
  [ 6  7  8]]

 [[12 13 14]
  [18 19 20]
  [21 22 23]]

 [[27 28 29]
  [30 31 32]
  [33 34 35]]

 [[36 37 38]
  [39 40 41]
  [45 46 47]]

 [[48 49 50]
  [54 55 56]
  [57 58 59]]]