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)