TF Lite口罩检测图像分类模型如何关联打印标签与预测结果

你只需要在获取result的代码后，添加标签映射和遍历打印的逻辑即可，修改后的完整代码如下：

import tensorflow as tf
import cv2
from pprint import pprint

interpret= tf.lite.Interpreter(model_path="<FILE_PATH>")
input_details = interpret.get_input_details()
output_details = interpret.get_output_details()

interpret.allocate_tensors()

pprint(input_details)
pprint(output_details)

data= cv2.imread("file.jpeg")
new_image= cv2.resize(data,(224,224))

interpret.resize_tensor_input(input_details[0]["index"],[1,224,224,3])
interpret.allocate_tensors()
interpret.set_tensor(input_details[0]["index"],[new_image])
interpret.invoke()
result= interpret.get_tensor(output_details[0]['index'])

# 新增标签映射打印逻辑
# 标签顺序与模型训练时的输出顺序保持一致
labels = ["mask", "no_mask"]
pred_scores = result[0]
for label, score in zip(labels, pred_scores):
    # 如果需要输出0-100区间的百分比置信度，可取消下行注释对原始logits做转换
    # score = round(tf.nn.softmax(pred_scores).numpy()[list(labels).index(label)] * 100)
    print(f"{label}: {score}")

补充说明：

• 你当前拿到的[[30 246]]是模型输出的原始logits值，如果需要得到和你示例一致的百分比格式输出，取消代码中softmax相关的注释即可。
• 务必保证labels列表的顺序和你训练模型时分类的输出顺序完全一致，否则会出现标签和数值匹配错误的问题。

内容的提问来源于stack exchange，提问作者Ankita Shrivastava