yolov5实现小数据集的目标检测--kolektor缺陷数据集

要求:

      Python>=3.7.0 

      PyTorch>=1.7

我直接在colab上运行,直接不用管这些。

数据集:

可以使用自己做的数据集,我用的是kolektor数据集,下载地址https://www.vicos.si/Downloads/KolektorSDD

KSDD里面只有50个缺陷数据,我做了旋转,差不多一百个缺陷数据

(自己想做其他数据的检测原理也是一样的)

标注数据集就用labelImg标注,效果如下:

yolov5实现小数据集的目标检测--kolektor缺陷数据集_第1张图片

yolov5实现小数据集的目标检测--kolektor缺陷数据集_第2张图片

 一、在colab上直接克隆源码,再下载requirements.txt文件

yolov5实现小数据集的目标检测--kolektor缺陷数据集_第3张图片

% cd /content/yolov5
!git clone https://github.com/ultralytics/yolov5
!pip install -r requirements.txt 

 二、将数据集与yolov5文件放置同级(根据官网)

 分别将图片和txt文件放到images和labels文件夹内。

然后根据自己数据集的类别,在路径yolov5/data/下写一个含有后缀.yaml的配置文件,我的只需要检测一类缺陷,nc为1,names为类名。

下面是我定义的,test.yaml:

yolov5实现小数据集的目标检测--kolektor缺陷数据集_第4张图片

三、进行训练(可以选择不同的版本:yolov5n/s/m/l/x

 

! python train.py --img 640 --batch 30 --epochs 100 --data test.yaml --weights yolov5s.pt

可以自己定义batch、epochs

注意:在我第一次训练的时候,用的官网的batch 16 、epochs 3,导致我val_batch_pred.jpg的图片并没有预测信息,当我调大之后就有了。

训练之后的结果:

yolov5实现小数据集的目标检测--kolektor缺陷数据集_第5张图片

yolov5实现小数据集的目标检测--kolektor缺陷数据集_第6张图片

 然后运行detect.py ,选择自己的图片验证:

yolov5实现小数据集的目标检测--kolektor缺陷数据集_第7张图片

!python detect.py --weights runs/train/exp4/weights/last.pt --source data/images/Part0.jpg --device 0 --save-txt

有缺陷的:

yolov5实现小数据集的目标检测--kolektor缺陷数据集_第8张图片

无缺陷的:

yolov5实现小数据集的目标检测--kolektor缺陷数据集_第9张图片

你可能感兴趣的:(人工智能,目标检测,人工智能,计算机视觉,pytorch)