yolov5训练自己的数据集(缺陷检测分类——看了就懂系列)

Posted 凌风探梅

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了yolov5训练自己的数据集(缺陷检测分类——看了就懂系列)相关的知识,希望对你有一定的参考价值。

这里记录一下将yolov5用于训练自己的数据集,这里v5的github地址

Environment

本文是直接 pip install -r requirement.txt直接安装成功,如果不容易成功,就一个个安装即可。如果出现报错,多关注一下torch的版本即可。

Data preparation

label的格式:
<object-class> <x_center> <y_center> <width> <height>

举个🌰:
414_735_0.txt

7 0.503 0.499 0.098 0.451
7 0.465 0.633 0.033 0.053

 

数据集格式
data
├──coco.yaml
├──coco128.yaml
├──defect_data
   ├──images
      ├──train
      ├──val
   ├──labels
      ├──train
      ├──val
├──…

然后制作自己的defect.yaml配置文件(直接复制coco.yaml修改或者新建一个文件将一下代码copy进去):

train: ./data/defect_data/images/train  # absolute_path_train_img
val: ./data/defect_data/images/val  # absolute_path_val_img

nc: 8 # 类别数量

names: [‘quexian1’, ‘quexian2’, ... ‘quexian8’]

修改models/xxx.yaml,有s, m, l, x四个模型供选择,需要修改类别nc数目;同时最好把anchors也修改为自己数据对应的anchors(需要自己提前聚类),不改也可train。

Pretrained Model

建议爬梯子下载最新的预训练权重,这个repo经常更新,可能和以前的预训练模型不适应。下载好放在weights下,在train.py里修改相应路径。

Training

配置train.py

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='weights/yolov5s.pt', help='initial weights path')
    parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model.yaml path')
    parser.add_argument('--data', type=str, default='data/defect.yaml', help='data.yaml path')
    parser.add_argument('--epochs', type=int, default=300)
    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')
    parser.add_argument('--device', default='0,1,2,3', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    opt = parser.parse_args()

 

训练好的模型会被保存在./runs/train/exp3/weights/last.pt和best.pt

Inference

配置detect.py

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default='runs/train/exp3/weights/last.pt', help='model.pt path(s)')
    parser.add_argument('--source', type=str, default='data/defect_data/images/val', help='source')  # file/folder, 0 for webcam
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--device', default='0,1,2,3', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    opt = parser.parse_args()

 

如果需要inference 自己数据集原图大小,需要修改utils/dataset.py中

class LoadImages:  # for inference
		# Padded resize
        img = letterbox(img0, new_shape=self.img_size)[0]

修改为:

class LoadImages:  # for inference
		# Padded resize
        img = letterbox(img0, new_shape=(math.ceil(img0.shape[0]/32)*32, math.ceil(img0.shape[1]/32)*32))[0]

Visualization of Training

利用tensorboard进行可视化,训练开始会在主目录生成一个runs文件,利用tensorboard打开:
tensorboard --logdir=./runs

Thanks

欢迎评论,每天回复☼

以上是关于yolov5训练自己的数据集(缺陷检测分类——看了就懂系列)的主要内容,如果未能解决你的问题,请参考以下文章

yolov5训练自己的数据集(缺陷检测分类——看了就懂系列)

yolov5训练自己的数据集(缺陷检测分类——看了就懂系列)

智能生活垃圾检测与分类系统(UI界面+YOLOv5+训练数据集)

笔记2:yolov5训练自己的目标检测模型_创建并划分数据集

深度学习目标检测---使用yolov5训练自己的数据集模型(Windows系统)

资源YOLOV5垃圾分类数据集(已标注好)