正在打印Tf-Lite模型图像分类的标签

Question

我正在开发一个Image Claasification TF Lite模型，以使用此链接检测人脸上的面具或无面具。我按照链接，在顶点AI中训练了一个图像多类分类，并下载了TF-lite模型。模型的标签是"；掩模"；以及"；no_mask"；。为了测试模型，我编写了以下代码：

interpret= tf.lite.Interpreter(model_path="<FILE_PATH>")
input= interpret.get_input_details()
output= interpret.get_output_details()
interpret.allocate_tensors()
pprint(input)
pprint(output)
data= cv2.imread("file.jpeg")
new_image= cv2.resize(data,(224,224))
interpret.resize_tensor_input(input[0]["index"],[1,224,224,3])
interpret.allocate_tensors()
interpret.set_tensor(input[0]["index"],[new_image])
interpret.invoke()
result= interpret.get_tensor(output[0]['index'])
print (" Prediction is - {}".format(result))

将此代码用于我的一张图像会给我以下结果：

[[30 246]]

现在我也想在结果中打印标签。例如：

掩码：30

no_mask:46

有什么办法可以实现吗？

请帮忙，因为我是TF Lite 的新手

Answer 1

我自己解决了。从Vertex AI下载的.tflite模型包含名为"dict.txt"的标签文件，该文件包含所有标签。请在此处查看GCP文档。要获得这个标签文件，我们首先需要解压缩.tflite文件，该文件将提供dict.txt。有关更多信息，请查看tflite文档以及如何从模型中读取关联文件。

之后，我参考github链接标签.py执行了以下代码：

import argparse
import time
import numpy as np
from PIL import Image
import tensorflow as tf
interpret= tf.lite.Interpreter(model_path="<FILE_PATH>")
input= interpret.get_input_details()
output= interpret.get_output_details()
interpret.allocate_tensors()
pprint(input)
pprint(output)
data= cv2.imread("file.jpeg")
new_image= cv2.resize(data,(224,224))
interpret.resize_tensor_input(input[0]["index"],[1,224,224,3])
interpret.allocate_tensors()
interpret.set_tensor(input[0]["index"],[new_image])
interpret.invoke()
floating_model= input[0]['dtype'] == np.float32
op_data= interpret.get_tensor(output[0]['index'])
result= np.squeeze(op_data)
top_k=result.agrsort()[-5:][::1]
labels=load_labels("dict.txt")
for i in top_k:
if floating_model:
print('{:08.6f}: {}'.format(float(result[i]), labels[i]))
else:
print('{:08.6f}: {}'.format(float(result[i] / 255.0), labels[i]))