目录
- 一、新建conda环境
- 二、安装yolov8环境
- 1.克隆安装包
- 2.安装依赖包
- 3.测试模型
- 任务2:基于RT-DETR实现PKU-PCB表面缺陷检测
- 数据准备
- 数据增强
- 测试
- 总结
一、新建conda环境
创建并激活conda环境:
在conda创建一个名为yolov8的新环境,并在其中安装python3.8,这个环境是独立的不会影响系统中的其他环境:
conda create -n PCBDet python=3.8
conda activate PCBDet
如何删除conda环境
conda remove --name yolov8cs --all
二、安装yolov8环境
1.克隆安装包
代码如下(示例):
git clone https://github.com/ultralytics/ultralytics.git
2.安装依赖包
pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple
国内的镜像源,不用翻墙,速度也非常快,达到2.2M,需要稍等一会儿下载与安装;
安装的依赖包如此之多:
Installing collected packages: pytz, py-cpuinfo, mpmath, zipp, urllib3, tzdata, typing-extensions, tqdm, sympy, six, pyyaml, pyparsing, psutil, pillow, packaging, nvidia-nvtx-cu12, nvidia-nvjitlink-cu12, nvidia-nccl-cu12, nvidia-curand-cu12, nvidia-cufft-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, numpy, networkx, MarkupSafe, kiwisolver, idna, fsspec, fonttools, filelock, cycler, charset-normalizer, certifi, triton, scipy, requests, python-dateutil, opencv-python, nvidia-cusparse-cu12, nvidia-cudnn-cu12, jinja2, importlib-resources, contourpy, pandas, nvidia-cusolver-cu12, matplotlib, torch, seaborn, ultralytics-thop, torchvision, ultralytics
Successfully installed MarkupSafe-2.1.5 certifi-2024.6.2 charset-normalizer-3.3.2 contourpy-1.1.1 cycler-0.12.1 filelock-3.15.3 fonttools-4.53.0 fsspec-2024.6.0 idna-3.7 importlib-resources-6.4.0 jinja2-3.1.4 kiwisolver-1.4.5 matplotlib-3.7.5 mpmath-1.3.0 networkx-3.1 numpy-1.24.4 nvidia-cublas-cu12-12.1.3.1 nvidia-cuda-cupti-cu12-12.1.105 nvidia-cuda-nvrtc-cu12-12.1.105 nvidia-cuda-runtime-cu12-12.1.105 nvidia-cudnn-cu12-8.9.2.26 nvidia-cufft-cu12-11.0.2.54 nvidia-curand-cu12-10.3.2.106 nvidia-cusolver-cu12-11.4.5.107 nvidia-cusparse-cu12-12.1.0.106 nvidia-nccl-cu12-2.20.5 nvidia-nvjitlink-cu12-12.5.40 nvidia-nvtx-cu12-12.1.105 opencv-python-4.10.0.84 packaging-24.1 pandas-2.0.3 pillow-10.3.0 psutil-6.0.0 py-cpuinfo-9.0.0 pyparsing-3.1.2 python-dateutil-2.9.0.post0 pytz-2024.1 pyyaml-6.0.1 requests-2.32.3 scipy-1.10.1 seaborn-0.13.2 six-1.16.0 sympy-1.12.1 torch-2.3.1 torchvision-0.18.1 tqdm-4.66.4 triton-2.3.1 typing-extensions-4.12.2 tzdata-2024.1 ultralytics-8.2.38 ultralytics-thop-2.0.0 urllib3-2.2.2 zipp-3.19.2
3.测试模型
需要下载RT-DETR.pt模型,直接放在ultralytics安装目录下;
模型下载链接:
https://github.com/ultralytics/assets/releases/download/v8.2.0/rtdetr-l.pt
第二步,新建一个运行脚本runRTDETR.py:
from ultralytics import RTDETR
from PIL import Image
import cv2
# 加载预训练的COCO RT-DETR-l模型
model = RTDETR('rtdetr-l.pt')
# 显示模型信息(可选)
model.info()
# from PIL
im1 = Image.open("figures/bus.jpg")
results = model.predict(source=im1, save=True) # save plotted images
通过以下命令行‘ python runRTDETR.py’测试模型是否运行成功:
python runRTDETR.py
rt-detr-l summary: 673 layers, 32970476 parameters, 0 gradients, 108.3 GFLOPs
0: 640x640 1 tv, 20.9ms
Speed: 2.9ms preprocess, 20.9ms inference, 1.5ms postprocess per image at shape (1, 3, 640, 640)
Results saved to /home/wqt/Projects/PCBDet/ultralytics/runs/detect/predict1
运行结果显示:
到这一步,表示整个安装流程成功!!!
任务2:基于RT-DETR实现PKU-PCB表面缺陷检测
数据准备
1、数据下载链接:https://pan.baidu.com/s/1hoPNd7_SAxOWa2XbBZZuTg
(1)训练集划分主要是将数据集分类成训练代码和测试代码,默认数据集会随机按比例8:1:1分类在train,val,test三个文本文件中,运行代码之后路径ImageSets/Main会出现四个文件,主要是写入的内容是训练数据集和测试数据集的图片名称
import os
import random
import argparse
parser = argparse.ArgumentParser()
#xml文件的地址,根据自己的数据进行修改 xml一般存放在Annotations下
parser.add_argument('--xml_path', default='Annotations', type=str, help='input xml label path')
#数据集的划分,地址选择自己数据下的ImageSets/Main
parser.add_argument('--txt_path', default='ImageSets/Main', type=str, help='output txt label path')
opt = parser.parse_args()
trainval_percent = 0.9
train_percent = 0.8
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)
if not os.path.exists(txtsavepath):
os.makedirs(txtsavepath)
num = len(total_xml)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)
file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')
for i in list_index:
name = total_xml[i][:-4] + '\n'
if i in trainval:
file_trainval.write(name)
if i in train:
file_train.write(name)
else:
file_val.write(name)
else:
file_test.write(name)
file_trainval.close()
file_train.close()
file_val.close()
file_test.close()
(2)用于yolo训练的txt格式代码,主要是将图片数据集标注后的xml文件中的标注信息读取出来并写入txt文件,运行后在labels文件夹中出现所有图片数据集的标注信息并且会在文件夹生成train.txt和val.txt,里面主要是训练集和测试集的图片路径。
# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import os
from os import getcwd
sets = ['train', 'val']
classes = ["missing_hole","mouse_bite","open_circuit","short","spur","spurious_copper"] # 改成自己的类别
abs_path = os.getcwd()
print(abs_path)
def convert(size, box):
dw = 1. / (size[0])
dh = 1. / (size[1])
x = (box[0] + box[1]) / 2.0 - 1
y = (box[2] + box[3]) / 2.0 - 1
w = box[1] - box[0]
h = box[3] - box[2]
x = x * dw
w = w * dw
y = y * dh
h = h * dh
return x, y, w, h
# Spur Spurious_copper
def conver_annotation(image_id):
in_file = open('Annotations/%s.xml' % (image_id), encoding='UTF-8')
out_file = open('labels/%s.txt' % (image_id), 'w')
tree = ET.parse(in_file)
root = tree.getroot()
size = root.find('size')
w = int(size.find('width').text)
h = int(size.find('height').text)
for obj in root.iter('object'):
difficult = obj.find('difficult').text
#difficult = obj.find('Difficult').text
cls = obj.find('name').text
if cls not in classes or int(difficult) == 1:
continue
cls_id = classes.index(cls)
xmlbox = obj.find('bndbox')
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
b1, b2, b3, b4 = b
# 标注越界修正
if b2 > w:
b2 = w
if b4 > h:
b4 = h
b = (b1, b2, b3, b4)
bb = convert((w, h), b)
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
wd = getcwd()
for image_set in sets:
if not os.path.exists('labels/'):
os.makedirs('labels/')
image_ids = open('ImageSets/Main/%s.txt' % (image_set)).read().strip().split()
list_file = open('%s.txt' % (image_set), 'w')
for image_id in image_ids:
list_file.write(abs_path + '/images/%s.jpg\n' % (image_id))
conver_annotation(image_id)
list_file.close()
2、工程怎么读:ultralytics/cfg/datasets下写一个读取数据集的yaml文件
train: /home/wqt/Projects/data/PCB_DATASET2YOLO/train.txt #训练数据集的路径
val: /home/wqt/Projects/data/PCB_DATASET2YOLO/val.txt #测试数据集的路经
nc: 6 #缺陷标签数量
names: ['missing_hole', 'mouse_bite', 'open_circuit', 'short', 'spur', 'spurious_copper'] #缺陷名称
3在根目录下建立一个运行脚本trainPCB.py:
from ultralytics import RTDETR
model = RTDETR('rtdetr-l.pt') # 加载预训练的COCO RT-DETR-l模型
model.train(model="rtdetr-l.yaml",
data='PKUPCB.yaml', #输入上面建立读取数据集的yaml文件
epochs=230, #循环230次
batch=16, #批处理,每次输入16个样本
imgsz=640) #输入图片归一化,统一输入尺度为640x640
开始训练:
通过以下命令行运行训练模型:
python trainPCB.py #这里写自己训练脚本的名字
训练结果,保存到文档;
Validating /home/wqt/Projects/PCBDet/ultralytics/runs/detect/train/weights/best.pt...
Ultralytics YOLOv8.2.38 🚀 Python-3.8.19 torch-2.3.1+cu121 CUDA:0 (NVIDIA GeForce RTX 4090, 24209MiB)
rt-detr-l summary: 498 layers, 31996070 parameters, 0 gradients, 103.5 GFLOPs
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 4/4 [00:00<00:00, 5.27it/s]
all 125 533 0.957 0.965 0.967 0.507
missing_hole 18 77 0.989 1 0.995 0.579
mouse_bite 19 84 0.951 0.976 0.969 0.49
open_circuit 22 92 0.966 0.978 0.971 0.528
short 21 86 0.942 0.951 0.949 0.47
spur 22 93 0.933 0.935 0.931 0.485
spurious_copper 23 101 0.96 0.948 0.984 0.489
Speed: 0.1ms preprocess, 3.0ms inference, 0.0ms loss, 0.1ms postprocess per image
Results saved to /home/wqt/Projects/PCBDet/ultralytics/runs/detect/train #这个就是训练结果的绝对路径
数据增强
数据增强的方法主要有以下几种:
而RT-DETR中要进行数据增强可以通过改变ultralytics/cfg/default.yaml中的数据来达到数据增强的目的:
# degrees: 0.0 # (float) image rotation (+/- deg)
degrees: 10.0 # (float) image rotation (+/- deg)
translate: 0.1 # (float) image translation (+/- fraction)
# scale: 0.5 # (float) image scale (+/- gain)
scale: 0.25 # (float) image scale (+/- gain)
shear: 0.0 # (float) image shear (+/- deg)
perspective: 0.0 # (float) image perspective (+/- fraction), range 0-0.001
# flipud: 0.0 # (float) image flip up-down (probability)
flipud: 0.2 # (float) image flip up-down (probability)
这里主要是对数据集图片进行了10度的旋转(degrees),缩小成原图的0.75倍(scale),最后是百分之二十的上下翻转的概率。
然后在运行训练脚本。
训练结果:
Validating /home/wqt/Projects/PCBDet/ultralytics/runs/detect/train/weights/best.pt...
Ultralytics YOLOv8.2.38 🚀 Python-3.8.19 torch-2.3.1+cu121 CUDA:0 (NVIDIA GeForce RTX 4090, 24209MiB)
rt-detr-l summary: 498 layers, 31996070 parameters, 0 gradients, 103.5 GFLOPs
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 4/4 [00:00<00:00, 5.43it/s]
all 125 533 0.965 0.965 0.977 0.52
missing_hole 18 77 0.977 1 0.995 0.594
mouse_bite 19 84 0.964 0.964 0.973 0.52
open_circuit 22 92 0.977 0.92 0.951 0.517
short 21 86 0.977 0.97 0.985 0.479
spur 22 93 0.939 0.957 0.971 0.488
spurious_copper 23 101 0.957 0.98 0.985 0.524
Speed: 0.1ms preprocess, 3.0ms inference, 0.0ms loss, 0.1ms postprocess per image
Results saved to /home/wqt/Projects/PCBDet/ultralytics/runs/detect/train
正常测试的数据:
没有使用数据增强的训练:
数据增强的数据:
测试
利用训练好的模型来测试单张图片:
from ultralytics import RTDETR
from PIL import Image
import cv2
modelPath = "runs/detect/train6/weights/best.pt"
# 加载预训练的COCO RT-DETR-l模型
model = RTDETR(modelPath)
# 显示模型信息(可选)
model.info()
imgPath = "resource/01_missing_hole_01.jpg"
# from PIL
im1 = Image.open(imgPath)
results = model.predict(source=im1, save=True) # save plotted images
总结
通过不用翻墙就能实现yolo平台的RT-DETR使用,并迁移学习PCB表面缺陷检测。