从零训练猫脸识别模型并部署到K230开发板
前言#
众所不周知,我最后以635分进入了河海大学的自动化专业,非常遗憾没能水进西电的管理学当跳板
然而事已至此,不如继续玩点新玩具吧(
正好搞到了立创的K230开发板,试验了里面的几个demo以后,想试试自己训练一个模型出来
嗯,那就训练一个猫脸识别模型吧
获取猫脸数据集#
此次训练采用了https://www.kaggle.com/datasets/georgemartvel/catflw的数据集(CC BY-NC 4.0,不能直接用于商业项目)
https://github.com/martvelge/CatFLW
对数据集进行处理#
由于该数据集有约 2,000 张猫脸图像,每张含一个猫脸框和 48 个关键点,因此已经省去了我们前期进行数据标注的操作,仅需要将 bounding_boxes 转成YOLO标注后再进行训练即可
使用如下代码进行转换
"""Convert CatFLW JSON labels to a one-class YOLO detection dataset."""
import argparse
import json
import random
import shutil
from pathlib import Path
from PIL import Image
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--source", type=Path, default=Path("CatFLW dataset"))
parser.add_argument("--output", type=Path, default=Path("catflw_yolo"))
args = parser.parse_args()
images_dir = args.source / "images"
labels_dir = args.source / "labels"
if not images_dir.is_dir() or not labels_dir.is_dir():
raise SystemExit(f"Expected images/ and labels/ under {args.source}")
if args.output.exists():
shutil.rmtree(args.output)
for split in ("train", "val", "test"):
(args.output / "images" / split).mkdir(parents=True, exist_ok=True)
(args.output / "labels" / split).mkdir(parents=True, exist_ok=True)
image_paths = sorted(images_dir.glob("*.png"))
random.Random(42).shuffle(image_paths)
train_end = int(len(image_paths) * 0.8)
val_end = int(len(image_paths) * 0.9)
counts = {"train": 0, "val": 0, "test": 0}
for index, image_path in enumerate(image_paths):
json_path = labels_dir / f"{image_path.stem}.json"
if not json_path.exists():
raise SystemExit(f"Missing annotation: {json_path}")
split = "train" if index < train_end else "val" if index < val_end else "test"
with Image.open(image_path) as image:
width, height = image.size
annotation = json.loads(json_path.read_text(encoding="utf-8"))
x1, y1, x2, y2 = annotation["bounding_boxes"]
x1 = max(0.0, min(float(x1), width))
y1 = max(0.0, min(float(y1), height))
x2 = max(x1, min(float(x2), width))
y2 = max(y1, min(float(y2), height))
if x2 <= x1 or y2 <= y1:
raise SystemExit(f"Invalid bounding box: {json_path}")
x_center = ((x1 + x2) / 2) / width
y_center = ((y1 + y2) / 2) / height
box_width = (x2 - x1) / width
box_height = (y2 - y1) / height
label = f"0 {x_center:.6f} {y_center:.6f} {box_width:.6f} {box_height:.6f}\n"
shutil.copy2(image_path, args.output / "images" / split / image_path.name)
(args.output / "labels" / split / f"{image_path.stem}.txt").write_text(
label, encoding="utf-8"
)
counts[split] += 1
yaml = """path: catflw_yolo\ntrain: images/train\nval: images/val\ntest: images/test\nnc: 1\nnames: [cat_face]\n"""
(args.output / "data.yaml").write_text(yaml, encoding="utf-8")
print(f"Created {args.output}: {counts}")
if __name__ == "__main__":
main()
最终会生成 YOLO 数据集:catflw_yolo
- 训练集:1663 张
- 验证集:208 张
- 测试集:208 张
- 类别:
cat_face
正式训练#
安装训练环境#
建议使用 Conda:
conda create -n catface python=3.10 -y
conda activate catface
pip install ultralytics pillow检查安装:
yolo checks如果有 NVIDIA GPU,需要根据 CUDA 版本安装对应的 PyTorch;没有 GPU 也可以训练,但会比较慢。4
开始训练 YOLOv8n#
在当前目录执行:
yolo detect train `
data=catflw_yolo\data.yaml `
model=yolov8n.pt `
epochs=100 `
imgsz=320 `
batch=16 `
device=0 `
project=runs `
name=catface_yolov8n| 参数 | 含义 |
|---|---|
| yolo | 调用 Ultralytics YOLO 命令行工具 |
| detect | 使用目标检测任务 |
| train | 执行训练 |
| data=catflw_yolo\data.yaml | 数据集配置文件路径 |
| model=yolov8n.pt | 使用 YOLOv8 Nano 预训练模型,体积小、速度快,适合 K230 |
| epochs=100 | 完整遍历训练集 100 次 |
| imgsz=320 | 将图片缩放到约 320×320,影响速度和精度 |
| batch=8 | 每次同时处理 8 张图片,显存不足可改成 4 或 2 |
| device=0 | 使用第 0 号 GPU,也就是你的 RTX 5070 Laptop GPU |
| project=runs | 训练结果保存到 runs 目录 |
| name=catface_yolov8n | 本次训练实验名称 |
| workers=0 | 使用 0 个数据加载进程,Windows 下更稳定 |
参数说明:
yolov8n.pt:Nano 版本,适合 K230imgsz=320:K230 上速度和精度的折中batch=16:显存不足时改成8、4或2device=0:使用第一张 NVIDIA GPU- 没有 GPU 时改为
device=cpu
训练完成后,模型位于:
runs\catface_yolov8n\weights\best.pt用yolo验证模型#
yolo detect val `
model=runs\catface_yolov8n\weights\best.pt `
data=catflw_yolo\data.yaml `
imgsz=320用测试图片查看效果:
yolo detect predict `
model=runs\catface_yolov8n\weights\best.pt `
source=catflw_yolo\images\test `
imgsz=320 `
conf=0.25 `
save=True结果会保存到类似:
runs\detect\predict用图形化界面和电脑摄像头进行验证#
"""CanMV K230 real-time cat-face detection demo."""
from libs.PipeLine import PipeLine, ScopedTiming
from libs.YOLO import YOLOv8
import gc
import os
if __name__ == "__main__":
kmodel_path = "/sdcard/catface_320.kmodel"
labels = ["cat_face"]
model_input_size = [320, 320]
confidence_threshold = 0.35
nms_threshold = 0.45
# Use "lcd" for the common 800x480 display. Change to "hdmi" if needed.
display_mode = "lcd"
rgb888p_size = [640, 360]
pipeline = PipeLine(
rgb888p_size=rgb888p_size,
display_mode=display_mode,
)
pipeline.create()
display_size = pipeline.get_display_size()
yolo = YOLOv8(
task_type="detect",
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,
nms_thresh=nms_threshold,
max_boxes_num=10,
debug_mode=1,
)
yolo.config_preprocess()
try:
print("Cat-face model loaded. Starting camera detection...")
while True:
os.exitpoint()
with ScopedTiming("total", True):
frame = pipeline.get_frame()
result = yolo.run(frame)
yolo.draw_result(result, pipeline.osd_img)
pipeline.show_image()
gc.collect()
finally:
yolo.deinit()
pipeline.destroy()
导出 ONNX#
K230 通常通过 ONNX -> kmodel 转换:
yolo export `
model=runs\catface_yolov8n\weights\best.pt `
format=onnx `
imgsz=320 `
opset=11 `
simplify=True生成:
runs\catface_yolov8n\weights\best.onnx建议使用 Netron 检查 ONNX 输入,一般应为:
[1, 3, 320, 320]安装 nncase#
K230 的 nncase 版本必须和开发板固件匹配。官方文档说明 Windows 下需要:
pip install nncase然后从 nncase Releases 下载对应版本的:
nncase_kpu-版本号-py2.py3-none-win_amd64.whl安装:
pip install nncase_kpu-版本号-py2.py3-none-win_amd64.whl官方 K230 文档:
重点注意:不要随便使用最新 nncase,先确认你的 K230 是 CanMV还是其他版本
ONNX 转 KModel#
执行:
Invoke-WebRequest `
-Uri "https://kendryte-download.canaan-creative.com/developer/k230/yolo_files/test_yolov8.zip" `
-OutFile ".\test_yolov8.zip"解压:
Expand-Archive `
-Path ".\test_yolov8.zip" `
-DestinationPath "." `
-Force检查检测转换脚本是否存在:
Test-Path ".\test_yolov8\detect\to_kmodel.py"如果返回:
True说明脚本准备成功。
再安装转换脚本依赖:
conda run -n k230 python -m pip install onnx onnxruntime onnxsim pillow numpy确认当前模型和校准集存在:
Test-Path ".\runs\detect\runs\catface_yolov8n\weights\best.onnx"
Test-Path ".\calibration"两个都应返回:
True注意,官方示例中的 test_yolov8\detect\to_kmodel.py 是目标检测转换脚本,不要使用 classify 目录中的脚本。
完成后执行:
Test-Path ".\test_yolov8\detect\to_kmodel.py"在项目根目录执行:
conda run -n k230 python ".\test_yolov8\detect\to_kmodel.py" `
--target k230 `
--model ".\runs\detect\runs\catface_yolov8n\weights\best.onnx" `
--dataset ".\calibration" `
--input_width 320 `
--input_height 320 `
--ptq_option 0参数含义:
--target k230 目标芯片为 K230
--model 输入 ONNX 模型
--dataset INT8 量化校准图片目录
--input_width 320 模型输入宽度
--input_height 320 模型输入高度
--ptq_option 0 使用 uint8 权重和 uint8 激活量化转换过程可能需要几分钟。成功后,通常会在 ONNX 所在目录生成:
runs\detect\runs\catface_yolov8n\weights\best.kmodel部署到 K230#
如果使用 CanMV:
1.将 catface_yolov8n.kmodel 复制到开发板 SD 卡。
2.参考官方 YOLOv8 K230 示例。
3.修改示例代码中的模型路径。
4.将类别名称改为:
labels = ["cat_face"]5.调整置信度阈值:
confidence_threshold = 0.256.示例代码:
from libs.PipeLine import PipeLine, ScopedTiming
from libs.YOLO import YOLOv8
import gc
import os
import sys
if __name__ == "__main__":
# 确认模型实际上传位置
kmodel_path = "/data/models/catface/best.kmodel"
labels = ["cat_face"]
model_input_size = [320, 320]
confidence_threshold = 0.35
nms_threshold = 0.45
# LCD 使用 lcd,HDMI 使用 hdmi
display_mode = "lcd"
# LCKFB-K230 通常可以使用 1920x1080 摄像头输入
rgb888p_size = [1920, 1080]
if display_mode == "lcd":
display_size = [800, 480]
else:
display_size = [1920, 1080]
pipeline = PipeLine(
rgb888p_size=rgb888p_size,
display_size=display_size,
display_mode=display_mode,
)
pipeline.create()
yolo = YOLOv8(
task_type="detect",
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,
nms_thresh=nms_threshold,
max_boxes_num=10,
debug_mode=0,
)
yolo.config_preprocess()
try:
print("Cat-face model loaded. Starting detection...")
while True:
os.exitpoint()
with ScopedTiming("total", 1):
frame = pipeline.get_frame()
result = yolo.run(frame)
yolo.draw_result(result, pipeline.osd_img)
pipeline.show_image()
gc.collect()
except Exception as e:
sys.print_exception(e)
finally:
yolo.deinit()
pipeline.destroy()