1 准备工作
下载yolov5项目代码,并在本机上配置环境:正常按照requirements安装依赖包,下载权重文件即可。
2 训练数据
2.1 创建文件目录
在data文件夹下新建数据文件fire_data,fire_data文件下新建images和labels两个文件夹,文件夹内容如下:
-
images
- train -----放训练图片
- val -----放验证图片
-
labels —图像标注后的结果存放在这里
- train -----放训练label
- val -----放验证label
2.2 图像标记
通过labelimg工具,对图片进行标注,值得注意的是yolov5用的是yolo格式的文件:
第一个数字为标签序号,后四个数值为坐标【(x,y,w,h),其中(x,y)为矩形中心坐标,w为矩形相对宽度,h为矩形相对高度。(x,y,w,h)均进行了归一化】。
在标注前将VOC格式转换为YOLO即可:
2.3 新建yaml配置文件
在data文件夹下新建一个fire.yaml配置文件(可参考coco128.yaml文件)
# train and val datasets (image directory or *.txt file with image paths)
train: ./data/fire_data/images/train/
val: ./data/fire_data/images/val/
# number of classes
nc: 2
# class names
names: ['fire', 'smog']
同时将models/yolov5s.yaml文件中的类别数nc修改为2(我标记了两类fire和smog),这里我用的是yolov5s模型,如果你选则使用其他如yolov5l、yolov5m、yolov5x则修改对应yaml文件中的nc值。
3 训练模型
3.1 train.py
修改参数为前面设置好的配置文件:
3.2 test.py
同样的,因为train.py中调用了test.py,因此也要在test.py中修改参数。
3.3 运行train.py
训练时间有点长,默默等待就好啦!
训练完成后,训练好的模型在weights文件夹下,分别是last.pt和best.pt,可运行detect.py验证啦~
效果如下,开心( * ^ ▽ ^ * )