从零训练猫脸识别模型并部署到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 版本，适合 K230
imgsz=320：K230 上速度和精度的折中
batch=16：显存不足时改成 8、4 或 2
device=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.25

6.示例代码:

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()