yolov3-tiny训练自己的数据集TF2实现
原创
©著作权归作者所有:来自51CTO博客作者wx63dcd9d7dd8a8的原创作品,请联系作者获取转载授权,否则将追究法律责任
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