Vitis AI 基本认知（Tiny-VGG 标签获取+预测后处理）

news2024/9/20 9:25:43

1. 简介

2. 解析

2.1 获取标签

2.1.1 载入数据集

2.1.2 标签-Index

2.1.3 保存和读取类别标签

2.2 读取单个图片

2.3 载入模型并推理

2.3.1 tiny-vgg 模型结构

2.3.2 运行推理

2.4 置信度柱状图

2.5 预测标签

3. 完整代码

4. 总结

1. 简介

本博文在《Vitis AI 基本认知（Tiny-VGG 项目代码详解）-CSDN博客》基础上，详细介绍如何使用TensorFlow框架进行单个图片的推理，从获取和处理数据集的标签开始，到模型的加载与推理，再到结果的可视化展示。关键信息如下：

获取数据集的标签
保存和读取类别标签
加载模型并推理
绘制图像
使用中文标签
置信度柱状图

2. 解析

2.1 获取标签

2.1.1 载入数据集

通过 image_dataset_from_directory 方法

vali_dataset = tf.keras.preprocessing.image_dataset_from_directory(
    './dataset/class_10_val/val_images/',
    image_size=(64, 64),
    batch_size=32)

取出一个图片，并查看其标签：

for images, labels in vali_dataset.take(1):
    # 取出第一个图片和标签
    image = images[0].numpy().astype("uint8")
    label = labels[0].numpy()

    # 显示图片
    plt.figure(figsize=(2, 2))
    plt.imshow(image)
    plt.title(f"Label: {label}")
    plt.axis('off')
    plt.show()

2.1.2 标签-Index

查看类别标签及其 Index：

class_names = vali_dataset.class_names

for i, class_name in enumerate(class_names):
    print(f"Class name: {class_name:<4}, Index: {i}")
---
Class name: 咖啡   , Index: 0
Class name: 小熊猫 , Index: 1
Class name: 披萨   , Index: 2
Class name: 救生艇 , Index: 3
Class name: 校车   , Index: 4
Class name: 橙子   , Index: 5
Class name: 灯笼椒 , Index: 6
Class name: 瓢虫   , Index: 7
Class name: 考拉   , Index: 8
Class name: 跑车   , Index: 9

类别标签对应的 one-hot 标签：

for index, class_name in enumerate(class_names):
    one_hot = tf.one_hot(index, len(class_names)).numpy()
    print(f"Class: {class_name}, One-hot: {one_hot}")
---
Class: 咖啡  , One-hot: [1. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
Class: 小熊猫, One-hot: [0. 1. 0. 0. 0. 0. 0. 0. 0. 0.]
Class: 披萨  , One-hot: [0. 0. 1. 0. 0. 0. 0. 0. 0. 0.]
Class: 救生艇, One-hot: [0. 0. 0. 1. 0. 0. 0. 0. 0. 0.]
Class: 校车  , One-hot: [0. 0. 0. 0. 1. 0. 0. 0. 0. 0.]
Class: 橙子  , One-hot: [0. 0. 0. 0. 0. 1. 0. 0. 0. 0.]
Class: 灯笼椒, One-hot: [0. 0. 0. 0. 0. 0. 1. 0. 0. 0.]
Class: 瓢虫  , One-hot: [0. 0. 0. 0. 0. 0. 0. 1. 0. 0.]
Class: 考拉  , One-hot: [0. 0. 0. 0. 0. 0. 0. 0. 1. 0.]
Class: 跑车  , One-hot: [0. 0. 0. 0. 0. 0. 0. 0. 0. 1.]

2.1.3 保存和读取类别标签

将类别标签写入文本文档：

with open('tiny_VGG_class_names.txt', 'w') as file:
    for class_name in class_names:
        file.write(f"{class_name}\n")

从文本文档中读取类别标签：

with open('tiny_VGG_class_names.txt', 'r') as file:
    class_names = [line.strip() for line in file]

print(class_names)
---
['咖啡', '小熊猫', '披萨', '救生艇', '校车', '橙子', '灯笼椒', '瓢虫', '考拉', '跑车']

2.2 读取单个图片

读取图片，并显示在 Jupyter Lab 中：

img = cv2.imread('./dataset/class_10_val/val_images/橙子/val_1067.JPEG')

plt.figure(figsize=(2, 2))
plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
plt.axis('off')
plt.show()

对图片归一化操作：

normalization_layer = tf.keras.layers.Rescaling(1./255)
img_norm = normalization_layer(img)
img_norm = np.expand_dims(img_norm, axis=0)
np.shape(img_norm)
---
(1, 64, 64, 3)

训练过程中，对数据集做了归一化处理，推理时也要做同样的处理。