YOLOv3-Pytorch:制作并训练自己的数据集(single class)

WPJAM TOC

一、数据准备与预处理

使用labelimg对图像中需要识别的物体进行标识。下载地址

这里说一下可以提高效率的方法:

  • 创建一个空文件夹Annotations,点击Change save dir选择这个文件夹。
  • 在界面中的View选项中,勾选Auto save mode,这样可自动将XML存在Annotations文件夹内;
  • 如果需要识别的类别只有一种,勾选Single class mode,这样不需要每次都确认标签名了。
  • 标记图片是个枯燥且废手的事情,建议注意放松手指。

打开:yolo3 in pytorch下载地址

这里请注意,一定要clone下来,别下载zip,不然训练的时候会报错:“subprocess.CalledProcessError: Command 'git status -uno' returned non-zero exit status 128”,不过这个bug可通过删除里面相关git的代码解决。

在data目录下创建下面几个文件夹:

  1. Annotations,这个直接复制labelimg保存XML的那个即可;
  2. JPEGImges,把图片放这
  3. images,复制JPEGImages并重命名
  4. 空文件夹labels
  5. 空文件夹ImageSets

在项目根目录下创建并依次运行makeTxt.py,和voc_label.py,在data文件夹下创建*.names和*.data文件,此处参考该博客:YOLOV3训练自己的数据集(PyTorch版本)

下面需要修改cfg文件,在cfg文件夹下选择你需要选用的训练网络,yolo3和yolo3-tiny的区别很大,yolo3对显卡要求比较高,本人1050ti跑不动,无法解决“CUDA out of memory”的bug,所以选择小一点的yolo3-tiny,但是tiny在卷积的时候会丢失很多信息,所以谨慎使用。

在你选用的cfg中,查找yolo层的位置,将yolo层中的class设为你要识别的类的数量n。在每个[yolo]上方的那个[convolutional]中的filters设置为(n+4+1)*3,由于我的识别的类只有一个,所以设置成18。

二、训练模型

在命令行中执行以下脚本:

python train.py --data data/stomata.data --cfg cfg/yolov3-tiny.cfg --epochs 300 --single-cls --weights weights/yolov3-tiny.weights --name data/stomata.names --batch-size 8

如果不是single class 请去除--single-cls这个选项;stomata是我的类名,所以以此命名*.data和*.names;迭代次数请多一点。而且提一点,在cfg中修改batch是木有用的,需要提供命令修改batch_size,.cfg、.weight、.data、.names文件亦是如此。

三、查看训练结果和预测

训练结束后,在项目根目录下会产生results.png、results.txt、test_batch0.png、train_batch0.png四个文件,在weights目录下回产生best.pt和last.pt(或yourFileName.pt)。可能训练结束后命令行里会出现无法将last.pt重命名的报错,导致result.png无法输出,没关系,只要有results.txt,在Python Console中输入from utils import utils; utils.plot_results(),即可查看结果统计。如下图:

可以看到问题不大之后,把需要预测的图像放在data/sample下,在命令行中输入:

python detect.py  --cfg cfg/yolov3-tiny.cfg --weights weights/best.pt --names data/stomata.names

成功之后会在output目录下看到预测结果。框上的小数为置信度。

预测结果

发表评论

电子邮件地址不会被公开。 必填项已用*标注