tensorflow - 用 CNN 预测图像的类别
问题描述
我试图预测训练模型上单个图像的类别,但我得到一个奇怪的输出,所以这是我的代码:
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
import matplotlib.pyplot as plt
import numpy as np
import os
def load_image(img_path, show=False):
img = image.load_img(img_path, target_size=(300, 300))
img_tensor = image.img_to_array(img)
img_tensor = np.expand_dims(img_tensor, axis=0)
img_tensor /= 255.
if show:
plt.imshow(img_tensor[0])
plt.axis('off')
plt.show()
return img_tensor
if __name__ == "__main__":
# load model
model = load_model("model1.h5")
# image path
img_path = 'dog.jpg' # dog
# load a single image
new_image = load_image(img_path, True)
# check prediction
pred = model.predict(new_image)
print(pred)
但我得到[[0.8189566 0.18104333]]
了输出。但是我有 0 类和 1 类。这可能是因为在任何地方都没有指定批量大小吗?
解决方案
与model.predict(new_image)
您获得每个测试图像属于特定类别的概率
要获得输出类,您需要选择预测的最大概率的类:
np.argmax(pred, axis=1)