EEPW论坛

简介

通过上次的发帖后https://forum.eepw.com.cn/thread/394354/1/#1 反响不错，看到论坛内的电子工程师对这个训练的过程还是蛮感兴趣的，于是决定再发一个帖子，这次使用卷积神经网络基于TensorFlow训练一个宝可梦的分类。其实是可以直接部署在ESP32的单片机上的，但是目前我驱动OV7670遇到了一点问题（部署的帖子直接看上述连接即可，通用）。 所以这篇文章着重讲解一下训练的过程。

首先让我介绍一下数据集，在精灵宝可梦绿的游戏中，一共是具有151个精灵。从妙蛙种子，一直到最后的梦幻。由于宝可梦总共就这么多，所以我们没有太大的必要来在这个有限的数据上划分训练数据集和测试数据集或者交叉验证等。无论训练的模型再怎么过拟合都对我们的分类没有任何的影响。

在我们获得数据集之后我们需要给每一个精灵宝可梦打进行数据的标注，从而告诉我们的模型我们期望的输出是什么，通产的标注办法比较麻烦。这里建议使用LableStudio进行数据的标注。

此时Lable studio已经启动了。 我们可以创建项目并且将宝可梦都导入进去。

然后在项目中手动对每一个图片进行标注 （这是一个麻烦的过程）

在标准完成之后点击Submit，然后点击右上角的export导出数据。选择导出JSON或者CSV都可以

之后进入到我们实现准备好的Tensorflow的环境中进行代码的编写。由于当前的数据结构都是被导入到了LableStudio中，所以我们需要找到上述的位置，并且拷贝到我们的项目目录中，或者修改JSON的配置文件指向LableStudio中。

接下来我们开始代码的编写。

import json
import pandas as pd

# JSON 文件名
json_file = "pokemon_labels.csv.json"

with open(json_file, "r", encoding="utf-8") as f:
    data = json.load(f)

rows = []
for task in data:
    # 当前目录下的实际路径
    image_path = f"data/upload/3/{task['file_upload']}"
    label = task["annotations"][0]["result"][0]["value"]["choices"][0]
    rows.append([image_path, label])

df = pd.DataFrame(rows, columns=["image", "label"])
df.to_csv("pokemon_labels.csv", index=False, encoding="utf-8")

print("CSV 已生成，路径指向当前目录的 data/upload/3")

首先我们使用Pandas来将JSON转换成CSV

import tensorflow as tf
from tensorflow.keras import layers, models

# -----------------------------
# 1. 读取 CSV 并构建 dataset
# -----------------------------
df = pd.read_csv("pokemon_labels.csv")

# 构建标签映射
labels = sorted(df["label"].unique())
label_to_index = {label: i for i, label in enumerate(labels)}
df["label_idx"] = df["label"].map(label_to_index)

# 加载图片函数
def load_image(path, label):
    img = tf.io.read_file(path)
    img = tf.image.decode_png(img, channels=3)
    img = tf.image.resize(img, [64, 64])  # 调整大小
    img = img / 255.0  # 归一化
    return img, label

# 创建 tf.data.Dataset
dataset = tf.data.Dataset.from_tensor_slices((df["image"], df["label_idx"]))
dataset = dataset.map(load_image)
dataset = dataset.shuffle(buffer_size=len(df)).batch(32).prefetch(tf.data.AUTOTUNE)

然后使用tf读取数据集，并且来进行图片大小的调整和进行Nomalization来进行归一化，使其图像的范围进行收缩。使其更加容易训练。

num_classes = len(labels)

model = models.Sequential([
    layers.Conv2D(32, (3,3), activation='relu', input_shape=(64,64,3)),
    layers.MaxPooling2D((2,2)),

    layers.Conv2D(64, (3,3), activation='relu'),
    layers.MaxPooling2D((2,2)),

    layers.Conv2D(128, (3,3), activation='relu'),
    layers.MaxPooling2D((2,2)),

    layers.Flatten(),
    layers.Dense(128, activation='relu'),
    layers.Dense(num_classes, activation='softmax')
])

然后就可以构建我们的卷积神经网络，首先第一层为输入层，创建了32个3*3 的卷积核心，使用relu function作为激活函数，并且将上述resize之后的图片的大小作为输入（3个通道，RGB）。之后对上述的卷积层进行池化，然后再创建卷积层，和更多的卷积核对上述的池化后的层进行特征提取。重复上述操作，然后将所有的输出进行展平，进入隐层层来提取抽象特征，最后来构建根据宝可梦的数量来构建输出层，使用softmax进行分类。

# 3. 编译模型
# -----------------------------
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# -----------------------------
# 4. 训练模型
# -----------------------------
model.fit(dataset, epochs=500)

然后就是对模型进行编译和训练，采用的是adam量化器和交叉熵损失函数，然后进行训练。

可以看到针对训练数据集的准确率基本上就是1. 100%能够识别正确。

model.save("pokemon_cnn_model.keras")

之后便是保存模型，方便进行二次调用或者微调，然后我们来进行一次测试。

我这里有一个卡比兽的图片，我们把它作为输入。查看模型的分类输出是什么。

import tensorflow as tf
from tensorflow.keras.models import load_model
import numpy as np
from tensorflow.keras.preprocessing import image
import pandas as pd

# -----------------------------
# 1. 加载模型
# -----------------------------
model = load_model("pokemon_cnn_model.keras")

# -----------------------------
# 2. 加载标签映射
# -----------------------------
df = pd.read_csv("pokemon_labels.csv")
labels = sorted(df["label"].unique())  # 确保顺序和训练时一致

# -----------------------------
# 3. 加载单张图片
# -----------------------------
img_path = "143.png"  
img = image.load_img(img_path, target_size=(64, 64))  # resize到训练尺寸
img_array = image.img_to_array(img)
img_array = img_array / 255.0  # 归一化
img_array = np.expand_dims(img_array, axis=0)  

# -----------------------------
# 4. 模型预测
# -----------------------------
pred = model.predict(img_array)
pred_label_idx = np.argmax(pred[0])
pred_label = labels[pred_label_idx]

print("预测标签:", pred_label)

分类非常正确，就是卡比兽！

后续

后续正在看怎么驱动这个OV7670单片机来读取图像的显示，ESP-IDF有正确的组件，但是不知道为什么烧录后显示不能识别到这个摄像头，ID不对或者是型号不支持。 还在研究中。那么等到成功驱动的话便会将模型部署到ESP32上在ESP32上进行推理。我已经给大家准备好了数据集，代码和对应的模型。还有打好标签的JSON文件，大家可以自行进行尝试。

tf.zip