python - Append 2D array to 3D array, extending third dimension

Question

I have an array A that has shape (480, 640, 3), and an array B with shape (480, 640).

How can I append these two as one array with shape (480, 640, 4)?

I tried np.append(A,B) but it doesn't keep the dimension, while the axis option causes the ValueError: all the input arrays must have same number of dimensions.

Answer 1

Use dstack:

>>> np.dstack((A, B)).shape
(480, 640, 4)

This handles the cases where the arrays have different numbers of dimensions and stacks the arrays along the third axis.

Otherwise, to use append or concatenate, you'll have to make B three dimensional yourself and specify the axis you want to join them on:

>>> np.append(A, np.atleast_3d(B), axis=2).shape
(480, 640, 4)