Keras - How to use argmax for predictions

Question

I have 3 categories of classes Tree, Stump, Ground. I've made a list for these categories:

CATEGORIES = ["Tree", "Stump", "Ground"]

When i print my prediction, it gives me the output of

[[0. 1. 0.]]

I've read up about numpy's Argmax but I'm not entirely sure how to use it in this case.

I've tried using

print(np.argmax(prediction))

But that gives me the output of 1. That's great but I would like to find out what's the index of 1 and then print out the Category instead of the highest value.

import cv2
import tensorflow as tf
import numpy as np

CATEGORIES = ["Tree", "Stump", "Ground"]


def prepare(filepath):
    IMG_SIZE = 150 # This value must be the same as the value in Part1
    img_array = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)
    new_array = cv2.resize(img_array, (IMG_SIZE, IMG_SIZE))
    return new_array.reshape(-1, IMG_SIZE, IMG_SIZE, 1)

# Able to load a .model, .h3, .chibai and even .dog
model = tf.keras.models.load_model("models/test.model")

prediction = model.predict([prepare('image.jpg')])
print("Predictions:")
print(prediction)
print(np.argmax(prediction))

I expect my prediction to show me:

Predictions:
[[0. 1. 0.]]
Stump

Thanks for reading :) I appreciate any help at all.

Answer 1

You just have to index categories with the result of np.argmax:

pred_name = CATEGORIES[np.argmax(prediction)]
print(pred_name)