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 主要修改两个文件
  1. /models/yolov5s.yaml,将nc = 80修改为nc = 4(本文为4类目标)
  2. 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

参考博客:
FLIR Thermal Starter 数据集详解

Logo

更多推荐