基于 YOLOV5 的 SAR 图像舰船检测

在 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     [email protected] [email protected]:.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 中查看结果