3、数据集准备与预处理：从原始图片到模型养料

好，咱们接着往下走。上一章聊了模型选型，这一章要干点「脏活累活」——数据集准备。说实话，很多做嵌入式AI的朋友，模型调了半天效果不好，最后发现是数据出了问题。我刚开始做STM32图像识别时也踩过这个坑，辛辛苦苦采集了1000张图片，训练出来准确率不到60%，后来才发现标注错了将近一半。

数据集准备，说白了就是给模型准备「教材」。教材质量不行，学生再聪明也学不好。咱们一步步来。

3.1 图像数据集采集：从哪里搞图片？

采集数据，我建议优先考虑这三种途径：

  • 自己拍：最靠谱，但最累。用手机或USB摄像头，在不同光照、角度下拍摄。我做过一个口罩识别项目，光采集就花了两天，拍了3000多张。
  • 公开数据集：比如ImageNet子集、COCO、Open Images。注意看许可证，商用项目要小心。
  • 网络爬虫：用Python脚本从搜索引擎爬。但记得人工筛选，网上图片质量参差不齐。
我的小建议：嵌入式项目，每类至少500张图片。太少的话，模型容易过拟合。我见过有人每类只拍30张，结果模型在STM32上跑起来，连自家猫都认不出来。

3.2 数据清洗：剔除「垃圾」图片

采集完图片，第一件事不是标注，而是清洗。你想想看，如果混进去几张模糊的、重复的、甚至完全不相关的图片，模型学到的特征就歪了。

清洗我一般做这几步：

  1. 去模糊：用OpenCV的拉普拉斯算子检测，方差低于某个阈值就删掉。
  2. 去重复：计算图片哈希值（pHash），相似度高于95%的只留一张。
  3. 去无关：人工过一遍，把拍错的、背景太乱的剔除。
# 用OpenCV检测模糊图片的示例
import cv2

def is_blurry(image_path, threshold=100):
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    laplacian_var = cv2.Laplacian(gray, cv2.CV_64F).var()
    return laplacian_var < threshold

# 调用
if is_blurry("cat_001.jpg"):
    print("这张太糊了，删掉！")
注意：阈值要根据你的场景调整。我做过一个工业质检项目，因为光照太暗，所有图片的拉普拉斯方差都偏低，差点把好图全删了。后来把阈值从100降到50才正常。

3.3 数据标注：给图片打标签

标注是体力活，但也是决定模型上限的关键。分类任务简单，每张图片对应一个类别标签就行。检测任务就麻烦些，要画边界框。

常用的标注工具有：

  • LabelImg：轻量级，适合小项目。我早期一直用这个。
  • Label Studio：功能强，支持分类、检测、分割多种任务。
  • CVAT：在线工具，适合团队协作。

标注格式，我建议统一用COCO格式或YOLO格式。STM32CubeAI对COCO格式支持得比较好。

避坑指南：我曾经在一个项目中，标注员把「猫」和「狗」的标签搞反了200多张。训练出来的模型，看到猫就识别成狗，看到狗就识别成猫。后来我加了一道交叉验证的流程——随机抽10%的标注结果，自己再过一遍。

3.4 数据增强：让模型见多识广

嵌入式设备上，摄像头角度、光照变化很大。如果训练数据太「干净」，模型到了真实场景就抓瞎。数据增强就是人为制造变化，让模型学会「不管怎么变，我都能认出来」。

常用的增强操作：

增强方式 参数范围 适用场景
随机旋转 -30° ~ +30° 摄像头安装角度不固定
随机缩放 0.8 ~ 1.2倍 物体距离摄像头远近不同
水平翻转 50%概率 左右对称的物体（如人脸）
亮度调整 0.7 ~ 1.3 光照变化大的场景
高斯噪声 方差0.01 传感器噪声较大的低端摄像头

用Python的imgaug库或albumentations库实现很方便：

import imgaug.augmenters as iaa

seq = iaa.Sequential([
    iaa.Affine(rotate=(-30, 30)),        # 旋转
    iaa.Affine(scale=(0.8, 1.2)),        # 缩放
    iaa.Fliplr(0.5),                     # 水平翻转
    iaa.Multiply((0.7, 1.3)),            # 亮度调整
    iaa.AdditiveGaussianNoise(scale=0.01) # 高斯噪声
])

# 对一张图片做增强
augmented_image = seq(image=original_image)
我的经验：增强不是越多越好。我见过有人把旋转角度设到±90°，结果模型把倒着的杯子也当成正着的，反而学歪了。一般来说，旋转不超过±30°，缩放不超过20%，就够用了。

3.5 数据集划分：训练集/验证集/测试集

数据准备好了，接下来要分成三份。为什么不是两份？你想想看，如果只用训练集和验证集，你调参时模型会「记住」验证集的特征，最后测试结果虚高。所以必须留一份「没见过」的测试集。

划分比例，我一般这样：

  • 训练集：70% —— 模型学习用
  • 验证集：15% —— 调参、选模型用
  • 测试集：15% —— 最终评估用

如果数据量少（比如每类只有500张），可以调整为80%、10%、10%。

划分时要注意：

  1. 分层采样：保证每类图片在三个集合中的比例一致。别让训练集全是猫，测试集全是狗。
  2. 随机打乱：避免连续拍摄的图片集中在同一个集合里。
  3. 固定种子：用random.seed(42)固定随机数，保证每次划分结果一致，方便复现。
from sklearn.model_selection import train_test_split
import os, random

# 假设image_paths是所有图片路径列表，labels是对应标签
random.seed(42)

# 先分出测试集
train_val_paths, test_paths, train_val_labels, test_labels = train_test_split(
    image_paths, labels, test_size=0.15, stratify=labels, random_state=42
)

# 再从剩余数据中分出验证集
train_paths, val_paths, train_labels, val_labels = train_test_split(
    train_val_paths, train_val_labels, test_size=0.176, stratify=train_val_labels, random_state=42
)
# 0.176 = 15% / 85%，这样验证集占总体的15%

关键提醒：测试集一旦划分好，就不要再碰了。我见过有人调参时偷偷看测试集的结果，最后模型在真实场景中表现很差。测试集是「期末考试」，不能提前泄露答案。

3.6 数据格式转换：喂给STM32CubeAI

STM32CubeAI支持的输入格式一般是：

  • 图片尺寸：统一缩放到模型输入大小（比如96x96、128x128）
  • 颜色通道：RGB三通道，或灰度单通道
  • 数据格式：.jpg或.png都可以，但建议用.jpg压缩存储，节省空间
  • 标签文件：一个.txt或.csv文件，每行是「图片路径,类别ID」

我习惯把所有图片统一放到一个文件夹，然后生成一个CSV清单：

import pandas as pd
import os

data = []
for img_name in os.listdir("dataset/train"):
    label = 0 if "cat" in img_name else 1  # 假设文件名包含类别
    data.append([f"train/{img_name}", label])

df = pd.DataFrame(data, columns=["image_path", "label"])
df.to_csv("train.csv", index=False)

嗯，到这里，数据集就准备好了。下一章咱们要把它导入STM32CubeAI，开始训练模型。到时候你会发现，前面这些准备工作做得越扎实，后面就越省心。

← 上一章 📖 返回目录 下一章 →