【YOLOv7】训练自己的数据集 - 目标检测

1. 准备深度学习环境

本人的笔记本环境 Ubuntu16.04, cuda-8.0, conda-4.3.8

1.1 代码克隆

git clone https://github.com/WongKinYiu/yolov7.git

1.2 环境安装

## 创建conda环境 
conda create -n yolov7 python=3.6

## 安装包
source activate yolov7

## 使用 conda 单独安装 torch 和 torchvision
conda install pytorch==1.7.0 torchvision==0.8.0 torchaudio==0.7.0 cudatoolkit=10.2 -c pytorch

## 打开 requirements.txt, 注释掉 torch 和 torchvision, 使用 pip 安装其他工具包
# torch>=1.7.0,!=1.12.0
# torchvision>=0.8.1,!=0.13.0
pip install -r requirements.txt

1.3 权重下载

yolov7预训练权重网址

 mkdir weights
 cd weights
 wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7.pt
 wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-tiny.pt

1.4 环境测试

## CPU测试
python detect.py --weights models/yolov7-tiny.pt --conf 0.25 --img-size 640 --source inference/images/horses.jpg --device cpu

## GPU测试
python detect.py --weights models/yolov7-tiny.pt --conf 0.25 --img-size 640 --source inference/images/horses.jpg --device 0

2. 准备数据

2.1 yolo格式的数据集准备

• 使用labelImg工具制作YOLO目标检测数据集

  • 一张图片对应一个标签文件：a.jpg -> a.txt。其中a.txt文件中内容如下：

    ## class  x  y  w  h
    0 0.421094 0.512500 0.051562 0.091667
    0 0.514062 0.508333 0.034375 0.087500
    0 0.591406 0.551042 0.082812 0.039583
    

  • 数据保存目录

    - /data/YOLOLable/
        - train/
    		- images/
    			- a.jpg
    			- ......
    		- labels/
    			- a.txt
    			- ...... 
    	- val/
    		- images/
    			- aa.jpg
    			- ......
    		- labels/
    			- aa.txt
    			- ...... 
    	- test/
    	    - images/
    			- aaa.jpg
    			- ......
    		- labels/
    			- aaa.txt
    			- ...... 
    

3. 配置训练的相关文件

3.1 修改数据配置文件

• 在 yolov7/data 目录下创建test.yaml （yolov7/data/test.yaml）

  # train and val data as
  train: ./datasets/train.txt
  val: ./datasets/val.txt
  test: ./datasets/test.txt
  
  # number of classes
  nc: 1
  
  # class names
  names: [ 'book']
  

• ./datasets/train.txt：每一行为训练图片的绝对路径

  /data/YOLOLable/train/images/a.jpg
  .......
  

3.2 修改模型配置文件

• 在 yolov7/cfg/training 目录下复制yolov7.yaml，并命名为yolov7-test.yaml （yolov7/data/yolov7-test），修改 nc 的值

   # parameters
   # nc: 80  # number of classes
   nc: 1     # 你的数据集中类别的个数
   ...
  

4. 模型训练和测试

模型训练和测试可以根据yolov7根目录中 readme.md 文件描述的进行。

4.1 模型训练

## p5 models
python train.py --weights weights/yolov7.pt --data data/test.yaml --epochs 100 --batch-size 8 --cfg cfg/training/yolov7-test.yaml --workers 0 --device 0 --img-size 640 640

4.2 模型测试

python test.py --data data/test.yaml --img 640 --batch 8 --conf 0.001 --iou 0.65 --device 0 --weights runs/train/exp/weights/best.pt