yolov3-tiny训练自己的数据集TF2实现

简介

YOLOv3是一种基于深度神经网络的对象识别和定位算法，其最大的特点是运行速度很快，可以用于实时系统。而YOLOv3-tiny是YOLOv3的简化版。

YOLOv3-tiny网络结构图

layer     filters    size              input                output
0 conv     16  3 x 3 / 1   416 x 416 x   3   ->   416 x 416 x  16  0.150 BFLOPs
1 max          2 x 2 / 2   416 x 416 x  16   ->   208 x 208 x  16
2 conv     32  3 x 3 / 1   208 x 208 x  16   ->   208 x 208 x  32  0.399 BFLOPs
3 max          2 x 2 / 2   208 x 208 x  32   ->   104 x 104 x  32
4 conv     64  3 x 3 / 1   104 x 104 x  32   ->   104 x 104 x  64  0.399 BFLOPs
5 max          2 x 2 / 2   104 x 104 x  64   ->    52 x  52 x  64
6 conv    128  3 x 3 / 1    52 x  52 x  64   ->    52 x  52 x 128  0.399 BFLOPs
7 max          2 x 2 / 2    52 x  52 x 128   ->    26 x  26 x 128
8 conv    256  3 x 3 / 1    26 x  26 x 128   ->    26 x  26 x 256  0.399 BFLOPs
9 max          2 x 2 / 2    26 x  26 x 256   ->    13 x  13 x 256
10 conv    512  3 x 3 / 1    13 x  13 x 256   ->    13 x  13 x 512  0.399 BFLOPs
11 max          2 x 2 / 1    13 x  13 x 512   ->    13 x  13 x 512
12 conv   1024  3 x 3 / 1    13 x  13 x 512   ->    13 x  13 x1024  1.595 BFLOPs
13 conv    256  1 x 1 / 1    13 x  13 x1024   ->    13 x  13 x 256  0.089 BFLOPs
14 conv    512  3 x 3 / 1    13 x  13 x 256   ->    13 x  13 x 512  0.399 BFLOPs
15 conv    255  1 x 1 / 1    13 x  13 x 512   ->    13 x  13 x 255  0.044 BFLOPs
16 yolo
17 route  13
18 conv    128  1 x 1 / 1    13 x  13 x 256   ->    13 x  13 x 128  0.011 BFLOPs
19 upsample            2x    13 x  13 x 128   ->    26 x  26 x 128
20 route  19 8
21 conv    256  3 x 3 / 1    26 x  26 x 384   ->    26 x  26 x 256  1.196 BFLOPs
22 conv    255  1 x 1 / 1    26 x  26 x 256   ->    26 x  26 x 255  0.088 BFLOPs
23 yolo

环境条件

• Window 10
• GTX1050 2G
• Python 3.6.7
• Opencv 3.4.2
• Tensorflow-gpu 2.4.1

源码获取

https://github.com/FreemanTang/yolov3_tiny_tf2
# 或者
https://gitee.com/freemantang/yolov3_tiny_tf2

项目结构

准备数据

将所有标注的图片放到/data/voc2012_raw/VOCdevkit/VOC2012/JPEGImages目录下，如下图。

将所有标注的XML文件放到/data/voc2012_raw/VOCdevkit/VOC2012/Annotations目录下，如下图。

运行/data/voc2012_raw/VOCdevkit/VOC2012/ImageName_to_txt.py代码，如下图。

运行成功，会在/data/voc2012_raw/VOCdevkit/VOC2012/ImageSets/Main目录下，生成4个TXT文件，如下图。

相关配置

将/data/voc2012.names用记事本打开，输入自己所标注的类名，如下图。

数据预处理

# 预处理训练数据集
python voc2012.py --data_dir ./VOCdevkit/VOC2012 --split train --output_file ./data/voc2012_train.tfrecord

运行成功，会生成一个voc2012_train.tfrecord文件

# 预处理验证数据集
python voc2012.py --data_dir ./VOCdevkit/VOC2012 --split val --output_file ./data/voc2012_val.tfrecord

运行成功，会生成一个voc2012_val.tfrecord文件

训练

运行yolov3_tiny_tf2/train.py代码，如下图。

python train.py --dataset ./data/voc2012_train.tfrecord --val_dataset ./data/voc2012_val.tfrecord --classes ./data/voc2012.names --num_classes 4 --batch_size 1 --epochs 20

训练完成后，会生成一个保存了权重的checkpoints文件夹，如下图。

测试

Window系统下测试

图片测试

python detect.py --classes ./data/voc2012.names --weights ./checkpoints/yolov3_tiny_train_5.tf  --tiny --image ./data/head.jpg --num_classes 4

视频测试

python detect_video.py --classes ./data/voc2012.names --weights ./checkpoints/yolov3_tiny_train_5.tf  --tiny --video test_video.mp4 --num_classes 4

摄像头测试

python detect_video.py --classes ./data/voc2012.names --weights ./checkpoints/yolov3_tiny_train_5.tf  --tiny --video 0 --num_classes 4

Jetson Nano下测试

将训练好的项目，直接用U盘打包移植到Jetson Nano，即可！

图片测试

python3 detect.py --classes ./data/voc2012.names --weights ./checkpoints/yolov3_tiny_train_5.tf  --tiny --image ./data/head.jpg --num_classes 4

视频测试

python3 detect_video.py --classes ./data/voc2012.names --weights ./checkpoints/yolov3_tiny_train_5.tf  --tiny --video test_video.mp4 --num_classes 4

摄像头测试

python3 detect_video.py --classes ./data/voc2012.names --weights ./checkpoints/yolov3_tiny_train_5.tf  --tiny --video CSI --num_classes 4