Keras-如何使用argmax进行预测

Question

我有3类类别Tree, Stump, Ground 。 我列出了以下类别的清单：

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

当我打印我的预测时，它给我输出

[[0. 1. 0.]]

我已经阅读了有关numpy的Argmax的信息，但是我不确定在这种情况下如何使用它。

我试过使用

print(np.argmax(prediction))

但这给了我1的输出。 很好，但是我想找出1的索引是什么，然后打印出Category而不是最大值。

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))

我希望我的预测能告诉我：

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

感谢您的阅读:)非常感谢您的帮助。

Answer 1

您只需要使用np.argmax的结果索引类别：

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