python - How to converse an array of n-column labels to 1-column?

Question

this is a label array that I use in my classification model:

array([[1, 0, 0],
   [0, 1, 0],
   [0, 0, 1]], dtype=uint8)

But I'd like to reverse it to one column, so it will look like this:

  array([[0],
   [1],
   [2]], dtype=uint8)

Any suggestions are appreciated.

question from:https://stackoverflow.com/questions/65869353/how-to-converse-an-array-of-n-column-labels-to-1-column

Answer 1

You could use np.argmax

np.argmax(arr, axis=1).reshape(arr.shape[0], 1).astype(np.int8)

# array([[0],
#        [1],
#        [2]], dtype=int8)

If you want to take always the position of the ones:

np.argmax(arr==1, axis=1)