- ubuntu 18.04
- python 3.6.9
- opencv 4.5.0
- pytorch 1.9.0
- torchvision 0.10.0
（1）解压数据集，放在yolov5-6.0根目录下
（2）删除ImageSets文件夹下的内容，在/ImageSets文件夹下新建Main文件夹
（3）数据集分类，运行test.py（注意路径正确）。会在/ImageSets/Main文件夹下生成四个txt文件，trainval.txt、test.txt、train.txt、val.txt。
import os
import random
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--xml_path', default='./Annotations', type=str, help='input xml label path')
parser.add_argument('--txt_path', default='./ImageSets/Main', type=str, help='output txt label path')
opt = parser.parse_args()
trainval_percent = 1.0
train_percent = 0.9
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()
（4）在data文件夹下新建labels文件夹，然后在yolov5-6.0文件夹下新建my_labels.py文件，并运行（注意路径，注意数据集类别）。会在data文件夹下生成三个txt文件，train.txt、val.txt、test.txt。
import xml.etree.ElementTree as ET
import os
from os import getcwd
sets = ['train', 'val', 'test']
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
def convert_annotation(image_id):
in_file = open ('/home/lyp/Downloads/yolo/yolov5-opencv-dnn-cpp-main/yolov5-6.0/data/Annotations/%s.xml' % (image_id), encoding='UTF-8')
out_file = open('/home/lyp/Downloads/yolo/yolov5-opencv-dnn-cpp-main/yolov5-6.0/data/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'):
if obj.find('difficult'):
difficult = float(obj.find('difficult').text)
else:
difficult = 0
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('/home/lyp/Downloads/yolo/yolov5-opencv-dnn-cpp-main/yolov5-6.0/data/labels/'):
os.makedirs('/home/lyp/Downloads/yolo/yolov5-opencv-dnn-cpp-main/yolov5-6.0/data/labels/')
image_ids = open('/home/lyp/Downloads/yolo/yolov5-opencv-dnn-cpp-main/yolov5-6.0/data/ImageSets/Main/%s.txt' % (image_set)).read().strip().split()
list_file = open('/home/lyp/Downloads/yolo/yolov5-opencv-dnn-cpp-main/yolov5-6.0/data/%s.txt' % (image_set), 'w')
for image_id in image_ids:
list_file.write(abs_path + '/data/images/%s.jpg\n' % (image_id))
convert_annotation(image_id)
list_file.close()
（5）在models文件夹下新建文件my_pcb.yaml（注意修改类别数量nc）
nc: 6
depth_multiple: 0.33
width_multiple: 0.50
anchors:
- [10,13, 16,30, 33,23]
- [30,61, 62,45, 59,119]
- [116,90, 156,198, 373,326]
backbone:
[[-1, 1, Conv, [64, 6, 2, 2]],
[-1, 1, Conv, [128, 3, 2]],
[-1, 3, C3, [128]],
[-1, 1, Conv, [256, 3, 2]],
[-1, 6, C3, [256]],
[-1, 1, Conv, [512, 3, 2]],
[-1, 9, C3, [512]],
[-1, 1, Conv, [1024, 3, 2]],
[-1, 3, C3, [1024]],
[-1, 1, SPPF, [1024, 5]],
]
head:
[[-1, 1, Conv, [512, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 6], 1, Concat, [1]],
[-1, 3, C3, [512, False]],
[-1, 1, Conv, [256, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 4], 1, Concat, [1]],
[-1, 3, C3, [256, False]],
[-1, 1, Conv, [256, 3, 2]],
[[-1, 14], 1, Concat, [1]],
[-1, 3, C3, [512, False]],
[-1, 1, Conv, [512, 3, 2]],
[[-1, 10], 1, Concat, [1]],
[-1, 3, C3, [1024, False]],
[[17, 20, 23], 1, Detect, [nc, anchors]],
]
（6）在data文件夹下新建my_data.yaml（注意train和val路径、nc和数据类别名称names）
train: /home/lyp/Downloads/yolo/yolov5-opencv-dnn-cpp-main/yolov5-6.0/data/train.txt
val: /home/lyp/Downloads/yolo/yolov5-opencv-dnn-cpp-main/yolov5-6.0/data/val.txt
nc: 6
names: ['missing_hole','mouse_bite', 'open_circuit', 'short', 'spur', 'spurious_copper']
（7）准备预训练文件yolov5s.pt，放在weights文件夹下。
（8）开始训练，在yolov5-6.0根目录下执行训练命令。
python3 train.py --data data/my_data.yaml --cfg models/my_pcb.yaml --weights weights/yolov5s.pt --batch-size 8 --epochs 100
训练完成后，可以在runs/train文件夹下找到训练的pt权重文件。
python3 detect.py --weights ./runs/train/exp4/weights/best.pt --source ./data/images/05_spur_07.jpg
pip install -U coremltools onnx scikit-learn==0.19.2
python3 export.py --weights ./runs/train/exp4/weights/best.pt --include onnx --dynamic
python3 export.py --weights ./runs/train/exp4/weights/best.pt --include onnx
在export.py中修改下参数（用命令行导出的请加上–opset 12，请注意不要带上dynamic参数，opencv对动态输入支持做的不好，应该是需要opencv4.5.3以上的版本，并且需要TBB支持才行来着）
python3 detect.py --weights ./runs/train/exp4/weights/best.onnx --source ./data/images/05_mouse_bite_08.jpg
参考链接：https://blog.csdn.net/magic_ll/article/details/121352907
参考链接：https://blog.csdn.net/qq_34124780/article/details/121079317
（1）将转换得到的best.onnx权重文件拷贝到yolov5-opencv-dnn-cpp-main文件夹下。
（2）修改yolo.h文件（修改成自己的数据类别名称）。
#pragma once
#include
#include
#define YOLO_P6 false
struct Output {
int id;
float confidence;
cv::Rect box;
};
class Yolo {
public:
Yolo() {
}
~Yolo() {}
bool readModel(cv::dnn::Net& net, std::string& netPath, bool isCuda);
bool Detect(cv::Mat& SrcImg, cv::dnn::Net& net, std::vector<Output>& output);
void drawPred(cv::Mat& img, std::vector<Output> result, std::vector<cv::Scalar> color);
private:
#if(defined YOLO_P6 && YOLO_P6==true)
const float netAnchors[4][6] = { { 19,27, 44,40, 38,94 },{ 96,68, 86,152, 180,137 },{ 140,301, 303,264, 238,542 },{ 436,615, 739,380, 925,792 } };
const int netWidth = 1280;
const int netHeight = 1280;
const int strideSize = 4;
#else
const float netAnchors[3][6] = { { 10,13, 16,30, 33,23 },{ 30,61, 62,45, 59,119 },{ 116,90, 156,198, 373,326 } };
const int netWidth = 640;
const int netHeight = 640;
const int strideSize = 3;
#endif
const float netStride[4] = { 8, 16.0,32,64 };
float boxThreshold = 0.25;
float classThreshold = 0.25;
float nmsThreshold = 0.45;
float nmsScoreThreshold = boxThreshold * classThreshold;
if 0
std::vector<std::string> className = { "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
"tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
"hair drier", "toothbrush" };
endif
std::vector<std::string> className = { "missing_hole", "mouse_bite", "open_circuit", "short", "spur", "spurious_copper"};
};
（3）修改main.cpp文件（注意修改成自己的权重文件路径和待检测图片）。
#include "yolo.h"
#include
#include
#include
using namespace std;
using namespace cv;
using namespace dnn;
int main()
{
string img_path = "/home/lyp/Downloads/yolo/yolov5-opencv-dnn-cpp-main/image/01_short_02.jpg";
string model_path = "/home/lyp/Downloads/yolo/yolov5-opencv-dnn-cpp-main/best.onnx";
Yolo test;
Net net;
if (test.readModel(net, model_path, false)) {
cout << "read net ok!" << endl;
}
else {
return -1;
}
vector<Scalar> color;
srand(time(0));
for (int i = 0; i < 80; i++) {
int b = rand() % 256;
int g = rand() % 256;
int r = rand() % 256;
color.push_back(Scalar(b, g, r));
}
vector<Output> result;
Mat img = imread(img_path);
if (test.Detect(img, net, result)) {
test.drawPred(img, result, color);
}
else {
cout << "Detect Failed!"<<endl;
}
system("pause");
return 0;
}
（4）新建在yolov5-opencv-dnn-cpp-main文件夹下build文件夹，然后在build文件夹下编译。
cd build
cmake ..
make
./main
参考链接：https://blog.csdn.net/weixin_43199832/article/details/122323192
Original: https://blog.csdn.net/qq_41821678/article/details/124827240
Author: 酸梅果茶
Title: 基于yolov5-6.0版本的PCB板缺陷检测（Python/C++部署）
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/711297/
转载文章受原作者版权保护。转载请注明原作者出处！