目录
- 1.项目克隆和环境配置
- 1.1 我这里使用的是v8.0.6版本
- 1.2 项目代码结构介绍
- 2.数据集介绍
- 2.1 数据集采集
- 2.2采集结果介绍
- 3.模型训练
- 4.pyqt界面设计
- 4.1 界面内容介绍
- 4.2 界面实现
- 5.操作中的逻辑实现
- 5.1 图片检测
- 5.2 文件夹检测
- 5.3 视频检测和摄像头检测
- 6. 效果展示
1.项目克隆和环境配置
1.1 我这里使用的是v8.0.6版本
这里为了方便学习,我使用的是旧的版本,先从官网上Tag下找到v8.0.6版本的安装包,然后下载解压,下面就是项目下载地方
1.2 项目代码结构介绍
将这个版本的yolov8下载到本地后整个项目结构和下面这张图是类似的,整体项目结构和v5有些不同,其中重要的代码都在ultralytics中
|——bolt_datasets:这个里面存放的是我用来进行螺栓松动检测的数据集。
|——docker:是一个应用容器引擎。
|——doc:这里面的东西不用管。
|——example:这里面是下载ultralytics这个python库,用这个来进行训练和推理的一 些例子。
|——Flowers_Dataset:这是另一个花朵数据集。
|——ultralytics:这里面包含模型配置,数据集配置yolov8模型的实现等。
ultralytics中代码结构如下图,主要使用的是ultralytics\yolo\v8\classify\predict.py这部分代码,如果想用自己的数据集训练模型可以看这篇文章,v8和v5项目结构有一些变化,不过数据集处理等是一样的,用python中labeling标注数据。
2.数据集介绍
2.1 数据集采集
我是用的自己手机进行采集的,用三脚架固定手机,通过调整手机高度、距离螺栓远近和拍摄角度来采集数据的。
总共拍摄了5个松动的10秒视频和5个拧紧的10秒视频,再将视频处理为一帧一帧的形式生成一个数据集,将每帧调整为640×640大小的图片,最后使用其中的200张图片训练,每张图片包含4个螺栓。
2.2采集结果介绍
上面两张是螺栓拧紧的图片,下面两张是螺栓松动的图片
下面是我用yolov8进行训练的的数据集结构,制作类似这样的数据集可以看这篇文章,该数据集共800个螺栓,图片大小为640×640,训练集和测试集比例为9:1。
3.模型训练
先修改ultralytics\yolo\configs\default.yaml这部分中的配置文件,我是用的是预训练yolov8n.pt模型进行训练,共训练100次,最后训练出的文件如下图
4.pyqt界面设计
4.1 界面内容介绍
界面中内容包括5个部分,分别是检测窗口、检测结果与位置信息、检测参数设置、检测结果和操作。检测窗口用来显示图片,检测结果与位置信息显示序号、检测图片文件路径、目标编号、类别、置信度和坐标位置这六个部分,然后检测参数包括阈值和交并比阈值,还有检测结果和操作等等。
4.2 界面实现
这里我用的是python带的pyqt5写的,用designer设计布局,再将ui文件转为py文件,这里注意一定要使用界面布局,这样整个界面布局更清晰并且也能放大放小,不容易混乱。
下面是我的界面图:
5.操作中的逻辑实现
我的实现方法是在detect\predict.py这个文件内容上新创建了一个类来实现操作中功能,这个类能初始化界面,实现按钮功能等。
5.1 图片检测
实现步骤大致为三步,打开要检测的图片,对图片进行检测,最后将检测结果显示在界面上。
用QFileDialog打开图片,打开图片后用yolo\engine\predictor中的类实例进行检测,检测过程中同过yield返回必要的数据,包括模型model、检测花费时间tm和原图im0s。之后根据返回的信息存储边界框、置信度、类别序号信息,向界面tableview和comboBox中插入信息等操作。
具体代码
# 文件中图片检测
def open_file(self):
folder_path = QFileDialog.getExistingDirectory(self, 'Select Folder', '')
# 清空表格数据
row_count = self.tb_model.rowCount()
self.tb_model.removeRows(0, row_count)
self.flag == 0
for filename in os.listdir(folder_path):
self.sum = 0
# 构造完整的文件路径
file_path = os.path.join(folder_path, filename)
# # 检查文件是否是图片文件,可以根据文件扩展名来判断
if file_path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
self.fileName1 = file_path
try:
for model, tm, im in self.predict():
print(self.save_dir)
self.model = model
det = self.predictor.output['det']
self.det = det
# 存储结果,向tableview和comboBox中添加信息
self.image_detections[self.fileName1] = det # 存储检测结果
self.image_path[filename] = file_path
self.image_tm[self.fileName1] = tm
self.image_allLabel[self.fileName1] = os.path.join(self.save_dir, filename)
box = det[:,:4] # 边界框
confidences = det[:,4] # 置信度
class_indices = det[:,5].astype(int) # 类序号
predicted_classes = [model.names[idx] for idx in class_indices] # 序号转成类别名
class_count = {}
new_predicted_classes = []
for box, confidences, class_name in zip(box,confidences, predicted_classes): # 修改重复的类别名
# 给相同类别的标签添加数字后缀
if class_name not in class_count:
class_count[class_name] = 0
class_count[class_name] += 1
#创建新标签
new_class_name = f"{class_name}{class_count[class_name] - 1}" if class_count[class_name] > 1 else class_name
new_predicted_classes.append(new_class_name)
pp = '(' + ', '.join([str(item) for item in box]) + ')'
self.sum += 1
# 向TableView添加信息
row = []
file=f'{self.save_dir}/crops/{new_class_name}/{new_class_name}.jpg'
row.append(QStandardItem(str(self.sum)))
row.append(QStandardItem(str(file)))
row.append(QStandardItem(str(self.sum)))
row.append(QStandardItem(str(new_class_name)))
row.append(QStandardItem(str(f'{confidences:.2f}')))
row.append(QStandardItem(str(pp)))
self.tb_model.appendRow(row)
self.comboBox.clear()
self.comboBox.addItem('全部')
self.comboBox.addItems(new_predicted_classes)
self.image_boBox[self.fileName1] = new_predicted_classes
self.image_sum[self.fileName1] = self.sum
self.image = self.predictor.annotator.result() # 标记所有的结果图
self.label_6.setText(str(self.sum)) # 总目标数
self.label_8.setText(str(round(tm*1000, 3))+'ms') # 用时
# 将numpy图片转成QImage
height, width, channels = self.image.shape
bytes_per_line = 3 * width
q_image = QImage(self.image.data, width, height, bytes_per_line, QImage.Format_BGR888)
# Convert QImage to QPixmap
pixmap = QPixmap.fromImage(q_image)
self.label_2.setPixmap(pixmap.scaled(self.label_2.size(), Qt.IgnoreAspectRatio, Qt.SmoothTransformation))
self.label_2.setText('') # Clear the label text
except Exception as e:
print(f"Error processing {filename}: {e}")
5.2 文件夹检测
这部分是对文件夹中存在的所有图片进行检测,并且鼠标点击tableview可以切换检测图片,同时能够切换检测图片中的单个检测对象。
实现流程跟上面类似,就是多了个读取文件夹中每张图片的过程,代码如下:
for filename in os.listdir(folder_path):
self.sum = 0
# 构造完整的文件路径
file_path = os.path.join(folder_path, filename)
# # 检查文件是否是图片文件,可以根据文件扩展名来判断
if file_path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
self.fileName1 = file_path
try:
for model, tm, im in self.predict():
print(self.save_dir)
self.model = model
det = self.predictor.output['det']
self.det = det
# 存储结果,向tableview和comboBox中添加信息
self.image_detections[self.fileName1] = det # 存储检测结果
self.image_path[filename] = file_path
self.image_tm[self.fileName1] = tm
self.image_allLabel[self.fileName1] = os.path.join(self.save_dir, filename)
box = det[:,:4] # 边界框
confidences = det[:,4] # 置信度
class_indices = det[:,5].astype(int) # 类序号
predicted_classes = [model.names[idx] for idx in class_indices] # 序号转成类别名
class_count = {}
new_predicted_classes = []
for box, confidences, class_name in zip(box,confidences, predicted_classes): # 修改重复的类别名
# 给相同类别的标签添加数字后缀
if class_name not in class_count:
class_count[class_name] = 0
class_count[class_name] += 1
#创建新标签
new_class_name = f"{class_name}{class_count[class_name] - 1}" if class_count[class_name] > 1 else class_name
new_predicted_classes.append(new_class_name)
pp = '(' + ', '.join([str(item) for item in box]) + ')'
self.sum += 1
# 向TableView添加信息
row = []
file=f'{self.save_dir}/crops/{new_class_name}/{new_class_name}.jpg'
row.append(QStandardItem(str(self.sum)))
row.append(QStandardItem(str(file)))
row.append(QStandardItem(str(self.sum)))
row.append(QStandardItem(str(new_class_name)))
row.append(QStandardItem(str(f'{confidences:.2f}')))
row.append(QStandardItem(str(pp)))
self.tb_model.appendRow(row)
self.comboBox.clear()
self.comboBox.addItem('全部')
self.comboBox.addItems(new_predicted_classes)
self.image_boBox[self.fileName1] = new_predicted_classes
self.image_sum[self.fileName1] = self.sum
self.image = self.predictor.annotator.result() # 标记所有的结果图
self.label_6.setText(str(self.sum)) # 总目标数
self.label_8.setText(str(round(tm*1000, 3))+'ms') # 用时
# 将numpy图片转成QImage
height, width, channels = self.image.shape
bytes_per_line = 3 * width
q_image = QImage(self.image.data, width, height, bytes_per_line, QImage.Format_BGR888)
# Convert QImage to QPixmap
pixmap = QPixmap.fromImage(q_image)
self.label_2.setPixmap(pixmap.scaled(self.label_2.size(), Qt.IgnoreAspectRatio, Qt.SmoothTransformation))
self.label_2.setText('') # Clear the label text
except Exception as e:
print(f"Error processing {filename}: {e}")
5.3 视频检测和摄像头检测
这部分我使用的是源代码对这两种流数据处理过程,通过__call__函数和yield返回程序处理过程中的中间数据,再对中间数据进行存储、显示等操作。代码与上面类似,注意这里没有使用Timer对数据帧定时处理。
6. 效果展示
对全部检测结果显示
对部分检测结果显示
到这里螺栓松动检测系统就完成了,如果有需要源代码的可以小刀。