【超详细】基于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_hole1: mouse_bite2: open_circuit3: short4: spur5: spurious_copper

4.3 创建一个训练脚本

在主目录下创建一个train.py，内容如下：

from ultralytics import YOLOif __name__ == '__main__':# Load a model# model = YOLO("yolo11s.yaml")  # build a new model from scratchmodel = YOLO("yolo11s.pt")  # load a pretrained model (recommended for training)# Use the modelmodel.train(data="PCB.yaml", imgsz=640, batch=16, workers=8, cache=True, epochs=100)  # train the modelmetrics = model.val()  # evaluate model performance on the validation set# results = model("ultralytics\\assets\\bus.jpg")  # predict on an imagepath = model.export(format="onnx", opset=13)  # export the model to ONNX format

问题1：若爆显存，降低batch和workers大小！

训练结果如下【比yolov8s略高】：
在这里插入图片描述

4.4 创建一个预测脚本

在主目录下创建一个detect.py，内容如下：

from ultralytics import YOLOif __name__ == '__main__':# Load a modelmodel = YOLO(r"runs\detect\train\weights\best.pt")  # load modelmodel.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 npimport 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 modelself.ndtype = np.half if self.ort_session.get_inputs()[0].type == 'tensor(float16)' else np.singleself.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-processt1 = 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))# 推理 inferencet2 = 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-processt3 = 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 -> BCHWdef 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 shapenew_shape = (self.model_height, self.model_width)r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])ratio = r, rnew_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 paddingif shape[::-1] != new_unpad:  # resizeimg = 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.0img_process = img[None] if len(img.shape) == 3 else imgreturn img_process, ratio, (pad_w, pad_h)# 后处理，包括：阈值过滤与NMSdef 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-thresholdx = 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_.htmlx = 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 -> xyxyx[..., [0, 1]] -= x[..., [2, 3]] / 2x[..., [2, 3]] += x[..., [0, 1]]# Rescales bounding boxes from model shape(model_height, model_width) to the shape of original imagex[..., :4] -= [pad_w, pad_h, pad_w, pad_h]x[..., :4] /= min(ratio)# Bounding boxes boundary clampx[..., [0, 2]] = x[:, [0, 2]].clip(0, im0.shape[1])x[..., [1, 3]] = x[:, [1, 3]].clip(0, im0.shape[0])return x[..., :6]  # boxeselse:return []if __name__ == '__main__':# Create an argument parser to handle command-line argumentsparser = 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 modeldet_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 OpenCVimg = cv2.imread(os.path.join(args.source, img_name))# 检测Inferenceboxes, (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)