前言
文章目录
项目环境安装
ESP-IDF安装
ESP-IDF 5.0+ 的版本有较大改动，在部署过程中会出现一些问题，建议使用 4.4 版本的进行安装。
基于 Windows 平台的软件安装，可以参考 https://dl.espressif.com/dl/esp-idf/. 按照流程完成安装即可。
开发环境
本项目整体开发环境主要基于训练框架，以及对应esp32的模型部署框架，具体如下：
- 训练、转换模型: Model Assistant
- 模型部署: sscma-example-esp32(1.0.0)
运行环境
python
python3.10 + CUDA11.7 + esp-idf 4.4
# 主要按照 ModelAssistant/requirements_cuda.txt 进行安装
torch 2.0.0+cu117
torchaudio 2.0.1+cu117
torchvision 0.15.1+cu117
yapf 0.40.2
typing_extensions 4.5.0
tensorboard 2.13.0
tensorboard-data-server 0.7.2
tensorflow 2.13.0
keras 2.13.1
tensorflow-estimator 2.13.0
tensorflow-intel 2.13.0
tensorflow-io-gcs-filesystem 0.31.0
sscma 2.0.0rc3
setuptools 60.2.0
rich 13.4.2
Pillow 9.4.0
mmcls 1.0.0rc6
mmcv 2.0.0
mmdet 3.0.0
mmengine 0.10.1
mmpose 1.2.0
mmyolo 0.5.0conda 的环境依赖主要见上面各种库的版本，其中
mmcv库安装
mmcv库的安装需要对应cuda版本、torch版本以及python版本，具体说明：cu117，torch2.0.0，python3.10可以参考
https://download.openmmlab.com/mmcv/dist/cu117/torch2.0.0/index.html，对应其中主要根据操作系统选择性安装，./mmcv-2.0.1-cp310-cp310-manylinux1_x86_64.whl和./mmcv-2.0.1-cp310-cp310-win_amd64.whl文件，具体如下，
训练数据集准备
添加自定义数据集
主要的数据集可以从开源的 Roboflow Universe 搜集，比如我们需要识别某些类别，可以在该网站上下载对应的数据集，下载格式选择 COCO 格式，如下图所示：
将下载的数据集压缩包放置在 data/collection 目录下面，对各个类别数据集加压并重命名为类别名称，例如 "face", "phone"，即类别标签。
挑选多个类别数据集后（这里选取 "face, phone"），需要对其进行合并，利用下述代码进行合并（主要合并 json 文件，并拷贝各类别数据集的图像文件）：
python
import os
import json
import shutil
# Save paths
train_path = 'datasets/collection/train'
valid_path = 'datasets/collection/valid'
# Create directories if not exist
if not os.path.exists(train_path):
os.makedirs(train_path)
if not os.path.exists(valid_path):
os.makedirs(valid_path)
def Add_Class_COCO_Data(dataset_path, classname, dataset, dataset_info):
with open(os.path.join(dataset_path, '_annotations.coco.json'), 'r') as file:
class_data = json.load(file)
# check if the class already exists in the train data
class_exist, class_id = False, None
for item in dataset['categories']:
if item['name'] == classname:
class_exist, class_id = True, item['id']
if class_exist is False:
class_id = dataset_info['category_id']
dataset['categories'].append({'id': class_id, 'name': classname})
dataset_info['category_id'] += 1
# add the class images for the train data
image_id = dataset_info['img_id']
for image in class_data['images']:
image_info = {'id': dataset_info['img_id'], 'file_name': image['file_name'], 'width': image['width'], 'height': image['height']}
dataset['images'].append(image_info)
dataset_info['img_id'] += 1
# add the class annotations for the train data
for ann in class_data['annotations']:
ann_info = {'id': dataset_info['ann_id'], 'image_id': image_id+ann['image_id'], 'category_id': class_id, 'bbox': ann['bbox'],
'area': ann['area'], 'iscrowd': ann['iscrowd']}
dataset['annotations'].append(ann_info)
dataset_info['ann_id'] += 1
def Copy_Files(src_dir, dst_dir, skip_files=['_annotations.coco.json']):
"""
Copy the files from source directory to the destination directory.
"""
os.makedirs(dst_dir, exist_ok=True)
for file_name in os.listdir(src_dir):
# skip some files
if file_name in skip_files:
continue
src_file = os.path.join(src_dir, file_name)
dst_file = os.path.join(dst_dir, file_name)
if os.path.isfile(src_file):
shutil.copy(src_file, dst_file)
print(f"Copied: {src_file} -> {dst_file}")
else:
print(f"Skipping directory: {src_file}")
# training classes
classes = ["face", "phone"]
# store the combined data
category_id = 0
train_data = {'categories': [], 'images': [], 'annotations': []}
train_info = {'category_id': 0, 'img_id': 0, 'ann_id': 0}
valid_data = {'categories': [], 'images': [], 'annotations': []}
valid_info = {'category_id': 0, 'img_id': 0, 'ann_id': 0}
class_data_root = './data/collection'
for cls_name in classes:
class_train_path = os.path.join(class_data_root, cls_name, 'train')
class_valid_path = os.path.join(class_data_root, cls_name, 'valid')
Add_Class_COCO_Data(class_train_path, cls_name, train_data, train_info)
Add_Class_COCO_Data(class_valid_path, cls_name, valid_data, valid_info)
Copy_Files(class_train_path, train_path)
Copy_Files(class_valid_path, valid_path)
# Save data to file
with open(os.path.join(train_path, '_annotations.coco.json'), 'w') as f:
json.dump(train_data, f, indent=4)
with open(os.path.join(valid_path, '_annotations.coco.json'), 'w') as f:
json.dump(valid_data, f, indent=4)
print(f">>> length of categories: {len(train_data['categories'])}")
print(f">>> length of train_data: {len(train_data['images'])}")
print(f">>> length of valid_data: {len(valid_data['images'])}")合并后的统一的训练和验证数据集全部放在了 datasets/collection 目录下面，以便于后续模型训练使用。
下载预训练模型
参考 Face Detection - Swift-YOLO 下载预训练模型权重文件 pretrain.pth，然后保存在 ModelAssistant/checkpoints 文件夹下。
训练 YOLO 模型
在 ModelAssistant 项目下，采用的是仓库提供的 Swift YOLO 模型，作者解释说明该模型具有优化的端侧运行性能："We implemented a lightweight object detection algorithm called Swift YOLO, which is designed to run on low-cost hardware with limited computing power. The visualization tool, model training and export command-line interface has refactored now".
采用 swift_yolo_tiny_1xb16_300e_coco 的配置文件进行训练，训练命令如下：
python
# training the yolo model
python tools/train.py configs/swift_yolo/swift_yolo_tiny_1xb16_300e_coco.py \
--cfg-options \
work_dir=work_dirs/collection \
num_classes=2 \
epochs=300 \
height=96 \
width=96 \
data_root=datasets/collection/ \
load_from=checkpoints/pretrain.pth模型量化和格式转换
训练完毕后，还需要对模型进行权重量化以及格式转换，这样才能够让模型成功在 ESP32S3 主板上运行。在工作目录 work_dirs/collection 下，找到最好的 bbox_mAP 的模型，例如这里是 best_coco_bbox_mAP_epoch_300.pth，采用以下命令导出模型：
python
# export the model
python tools/export.py configs/swift_yolo/swift_yolo_tiny_1xb16_300e_coco.py ./work_dirs/collection/best_coco_bbox_mAP_epoch_300.pth --cfg-options \
work_dir=work_dirs/collection \
num_classes=2 \
epochs=300 \
height=96 \
width=96 \
data_root=datasets/collection/ \
load_from=checkpoints/pretrain.pth导出的模型会保存在 work_dirs/collection 文件夹下，生成 best_coco_bbox_mAP_epoch_300_int8.tflite 文件，这是量化到 Int8 格式的 tflite 文件，可以用于后续模型的部署。
模型结果评估
训练损失
在整个 300 epoches 的训练过程中，对应的类别损失以及目标检测损失的变化如下图所示：
评估指标
在该项目中，主要测定了 Swift YOLO Tiny 结构以及 Swift YOLO Tiny Nano 结构（一种高度紧凑的深卷积神经网络，用于使用人机协同设计策略设计的嵌入式目标检测），主要的评估指标如下所示：
考虑到部署到端侧设备上时，更关注于目标检测的置信度，因此下面主要从置信度、平均推理时间的角度进行评估。
| Model | Precision | Class | Confidence | Infer_Time(ms) | Size(MB) |
|---|---|---|---|---|---|
| Swift YOLO Tiny | Float32 | Face | 69.57 % | 6.44 | 3.63 |
| Swift YOLO Tiny | Float32 | Phone | 54.86 % | 6.44 | 3.63 |
| Swift YOLO Tiny | Float32 | [Face, Phone] | 62.21 % | 6.44 | 3.63 |
| Swift YOLO Tiny | Int8 | Face | 68.75 % | 6.69 | 1.05 |
| Swift YOLO Tiny | Int8 | Phone | 55.18 % | 6.69 | 1.05 |
| Swift YOLO Tiny | Int8 | [Face, Phone] | 61.97 % | 6.69 | 1.05 |
| Swift YOLO Tiny Nano | Float32 | Face | 73.62 % | 8.02 | 9.13 |
| Swift YOLO Tiny Nano | Float32 | Phone | 55.76 % | 8.02 | 9.13 |
| Swift YOLO Tiny Nano | Float32 | [Face, Phone] | 64.69 % | 8.02 | 9.13 |
| Swift YOLO Tiny Nano | Int8 | Face | 74.62 % | 12.86 | 2.49 |
| Swift YOLO Tiny Nano | Int8 | Phone | 55.62 % | 12.86 | 2.49 |
| Swift YOLO Tiny Nano | Int8 | [Face, Phone] | 65.12 % | 12.86 | 2.49 |
- 可以看到，模型量化压缩后在能够维持较高的精度的情况下，模型大小显著减小，但是推理时间并没有降低，甚至约有增加；
- Swift YOLO Tiny Nano 的精度更高，但是模型大小更大，推理时间更长，在 ESP32S3-EYE 设备上牺牲的代价就是帧率较低；
- 整体而言，模型的置信度都已经比较高以及推理速度能比较不错，能够满足实际应用需求。
此外，我还尝试了保持原图像尺寸，即对于Swift YOLO Tiny配置而言，设置模型处理的图像宽高都是640，具体如下：
python
# training the yolo model
python tools/train.py configs/swift_yolo/swift_yolo_tiny_1xb16_300e_coco.py \
--cfg-options \
work_dir=work_dirs/collection_640 \
num_classes=2 \
epochs=300 \
height=640 \
width=640 \
data_root=datasets/collection/ \
load_from=checkpoints/pretrain.pth训练后模型的评估指标如下所示：
| Model | Precision | Class | Confidence | Infer_Time(ms) | Size(MB) |
|---|---|---|---|---|---|
| Swift YOLO Tiny WH640 | Float32 | Face | 74.83 % | 54.08 | 3.88 |
| Swift YOLO Tiny WH640 | Float32 | Phone | 65.47 % | 54.08 | 3.88 |
| Swift YOLO Tiny WH640 | Float32 | [Face, Phone] | 70.15 % | 54.08 | 3.88 |
| Swift YOLO Tiny WH640 | Int8 | Face | 75.52 % | 71.87 | 1.20 |
| Swift YOLO Tiny WH640 | Int8 | Phone | 65.46 % | 71.87 | 1.20 |
| Swift YOLO Tiny WH640 | Int8 | [Face, Phone] | 70.49 % | 71.87 | 1.20 |
虽然整体的精度有所提升，但是推理时间显著增加，模型虽然大小基本维持不变，但是推理时间增大了将近10倍，处理图像的分辨率所带来的开销远远超过了模型识别的精度。事实上，Swift YOLO Tiny WH640 在 ESP32S3-EYE 设备上已经没办法运行，实际中会出现数据存储栈溢出的问题，摄像头采集的图像分辨率太高导致栈空间不足。
模型推理
主要是观测量化后模型对验证集的推理结果，具体如下：
模型部署
部署环境
部署环境为 ESP32-S3 EYE 开发板，没判断错的话，它有 4MB Flash，我们烧录的模型也主要存储在这个区域，当然它也附带了 SD 卡功能，可以从 SD 卡中加载模型。4 MB 的 Flash 分配主要在 partitions.csv 文件中，具体如下：
csv
# Name Type SubType Offset Size Flags
# Note: if you change the phy_init or app partition offset make sure to change the offset in Kconfig.projbuild
factory app factory 0x010000 2048K
nvs data nvs 0x3D0000 64K
fr data 0x3E0000 128K默认 app 分区最多有 2048K 即 2MB 的大小，但是由于地址空间还很充裕，例如从 0x010000 到 0x3D0000 的大小，最多可以分配 3MB 的空间，足够上述量化后的模型存储。此外，分区表主要以 64KB 为单位分配的，所以偏移地址后四位都为0. 可以在项目中通过 idf.py build 查看输出信息，其中包含了对分区大小是否合适以及剩余空间的判断。
对于 Swift YOLO Tiny Nano 配置，需要将分区表中的 app 分区增大到 3MB 即 3072K.
导入模型
对于烧录程序而言，需要将 tflite 格式模型转换为 c 语言格式，具体在项目 sscma-example-esp32-1.0.0 中，通过 tools/tflite2c.py 文件进行转换，但是为了能够在显示屏上同步显示检测的类别名称，需要修改其中的代码，即将 classes 换成字符串，然后通过分词找到各个类别，并将其转换为字符串列表，其中每个字符串都是一个类别名称，具体如下：
python
import sys
import os
import binascii
import argparse
def parse_args():
parser = argparse.ArgumentParser(
description='Convert tflite to c or cpp file')
parser.add_argument('--input', help='input tflite file')
parser.add_argument('--output_dir', help='output directory')
parser.add_argument('--name', help='model name')
parser.add_argument('--cpp', action='store_true',
default=True, help='output cpp file')
parser.add_argument('--classes', type=str, help='classes name')
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
input = args.input
name = args.name
output_dir = args.output_dir
classes = args.classes
if classes != None:
classes = list(classes.split(','))
if not os.path.exists(input):
print('input file not exist')
sys.exit(1)
if name == None:
name = input.split('/')[-1].split('.')[0]
output_h = os.path.join(output_dir, name + '_model_data.h')
if args.cpp:
output_c = os.path.join(output_dir, name + '_model_data.cpp')
else:
output_c = os.path.join(output_dir, name + '_model_data.c')
with open(input, 'rb') as f_input:
data = f_input.read()
if data[4:8] != b'TFL3':
print('input file is not tflite')
sys.exit(1)
data = binascii.hexlify(data)
data = data.decode('utf-8')
with open(output_h, 'w') as f_output_h:
f_output_h.write('#ifndef __%s_MODEL_DATA_H__\r\n' % name.upper())
f_output_h.write('#define __%s_MODEL_DATA_H__\r\n' % name.upper())
f_output_h.write('\r\n//this file is generated by tflite2c.py\r\n')
f_output_h.write('\r\n#include <stdint.h>\r\n')
f_output_h.write(
'extern const unsigned char g_%s_model_data[];\r\n' % name)
f_output_h.write(
'extern const unsigned int g_%s_model_data_len;\r\n' % name)
if classes != None:
f_output_h.write(
'extern const char* g_%s_model_classes[];\r\n' % name)
f_output_h.write(
'extern const unsigned int g_%s_model_classes_num;\r\n' % name)
f_output_h.write('\r\n#endif\r\n')
f_output_h.close()
with open(output_c, 'w') as f_output_c:
f_output_c.write('#include <stdint.h>\r\n')
f_output_c.write('\r\n#include "%s_model_data.h"\r\n\r\n' % name)
f_output_c.write(
'const unsigned char g_%s_model_data[] = {\r\n' % name)
for i in range(0, len(data), 2):
f_output_c.write('0x')
f_output_c.write(data[i])
f_output_c.write(data[i+1])
f_output_c.write(', ')
if i % 36 == 34:
f_output_c.write('\r\n')
f_output_c.write('};\r\n\r\n')
f_output_c.write(
'const unsigned int g_%s_model_data_len = %d;\r\n' % (name, len(data) // 2))
if classes != None:
f_output_c.write(
'const char* g_%s_model_classes[] = {' % name)
for i in range(len(classes)):
f_output_c.write('"%s", ' % classes[i])
f_output_c.write('};\r\n\r\n')
f_output_c.write(
'const unsigned int g_%s_model_classes_num = %d;\r\n' % (name, len(classes)))
else:
f_output_c.write(
'const char* g_%s_model_classes[] = {};\r\n' % name)
f_output_c.write(
'const unsigned int g_%s_model_classes_num = 0;\r\n' % name)
f_output_c.close()
f_input.close()同时为了能够在显示屏上显示置信度，还需要在 components/modules/algorithm/algo_yolo.cpp 文件中修改相应的代码，具体如下所示：
cpp
if (std::distance(_yolo_list.begin(), _yolo_list.end()) > 0)
{
int index = 0;
found = true;
printf(" Objects found: %d\n", std::distance(_yolo_list.begin(), _yolo_list.end()));
printf(" Objects:\n");
printf(" [\n");
for (auto &yolo : _yolo_list)
{
yolo.x = uint16_t(float(yolo.x) / float(w) * float(frame->width));
yolo.y = uint16_t(float(yolo.y) / float(h) * float(frame->height));
yolo.w = uint16_t(float(yolo.w) / float(w) * float(frame->width));
yolo.h = uint16_t(float(yolo.h) / float(h) * float(frame->height));
fb_gfx_drawRect2(frame, yolo.x - yolo.w / 2, yolo.y - yolo.h / 2, yolo.w, yolo.h, box_color[index % (sizeof(box_color) / sizeof(box_color[0]))], 4);
// fb_gfx_printf(frame, yolo.x - yolo.w / 2, yolo.y - yolo.h/2 - 5, 0x1FE0, 0x0000, "%s", g_yolo_model_classes[yolo.target]);
fb_gfx_printf(frame, yolo.x - yolo.w / 2, yolo.y - yolo.h/2 - 5, 0x1FE0, "%s:%d", g_yolo_model_classes[yolo.target], yolo.confidence);
printf(" {\"class\": \"%d\", \"x\": %d, \"y\": %d, \"w\": %d, \"h\": %d, \"confidence\": %d},\n", yolo.target, yolo.x, yolo.y, yolo.w, yolo.h, yolo.confidence);
index++;
}
printf(" ]\n");
}然后具体的转换命令如下所示，这将会在文件夹 components/modules/model 中生成两个文件 yolo_model_data.h 和 yolo_model_data.cpp，其中 yolo_model_data.h 中包含了模型数据的声明，yolo_model_data.cpp 中包含了模型数据的定义。
python
python tools/tflite2c.py --input ./model_zoo/facephone_96/best_coco_bbox_mAP_epoch_300_int8.tflite --name yolo --output_dir ./components/modules/model --classes "person,phone"烧录模型
模型准备完毕后，就可以利用乐鑫的 IDF 开发工具链将模型烧录进开发板中，具体过程主要是：
- 使用
ESP-IDF 4.4 CMD进入sscma-example-esp32-1.0.0项目的examples/yolo目录下； - 使用
idf.py set-target esp32s3设置目标芯片为esp32s3； - 使用
idf.py build命令编译项目； - 使用
idf.py flash monitor命令烧录项目。
结果展示
Swift YOLO Tiny Nano
在 ESP32S3 设备上平均每张图像的处理时间是 628 ms，在端侧设备上出现帧率较低，稍微偏卡顿的效果。但是准确度提高的也不是很多，因此首选还是 Swift YOLO Tiny。
Swift YOLO Tiny
在 ESP32S3 设备上平均每张图像的处理时间是 268 ms，实时监测的帧率非常不错，已经很流畅，准确性而言也是较为不错的，具体结果如下所示：
实际的开发板检测结果如下所示：
总结
本文主要介绍了如何利用 Model Assistant、sscma-example-esp32-1.0.0 项目将 YOLO 模型部署到 ESP32S3 EYE 设备上，并且通过自定义数据集实现训练、导出、转换、导入、部署模型整个流程，也对实际开发板进行了测试，最终得到了较为理想的结果。
本文主要选取了 "face, phone" 两个类别进行训练，如果需要训练更多的类别，可以添加更多的类别，但是会存在目标检测上限以及端侧设备计算性能的限制，需要进一步平衡。