Convert probability from dummy variables to labels

Question

I have trained a CNN where I transformed my variables (5) into dummy variables, and now I would like to convert the output (probabilities) to the original label. my output is a matrix NxM for example

                                      predictions =  [0.1,0.0,0.1,0.6,0.2]
                                                     [0.2,0.5,0.1,0.1,0.1]
                                                     [0.4,0.4,0.1,0.1,0.0]
                                                              .
                                                              .
                                                              .
                                                     [0.1,0.0,0.1,0.6,0.2]

My labels would be classes labels = [1,2,3,4,5] and what I did to convert the output to my labels is:

i, j = np.where(predictions >= 0.5) # find the index where the probability is larger than 0.5
Y_pred =[]
for elem, value in enumerate(j): # loop through the lines
    x = labels[value]
    Y_pred.append(x)

But this code fails when I don't have a variable with a probability larger than 0.5, like in the 3rd line from my predictions, or when I have 2 variables with the same larger probability. Some idea how to overcome this issue? which condition I could apply? I can't provide the data, and with synthetic data, the CNN is not working well.

Answer 1

Well, the solution was to convert my output to pandas dataframe, then I found the index of the maximum value at each row: Y_pred = label[df.idxmax(axis=1)]

The point is to convert the result to pandas dataframe, then everything is easier.