Yolov5代码:https://github.com/ultralytics/yolov5
数据集:https://www.flir.com/oem/adas/adas-dataset-form/
运行环境:Ubuntu 16.04+CUDA10.1+Python3.6
1、创建环境
1.1 创建虚拟环境yolo
conda create -n yolo python=3.6
source activate yolo
1.2 解压yolov5-master,并进入文件夹
安装需要的各种库
pip install -r requirements.txt
默认装的torch版本应该是最新版本的1.9,torchvision是0.10,低版本CUDA可能无法支持,这会导致训练的时候默认调用CPU而不是GPU。
2、数据集制作
2.1 文件夹路径如下
--dataset
--train
--images
--labels
--val
--images
--labels
值得注意的是,images底下的图片名字(00001.jpg)和labels的groundtruth(00001.txt)名字是一一对应的。
原始的数据集thermal_annotations.json文件中包含了全部训练集和验证集的标注信息,但yolo是一个文件对应一张图片,需要自己生成上图所示的标签文件。
另外,thermal_annotations.json中的”image_id“与图片id不对应
train图片是从00001.jpg开始,而对应的thermal_annotations.json中的”image_id“是从00000开始
val图片是从08863.jpg开始,而对应的thermal_annotations.json中的”image_id“是从00000开始
需要矫正一下对应关系!!!
2.2 labels标签的txt文件生成
这一步是训练数据准备的关键,内容如下:
每一行有五个数字组成,类别、x、y,w,h。类别由0,1,2,3,…等组成(如果是4类,应该取值为0~3),xywh是归一化后的数值,且xy是目标中心坐标,假设bbox保存的是目标左上角的坐标和宽高,转化过程如下:
def convert(size, box):
dw = 1./(size[0])
dh = 1./(size[1])
x = (box[0] + box[1])/2.0 - 1
y = (box[2] + box[3])/2.0 - 1
w = box[1] - box[0]
h = box[3] - box[2]
x = x*dw
w = w*dw
y = y*dh
h = h*dh
return (x,y,w,h)
3、训练
3.1 主要修改两个文件
- /models/yolov5s.yaml,将nc = 80修改为nc = 4(本文为4类目标)
- data文件夹下新建一个xx.yaml(数据集路径),本文叫做FLIR.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: /data3/JJChen/YOLOtrainDataset/Flir/FLIR_ADAS_1_3/ # dataset root dir
train: train/images/ # train images (relative to 'path') 128 images
val: val/images/ # val images (relative to 'path') 128 images
test: # test images (optional)
# Classes
nc: 4 # number of classes
names: ['Dogs', 'People', 'Bicycles', 'Cars'] # class names
path+train即训练集图片的路径
3.2 开始训练
从零开始训练:python train.py --img 640 --epochs 300 --data data/FLIR.yaml --cfg models/yolov5s.yaml --weight " "
从预训练模型进行微调:python train.py --img 640 --epochs 300 --data data/FLIR.yaml --cfg models/yolov5s.yaml --weight weights/yolov5s.pt
开始训练界面如下:
可能出现的错误:无法调用GPU
运行时gpu_mem为0G,这是用CPU在训练,所以一开始速度贼慢。打开train.py,修改436‘device’为0,1,2,3,结果却报错:
这主要原因是cuda和torch版本不匹配,直接安装的torch是1.9的,所以得降版本
Pytorch1.7.1+torchvision0.8.2下载地址:https://download.pytorch.org/whl/torch_stable.html
卸载torch:pip uninstall torch
卸载torchvision:pip uninstall torchvison
本地安装:
pip install torch-1.7.1+cu101-cp36-cp36m-linux_x86_64.whl
pip install torchvision-0.8.2+cu101-cp36-cp36m-linux_x86_64.whl
重新训练:
4、测试
下载一个红外视频python detect.py --source xx.mp4 --weights runs/train/exp7/weights/best.pt