零基础实战：手把手教你用YOLOv5-Lite在树莓派上搭建数字识别系统（附完整代码与避坑指南）

1. 环境准备与硬件选择

在开始搭建数字识别系统之前，我们需要准备好合适的硬件和软件环境。树莓派作为一款小巧而强大的单板计算机，非常适合运行轻量化的深度学习模型。我推荐使用树莓派4B 4GB版本，它的处理能力足够运行YOLOv5-Lite这样的轻量化模型。

硬件方面，除了树莓派主板，你还需要一个稳定的电源适配器（至少5V/3A），一个高质量的MicroSD卡（建议32GB以上），以及一个散热外壳。树莓派在运行时会产生一定的热量，良好的散热能保证系统稳定运行。摄像头方面，我强烈建议使用免驱USB摄像头而不是树莓派专用的CSI摄像头。USB摄像头不仅线材更长、更灵活，而且方便在电脑上进行测试，排除了硬件兼容性问题。

软件环境准备分为电脑端和树莓派端。电脑端我们主要用来训练模型，树莓派端则用来部署和运行训练好的模型。在电脑上，我们需要安装Python 3.8-3.12之间的版本，推荐使用3.10版本，这个版本在兼容性和稳定性方面表现都很不错。开发工具推荐使用PyCharm社区版，它是完全免费的，而且功能足够我们使用。

# 树莓派系统初始化
sudo apt update
sudo apt upgrade -y
sudo apt install python3-pip python3-venv

2. 数据集制作与增强

数据集的质量直接决定了最终模型的性能。对于数字识别这种相对简单的任务，我们可以通过数据增强技术来大幅扩充数据集规模。我在这里分享一个实用的数据增强方法，可以将单张图片生成256张不同的变体。

首先，我们需要收集基础数字图片。你可以在白纸上用黑色马克笔写下数字1-8，然后拍照保存。每张图片最好背景简洁，数字清晰。有了基础图片后，我们就可以使用下面的代码进行数据增强了：

import cv2
import numpy as np
import os

def augment_digits(image_path, output_dir, num_augment=256):
    """
    数字图像数据增强函数
    image_path: 原始图像路径
    output_dir: 输出目录
    num_augment: 每张图像增强数量
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    image = cv2.imread(image_path)
    if image is None:
        print(f"无法读取图像: {image_path}")
        return
    
    h, w = image.shape[:2]
    center = (w // 2, h // 2)
    
    for i in range(num_augment):
        # 随机旋转
        angle = np.random.uniform(-30, 30)
        M_rotate = cv2.getRotationMatrix2D(center, angle, 1.0)
        rotated = cv2.warpAffine(image, M_rotate, (w, h), 
                                borderValue=(255, 255, 255))
        
        # 随机仿射变换
        pts1 = np.float32([[w//4, h//4], [3*w//4, h//4], [w//4, 3*h//4]])
        pts2 = np.float32([
            [w//4 + np.random.uniform(-20, 20), h//4 + np.random.uniform(-20, 20)],
            [3*w//4 + np.random.uniform(-20, 20), h//4 + np.random.uniform(-20, 20)],
            [w//4 + np.random.uniform(-20, 20), 3*h//4 + np.random.uniform(-20, 20)]
        ])
        M_affine = cv2.getAffineTransform(pts1, pts2)
        transformed = cv2.warpAffine(rotated, M_affine, (w, h),
                                    borderValue=(255, 255, 255))
        
        # 随机亮度调整
        brightness = np.random.uniform(0.8, 1.2)
        transformed = np.clip(transformed * brightness, 0, 255).astype(np.uint8)
        
        # 保存增强后的图像
        output_path = os.path.join(output_dir, f"aug_{i:03d}.jpg")
        cv2.imwrite(output_path, transformed)
    
    print(f"已完成 {num_augment} 张图像增强")

# 使用示例
augment_digits("digit_3.jpg", "augmented_digits/3", 256)

这个增强方法包含了旋转、仿射变换和亮度调整，能够很好地模拟实际场景中数字可能出现的变化。记得为每个数字创建单独的文件夹存放增强后的图片，这样便于后续的标注工作。

3. 自动化标注技巧

手动标注图像既耗时又容易出错，特别是当我们有大量数据时。我开发了一个自动化标注方案，可以大大减少标注工作量。这个方案基于一个简单的观察：在数字识别任务中，数字通常是图像中最暗的区域。

import cv2
import numpy as np
import os
import json

def auto_annotate(image_dir, output_dir, label):
    """
    自动标注数字图像
    image_dir: 图像目录
    output_dir: 输出目录
    label: 数字标签
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    image_files = [f for f in os.listdir(image_dir) if f.endswith('.jpg')]
    
    for image_file in image_files:
        image_path = os.path.join(image_dir, image_file)
        img = cv2.imread(image_path)
        
        if img is None:
            continue
        
        # 转换为灰度图并二值化
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        _, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY_INV)
        
        # 寻找轮廓
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, 
                                     cv2.CHAIN_APPROX_SIMPLE)
        
        if contours:
            # 合并所有轮廓并计算边界框
            all_contours = np.vstack(contours)
            x, y, w, h = cv2.boundingRect(all_contours)
            
            # 创建LabelMe格式的标注
            annotation = {
                "version": "5.4.1",
                "flags": {},
                "shapes": [{
                    "label": str(label),
                    "points": [
                        [x, y],
                        [x + w, y],
                        [x + w, y + h],
                        [x, y + h]
                    ],
                    "group_id": None,
                    "shape_type": "polygon",
                    "flags": {}
                }],
                "imagePath": image_file,
                "imageData": None,
                "imageHeight": img.shape[0],
                "imageWidth": img.shape[1]
            }
            
            # 保存标注文件
            json_path = os.path.join(output_dir, 
                                   os.path.splitext(image_file)[0] + '.json')
            with open(json_path, 'w') as f:
                json.dump(annotation, f, indent=2)

# 为每个数字文件夹执行自动标注
digits = ['1', '2', '3', '4', '5', '6', '7', '8']
for digit in digits:
    image_dir = f"augmented_digits/{digit}"
    output_dir = f"annotations/{digit}"
    auto_annotate(image_dir, output_dir, digit)

这个自动标注方法在处理简单背景下的数字图像时效果很好。如果图像背景复杂，可能还需要手动调整一些标注。标注完成后，我们需要将LabelMe的JSON格式转换为YOLO所需的TXT格式。

4. 模型训练与优化

现在来到最重要的环节——模型训练。我们使用YOLOv5-Lite v1.4版本，这个版本在树莓派上的表现更加稳定。首先需要从GitHub下载源码：

git clone https://github.com/ppogg/YOLOv5-Lite
cd YOLOv5-Lite
git checkout v1.4

安装所需的依赖包：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

接下来配置数据集文件。在data目录下创建mydata.yaml文件：

# 数据集配置文件
path: ../datasets
train: images/train
val: images/val

# 类别数量
nc: 8

# 类别名称
names: ['1', '2', '3', '4', '5', '6', '7', '8']

修改train.py中的训练参数：

# 在parse_opt函数中修改以下参数
parser.add_argument('--weights', type=str, default='v5lite-s.pt', help='初始权重路径')
parser.add_argument('--cfg', type=str, default='models/v5lite-s.yaml', help='模型配置文件路径')
parser.add_argument('--data', type=str, default='data/mydata.yaml', help='数据集配置文件路径')
parser.add_argument('--epochs', type=int, default=300, help='训练轮数')
parser.add_argument('--batch-size', type=int, default=8, help='批次大小')
parser.add_argument('--imgsz', '--img-size', type=int, default=320, help='输入图像大小')
parser.add_argument('--device', default='cpu', help='使用设备')

开始训练：

python train.py

训练过程中要密切关注损失函数的变化和验证集上的表现。如果发现过拟合（训练损失持续下降但验证损失开始上升），可以提前停止训练或者增加数据增强的强度。

5. 模型转换与部署

训练完成后，我们需要将PyTorch模型转换为ONNX格式，以便在树莓派上高效运行。使用以下代码进行转换：

import torch
from models.experimental import attempt_load

def convert_to_onnx(weights_path, img_size=320):
    """
    将PyTorch模型转换为ONNX格式
    weights_path: 权重文件路径
    img_size: 输入图像大小
    """
    device = torch.device('cpu')
    model = attempt_load(weights_path, map_location=device)
    model.eval()
    
    # 创建示例输入
    dummy_input = torch.randn(1, 3, img_size, img_size, device=device)
    
    # 导出ONNX模型
    onnx_path = weights_path.replace('.pt', '.onnx')
    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        opset_version=12,
        input_names=['images'],
        output_names=['output'],
        dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}}
    )
    
    print(f"模型已转换为: {onnx_path}")

# 转换最佳权重
convert_to_onnx('runs/train/exp/weights/best.pt')

现在将转换好的ONNX模型和推理代码部署到树莓派上。在树莓派上需要安装以下依赖：

# 安装必要的依赖
sudo apt update
sudo apt install libatlas-base-dev libopenblas-dev libblas-dev
pip install opencv-python-headless onnxruntime numpy

创建树莓派推理脚本：

import cv2
import numpy as np
import onnxruntime as ort
import time

class DigitDetector:
    def __init__(self, model_path, conf_threshold=0.5):
        self.conf_threshold = conf_threshold
        self.session = ort.InferenceSession(model_path)
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
        
        # 类别映射
        self.class_names = ['1', '2', '3', '4', '5', '6', '7', '8']
    
    def preprocess(self, image):
        """预处理输入图像"""
        img = cv2.resize(image, (320, 320))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = img.astype(np.float32) / 255.0
        img = np.transpose(img, (2, 0, 1))
        img = np.expand_dims(img, axis=0)
        return img
    
    def detect(self, image):
        """执行数字检测"""
        input_img = self.preprocess(image)
        outputs = self.session.run([self.output_name], 
                                 {self.input_name: input_img})[0]
        
        detections = self.postprocess(outputs, image.shape)
        return detections
    
    def postprocess(self, outputs, orig_shape):
        """后处理检测结果"""
        # 实现后处理逻辑，包括NMS和置信度过滤
        # 这里简化处理，实际需要根据模型输出格式调整
        detections = []
        # 具体的后处理实现...
        return detections

# 使用示例
detector = DigitDetector('best.onnx')
cap = cv2.VideoCapture(0)

while True:
    ret, frame = cap.read()
    if not ret:
        break
    
    detections = detector.detect(frame)
    
    # 在图像上绘制检测结果
    for det in detections:
        x1, y1, x2, y2, conf, cls_id = det
        label = f'{detector.class_names[cls_id]}: {conf:.2f}'
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(frame, label, (x1, y1-10),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    
    cv2.imshow('Digit Detection', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

6. 性能优化技巧

在树莓派上运行深度学习模型时，性能优化至关重要。以下是我总结的几个实用优化技巧：

模型量化：将FP32模型转换为INT8格式，可以显著减少模型大小和推理时间：

def quantize_model(onnx_path):
    """量化ONNX模型"""
    from onnxruntime.quantization import quantize_dynamic, QuantType
    
    quantized_path = onnx_path.replace('.onnx', '_quantized.onnx')
    quantize_dynamic(onnx_path, quantized_path, weight_type=QuantType.QUInt8)
    return quantized_path

多线程处理：使用多线程来并行处理图像采集和推理：

import threading
from queue import Queue

class ProcessingPipeline:
    def __init__(self, model_path):
        self.detector = DigitDetector(model_path)
        self.frame_queue = Queue(maxsize=2)
        self.result_queue = Queue(maxsize=2)
        
    def capture_thread(self):
        """图像采集线程"""
        cap = cv2.VideoCapture(0)
        while True:
            ret, frame = cap.read()
            if ret:
                if self.frame_queue.full():
                    self.frame_queue.get()
                self.frame_queue.put(frame)
    
    def process_thread(self):
        """处理线程"""
        while True:
            if not self.frame_queue.empty():
                frame = self.frame_queue.get()
                results = self.detector.detect(frame)
                self.result_queue.put((frame, results))
    
    def display_thread(self):
        """显示线程"""
        while True:
            if not self.result_queue.empty():
                frame, results = self.result_queue.get()
                # 显示处理结果...

内存优化：合理管理内存使用，避免不必要的内存分配：

def optimize_memory_usage():
    """优化内存使用"""
    import gc
    import onnxruntime as ort
    
    # 配置ONNX Runtime选项
    options = ort.SessionOptions()
    options.enable_mem_pattern = False
    options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
    
    # 定期垃圾回收
    gc.collect()

7. 常见问题解决

在实际部署过程中，你可能会遇到各种问题。这里我总结了一些常见问题及其解决方案：

问题1：模型推理速度慢 解决方案：使用模型量化、减少输入图像尺寸、启用ONNX Runtime优化：

# 启用ONNX Runtime优化
options = ort.SessionOptions()
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
session = ort.InferenceSession('model.onnx', sess_options=options)

问题2：内存不足 解决方案：使用内存映射方式加载模型，减少批量大小：

# 使用内存映射加载大模型
with open('model.onnx', 'rb') as f:
    model_data = f.read()
session = ort.InferenceSession(model_data)

问题3：摄像头帧率低 解决方案：降低分辨率，使用硬件加速：

# 设置摄像头参数
cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
cap.set(cv2.CAP_PROP_FPS, 15)

问题4：误检和漏检 解决方案：调整置信度阈值，增加后处理过滤：

def filter_detections(detections, conf_threshold=0.5, iou_threshold=0.4):
    """过滤检测结果"""
    # 按置信度过滤
    detections = [d for d in detections if d[4] > conf_threshold]
    
    # NMS处理
    boxes = [d[:4] for d in detections]
    scores = [d[4] for d in detections]
    indices = cv2.dnn.NMSBoxes(boxes, scores, conf_threshold, iou_threshold)
    
    return [detections[i] for i in indices]

问题5：树莓派过热降频 解决方案：改善散热，添加散热片和风扇，监控温度：

# 监控树莓派温度
vcgencmd measure_temp
# 设置温度警告
echo "Temp: $(vcgencmd measure_temp)" >> /var/log/temp.log

通过以上优化和问题解决方案，你应该能够在树莓派上稳定运行数字识别系统。记得在实际部署前充分测试每个组件，确保系统的稳定性和可靠性。