《2025 DigiKey AI应用创意挑战赛》基于Yolo的PCB缺陷检测

eefocus_3956818

本帖最后由 eefocus_3956818 于 2026-1-25 02:38 编辑

项目名称：基于Yolo的PCB缺陷检测

项目概述

本项目为基于YOLO系列模型的缺陷检测系统，支持对PCB等工业产品的缺陷自动识别。系统包含数据集处理、模型训练、推理检测、结果可视化等完整流程，并集成了Web前端界面，方便用户上传图片并查看检测结果。项目适配树莓派5环境，便于在边缘设备上部署和运行。

主要功能模块：

数据集处理与转换（convert.ipynb, dataset.ipynb）
模型训练与测试（training.ipynb, testing.ipynb，支持yolov8n/8s/11n等权重）
缺陷检测API与Web服务（app.py, app_raw.py，Flask框架）
前端页面与交互（templates/index.html, static/app.js, static/style.css）
Docker容器化部署支持（Dockerfile, docker-build.sh, docker-run.sh）

部署环境要求（树莓派5）：

操作系统：Raspberry Pi OS 64位
Python 3.8及以上
依赖库：requirements.txt中已列出（如Flask、torch、ultralytics等）
Docker（可选，推荐用于快速部署）
推荐使用硬件加速（如树莓派5的NPU/GPU，当前未使用）以提升推理速度

作品实物图

本次的软件开发，AI的应用多集中于软件层面，对硬件并无实际的应用控制，仅仅是使用USB摄像头作为输入来到树莓派进行推理。因此故第一张图片展示树莓派和摄像头的外观，其他界面则展示对应的功能。

完整的系统正面所示

PCB缺陷检测效果1 ：文件上传推理

PCB缺陷检测效果1 ：摄像头数据流实时推理。

项目软件设计

数据集选择：项目的数据集采用的是飞浆 AI studio中北京大学开源的PCB缺陷数据集，其中可以检测：缺失孔，鼠标咬伤，开路，短路，杂散，伪铜

原生数据集采用的是JSON格式或者VOC格式，这里对VOC数据格式进行清洗和转换。

数据清洗：数据清洗的目的是主要将VOC的格式转换成Yolo支持的格式

import os
import xml.etree.ElementTree as ET
# VOC XML 文件夹
annotations_dir = "/Users/wangchong/Downloads/VOCdevkit/VOC2007/Annotations"
# 输出 YOLO 格式 TXT 文件夹
labels_dir = "/Users/wangchong/Downloads/VOCdevkit/VOC2007/labels"
# YOLO 类别列表（顺序决定 class_id）
classes = ["missing_hole", "mouse_bite", "short", "spur", "spurious_copper", "open_circuit"]
os.makedirs(labels_dir, exist_ok=True)
# 遍历所有 XML 文件
for xml_file in os.listdir(annotations_dir):
if not xml_file.endswith(".xml"):
continue
xml_path = os.path.join(annotations_dir, xml_file)
tree = ET.parse(xml_path)
root = tree.getroot()
# 图片大小
size = root.find("size")
img_width = int(size.find("width").text)
img_height = int(size.find("height").text)
yolo_lines = []
# 遍历所有对象
for obj in root.findall("object"):
cls_name = obj.find("name").text.strip() # 去掉空格
if cls_name not in classes:
print(f"未匹配类别: {cls_name}") # 调试用
continue
class_id = classes.index(cls_name)
bbox = obj.find("bndbox")
xmin = int(bbox.find("xmin").text)
ymin = int(bbox.find("ymin").text)
xmax = int(bbox.find("xmax").text)
ymax = int(bbox.find("ymax").text)
# 转换为 YOLO 格式（归一化）
x_center = ((xmin + xmax) / 2) / img_width
y_center = ((ymin + ymax) / 2) / img_height
width = (xmax - xmin) / img_width
height = (ymax - ymin) / img_height
yolo_lines.append(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
# 保存 TXT 文件
txt_filename = os.path.splitext(xml_file)[0] + ".txt"
txt_path = os.path.join(labels_dir, txt_filename)
with open(txt_path, "w") as f:
f.write("\n".join(yolo_lines))
print("转换完成！TXT 文件已生成到", labels_dir)

复制代码

划分训练集和数据集

import os
import shutil
from sklearn.model_selection import train_test_split
# 原始文件夹
images_dir = "/Users/wangchong/Downloads/VOCdevkit/VOC2007/JPEGImages"
labels_dir = "/Users/wangchong/Downloads/VOCdevkit/VOC2007/labels"
# 输出目录：dataset 根目录
dataset_dir = "/Users/wangchong/Downloads/VOCdevkit/VOC2007/dataset"
# 创建目录结构
train_img_dir = os.path.join(dataset_dir, "images/train")
val_img_dir = os.path.join(dataset_dir, "images/val")
train_label_dir = os.path.join(dataset_dir, "labels/train")
val_label_dir = os.path.join(dataset_dir, "labels/val")
for d in [train_img_dir, val_img_dir, train_label_dir, val_label_dir]:
os.makedirs(d, exist_ok=True)
# 获取所有图片文件
all_images = [f for f in os.listdir(images_dir) if f.endswith((".jpg", ".png"))]
# 划分训练集和验证集（80%训练, 20%验证）
train_imgs, val_imgs = train_test_split(all_images, test_size=0.2, random_state=42)
# 复制训练集图片和标签
for img_file in train_imgs:
# 图片
shutil.copy2(os.path.join(images_dir, img_file),
os.path.join(train_img_dir, img_file))
# 对应标签
label_file = os.path.splitext(img_file)[0] + ".txt"
src_label = os.path.join(labels_dir, label_file)
if os.path.exists(src_label):
shutil.copy2(src_label, os.path.join(train_label_dir, label_file))
# 复制验证集图片和标签
for img_file in val_imgs:
# 图片
shutil.copy2(os.path.join(images_dir, img_file),
os.path.join(val_img_dir, img_file))
# 对应标签
label_file = os.path.splitext(img_file)[0] + ".txt"
src_label = os.path.join(labels_dir, label_file)
if os.path.exists(src_label):
shutil.copy2(src_label, os.path.join(val_label_dir, label_file))
print(f"训练集: {len(train_imgs)} 张图片")
print(f"验证集: {len(val_imgs)} 张图片")
print("已完成目录整理为 YOLOv8 官方推荐结构：")
print(dataset_dir)

复制代码

模型训练

from ultralytics import YOLO
import os
# 类别列表
class_names = ["missing_hole", "mouse_bite", "short", "spur", "spurious_copper", "open_circuit"]
model = YOLO("yolov8s.pt")
model.train(
data="/Users/wangchong/Downloads/VOCdevkit/VOC2007/dataset/pcb.yaml",
epochs=100, # 训练轮数
imgsz=640, # 输入图片尺寸
batch=8, # 批量大小
device="mps", # Mac CPU 训练
name="pcb_train" # 保存训练结果的目录名 runs/detect/pcb_train
)

复制代码

一共在Mac上耗费了两个多小时，跑了100轮的训练。最终的精度到达了96% ，召回率在92%。模型表现很好。

100 epochs completed in 2.366 hours.
Optimizer stripped from /Users/wangchong/DataspellProjects/defect_detection/runs/detect/pcb_train4/weights/last.pt, 22.5MB
Optimizer stripped from /Users/wangchong/DataspellProjects/defect_detection/runs/detect/pcb_train4/weights/best.pt, 22.5MB
Validating /Users/wangchong/DataspellProjects/defect_detection/runs/detect/pcb_train4/weights/best.pt...
Ultralytics 8.3.244 🚀 Python-3.10.19 torch-2.5.1 MPS (Apple M4)
Model summary (fused): 72 layers, 11,127,906 parameters, 0 gradients, 28.4 GFLOPs
Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 9/9 5.1s/it 45.7s4.8ss
all 139 603 0.964 0.916 0.946 0.521
missing_hole 20 92 0.987 0.989 0.986 0.612
mouse_bite 31 136 0.961 0.919 0.942 0.513
short 26 113 0.966 0.956 0.97 0.547
spur 19 88 0.962 0.852 0.89 0.464
spurious_copper 27 111 0.953 0.908 0.948 0.521
open_circuit 16 63 0.955 0.873 0.942 0.469
Speed: 4.9ms preprocess, 154.0ms inference, 0.0ms loss, 9.3ms postprocess per image

复制代码

树莓派推理代码

import os
import cv2
import base64
import numpy as np
from flask import Flask, render_template, request, jsonify, Response
from ultralytics import YOLO
from datetime import datetime
from werkzeug.utils import secure_filename
app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'static/uploads'
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024 # 16MB max file size
app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg', 'bmp'}
# 确保上传目录存在
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
# 类别名称
class_names = [
"missing_hole",
"mouse_bite",
"short",
"spur",
"spurious_copper",
"open_circuit"
]
# 加载模型
model = YOLO("runs/detect/pcb_train4/weights/best.pt")
# 性能优化配置
PERFORMANCE_CONFIG = {
'frame_skip': 2, # 每N帧处理一次（树莓派推荐2-3）
'video_width': 640, # 视频分辨率宽度
'video_height': 480, # 视频分辨率高度
'inference_size': 416, # 推理图像尺寸（越小越快，但精度可能降低）
'jpeg_quality': 70, # JPEG压缩质量（1-100）
}
# 全局变量用于摄像头控制
camera = None
camera_active = False
camera_detections = []
frame_counter = 0
last_annotated_frame = None
def allowed_file(filename):
"""检查文件扩展名是否允许"""
return '.' in filename and filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']
def detect_defects(image, conf=0.3, iou=0.45, img_size=None):
"""
对图像进行缺陷检测
返回：标注后的图像和检测结果列表
"""
# 如果指定了推理尺寸，则调整图像大小
if img_size:
h, w = image.shape[:2]
scale = min(img_size / w, img_size / h)
if scale < 1:
new_w, new_h = int(w * scale), int(h * scale)
resized_image = cv2.resize(image, (new_w, new_h))
else:
resized_image = image
else:
resized_image = image
results = model(
resized_image,
conf=conf,
iou=iou,
max_det=50, # 减少最大检测数量
device="cpu",
verbose=False,
half=False
)
detections = []
annotated_image = image.copy()
# 如果图像被缩放，计算缩放比例用于坐标还原
if img_size and resized_image.shape != image.shape:
scale_x = image.shape[1] / resized_image.shape[1]
scale_y = image.shape[0] / resized_image.shape[0]
else:
scale_x = scale_y = 1.0
boxes = results[0].boxes
if boxes is not None and len(boxes) > 0:
for box in boxes:
cls_id = int(box.cls.item())
conf_score = float(box.conf.item())
x1, y1, x2, y2 = map(int, box.xyxy[0])
# 还原到原始图像坐标
x1, y1 = int(x1 * scale_x), int(y1 * scale_y)
x2, y2 = int(x2 * scale_x), int(y2 * scale_y)
# 记录检测结果
detections.append({
'class': class_names[cls_id],
'confidence': conf_score,
'bbox': [x1, y1, x2, y2]
})
# 绘制边界框
label = f"{class_names[cls_id]} {conf_score:.2f}"
cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
# 绘制标签背景
(label_width, label_height), _ = cv2.getTextSize(
label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2
)
cv2.rectangle(
annotated_image,
(x1, y1 - label_height - 10),
(x1 + label_width, y1),
(0, 255, 0),
-1
)
# 绘制标签文本
cv2.putText(
annotated_image,
label,
(x1, y1 - 6),
cv2.FONT_HERSHEY_SIMPLEX,
0.5,
(0, 0, 0),
2
)
return annotated_image, detections
@app.route('/')
def index():
"""主页"""
return render_template('index.html')
@app.route('/upload', methods=['POST'])
def upload_file():
"""处理图片上传和检测"""
if 'file' not in request.files:
return jsonify({'error': '没有文件上传'}), 400
file = request.files['file']
if file.filename == '':
return jsonify({'error': '没有选择文件'}), 400
if file and allowed_file(file.filename):
# 保存原始文件
filename = secure_filename(file.filename)
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
filename = f"{timestamp}_{filename}"
filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
file.save(filepath)
# 读取图像
image = cv2.imread(filepath)
if image is None:
return jsonify({'error': '无法读取图像'}), 400
# 获取置信度阈值
conf_threshold = float(request.form.get('confidence', 0.3))
# 进行检测
annotated_image, detections = detect_defects(image, conf=conf_threshold)
# 保存标注后的图像
annotated_filename = f"annotated_{filename}"
annotated_filepath = os.path.join(app.config['UPLOAD_FOLDER'], annotated_filename)
cv2.imwrite(annotated_filepath, annotated_image)
return jsonify({
'success': True,
'original_image': f'/static/uploads/{filename}',
'annotated_image': f'/static/uploads/{annotated_filename}',
'detections': detections,
'total_defects': len(detections)
})
return jsonify({'error': '不支持的文件格式'}), 400
def generate_frames():
"""生成摄像头帧流"""
global camera, camera_active, camera_detections, frame_counter, last_annotated_frame
while camera_active:
if camera is None:
break
success, frame = camera.read()
if not success:
break
frame_counter += 1
# 帧跳过优化：每N帧才进行一次检测
if frame_counter % PERFORMANCE_CONFIG['frame_skip'] == 0:
# 进行检测（使用较小的推理尺寸）
annotated_frame, detections = detect_defects(
frame,
img_size=PERFORMANCE_CONFIG['inference_size']
)
# 保存检测结果和标注帧
camera_detections = detections
last_annotated_frame = annotated_frame
else:
# 使用上一帧的标注结果或原始帧
annotated_frame = last_annotated_frame if last_annotated_frame is not None else frame
# 编码为JPEG（使用压缩质量设置）
encode_params = [cv2.IMWRITE_JPEG_QUALITY, PERFORMANCE_CONFIG['jpeg_quality']]
ret, buffer = cv2.imencode('.jpg', annotated_frame, encode_params)
frame_bytes = buffer.tobytes()
yield (b'--frame\r\n'
b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')
@app.route('/video_feed')
def video_feed():
"""视频流路由"""
return Response(
generate_frames(),
mimetype='multipart/x-mixed-replace; boundary=frame'
)
@app.route('/start_camera', methods=['POST'])
def start_camera():
"""启动摄像头"""
global camera, camera_active, frame_counter, last_annotated_frame
if not camera_active:
camera = cv2.VideoCapture(0)
# 设置摄像头分辨率（降低分辨率提升性能）
camera.set(cv2.CAP_PROP_FRAME_WIDTH, PERFORMANCE_CONFIG['video_width'])
camera.set(cv2.CAP_PROP_FRAME_HEIGHT, PERFORMANCE_CONFIG['video_height'])
camera.set(cv2.CAP_PROP_FPS, 15) # 限制帧率到15fps
# 重置计数器
frame_counter = 0
last_annotated_frame = None
camera_active = True
return jsonify({'success': True, 'message': '摄像头已启动'})
return jsonify({'success': False, 'message': '摄像头已在运行'})
@app.route('/stop_camera', methods=['POST'])
def stop_camera():
"""停止摄像头"""
global camera, camera_active
if camera_active:
camera_active = False
if camera is not None:
camera.release()
camera = None
return jsonify({'success': True, 'message': '摄像头已停止'})
return jsonify({'success': False, 'message': '摄像头未运行'})
@app.route('/camera_detect', methods=['POST'])
def camera_detect():
"""摄像头单帧检测"""
global camera, camera_detections
if camera is None or not camera_active:
return jsonify({'error': '摄像头未启动'}), 400
success, frame = camera.read()
if not success:
return jsonify({'error': '无法读取摄像头帧'}), 400
# 获取置信度阈值
conf_threshold = float(request.form.get('confidence', 0.3))
# 进行检测
annotated_frame, detections = detect_defects(frame, conf=conf_threshold)
# 编码为base64
_, buffer = cv2.imencode('.jpg', annotated_frame)
img_base64 = base64.b64encode(buffer).decode('utf-8')
return jsonify({
'success': True,
'image': f'data:image/jpeg;base64,{img_base64}',
'detections': detections,
'total_defects': len(detections)
})
@app.route('/camera_detections', methods=['GET'])
def get_camera_detections():
"""获取最新的摄像头检测结果"""
global camera_detections
return jsonify({
'detections': camera_detections,
'total_defects': len(camera_detections)
})
if __name__ == '__main__':
app.run(debug=True, host='0.0.0.0', port=5000)

复制代码

这里主要是对我们训练出的最好的模型进行加载，并且使用Flask应用将其运行。前端页面通过Ajax请求来访问后端的推理接口。这里启动的话比较麻烦，因此我提供了一个Docker file 用来下载依赖和打包应用。

# 使用适用于ARM64架构的Python基础镜像
FROM python:3.10-slim-bullseye
# 设置工作目录
WORKDIR /app
# 安装系统依赖
RUN apt-get update && apt-get install -y \
libgl1-mesa-glx \
libglib2.0-0 \
libsm6 \
libxext6 \
libxrender-dev \
libgomp1 \
libgstreamer1.0-0 \
libavcodec58 \
libavformat58 \
libswscale5 \
&& rm -rf /var/lib/apt/lists/*
# 复制requirements文件
COPY requirements.txt ./
# 安装Python依赖
RUN pip install --no-cache-dir -r requirements.txt
# 复制应用程序文件（只复制必要的）
COPY app.py ./
COPY templates/ ./templates/
COPY static/style.css ./static/
COPY static/app.js ./static/
# 复制模型文件
COPY runs/detect/pcb_train4/weights/best.pt ./runs/detect/pcb_train4/weights/
# 创建必要的目录
RUN mkdir -p static/uploads
# 暴露端口
EXPOSE 5000
# 设置环境变量
ENV FLASK_APP=app.py
ENV PYTHONUNBUFFERED=1
# 启动命令
CMD ["python", "app.py"]

复制代码

并且提供了运行和部署的脚本。

#!/bin/bash
# 构建Docker镜像的脚本
# 使用方法：chmod +x docker-build.sh && ./docker-build.sh
IMAGE_NAME="pcb-defect-detection"
IMAGE_TAG="latest"
echo "开始构建Docker镜像: ${IMAGE_NAME}:${IMAGE_TAG}"
echo "目标平台: linux/arm64 (树莓派5)"
# 构建镜像（针对ARM64架构）
docker build \
--platform linux/arm64 \
-t ${IMAGE_NAME}:${IMAGE_TAG} \
.
if [ $? -eq 0 ]; then
echo "✅ 镜像构建成功！"
echo ""
echo "保存镜像到文件（用于传输到树莓派）："
echo " docker save ${IMAGE_NAME}:${IMAGE_TAG} | gzip > pcb-defect-detection.tar.gz"
echo ""
echo "在树莓派上加载镜像："
echo " gunzip -c pcb-defect-detection.tar.gz | docker load"
echo ""
echo "运行容器："
echo " docker run -d -p 5000:5000 --name pcb-detector ${IMAGE_NAME}:${IMAGE_TAG}"
echo ""
echo "如果需要使用摄像头，添加设备映射："
echo " docker run -d -p 5000:5000 --device=/dev/video0:/dev/video0 --name pcb-detector ${IMAGE_NAME}:${IMAGE_TAG}"
else
echo "❌ 镜像构建失败"
exit 1
fi

复制代码

部署脚本如下

#!/bin/bash
# 运行Docker容器的脚本（在树莓派上使用）
# 使用方法：chmod +x docker-run.sh && ./docker-run.sh
IMAGE_NAME="pcb-defect-detection"
CONTAINER_NAME="pcb-detector"
PORT=5000
echo "启动PCB缺陷检测容器..."
# 停止并删除已存在的容器
docker stop ${CONTAINER_NAME} 2>/dev/null
docker rm ${CONTAINER_NAME} 2>/dev/null
# 运行新容器
# --restart=unless-stopped: 自动重启（除非手动停止）
# -p 5000:5000: 端口映射
# --device=/dev/video0: 摄像头设备映射（如果使用摄像头功能）
# -v ./static/uploads:/app/static/uploads: 持久化上传文件
docker run -d \
--name ${CONTAINER_NAME} \
--restart=unless-stopped \
-p ${PORT}:5000 \
--device=/dev/video0:/dev/video0 \
-v "$(pwd)/static/uploads:/app/static/uploads" \
${IMAGE_NAME}:latest
if [ $? -eq 0 ]; then
echo "✅ 容器启动成功！"
echo ""
echo "访问地址: http://树莓派IP:${PORT}"
echo "本地访问: http://localhost:${PORT}"
echo ""
echo "查看日志: docker logs -f ${CONTAINER_NAME}"
echo "停止容器: docker stop ${CONTAINER_NAME}"
echo "重启容器: docker restart ${CONTAINER_NAME}"
else
echo "❌ 容器启动失败"
exit 1
fi

复制代码

部署方式

安装Python及依赖库：pip install -r requirements.txt
（可选）使用Docker一键部署：sh docker-build.sh && sh docker-run.sh
运行app.py启动Web服务，访问本地或树莓派IP的5000端口即可使用系统

项目文档

通过网盘分享的文件：defect_detection.zip
链接: https://pan.baidu.com/s/1RLiC4-bhgGXoWSvj0dDo0w?pwd=qavh 提取码: qavh 复制这段内容后打开百度网盘手机App，操作更方便哦

视频说明和功能演示

点击我观看

项目总结

本项目实现了一套基于 YOLO 的 PCB 缺陷检测系统，从数据处理、模型训练到实际部署进行了完整实现。项目使用北京大学开源的 PCB 缺陷数据集，对原始 VOC 数据进行清洗并转换为 YOLO 格式，训练了 YOLOv8 模型，最终在验证集上取得约 96% 的检测精度和 92% 的召回率。

训练完成后，将模型部署到树莓派 5，并基于 Flask 开发了 Web 检测系统，支持图片上传和摄像头实时检测。同时通过 Docker 实现一键部署，方便在边缘设备上快速运行。该项目验证了 YOLO 在 PCB 工业缺陷检测场景下的可行性和实际应用价值。

[项目提交] 《2025 DigiKey AI应用创意挑战赛》基于Yolo的PCB缺陷检测

eefocus_3956818 LV1

站长推荐 /2