Convert a 2D numpy array into a hot-encoded 3D numpy array, with same values in the same plane

Question

Suppose I have a Numpy array:

[
  [0, 1, 0],
  [0, 1, 4],
  [2, 0, 0],
]

How can I turn this into a "hot encoded" 3D array? something like this:

[
  # Group of 0's
  [[1, 0, 1],
   [1, 0, 0],
   [0, 1, 1]],
  # Group of 1's
  [[0, 1, 0],
   [0, 1, 0],
   [0, 0, 0]],
  # Group of 2's
  [[0, 0, 0],
   [0, 0, 0],
   [1, 0, 0]],
  # Group of 3's
  # the group is still here, even though there are no threes
  [[0, 0, 0],
   [0, 0, 0],
   [0, 0, 0]],
  # Group of 4's
  [[0, 0, 0],
   [0, 0, 1],
   [0, 0, 0]]
]

That is, how can I take each occurrence of a number in the array and "group" them into their own plane in a 3D matrix? As shown in the example, even the "gap" in numbers (i.e. the 3) should still appear. In my case, I know the range of the data beforehand (range (0, 6]), so that should make it easier.

BTW, I need this because I have a chessboard represented by numbers, but need it in this form to pass into a 2d convolutional neural network (different "channels" for different pieces).

I've seen Convert a 2d matrix to a 3d one hot matrix numpy, but that has a one-hot encoding for every value, which isn't what I'm looking for.

anky anky · Accepted Answer · 2021-04-25T04:32:18

Create the desired array (arr.max()+1 here) and then reshape it to compare to the original array:

Setup:

arr = np.array([
  [0, 1, 0],
  [0, 1, 4],
  [2, 0, 0],
])

u = np.arange(arr.max()+1)
(u[:,np.newaxis,np.newaxis]==arr).astype(int)

array([[[1, 0, 1],
        [1, 0, 0],
        [0, 1, 1]],

       [[0, 1, 0],
        [0, 1, 0],
        [0, 0, 0]],

       [[0, 0, 0],
        [0, 0, 0],
        [1, 0, 0]],

       [[0, 0, 0],
        [0, 0, 0],
        [0, 0, 0]],

       [[0, 0, 0],
        [0, 0, 1],
        [0, 0, 0]]])

Convert a 2D numpy array into a hot-encoded 3D numpy array, with same values in the same plane

1 Answers