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]