yolov5 环境搭建 完成后,就能自己训练模型了,这篇文章主要讲解使用 yolov5 实现SAR图像中舰船的检测。

1. 数据集制作

1.1 下载数据集

SSDD 的数据集可以在百度网盘,数据集比较下,不需要等待太久 👀
链接:https://pan.baidu.com/s/1aXeDJ1-bd3-Wftgk4_7glg
提取码:2022

1.2 数据集格式

我们手头上的 SSDD 数据集格式如下:

- SSDD
	- Annotations
		- 000001.xml
		- 000002.xml
		- ...
	- JPEGImages
		- 000001.JPG
		- 000002.JPG
		- ...

其中,Annotations为所有图片的标签,每一个 xml 标签文件内包含了目标类型,目标的位置(框的中心位置和宽高)


YOLOV5 不支持这种格式的数据集,我们可以把数据集修改为 COCO 格式:

- ships
	- images
		- train (JPG images)
			- 000341.JPG
			- 000024.JPG
			- ...
		- test (JPG images)
		- val (JPG images) 
	- labels
		- train (txt files)
			- 000341.txt
			- 000024.txt
			- ...
		- test (txt files)
		- val (txt files)
		- ...

可以看到,数据集包含 /images 和 /labels,images 中 包含三个文件夹分别存放三类数据集,对应的 labels 中也存放三类标签, 值得注意的是,labels 中不在是 xml 文件,而是 txt 文件,每一个 txt 文件内容如下:

格式:<object-class> <x> <y> <width> <height>
0  0.41416  0.61538  0.31607  0.11242

1.3 格式转换

处理起来比较麻烦,我写了一个函数,可以放到与原始的 SSDD 文件夹同一目录下,运行后自动生成对应 COCO 格式的数据集 ships_dataset,代码下载地址:

链接:https://pan.baidu.com/s/1GMODNViP0WaquqJUyme_YA
提取码:2022

运行后效果如下:

可以把生成的 ships_dataset 丢到 yolov5 目录内,待后续使用。


1.4 创建 yaml 文件

在yolov5源代码目录中可以找到 coco 数据集的 yaml文件,它用来说明数据集的位置和包含的样本类型数目:

内容如下:

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
# COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
# Example usage: python train.py --data coco128.yaml
# parent
# ├── yolov5
# └── datasets
#     └── coco128  ← downloads here (7 MB)


# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/coco128  # dataset root dir
train: images/train2017  # train images (relative to 'path') 128 images
val: images/train2017  # val images (relative to 'path') 128 images
test:  # test images (optional)

# Classes
nc: 80  # number of classes
names: ['person', 'bicycle', 'car', ...]  # class names

其中,path 指出了数据集存放的目录,train 指出了训练集图片 所在位置,val 指出了验证集图片 所在位置,test 不是必须的。

对应的,我们创建 ships.yaml,内容如下:

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ./ships_dataset  # dataset root dir
train: images/train  # train images (relative to 'path') 128 images
val: images/val  # val images (relative to 'path') 128 images
test: images/test  # test images (optional)

# Classes
nc: 1  # number of classes
names: ['ship']  # class names

2. 调试

使用下面的命令训练模型:

python train.py --img 320 --batch-size 8 --epochs 10 --data data/ships.yaml --cfg models/yolov5s.yaml --weights yolov5s.pt --device 0

命令中的参数说明可以参考官方解释这篇博客,注意此处指定了我们刚刚创建的 yaml 文件 (这些参数可以在 train.py 中设定默认值,也可以在命令中指定

第一次训练,需要下载一些小文件,代码自动通过 Github 官网下载,有可能因为无法访问 Github 报错,直接改代码:

yolov5/utils/downloads.py 中的所有 github.com 改为镜像网址 hub.xn--p8jhe.tw

response = requests.get(f'https://hub.xn--p8jhe.tw/repos/

这样就能正常执行了,跑 10 个epoch 可以检测到一些舰船,但精度不高,100 epoches 基本能达到 94% 左右的检测率。


3. 优化

由于还是小白,调参还没有经验,目前先通过图像增强来提高精度,参考博客

在 yolov5\data\hyps 中可以看到有很多超参数设置相关文件,创建自己的 hyp.ships.yaml,先复制 hyp.scratch-low.yaml 的全部内容,值修改图像增强部分:

hsv_h: 0.015  # image HSV-Hue augmentation (fraction)
hsv_s: 0.5  # image HSV-Saturation augmentation (fraction)
hsv_v: 0.3  # image HSV-Value augmentation (fraction)
degrees: 0  # image rotation (+/- deg)
translate: 0.1  # image translation (+/- fraction)
scale: 0.75  # image scale (+/- gain)
shear: 0.0  # image shear (+/- deg)
perspective: 0.003  # image perspective (+/- fraction), range 0-0.001
flipud: 0.4  # 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)

使用了图像增强后需要在 train.py 中把参数值设置为

parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.ships.yaml', help='hyperparameters path')

3. 训练与结果

使用了图像增强后,重新训练 150 个epoches(等待 3 小时💤):

150 epochs completed in 3.510 hours.
Optimizer stripped from runs\train\exp40\weights\last.pt, 14.3MB
Optimizer stripped from runs\train\exp40\weights\best.pt, 14.3MB

Validating runs\train\exp40\weights\best.pt...
Fusing layers...
YOLOv5s summary: 213 layers, 7012822 parameters, 0 gradients, 15.8 GFLOPs
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100%|██████████| 15/15 [00:04<0
                 all        232        578       0.95      0.865      0.947      0.526
Results saved to runs\train\exp40

可以在 yolov5/runs/train/exp40 中查看结果:

也可以拿训练好的 weight/best.pt 文件进行预测(下面预测所有测试集中的图片):

python detect.py --iou-thres 0.3 --weights ./runs/train/exp40/weights/best.pt --source ./ships_dataset/images/test --save-txt --save-crop

可以在 yolov5/runs/detect 中查看结果

Logo

为开发者提供学习成长、分享交流、生态实践、资源工具等服务,帮助开发者快速成长。

更多推荐