使用深度学习目标检测框架yolov5训练使用电力巡线安全帽检测数据集
2分类包含未带安全帽，带安全帽。共计图片295张2分类。
附划分数据集的脚本。未对数据进行增强，包含xml和txt。
数据集介绍：目前安全帽检测大多针对工地和大型机械，电力行业的安全帽检测数据集较少，这里挑选了一个安全帽检测数据集，大部分图片来自电力行业，包括电杆安装、电力高架施工、电力抢修等场景；小部分来自其他行业
包括电杆安装、电力高架施工、电力抢修等场景；小部分来自其他行业
以下代码仅供参考：
基于电力巡线安全帽检测数据集进行目标检测任务。数据集包含2个类别（未带安全帽和带安全帽），共计295张图片，并且已经划分了训练集、验证集和测试集。数据集提供了XML和TXT格式的标注文件。
环境准备
确保您已经安装了以下软件和库：
- Python 3.8 或更高版本
- PyTorch 1.9 或更高版本
- torchvision 0.10 或更高版本
- OpenCV
- numpy
- pandas
- matplotlib
- lxml（用于解析XML文件）
- albumentations（用于数据增强）
您可以使用以下命令安装所需的Python库：
pipinstalltorch torchvision opencv-python numpy pandas matplotlib lxml albumentations数据集准备
假设您的数据集已经按照YOLO格式组织好，并且包含训练集、验证集和测试集。以下是数据集的预期结构：
datasets/ └── power_line_helmet_detection/ ├── images/ │ ├── train/ │ ├── val/ │ └── test/ ├── labels/ │ ├── train/ │ ├── val/ │ └── test/ └── annotations/ ├── train/ ├── val/ └── test/同时，有一个classes.txt文件包含类别名称，每行一个类别名称。
类别文件 (classes.txt)
no_helmet helmet划分数据集的脚本
假设您已经有一个脚本来划分数据集，我们将其命名为split_dataset.py。以下是一个示例脚本：
[<title="Data Splitting Script">]importosimportshutilimportrandomfromsklearn.model_selectionimporttrain_test_splitdefcreate_directory_structure(root_dir):forsplitin['train','val','test']:os.makedirs(os.path.join(root_dir,'images',split),exist_ok=True)os.makedirs(os.path.join(root_dir,'labels',split),exist_ok=True)os.makedirs(os.path.join(root_dir,'annotations',split),exist_ok=True)defmove_files(source_dir,dest_dir,files):forfileinfiles:shutil.move(os.path.join(source_dir,file),os.path.join(dest_dir,file))defsplit_dataset(data_root,train_ratio=0.7,val_ratio=0.15,test_ratio=0.15):image_files=[fforfinos.listdir(os.path.join(data_root,'images'))iff.endswith('.jpg')]# Shuffle the filesrandom.shuffle(image_files)# Split the filestrain_files,temp_files=train_test_split(image_files,test_size=(1-train_ratio))val_files,test_files=train_test_split(temp_files,test_size=test_ratio/(val_ratio+test_ratio))# Move files to respective directoriesmove_files(os.path.join(data_root,'images'),os.path.join(data_root,'images/train'),train_files)move_files(os.path.join(data_root,'images'),os.path.join(data_root,'images/val'),val_files)move_files(os.path.join(data_root,'images'),os.path.join(data_root,'images/test'),test_files)move_files(os.path.join(data_root,'labels'),os.path.join(data_root,'labels/train'),[f.replace('.jpg','.txt')forfintrain_files])move_files(os.path.join(data_root,'labels'),os.path.join(data_root,'labels/val'),[f.replace('.jpg','.txt')forfinval_files])move_files(os.path.join(data_root,'labels'),os.path.join(data_root,'labels/test'),[f.replace('.jpg','.txt')forfintest_files])move_files(os.path.join(data_root,'annotations'),os.path.join(data_root,'annotations/train'),[f.replace('.jpg','.xml')forfintrain_files])move_files(os.path.join(data_root,'annotations'),os.path.join(data_root,'annotations/val'),[f.replace('.jpg','.xml')forfinval_files])move_files(os.path.join(data_root,'annotations'),os.path.join(data_root,'annotations/test'),[f.replace('.jpg','.xml')forfintest_files])if__name__=='__main__':data_root='../datasets/power_line_helmet_detection/'create_directory_structure(data_root)split_dataset(data_root)转换XML到YOLO格式
如果您的数据集是XML格式，需要将其转换为YOLO格式。以下是一个示例脚本convert_xml_to_yolo.py:
[<title="Convert XML to YOLO Format">]importosimportxml.etree.ElementTreeasETimportcv2importnumpyasnpdefconvert(size,box):dw=1./size[0]dh=1./size[1]x=(box[0]+box[1])/2.0-1y=(box[2]+box[3])/2.0-1w=box[1]-box[0]h=box[3]-box[2]x=x*dw w=w*dw y=y*dh h=h*dhreturn(x,y,w,h)defconvert_annotation(xml_file,output_path,class_mapping):tree=ET.parse(xml_file)root=tree.getroot()size=root.find('size')w=int(size.find('width').text)h=int(size.find('height').text)withopen(output_path,'w')asout_file:forobjinroot.iter('object'):difficult=obj.find('difficult').text cls=obj.find('name').textifclsnotinclass_mappingorint(difficult)==1:continuecls_id=class_mapping[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))bb=convert((w,h),b)out_file.write(str(cls_id)+" "+" ".join([str(a)forainbb])+'\n')defmain(input_dir,output_dir,class_mapping):forfilenameinos.listdir(input_dir):iffilename.endswith('.xml'):xml_path=os.path.join(input_dir,filename)txt_filename=filename.replace('.xml','.txt')txt_path=os.path.join(output_dir,txt_filename)convert_annotation(xml_path,txt_path,class_mapping)if__name__=="__main__":input_dir='../datasets/power_line_helmet_detection/annotations/train/'output_dir='../datasets/power_line_helmet_detection/labels/train/'class_mapping={'no_helmet':0,'helmet':1}main(input_dir,output_dir,class_mapping)input_dir='../datasets/power_line_helmet_detection/annotations/val/'output_dir='../datasets/power_line_helmet_detection/labels/val/'main(input_dir,output_dir,class_mapping)input_dir='../datasets/power_line_helmet_detection/annotations/test/'output_dir='../datasets/power_line_helmet_detection/labels/test/'main(input_dir,output_dir,class_mapping)模型训练
我们将使用YOLOv5进行训练。首先，克隆YOLOv5仓库并设置环境。
gitclone https://github.com/ultralytics/yolov5.gitcdyolov5 pipinstall-rrequirements.txt准备配置文件
创建一个hyp.scratch.yaml文件来定义超参数：
# Hyperparameters for YOLOv5 training from scratchlr0:0.01# initial learning rate (SGD=1E-2, Adam=1E-3)lrf:0.1# final OneCycleLR learning rate (lr0 * lrf)momentum:0.937# SGD momentum/Adam beta1weight_decay:0.0005# optimizer weight decay 5e-4warmup_epochs:3.0# warmup epochs (fractions ok)warmup_momentum:0.8# warmup initial momentumwarmup_bias_lr:0.1# warmup initial bias lrbox:0.05# box loss gaincls:0.5# cls loss gaincls_pw:1.0# cls BCELoss positive_weightobj:1.0# obj loss gain (scale with pixels)obj_pw:1.0# obj BCELoss positive_weightiou_t:0.20# IoU training thresholdanchor_t:4.0# anchor-multiple thresholdfl_gamma:0.0# focal loss gamma (efficientDet default gamma=1.5)hsv_h:0.015# image HSV-Hue augmentation (fraction)hsv_s:0.7# image HSV-Saturation augmentation (fraction)hsv_v:0.4# image HSV-Value augmentation (fraction)degrees:0.0# image rotation (+/- deg)translate:0.1# image translation (+/- fraction)scale:0.5# image scale (+/- gain)shear:0.0# image shear (+/- deg)perspective:0.0# image perspective (+/- fraction), range 0-0.001flipud:0.0# image flip up-down (probability)fliplr:0.5# image flip left-right (probability)mosaic:1.0# image mosaic (probability)mixup:0.0# image mixup (probability)copy_paste:0.0# segment copy-paste (probability)paste_in:0.0# segment paste-in (probability)rect:0# rectangular trainingresume:false# resume training from last checkpointnosave:false# only save final checkpointnoval:false# only validate final epochnoautoanchor:true# disable AutoAnchorevolve:false# evolve hyperparametersbucket:''# gsutil bucketcache_images:false# cache images for faster trainingimage_weights:false# use weighted image selection for trainingsingle_cls:false# train multi-class data as single-classoptimizer:SGD# optimizer: SGD or AdamWsync_bn:false# use SyncBatchNorm, only available in DDP modeworkers:8# dataloader workers (max is number of CPU cores)freeze:0# freeze first n layersv5_metric:true# assume maximum recall as 1.0 in AP calculationmulti_scale:true# vary input size between 320 and 640 pixelsrect_training:false# rectangular trainingcos_lr:false# cosine LR schedulerclose_mosaic:1000# close mosaic borderscales:[0.5,1.5]# image size scalesaugment:true# augment dataverbose:false# verbose printseed:0# random seed for reproducibilitylocal_rank:-1# ddp device id (-1 for single gpu train)entity:null# W&B entityupload_dataset:False# upload dataset as W&B artifact tablebbox_interval:-1# W&B bounding box logging intervalartifact_alias:latest# version of dataset artifact to useproject:runs/train# save results to project/nameexist_ok:false# existing project/name ok, do not incrementquad:false# quadrilateral anchorslinear_assignment:false# use linear assignment for NMS创建一个data.yaml文件来定义数据集路径和类别：
train:../datasets/power_line_helmet_detection/images/train/val:../datasets/power_line_helmet_detection/images/val/nc:2# number of classesnames:['no_helmet','helmet']# list of class names训练模型
使用以下命令开始训练：
python train.py--img640--batch16--epochs50--datadata.yaml--cfgyolov5s.yaml--weightsyolov5s.pt--hyphyp.scratch.yaml结果评估
训练完成后，可以使用以下命令评估模型性能：
python val.py--datadata.yaml--weightsruns/train/exp/weights/best.pt--tasktest使用说明
配置路径：
- 确保
datasets/power_line_helmet_detection/目录结构正确。 - 确保
data.yaml中的路径和类别名称正确。
- 确保
运行脚本：
- 在终端中依次运行数据分割脚本、XML转YOLO脚本、训练脚本和评估脚本。
注意事项：
- 根据需要调整超参数和训练设置。
- 可以通过修改
data.yaml中的cfg参数来选择不同的YOLOv5模型架构（如yolov5m.yaml,yolov5l.yaml,yolov5x.yaml）。
示例
假设您的数据集文件夹结构如下：
datasets/ └── power_line_helmet_detection/ ├── images/ │ ├── all/ │ ├── train/ │ ├── val/ │ └── test/ ├── labels/ │ ├── train/ │ ├── val/ │ └── test/ └── annotations/ ├── all/ ├── train/ ├── val/ └── test/并且all/目录包含所有的图片和对应的XML标注文件。运行上述脚本后，您可以查看训练日志和最终的模型权重文件。
总结
这样的化，我们可以构建一个全面的安全帽检测系统，包括数据集准备、数据转换、模型训练和结果评估。以下是所有相关的代码文件：
- 数据分割脚本(
split_dataset.py) - XML转YOLO格式脚本(
convert_xml_to_yolo.py) - YOLOv5超参数配置文件(
hyp.scratch.yaml) - YOLOv5数据集配置文件(
data.yaml)
