【超详细】基于YOLOv11的PCB缺陷检测
主要内容如下：
1、数据集介绍
2、下载PCB数据集
3、不同格式数据集预处理（Json/xml），制作YOLO格式训练集
4、模型训练及可视化
5、Onnxruntime推理
运行环境：Python=3.8（要求>=3.8），torch1.12.0+cu113（要求>=1.8），onnxruntime-gpu=1.12.0
YOLO格式下载链接【可直接跳过步骤123】：https://aistudio.baidu.com/datasetdetail/297149
往期内容：
【超详细】跑通YOLOv8之深度学习环境配置1-Anaconda安装
【超详细】跑通YOLOv8之深度学习环境配置2-CUDA安装
【超详细】跑通YOLOv8之深度学习环境配置3-YOLOv8安装
【超详细】基于YOLOv8的PCB缺陷检测
0 YOLOv11
代码地址：https://github.com/ultralytics/ultralytics/tree/main/ultralytics/cfg/models/11
1 数据集介绍
1.1 简介
印刷电路板（PCB）瑕疵数据集：是一个公共的合成PCB数据集，由北京大学发布，其中包含1386张图像以及6种缺陷（缺失孔，鼠标咬伤，开路，短路，杂散，伪铜），用于检测，分类和配准任务。本文我们选取了其中适用与检测任务的693张图像，随机选择593张图像作为训练集，100张图像作为验证集。
1.2 示例
2 下载数据集
官方链接：https://robotics.pkusz.edu.cn/resources/dataset/
注意：百度网盘下载，速度很慢，不推荐！推荐去百度AI stduio数据集下载，速度快！
百度AI stduio下载链接：https://aistudio.baidu.com/datasetdetail/272346
3 制作YOLO格式训练集
具体见第3节3.2：【超详细】基于YOLOv8的PCB缺陷检测
YOLO格式下载链接【直接拿来训练Aistudio快速下载链接】：https://aistudio.baidu.com/datasetdetail/297149
4 模型训练及可视化
4.1 更新ultralytics
pip install --upgrade ultralytics
4.2 创建数据集yaml文件
注意：路径一定填对，类别与id一定要对应！！！
创建ultralytics\cfg\datasets\PCB.yaml文件，内容如下：
path: E:\\datasets\\PCB\\PCB_DATASET_YOLO # dataset root dir
train: images/train2017 # train images (relative to 'path') 4 images
val: images/val2017 # val images (relative to 'path') 4 images
# Classes for DOTA 1.0
names:
0: missing_hole
1: mouse_bite
2: open_circuit
3: short
4: spur
5: spurious_copper
4.3 创建一个训练脚本
在主目录下创建一个train.py，内容如下：
from ultralytics import YOLO
if __name__ == '__main__':
# Load a model
# model = YOLO("yolo11s.yaml") # build a new model from scratch
model = YOLO("yolo11s.pt") # load a pretrained model (recommended for training)
# Use the model
model.train(data="PCB.yaml", imgsz=640, batch=16, workers=8, cache=True, epochs=100) # train the model
metrics = model.val() # evaluate model performance on the validation set
# results = model("ultralytics\\assets\\bus.jpg") # predict on an image
path = model.export(format="onnx", opset=13) # export the model to ONNX format
问题1：若爆显存，降低batch和workers大小！
训练结果如下【比yolov8s略高】：
4.4 创建一个预测脚本
在主目录下创建一个detect.py，内容如下：
from ultralytics import YOLO
if __name__ == '__main__':
# Load a model
model = YOLO(r"runs\detect\train\weights\best.pt") # load model
model.predict(source=r"E:\datasets\PCB\PCB_DATASET_YOLO\images\val2017\01_missing_hole_07.jpg", save=True, save_conf=True, save_txt=True, name='output')
预测结果：
5 Onnxruntime推理及可视化
5.1 onnx推理
在主目录下创建一个onnx_infer.py，内容如下：
import argparse
import time
import cv2
import numpy as np
import onnxruntime as ort # 使用onnxruntime推理用上，pip install onnxruntime-gpu==1.12.0 -i https://pypi.tuna.tsinghua.edu.cn/simple，默认安装CPU
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
class YOLOv11:
"""YOLOv11 object detection model class for handling inference and visualization."""
def __init__(self, onnx_model, imgsz=(640, 640)):
"""
Initialization.
Args:
onnx_model (str): Path to the ONNX model.
"""
# 构建onnxruntime推理引擎
self.ort_session = ort.InferenceSession(onnx_model,
providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
if ort.get_device() == 'GPU' else ['CPUExecutionProvider'])
print(ort.get_device())
# Numpy dtype: support both FP32 and FP16 onnx model
self.ndtype = np.half if self.ort_session.get_inputs()[0].type == 'tensor(float16)' else np.single
self.model_height, self.model_width = imgsz[0], imgsz[1] # 图像resize大小
def __call__(self, im0, conf_threshold=0.4, iou_threshold=0.45):
"""
The whole pipeline: pre-process -> inference -> post-process.
Args:
im0 (Numpy.ndarray): original input image.
conf_threshold (float): confidence threshold for filtering predictions.
iou_threshold (float): iou threshold for NMS.
Returns:
boxes (List): list of bounding boxes.
"""
# 前处理Pre-process
t1 = time.time()
im, ratio, (pad_w, pad_h) = self.preprocess(im0)
pre_time = round(time.time() - t1, 3)
# print('det预处理时间：{:.3f}s'.format(time.time() - t1))
# 推理 inference
t2 = time.time()
preds = self.ort_session.run(None, {self.ort_session.get_inputs()[0].name: im})[0]
# print('det推理时间：{:.2f}s'.format(time.time() - t2))
det_time = round(time.time() - t2, 3)
# 后处理Post-process
t3 = time.time()
boxes = self.postprocess(preds,
im0=im0,
ratio=ratio,
pad_w=pad_w,
pad_h=pad_h,
conf_threshold=conf_threshold,
iou_threshold=iou_threshold,
)
# print('det后处理时间：{:.3f}s'.format(time.time() - t3))
post_time = round(time.time() - t3, 3)
return boxes, (pre_time, det_time, post_time)
# 前处理，包括：resize, pad, HWC to CHW，BGR to RGB，归一化，增加维度CHW -> BCHW
def preprocess(self, img):
"""
Pre-processes the input image.
Args:
img (Numpy.ndarray): image about to be processed.
Returns:
img_process (Numpy.ndarray): image preprocessed for inference.
ratio (tuple): width, height ratios in letterbox.
pad_w (float): width padding in letterbox.
pad_h (float): height padding in letterbox.
"""
# Resize and pad input image using letterbox() (Borrowed from Ultralytics)
shape = img.shape[:2] # original image shape
new_shape = (self.model_height, self.model_width)
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
ratio = r, r
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
pad_w, pad_h = (new_shape[1] - new_unpad[0]) / 2, (new_shape[0] - new_unpad[1]) / 2 # wh padding
if shape[::-1] != new_unpad: # resize
img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
top, bottom = int(round(pad_h - 0.1)), int(round(pad_h + 0.1))
left, right = int(round(pad_w - 0.1)), int(round(pad_w + 0.1))
img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114)) # 填充
# Transforms: HWC to CHW -> BGR to RGB -> div(255) -> contiguous -> add axis(optional)
img = np.ascontiguousarray(np.einsum('HWC->CHW', img)[::-1], dtype=self.ndtype) / 255.0
img_process = img[None] if len(img.shape) == 3 else img
return img_process, ratio, (pad_w, pad_h)
# 后处理，包括：阈值过滤与NMS
def postprocess(self, preds, im0, ratio, pad_w, pad_h, conf_threshold, iou_threshold):
"""
Post-process the prediction.
Args:
preds (Numpy.ndarray): predictions come from ort.session.run().
im0 (Numpy.ndarray): [h, w, c] original input image.
ratio (tuple): width, height ratios in letterbox.
pad_w (float): width padding in letterbox.
pad_h (float): height padding in letterbox.
conf_threshold (float): conf threshold.
iou_threshold (float): iou threshold.
Returns:
boxes (List): list of bounding boxes.
"""
x = preds # outputs: predictions (1, 84, 8400)
# Transpose the first output: (Batch_size, xywh_conf_cls, Num_anchors) -> (Batch_size, Num_anchors, xywh_conf_cls)
x = np.einsum('bcn->bnc', x) # (1, 8400, 84)
# Predictions filtering by conf-threshold
x = x[np.amax(x[..., 4:], axis=-1) > conf_threshold]
# Create a new matrix which merge these(box, score, cls) into one
# For more details about `numpy.c_()`: https://numpy.org/doc/1.26/reference/generated/numpy.c_.html
x = np.c_[x[..., :4], np.amax(x[..., 4:], axis=-1), np.argmax(x[..., 4:], axis=-1)]
# NMS filtering
# 经过NMS后的值, np.array([[x, y, w, h, conf, cls], ...]), shape=(-1, 4 + 1 + 1)
x = x[cv2.dnn.NMSBoxes(x[:, :4], x[:, 4], conf_threshold, iou_threshold)]
# 重新缩放边界框，为画图做准备
if len(x) > 0:
# Bounding boxes format change: cxcywh -> xyxy
x[..., [0, 1]] -= x[..., [2, 3]] / 2
x[..., [2, 3]] += x[..., [0, 1]]
# Rescales bounding boxes from model shape(model_height, model_width) to the shape of original image
x[..., :4] -= [pad_w, pad_h, pad_w, pad_h]
x[..., :4] /= min(ratio)
# Bounding boxes boundary clamp
x[..., [0, 2]] = x[:, [0, 2]].clip(0, im0.shape[1])
x[..., [1, 3]] = x[:, [1, 3]].clip(0, im0.shape[0])
return x[..., :6] # boxes
else:
return []
if __name__ == '__main__':
# Create an argument parser to handle command-line arguments
parser = argparse.ArgumentParser()
parser.add_argument('--det_model', type=str, default=r"E:\Code\yolov11\ultralytics-main\runs\detect\train\weights\best.onnx", help='Path to ONNX model')
parser.add_argument('--source', type=str, default=str(r'E:\datasets\PCB\PCB_DATASET_YOLO\images\val2017'), help='Path to input image')
parser.add_argument('--out_path', type=str, default=str(r'E:\Code\yolov11\ultralytics-main\runs\detect\res'), help='结果保存文件夹')
parser.add_argument('--imgsz_det', type=tuple, default=(640, 640), help='Image input size')
parser.add_argument('--classes', type=list, default=['missing_hole', 'mouse_bite', 'open_circuit', 'short', 'spur', 'spurious_copper'], help='类别')
parser.add_argument('--conf', type=float, default=0.25, help='Confidence threshold')
parser.add_argument('--iou', type=float, default=0.6, help='NMS IoU threshold')
args = parser.parse_args()
if not os.path.exists(args.out_path):
os.mkdir(args.out_path)
print('开始运行：')
# Build model
det_model = YOLOv11(args.det_model, args.imgsz_det)
color_palette = np.random.uniform(0, 255, size=(len(args.classes), 3)) # 为每个类别生成调色板
for i, img_name in enumerate(os.listdir(args.source)):
try:
start_time = time.time()
# Read image by OpenCV
img = cv2.imread(os.path.join(args.source, img_name))
# 检测Inference
boxes, (pre_time, det_time, post_time) = det_model(img, conf_threshold=args.conf, iou_threshold=args.iou)
print('{}/{} ==>总耗时间: {:.3f}s, 其中, 预处理: {:.3f}s, 推理: {:.3f}s, 后处理: {:.3f}s, 识别{}个目标'.format(i+1, len(os.listdir(args.source)), time.time() - start_time, pre_time, det_time, post_time, len(boxes)))
for (*box, conf, cls_) in boxes:
cv2.rectangle(img, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])),
color_palette[int(cls_)], 2, cv2.LINE_AA)
cv2.putText(img, f'{args.classes[int(cls_)]}: {conf:.3f}', (int(box[0]), int(box[1] - 9)),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)
cv2.imwrite(os.path.join(args.out_path, img_name), img)
except Exception as e:
print(e)
资源消耗：显存不到1.5G，推理速度5ms.
5.2 结果可视化
