voc怎么转yolo,如何分割数据集为验证集,怎样检测CUDA可用性 并使用yolov8训练安全帽数据集且 构建基于yolov8深度学习 安全帽检测系统 数据格式转化
voc怎么转yolo,如何分割数据集为验证集,怎样检测CUDA可用性
安全帽数据集,5000张图片和对应的xml标签,
五千个yolo标签,到手即可训练。另外附四个常用小脚本,非常实用voc转yolo代码.py
分割数据集为验证集.py
检测cuda.py
批量重命名.py。继续
以下代码仅供参考!
5000张图片和对应的YOLO格式标签,并且还附带了一些实用的小脚本,我们可以直接进行模型训练、评估以及可视化。以下是详细的步骤:
当然可以。根据您的描述,您提到的四个实用小脚本分别是:
- VOC转YOLO代码
- 分割数据集为验证集
- 检测CUDA可用性
- 批量重命名
下面是每个脚本的详细实现。
1. VOC转YOLO代码
这个脚本用于将Pascal VOC格式的标注文件转换为YOLO格式。
[<title="Convert VOC Annotations to YOLO TXT Format">]importosimportxml.etree.ElementTreeasETfrompathlibimportPath# Define pathsbase_path=Path('datasets/helmet')annotations_dir=base_path/'Annotations'images_dir=base_path/'JPEGImages'output_labels_dir=base_path/'labels'# Create output directory if it doesn't existos.makedirs(output_labels_dir,exist_ok=True)# Class names and their corresponding IDsclass_names={'helmet':0}defconvert_annotation(xml_file):tree=ET.parse(xml_file)root=tree.getroot()image_width=int(root.find('size/width').text)image_height=int(root.find('size/height').text)label_lines=[]forobjinroot.findall('object'):class_name=obj.find('name').text bbox=obj.find('bndbox')xmin=float(bbox.find('xmin').text)ymin=float(bbox.find('ymin').text)xmax=float(bbox.find('xmax').text)ymax=float(bbox.find('ymax').text)# Convert bounding box to YOLO format (center_x, center_y, width, height)center_x=(xmin+xmax)/2.0/image_width center_y=(ymin+ymax)/2.0/image_height width=(xmax-xmin)/image_width height=(ymax-ymin)/image_height class_id=class_names[class_name]label_line=f"{class_id}{center_x}{center_y}{width}{height}\n"label_lines.append(label_line)returnlabel_linesforannotation_fileinannotations_dir.glob('*.xml'):image_name=annotation_file.stem+'.jpg'# Assuming images are in JPEG formatlabel_file=output_labels_dir/(annotation_file.stem+'.txt')label_lines=convert_annotation(annotation_file)withopen(label_file,'w')asf:f.writelines(label_lines)print("Conversion completed.")2. 分割数据集为验证集
这个脚本用于将数据集划分为训练集和验证集。
[<title="Split Dataset into Train and Validation Sets">]importosimportrandomfromsklearn.model_selectionimporttrain_test_splitfrompathlibimportPath# Define pathsbase_path=Path('datasets/helmet')images_dir=base_path/'JPEGImages'annotations_dir=base_path/'labels'train_images_dir=base_path/'images/train'train_labels_dir=base_path/'labels/train'val_images_dir=base_path/'images/val'val_labels_dir=base_path/'labels/val'# Create directories if they don't existos.makedirs(train_images_dir,exist_ok=True)os.makedirs(train_labels_dir,exist_ok=True)os.makedirs(val_images_dir,exist_ok=True)os.makedirs(val_labels_dir,exist_ok=True)# List all image filesimage_files=list(images_dir.glob('*.jpg'))# Adjust extension if necessary# Shuffle the image filesrandom.shuffle(image_files)# Split ratiostrain_ratio=0.8val_ratio=0.2# Calculate split indicesnum_images=len(image_files)train_split=int(num_images*train_ratio)# Split images and labelstrain_images=image_files[:train_split]val_images=image_files[train_split:]defcopy_files(source_images,dest_images_dir,dest_labels_dir):forimg_fileinsource_images:label_file=annotations_dir/(img_file.stem+'.txt')iflabel_file.exists():os.symlink(img_file,dest_images_dir/img_file.name)os.symlink(label_file,dest_labels_dir/label_file.name)copy_files(train_images,train_images_dir,train_labels_dir)copy_files(val_images,val_images_dir,val_labels_dir)print("Dataset splitting completed.")3. 检测CUDA可用性
这个脚本用于检测CUDA是否可用。
[<title="Check CUDA Availability">]importtorchdefcheck_cuda_availability():cuda_available=torch.cuda.is_available()ifcuda_available:print(f"CUDA is available. Device count:{torch.cuda.device_count()}")print(f"Current device:{torch.cuda.current_device()}")print(f"Device name:{torch.cuda.get_device_name(torch.cuda.current_device())}")else:print("CUDA is not available.")if__name__=="__main__":check_cuda_availability()4. 批量重命名
这个脚本用于批量重命名文件夹中的所有图像文件。
[<title="Batch Rename Images">]importosfrompathlibimportPathdefbatch_rename_images(directory,prefix='image'):files=sorted(os.listdir(directory))foridx,filenameinenumerate(files):iffilename.endswith(('.png','.jpg','.jpeg')):new_filename=f"{prefix}{idx}.jpg"old_file=directory/filename new_file=directory/new_filename os.rename(old_file,new_file)print(f'Renamed:{old_file}->{new_file}')if__name__=="__main__":directory=Path('datasets/helmet/JPEGImages')# Change this path to your images directorybatch_rename_images(directory)使用说明
VOC转YOLO代码(
convert_voc_to_yolo.py):- 确保
datasets/helmet/Annotations和datasets/helmet/JPEGImages目录存在,并且包含相应的XML标注文件和图像文件。 - 运行此脚本后,会在
datasets/helmet/labels目录下生成对应的TXT标签文件。
- 确保
分割数据集为验证集(
split_dataset.py):- 确保
datasets/helmet/JPEGImages和datasets/helmet/labels目录存在,并且包含相应的图像文件和TXT标签文件。 - 运行此脚本后,会在
datasets/helmet/images/train,datasets/helmet/images/val,datasets/helmet/labels/train, 和datasets/helmet/labels/val目录下分别存储训练集和验证集的数据。
- 确保
检测CUDA可用性(
check_cuda.py):- 运行此脚本以检查系统中CUDA是否可用以及相关信息。
批量重命名(
batch_rename.py):- 修改
directory变量指向包含图像文件的目录。 - 运行此脚本后,目录中的所有图像文件将被重命名为
image0.jpg,image1.jpg, …。
- 修改
这些脚本应该能满足您处理安全帽数据集的需求。如果有任何问题或需要进一步的帮助,请告诉我!
环境准备:
- 安装必要的库。
数据集组织:
- 确认数据集的结构和路径。
模型定义与训练:
- 使用YOLOv8进行目标检测。
评估与可视化:
- 评估模型性能。
- 可视化结果。
环境准备
首先,我们需要安装必要的库。您可以使用以下命令来设置环境:
pipinstallultralytics opencv-python-headless pandas scikit-learn数据集组织
假设您的数据集已经按照YOLO格式组织好,并且包含以下文件夹结构:
datasets/helmet/ ├── images/ │ ├── train/ │ │ ├── image0.jpg │ │ ├── image1.jpg │ │ └── ... │ └── val/ │ ├── image0.jpg │ ├── image1.jpg │ └── ... └── labels/ ├── train/ │ ├── image0.txt │ ├── image1.txt │ └── ... └── val/ ├── image0.txt ├── image1.txt └── ...并且每个图像都有一个对应的TXT标签文件。
创建YAML配置文件
确保您有一个正确的YAML配置文件helmet.yaml来描述数据集。以下是示例配置文件:
[<title="YOLOv8 Configuration File for Helmet Detection">]train:../datasets/helmet/images/trainval:../datasets/helmet/images/valnc:1names:['helmet']将上述内容保存为datasets/helmet/helmet.yaml。
模型定义与训练
我们将使用YOLOv8进行目标检测。以下是训练脚本train_detection.py:
[<title="Training Script for Helmet Detection using YOLOv8">]fromultralyticsimportYOLO# Load a modelmodel=YOLO('yolov8n.pt')# load a pretrained model (recommended for training)# Train the modelresults=model.train(data='../datasets/helmet/helmet.yaml',epochs=50,imgsz=640,batch=16,project='../runs/train',name='helmet_detection')# Evaluate the modelmetrics=model.val()results=model.export(format='onnx')# export the trained model to ONNX format评估与可视化
使用YOLOv8自带的评估脚本来评估目标检测模型。
[<title="Evaluation Script for Helmet Detection using YOLOv8">]fromultralyticsimportYOLO# Load the best modelbest_model=YOLO('../runs/train/helmet_detection/weights/best.pt')# Evaluate the model on the validation datasetmetrics=best_model.val(data='../datasets/helmet/helmet.yaml',conf=0.5,iou=0.45)print(metrics)用户界面
我们将使用 PyQt5 创建一个简单的 GUI 来加载和运行模型进行实时预测。以下是用户界面脚本ui.py:
[<title="PyQt5 Main Window for Helmet Detection">]importsysimportcv2importnumpyasnpfromPyQt5.QtWidgetsimportQApplication,QMainWindow,QLabel,QPushButton,QVBoxLayout,QWidget,QFileDialogfromPyQt5.QtGuiimportQImage,QPixmapfromPyQt5.QtCoreimportQt,QTimerfromultralyticsimportYOLO# Load modeldetection_model=YOLO('../runs/train/helmet_detection/weights/best.pt')classMainWindow(QMainWindow):def__init__(self):super().__init__()self.setWindowTitle("安全帽检测系统")self.setGeometry(100,100,800,600)self.initUI()definitUI(self):self.central_widget=QWidget()self.setCentralWidget(self.central_widget)self.layout=QVBoxLayout()self.image_label=QLabel(self)self.image_label.setAlignment(Qt.AlignCenter)self.layout.addWidget(self.image_label)self.load_image_button=QPushButton("加载图像",self)self.load_image_button.clicked.connect(self.load_image)self.layout.addWidget(self.load_image_button)self.start_prediction_button=QPushButton("开始预测",self)self.start_prediction_button.clicked.connect(self.start_prediction)self.layout.addWidget(self.start_prediction_button)self.stop_prediction_button=QPushButton("停止预测",self)self.stop_prediction_button.clicked.connect(self.stop_prediction)self.layout.addWidget(self.stop_prediction_button)self.central_widget.setLayout(self.layout)self.image_path=Noneself.timer=QTimer()self.timer.timeout.connect(self.update_frame)defload_image(self):options=QFileDialog.Options()file_name,_=QFileDialog.getOpenFileName(self,"选择图像文件","","Images (*.png *.jpg *.jpeg);;All Files (*)",options=options)iffile_name:self.image_path=file_name self.display_image(file_name)defdisplay_image(self,path):pixmap=QPixmap(path)scaled_pixmap=pixmap.scaled(self.image_label.width(),self.image_label.height(),Qt.KeepAspectRatio)self.image_label.setPixmap(scaled_pixmap)defstart_prediction(self):ifself.image_pathisnotNoneandnotself.timer.isActive():self.timer.start(30)# Update frame every 30 msdefstop_prediction(self):ifself.timer.isActive():self.timer.stop()self.image_label.clear()defupdate_frame(self):original_image=cv2.imread(self.image_path)image_rgb=cv2.cvtColor(original_image,cv2.COLOR_BGR2RGB)# Detectionresults=detection_model.predict(image_rgb,size=640,conf=0.5,iou=0.45)[0]forboxinresults.boxes.cpu().numpy():r=box.xyxy[0].astype(int)cls=int(box.cls[0])conf=box.conf[0]# Map class ID to nameclass_names=['安全帽']class_name=class_names[cls]# Draw bounding boxcv2.rectangle(image_rgb,(r[0],r[1]),(r[2],r[3]),(0,255,0),2)# Put textfont=cv2.FONT_HERSHEY_SIMPLEX cv2.putText(image_rgb,f'{class_name}({conf:.2f})',(r[0],r[1]-10),font,0.9,(0,255,0),2)h,w,ch=image_rgb.shape bytes_per_line=ch*w qt_image=QImage(image_rgb.data,w,h,bytes_per_line,QImage.Format_RGB888)pixmap=QPixmap.fromImage(qt_image)scaled_pixmap=pixmap.scaled(self.image_label.width(),self.image_label.height(),Qt.KeepAspectRatio)self.image_label.setPixmap(scaled_pixmap)if__name__=="__main__":app=QApplication(sys.argv)window=MainWindow()window.show()sys.exit(app.exec_())请确保将路径替换为您实际的路径。
使用说明
配置路径:
- 确保
datasets/helmet目录结构正确,并且包含images和labels子目录。 - 确保
runs/train/helmet_detection/weights/best.pt是训练好的 YOLOv8 模型权重路径。
- 确保
运行脚本:
- 在终端中运行
train_detection.py脚本来训练目标检测模型。 - 在终端中运行
evaluate_detection.py来评估目标检测模型性能。 - 在终端中运行
ui.py来启动 GUI 应用程序。
- 在终端中运行
注意事项:
- 确保所有必要的工具箱已安装,特别是
ultralytics和PyQt5。 - 根据需要调整参数,如
epochs和batch_size。
- 确保所有必要的工具箱已安装,特别是
示例
假设您的数据文件夹结构如下:
datasets/ └── helmet/ ├── images/ │ ├── train/ │ │ ├── image0.jpg │ │ ├── image1.jpg │ │ └── ... │ └── val/ │ ├── image0.jpg │ ├── image1.jpg │ └── ... └── labels/ ├── train/ │ ├── image0.txt │ ├── image1.txt │ └── ... └── val/ ├── image0.txt ├── image1.txt └── ...并且每个图像都有一个对应的TXT标签文件。运行ui.py后,您可以点击按钮来加载图像并进行安全帽检测。