问题描述
图像分类作为计算机视觉领域的关键应用，在工业生产、物流管理、安防监控等多个场景发挥着重要作用。嘉楠科技的勘智 K230 芯片，凭借出色的运算性能与低功耗优势，为部署 Yolov5 图像分类模型提供了优质选择。接下来，我们将以数字分类为例，从环境搭建、数据采集、数据标注、模型训练、模型移植，到代码优化，一步步揭开图像分类的技术面纱，助力开发者在实际应用中落地创新。
一、训练环境准备
- 硬件：PC（考虑到训练需求，需要使用一台带GPU的电脑）
- 软件：Linux操作系统/Window安装WSL
- 网络：由于需要获取github的资源，最好是能够连接外网
1、搭建 YOLOv5 训练环境
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt
2、数据采集
我们已经提供 Micropython 脚本（请参考 K230 SDK 中代码 DataCollectionCamera.py），使用 K230 开发板就能完成数据采集的工作（数据采集脚本使用方法请参考文末 B站视频教程）
aaa@DESKTOP-OSN5BJK:~/canmv_k230_gitee_micropython$ cd src/canmv/resources/examples/16-AI-Cube/
aaa@DESKTOP-OSN5BJK:~/canmv_k230_gitee_micropython/src/canmv/resources/examples/16-AI-Cube$ ls
ClassificationApp.py DataCollectionCamera.py DetectionApp.py MultiLabelApp.py OCR_Det.py SegmentationApp.py SelfLearningApp.py
demo 运行起来如上图所示，中间会有一个红色提示框，我们在抓图的时候，让我们的目标物体尽量充满这个红色框框。
考虑到方便操作，推荐把这个脚本保存成 main.py 运行（具体操作办法见文末视频教程）。
3、数据标注
图像分类的数据集，只需要按照类别建立文件夹，把采集好的图片按照对应的目录放置进去就可以了。
tree -L 2 .
├── qual //量化数据集
├── test //测试数据集
│ ├── 0
│ ├── 1
│ ├── 2
│ └── 3
├── train //训练数据集
│ ├── 0
│ ├── 1
│ ├── 2
│ └── 3
└── val //验证数据集
├── 0
├── 1
├── 2
└── 3
二、 模型训练
一切准备就绪，让模型开始学习！在 yolov5 目录下执行如下指令：
python classify/train.py --model yolov5n-cls.pt --data datasets/num_cls --epochs 100 --batch-size 8 --imgsz 224 --device '0'
注意：这个命令的参数 --device '0' 是表示使用 GPU 资源训练，如果没有 GPU 的情况下，需要切换成 --device 'cpu'
训练完成显示如下（本地使用CPU电脑）：
Training complete (0.207 hours)
Results saved to runs/train-cls/exp3
Predict: python classify/predict.py --weights runs/train-cls/exp3/weights/best.pt --source im.jpg
Validate: python classify/val.py --weights runs/train-cls/exp3/weights/best.pt --data datasets/num_cls
Export: python export.py --weights runs/train-cls/exp3/weights/best.pt --include onnx
PyTorch Hub: model = torch.hub.load('ultralytics/yolov5', 'custom', 'runs/train-cls/exp3/weights/best.pt')
Visualize: https://netron.app
三、模型转换
1、安装 nncase 和 onnx 转换工具
# linux平台/Linux虚拟机：nncase和nncase-kpu可以在线安装，nncase-2.x 需要安装 dotnet-7
sudo apt-get install -y dotnet-sdk-7.0
pip install --upgrade pip
pip install nncase==2.9.0
pip install nncase-kpu==2.9.0
# windows平台：请自行安装dotnet-7并添加环境变量,支持使用pip在线安装nncase，但是nncase-kpu库需要离线安装，在https://github.com/kendryte/nncase/releases下载nncase_kpu-2.*-py2.py3-none-win_amd64.whl# 进入对应的python环境，在nncase_kpu-2.*-py2.py3-none-win_amd64.whl下载目录下使用pip安装
pip install nncase_kpu-2.*-py2.py3-none-win_amd64.whl
# 除nncase和nncase-kpu外，脚本还用到的其他库包括：
pip install onnx
pip install onnxruntime
pip install onnxsim
2、下载转换脚本工具 test_yolov5.zip 并解压到 yolov5 目录下
wget https://kendryte-download.canaan-creative.com/developer/k230/yolo_files/test_yolov5.zip
unzip test_yolov5.zip
3、模型转换
按照指令先将训练好的 pt 模型导出为 onnx 模型，再转换为 kmodel 模型
NNCASE 应用指南：https://www.kendryte.com/k230_rtos/zh/main/app_develop_guide/ai/nncase.html
量化数据集：我们在做模型转换的时候，需要使用训练数据的一个子集（推荐取训练数据的 20%）做量化，在模型转换的时候，需要指定到 --dataset 参数
# 导出onnx，pt模型路径请自行选择
python export.py --weight ./runs/train-cls/exp3/weights/best.pt --imgsz 224 --batch 1 --include onnx
cd test_yolov5/classify
# 转换kmodel,onnx模型路径请自行选择，生成的kmodel在onnx模型同级目录下
python to_kmodel.py --target k230 --model ../../runs/train-cls/exp3/weights/best.onnx --dataset ../../datasets/num_cls/qual --input_width 224 --input_height 224 --ptq_option 1
cd ../../
4、在文件下面，发现kmodel表示转换成功：
aaa@DESKTOP-OSN5BJK:~/yolov5/yolov5$ find . -name *.kmodel
./runs/train-cls/exp3/weights/best.kmodel
四、部署实战
1、准备开发环境
在 Canaan K230开发板上部署模型，首先要烧录镜像并安装 CanMV IDE。按照开发板的下载链接获取镜像并烧录，再安装好 CanMV IDE。
镜像文件下载地址：https://github.com/kendryte/canmv_k230/releases/tag/PreRelease
CanMV IDE 下载地址：https://www.kendryte.com/zh/resource/ide,k230
2、拷贝文件
连接 IDE ，把转换好的模型和测试图片拷贝到开发板目录下（路径可自定义，代码中记得同步修改）。
3、代码实现推理
图片推理：编写代码从本地读取图片，并进行格式转换。通过 YOLOv5 类初始化实例，设置好任务类型、推理模式、模型路径、标签列表等参数。例如：
from libs.YOLO import YOLOv5
import os,sys,gc
import ulab.numpy as np
import image
# 从本地读入图片，并实现HWC转CHW
def read_img(img_path):
img_data = image.Image(img_path)
img_data_rgb888=img_data.to_rgb888()
img_hwc=img_data_rgb888.to_numpy_ref()
shape=img_hwc.shape
img_tmp = img_hwc.reshape((shape[0] * shape[1], shape[2]))
img_tmp_trans = img_tmp.transpose()
img_res=img_tmp_trans.copy()
img_return=img_res.reshape((shape[2],shape[0],shape[1]))
return img_return,img_data_rgb888
if __name__=="__main__":
img_path="/data/test.jpg"
kmodel_path="/sdcard/examples/kmodel/best.kmodel"
labels = ["0","1","2"，"3"]
confidence_threshold = 0.5
model_input_size=[224,224]
img,img_ori=read_img(img_path)
rgb888p_size=[img.shape[2],img.shape[1]]
# 初始化YOLOv5实例
yolo=YOLOv5(task_type="classify",mode="image",kmodel_path=kmodel_path,labels=labels,rgb888p_size=rgb888p_size,model_input_size=model_input_size,conf_thresh=confidence_threshold,debug_mode=0)
yolo.config_preprocess()
try:
res=yolo.run(img)
yolo.draw_result(res,img_ori)
gc.collect()
except Exception as e:
sys.print_exception(e)
finally:
yolo.deinit()
运行代码，可以看到已经识别到这张图片是0的概率为0.83，并把结果绘制到左上角。
视频推理：从视频流中逐帧获取视频数据进行推理，实时展示识别结果。比如设置好显示模式、分辨率、模型路径等参数后，通过不断循环执行推理和显示操作：
from libs.PipeLine import PipeLine, ScopedTiming
from libs.YOLO import YOLOv5
import os,sys,gc
import ulab.numpy as np
import image
if __name__=="__main__":
# 显示模式，默认"hdmi",可以选择"hdmi"和"lcd"
display_mode="hdmi"
#需要根据输入图片的size修改rgb888p_size
rgb888p_size=[800，480]
if display_mode=="hdmi":
display_size=[1920,1080]
else:
display_size=[800,480]
# 模型路径
kmodel_path="/sdcard/examples/kmodel/yolov5n-cls.kmodel"
labels = ["0","1","2","3"]
confidence_threshold = 0.5
model_input_size=[224,224]
# 初始化PipeLine
pl=PipeLine(rgb888p_size=rgb888p_size,display_size=display_size,display_mode=display_mode)
pl.create()
# 初始化YOLOv5实例
yolo=YOLOv5(task_type="classify",mode="video",kmodel_path=kmodel_path,labels=labels,rgb888p_size=rgb888p_size,model_input_size=model_input_size,display_size=display_size,conf_thresh=confidence_threshold,debug_mode=0)
yolo.config_preprocess()
try:
while True:
os.exitpoint()
with ScopedTiming("total",1):
# 逐帧推理
img=pl.get_frame()
res=yolo.run(img)
yolo.draw_result(res,pl.osd_img)
pl.show_image()
gc.collect()
except Exception as e:
sys.print_exception(e)
finally:
yolo.deinit()
pl.destroy()
这样，无论是静态图片还是动态视频中的数字，都能被精准识别。
五、常见问题
1、在运行 python classify/train.py --model yolov5n-cls.pt --data datasets/num_cls --epochs 100 --batch-size 8 --imgsz 224 --device '0' 时报如下错误：
AttributeError: module 'importlib_metadata' has no attribute 'EntryPoints'
解决办法：
是因为 importlib_metadata 版本不兼容所导致的。 setuptools 依赖于 importlib_metadata，并且不同版本的 setuptools 对 importlib_metadata 版本的要求也不同。可以运行如下命令解决：
pip install --upgrade importlib_metadata
2、在安装模型转换工具运行 sudo apt-get install -y dotnet-sdk-7.0
报错：
E: Unable to locate package dotnet-sdk-7.0
E: Couldn't find any package by glob 'dotnet-sdk-7.0'
E: Couldn't find any package by regex 'dotnet-sdk-7.0'
解决办法：
请参考微软官方文档 （https://learn.microsoft.com/en-us/dotnet/core/install/linux），根据Linux操作系统的版本，选择相应的安装命令。
通过 YOLOv5 和 Canaan K230 开发板实现数字分类，不仅是技术的有趣应用，也展示了人工智能在日常生活中的巨大潜力。感兴趣的小伙伴不妨动手试试，探索更多人工智能与实际应用结合的可能！
立即解锁 K230 图像分类
官方资料:
https://www.kendryte.com/k230_canmv/zh/main/zh/example/ai/
Micropython AI 手册：
https://www.kendryte.com/k230_canmv/zh/main/zh/api/aidemo_libs.html
问答社区（请增加标签 [YOLO] ）：
https://www.kendryte.com/answer/