How to take n block of a 3d array in python?

I had 328 grayscale images with the size 128*128 and I convert all of them into a 3D array with the shape (128,128,328). Now I want to convert it into 5 separated 3d arrays with the shape of (128,128,64) without changing in sequence.

As you can see 328 is not divisible by 64 and using dsplit function is not working.

Is there any way to slice the 3d array on depth axes dynamically?

Answer

import numpy as np

arr = np.zeros(shape=(128,128,328))

# make aray divisible by 5 (trim it to the depth of 320)
arr = arr[:, :, :320]

# Split array
arrays = np.dsplit(arr, 5)

for array in arrays:
    print(array.shape)

Output:

(128, 128, 64)
(128, 128, 64)
(128, 128, 64)
(128, 128, 64)
(128, 128, 64)

EDIT: Here is the same thing written in a dynamic way.

import numpy as np

num_subarrays = 5
subarray_depth = 64

# Initialize array
arr = np.zeros(shape=(128,128,328))

# make aray divisible by subarray_depth
arr = arr[:, :, :(arr.shape[2] // subarray_depth) * subarray_depth]

# Split array
arrays = np.dsplit(arr, num_subarrays)

for array in arrays:
    print(array.shape)