如何使用Yolov8训练使用轴承缺陷数据集，3659张，轴承制造过程中的表面与结构缺陷

轴承缺陷数据集，3659张，提供yolo和voc两种所有代码仅供参考
轴承缺陷数据集，3659张，提供yolo和voc两种标注方式
8类，标注数量：
Casting_burr：铸造毛刺，750
crack：裂纹，1675
scratch：划痕，2446
pit：凹坑，843
Polished_casting：抛光铸件，2180
strain：应变，94
unpolished_casting：未抛光铸件，742
burr：毛刺，3
image num：图像数量，3659

数据准备、模型训练、评估和推理。整个项目结构和代码

完整项目结构

bearing_defect_detection/ ├── main.py ├── train.py ├── evaluate.py ├── infer.py ├── datasets/ │ ├── bearing_defects/ │ │ ├── Annotations/ │ │ ├── ImageSets/ │ │ │ └── Main/ │ │ │ ├── train.txt │ │ │ └── val.txt │ │ └── JPEGImages/ ├── best_bearing_defects.pt ├── requirements.txt └── data.yaml

文件内容

requirements.txt

opencv-python torch==1.9 ultralytics PyQt5

data.yaml

train:./datasets/bearing_defects/JPEGImages/trainval:./datasets/bearing_defects/JPEGImages/valtest:./datasets/bearing_defects/JPEGImages/testnc:8names:['Casting_burr','crack','scratch','pit','Polished_casting','strain','unpolished_casting','burr']

convert_voc_to_yolo.py

importosimportxml.etree.ElementTreeasETimportshutilimportcv2defxml_to_yolo(xml_file,image_width,image_height):yolo_labels=[]tree=ET.parse(xml_file)root=tree.getroot()forobjinroot.findall('object'):label=obj.find('name').text bbox=obj.find('bndbox')xmin=int(bbox.find('xmin').text)ymin=int(bbox.find('ymin').text)xmax=int(bbox.find('xmax').text)ymax=int(bbox.find('ymax').text)x_center=(xmin+xmax)/2.0/image_width y_center=(ymin+ymax)/2.0/image_height width=(xmax-xmin)/image_width height=(ymax-ymin)/image_height class_id={'Casting_burr':0,'crack':1,'scratch':2,'pit':3,'Polished_casting':4,'strain':5,'unpolished_casting':6,'burr':7}[label]yolo_labels.append(f"{class_id}{x_center}{y_center}{width}{height}")return'\n'.join(yolo_labels)defsplit_dataset(image_dir,annotations_dir,output_dir,train_ratio=0.8):images=[fforfinos.listdir(image_dir)iff.endswith('.jpg')]num_train=int(len(images)*train_ratio)train_images=images[:num_train]val_images=images[num_train:]withopen(os.path.join(output_dir,'ImageSets/Main/train.txt'),'w')asf:f.write('\n'.join([os.path.splitext(img)[0]forimgintrain_images]))withopen(os.path.join(output_dir,'ImageSets/Main/val.txt'),'w')asf:f.write('\n'.join([os.path.splitext(img)[0]forimginval_images]))defconvert_dataset(voc_dir,yolo_dir):annotations_dir=os.path.join(voc_dir,'Annotations')images_dir=os.path.join(voc_dir,'JPEGImages')yolo_labels_dir=os.path.join(yolo_dir,'labels')os.makedirs(yolo_labels_dir,exist_ok=True)os.makedirs(os.path.join(yolo_dir,'images'),exist_ok=True)os.makedirs(os.path.join(yolo_dir,'images/train'),exist_ok=True)os.makedirs(os.path.join(yolo_dir,'images/val'),exist_ok=True)os.makedirs(os.path.join(yolo_dir,'ImageSets/Main'),exist_ok=True)split_dataset(images_dir,annotations_dir,yolo_dir)forfilenameinos.listdir(annotations_dir):iffilename.endswith('.xml'):xml_file=os.path.join(annotations_dir,filename)image_filename=os.path.splitext(filename)[0]+'.jpg'image_path=os.path.join(images_dir,image_filename)image=cv2.imread(image_path)image_height,image_width,_=image.shape yolo_label=xml_to_yolo(xml_file,image_width,image_height)txt_filename=os.path.splitext(filename)[0]+'.txt'txt_file=os.path.join(yolo_labels_dir,txt_filename)withopen(txt_file,'w')asf:f.write(yolo_label)# Copy image to YOLO directorybase_image_dir=os.path.join(yolo_dir,'images')ifimage_filename.split('.')[0]in[line.strip()forlineinopen(os.path.join(yolo_dir,'ImageSets/Main/train.txt'))]:target_image_dir=os.path.join(base_image_dir,'train')else:target_image_dir=os.path.join(base_image_dir,'val')shutil.copy(image_path,target_image_dir)# 示例用法convert_dataset('./datasets/bearing_defects','./datasets/bearing_defects_yolo')

train.py

importtorchfromultralyticsimportYOLO# 设置随机种子以保证可重复性torch.manual_seed(42)# 定义数据集路径dataset_config='data.yaml'# 加载预训练的YOLOv8n模型model=YOLO('yolov8n.pt')# 训练模型results=model.train(data=dataset_config,epochs=50,imgsz=640,batch=16,name='bearing_defects',project='runs/train')# 评估模型metrics=model.val()# 保存最佳模型权重best_model_weights='runs/train/bearing_defects/weights/best.pt'print(f"最佳模型权重已保存到{best_model_weights}")

evaluate.py

fromultralyticsimportYOLO# 初始化YOLOv8模型model=YOLO('runs/train/bearing_defects/weights/best.pt')# 评估模型metrics=model.val()# 打印评估结果print(metrics)

infer.py

importsysimportcv2importnumpyasnpfromultralyticsimportYOLOfromPyQt5.QtWidgetsimportQApplication,QMainWindow,QFileDialog,QMessageBox,QLabel,QPushButtonfromPyQt5.QtGuiimportQImage,QPixmapfromPyQt5.QtCoreimportQTimerclassMainWindow(QMainWindow):def__init__(self):super(MainWindow,self).__init__()self.setWindowTitle("轴承缺陷检测")self.setGeometry(100,100,800,600)# 初始化YOLOv8模型self.model=YOLO('runs/train/bearing_defects/weights/best.pt')# 设置类别名称self.class_names=['Casting_burr','crack','scratch','pit','Polished_casting','strain','unpolished_casting','burr']# 创建界面元素self.label_display=QLabel(self)self.label_display.setGeometry(10,10,780,400)self.button_select_image=QPushButton("选择图片",self)self.button_select_image.setGeometry(10,420,150,30)self.button_select_image.clicked.connect(self.select_image)self.button_select_video=QPushButton("选择视频",self)self.button_select_video.setGeometry(170,420,150,30)self.button_select_video.clicked.connect(self.select_video)self.button_start_camera=QPushButton("开始摄像头",self)self.button_start_camera.setGeometry(330,420,150,30)self.button_start_camera.clicked.connect(self.start_camera)self.button_stop_camera=QPushButton("停止摄像头",self)self.button_stop_camera.setGeometry(490,420,150,30)self.button_stop_camera.clicked.connect(self.stop_camera)self.timer=QTimer()self.timer.timeout.connect(self.update_frame)self.cap=Noneself.results=[]defselect_image(self):options=QFileDialog.Options()file_path,_=QFileDialog.getOpenFileName(self,"选择图片","","图片 (*.jpg *.jpeg *.png);;所有文件 (*)",options=options)iffile_path:self.process_image(file_path)defprocess_image(self,image_path):frame=cv2.imread(image_path)results=self.model(frame)annotated_frame=self.draw_annotations(frame,results)self.display_image(annotated_frame)self.results.append((image_path,annotated_frame))defselect_video(self):options=QFileDialog.Options()file_path,_=QFileDialog.getOpenFileName(self,"选择视频","","视频 (*.mp4 *.avi);;所有文件 (*)",options=options)iffile_path:self.process_video(file_path)defprocess_video(self,video_path):self.cap=cv2.VideoCapture(video_path)whileself.cap.isOpened():ret,frame=self.cap.read()ifnotret:breakresults=self.model(frame)annotated_frame=self.draw_annotations(frame,results)self.display_image(annotated_frame)self.results.append((video_path,annotated_frame))ifcv2.waitKey(1)&0xFF==ord('q'):breakself.cap.release()defstart_camera(self):self.cap=cv2.VideoCapture(0)self.timer.start(30)defstop_camera(self):self.timer.stop()ifself.capisnotNone:self.cap.release()self.label_display.clear()defupdate_frame(self):ret,frame=self.cap.read()ifnotret:returnresults=self.model(frame)annotated_frame=self.draw_annotations(frame,results)self.display_image(annotated_frame)self.results.append(('camera',annotated_frame))defdraw_annotations(self,frame,results):forresultinresults:boxes=result.boxes.cpu().numpy()forboxinboxes:r=box.xyxy[0].astype(int)cls=int(box.cls[0])conf=box.conf[0]label=f"{self.class_names[cls]}{conf:.2f}"color=(0,255,0)cv2.rectangle(frame,(r[0],r[1]),(r[2],r[3]),color,2)cv2.putText(frame,label,(r[0],r[1]-10),cv2.FONT_HERSHEY_SIMPLEX,0.9,color,2)returnframedefdisplay_image(self,frame):rgb_image=cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)h,w,ch=rgb_image.shape bytes_per_line=ch*w qt_image=QImage(rgb_image.data,w,h,bytes_per_line,QImage.Format_RGB888)pixmap=QPixmap.fromImage(qt_image)self.label_display.setPixmap(pixmap.scaled(self.label_display.width(),self.label_display.height()))if__name__=="__main__":app=QApplication(sys.argv)window=MainWindow()window.show()sys.exit(app.exec_())

运行步骤总结

1. 克隆项目仓库（如果有的话）：

   gitclone https://github.com/yourusername/bearing_defect_detection.gitcdbearing_defect_detection

2. 安装依赖项：

   pipinstall-r requirements.txt

3. 转换数据集格式：

   python convert_voc_to_yolo.py

4. 训练模型：

   python train.py

5. 评估模型：

   python evaluate.py

6. 运行推理界面：

   python infer.py

操作界面

• 选择图片进行检测。
• 选择视频进行检测。
• 使用摄像头进行实时检测。
• 结果展示

你可以通过以下方式查看演示视频：

• 用上述步骤运行infer.py并按照界面上的按钮操作。

希望这些详细的信息和代码能够帮助你顺利实施和优化你的轴承缺陷检测系统。如果有其他需求或问题，请随时告知！

详细解释

requirements.txt

列出项目所需的所有Python包及其版本。

data.yaml

配置数据集路径和类别信息，用于YOLOv8模型训练。

convert_voc_to_yolo.py

将VOC格式的数据集转换为YOLO格式。读取XML标注文件并将其转换为YOLO所需的TXT标签格式。同时，将数据集分为训练集和验证集。

train.py

加载预训练的YOLOv8模型并使用自定义数据集进行训练。训练完成后评估模型并保存最佳模型权重。

evaluate.py

加载训练好的YOLOv8模型并对验证集进行评估，打印评估结果。

infer.py

创建一个GUI应用程序，支持选择图片、视频或使用摄像头进行实时检测，并显示检测结果。