YOLOV5 布料瑕疵检测

红亮

已于 2023-11-16 15:58:19 修改

阅读量9.5k

点赞数 13

文章标签： python

于 2022-04-11 17:22:45 首次发布

本文链接：https://blog.csdn.net/YEHUOLIANG/article/details/124101787

版权

1.前言

1.YOLOV5程序：https://github.com/ultralytics/yolov5

2.数据集convertor程序（后面附了网盘链接下载）https://github.com/datawhalechina/team-learning-cv/tree/master/DefectDetection

天池官方就是用的YOLOV5，但是里面yolov5版本老，这里只引用项目里的convertTrainLabel.py程序，下文叙述这个程序的修改。

import numpy as np # linear algebra
import os
import json
from tqdm.auto import tqdm
import shutil as sh
import cv2

josn_path = "./train_data/guangdong1_round2_train2_20191004_Annotations/Annotations/anno_train.json"
image_path = "./train_data/guangdong1_round2_train2_20191004_images/defect/"

name_list = []
image_h_list = []
image_w_list = []
c_list = []
w_list = []
h_list = []
x_center_list = []
y_center_list = []

with open(josn_path, 'r') as f:
    temps = tqdm(json.loads(f.read()))
    for temp in temps:
        # image_w = temp["image_width"]
        # image_h = temp["image_height"]
        name = temp["name"].split('.')[0]
        path = os.path.join(image_path, name, temp["name"])
        # print('path: ',path)
        im = cv2.imread(path)
        sp = im.shape
        image_h, image_w = sp[0], sp[1]
        # print("image_h, image_w: ", image_h, image_w)
        # print("defect_name: ",temp["defect_name"])
        #bboxs
        x_l, y_l, x_r, y_r = temp["bbox"]
        # print(temp["name"], temp["bbox"])
        if temp["defect_name"]=="沾污":
            defect_name = '0'
        elif temp["defect_name"]=="错花":
            defect_name = '1'
        elif temp["defect_name"] == "水印":
            defect_name = '2'
        elif temp["defect_name"] == "花毛":
            defect_name = '3'
        elif temp["defect_name"] == "缝头":
            defect_name = '4'
        elif temp["defect_name"] == "缝头印":
            defect_name = '5'
        elif temp["defect_name"] == "虫粘":
            defect_name = '6'
        elif temp["defect_name"] == "破洞":
            defect_name = '7'
        elif temp["defect_name"] == "褶子":
            defect_name = '8'
        elif temp["defect_name"] == "织疵":
            defect_name = '9'
        elif temp["defect_name"] == "漏印":
            defect_name = '10'
        elif temp["defect_name"] == "蜡斑":
            defect_name = '11'
        elif temp["defect_name"] == "色差":
            defect_name = '12'
        elif temp["defect_name"] == "网折":
            defect_name = '13'
        elif temp["defect_name"] == "其他":
            defect_name = '14'
        else:
            defect_name = '15'
            print("----------------------------------error---------------------------")
            raise("erro")
        # print(image_w, image_h)
        # print(defect_name)
        x_center = (x_l + x_r)/(2*image_w)
        y_center = (y_l + y_r)/(2*image_h)
        w = (x_r - x_l)/(image_w)
        h = (y_r - y_l)/(image_h)
        # print(x_center, y_center, w, h)
        name_list.append(temp["name"])
        c_list.append(defect_name)
        image_h_list.append(image_w)
        image_w_list.append(image_h)
        x_center_list.append(x_center)
        y_center_list.append(y_center)
        w_list.append(w)
        h_list.append(h)

    index = list(set(name_list))
    print(len(index))
    for fold in [0]:
        val_index = index[len(index) * fold // 5:len(index) * (fold + 1) // 5]
        print(len(val_index))
        for num, name in enumerate(name_list):
            print(c_list[num], x_center_list[num], y_center_list[num], w_list[num], h_list[num])
            row = [c_list[num], x_center_list[num], y_center_list[num], w_list[num], h_list[num]]
            if name in val_index:
                path2save = 'val/'
            else:
                path2save = 'train/'
            # print('convertor\\fold{}\\labels\\'.format(fold) + path2save)
            # print('convertor\\fold{}/labels\\'.format(fold) + path2save + name.split('.')[0] + ".txt")
            # print("{}/{}".format(image_path, name))
            # print('convertor\\fold{}\\images\\{}\\{}'.format(fold, path2save, name))
            if not os.path.exists('convertor/fold{}/labels/'.format(fold) + path2save):
                os.makedirs('convertor/fold{}/labels/'.format(fold) + path2save)
            with open('convertor/fold{}/labels/'.format(fold) + path2save + name.split('.')[0] + ".txt", 'a+') as f:
                for data in row:
                    f.write('{} '.format(data))
                f.write('\n')
                if not os.path.exists('convertor/fold{}/images/{}'.format(fold, path2save)):
                    os.makedirs('convertor/fold{}/images/{}'.format(fold, path2save))
                sh.copy(os.path.join(image_path, name.split('.')[0], name),
                        'convertor/fold{}/images/{}/{}'.format(fold, path2save, name))

3.数据集：数据集来源于天池2019广东工业智造创新大赛【赛场一】

数据集-阿里云天池

4.借鉴了这篇文章从零开始手把手教你利用yolov5训练自己的数据集(含coco128数据集/yolov5权重文件国内下载）更新于20200728_orangezs的博客-CSDN博客_coco128数据集

天池官方的的数据集和代码

链接：https://pan.baidu.com/s/1OgTERkCMCpVCkW5Doux2gg?pwd=onmo
提取码：onmo

2.数据处理

1.这里数据集只用到了天池提供的train2数据集，在yolov5项目文件夹中新建train_data文件夹，将train2数据集放在train_data文件夹内。

2.数据集转换为coco格式：

这里需要将天池官方程序中的convertTrainLabel.py放入yolov5文件夹中，并做以下修改：

(1).8，9行：path位置修改为:(这里我修改了数据集文件夹的名字)

josn_path = "./train_data/round2_train2/Annotations/anno_train.json"
image_path = "./train_data/round2_train2/defect_Images/"

(2)26行：

       path = os.path.join(image_path,temp["name"])

(2).36-69

        if temp["defect_name"]=="无疵点":
            defect_name = '0'
        elif temp["defect_name"]=="破洞":
            defect_name = '1'
        elif temp["defect_name"] == "水渍":
            defect_name = '2'
        elif temp["defect_name"] == "油渍":
            defect_name = '2'
        elif temp["defect_name"] == "污渍":
            defect_name = '2'
        elif temp["defect_name"] == "三丝":
            defect_name = '3'
        elif temp["defect_name"] == "结头":
            defect_name = '4'
        elif temp["defect_name"] == "花板跳":
            defect_name = '5'
        elif temp["defect_name"] == "百脚":
            defect_name = '6'
        elif temp["defect_name"] == "毛粒":
            defect_name = '7'
        elif temp["defect_name"] == "粗经":
            defect_name = '8'
        elif temp["defect_name"] == "松经":
            defect_name = '9'
        elif temp["defect_name"] == "断经":
            defect_name = '10'
        elif temp["defect_name"] == "吊经":
            defect_name = '11'
        elif temp["defect_name"] == "粗维":
            defect_name = '12'
        elif temp["defect_name"] == "纬缩":
            defect_name = '13'
        elif temp["defect_name"] == "浆斑":
            defect_name = '14'
        elif temp["defect_name"] == "整经结":
            defect_name = '15'
        elif temp["defect_name"] == "星跳":
            defect_name = '16'
        elif temp["defect_name"] == "跳花":
            defect_name = '16'
        elif temp["defect_name"] == "断氨纶":
            defect_name = '17'
        elif temp["defect_name"] == "稀密档":
            defect_name = '18'
        elif temp["defect_name"] == "浪纹档":
            defect_name = '18'
        elif temp["defect_name"] == "色差档":
            defect_name = '18'
        elif temp["defect_name"] == "磨痕":
            defect_name = '19'
        elif temp["defect_name"] == "轧痕":
            defect_name = '19'
        elif temp["defect_name"] == "修痕":
            defect_name = '19'
        elif temp["defect_name"] == "烧毛痕":
            defect_name = '19'
        elif temp["defect_name"] == "死皱":
            defect_name = '20'
        elif temp["defect_name"] == "云织":
            defect_name = '20'
        elif temp["defect_name"] == "双维":
            defect_name = '20'
        elif temp["defect_name"] == "双经":
            defect_name = '20'
        elif temp["defect_name"] == "跳纱":
            defect_name = '20'
        elif temp["defect_name"] == "筘路":
            defect_name = '20'
        elif temp["defect_name"] == "纬纱不良":
            defect_name = '20'

147行

                sh.copy(os.path.join(image_path, name),
                        'convertor/fold{}/images/{}/{}'.format(fold, path2save, name))

4.终端运行：python convertTrainLabel.py

5.把

3.对coco128.yaml的修改：

1）训练图像目录的路径（或带有训练图像列表的* .txt文件的路径）

2）与我们的验证图像相同的路径

3）类数

4）类列表名称：这里用21个数字代表了不同的瑕疵，后续再改。

path: ../coco128  # dataset root dir  /home/csim/HL/yolov5/coco
train: images/train  # train images (relative to 'path') 128 images
val: images/train  # val images (relative to 'path') 128 images
test:  # test images (optional)

# Classes
nc: 21  # number of classes
names: ['0','1', '2', '3', '4', '5', '6', '7', '8', '9', '10',
        '11', '12', '13', '14', '15', '16', '17', '18', '19', '20']  # class names

3.requirements

安装包：pip install -U -r requirements.txt

4.模型

从./models文件夹中，选择模型。yolov5提供了五个模型，这里我们以yolov5s.yaml为例子，双击打开，更新其中参数适应2.1中定义的类别。这里一般不用动更新nc就可以了。就是匹配你自己数据集的类别就可以了，其他的参数一般不做优化不需要动。

#parameters
nc: 21
#number of classes 更新匹配自己的数据集
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple

5.训练

可以从0开始训练也可以加载–cfg yolov5s.yaml --weights 通过传递匹配的权重文件从预训练的检查点进行训练：–cfg yolov5s.yaml --weights yolov5s.pt

$ python train.py --img 640 --batch 16 --epochs 5 --data ./data/coco128.yaml --cfg ./models/yolov5s.yaml --weights ''

train.py：

def parse_opt(known=False):
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
    parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
    parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml', help='hyperparameters 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, -1 for autobatch')
    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
    parser.add_argument('--rect', action='store_true', help='rectangular training')
    parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
    parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
    parser.add_argument('--noval', action='store_true', help='only validate final epoch')
    parser.add_argument('--noautoanchor', action='store_true', help='disable AutoAnchor')
    parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
    parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
    parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"')
    parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
    parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
    parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer')
    parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
    parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
    parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name')
    parser.add_argument('--name', default='exp', help='save to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--quad', action='store_true', help='quad dataloader')
    parser.add_argument('--cos-lr', action='store_true', help='cosine LR scheduler')
    parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
    parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')
    parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone=10, first3=0 1 2')
    parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
    parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')

    # Weights & Biases arguments
    parser.add_argument('--entity', default=None, help='W&B: Entity')
    parser.add_argument('--upload_dataset', nargs='?', const=True, default=False, help='W&B: Upload data, "val" option')
    parser.add_argument('--bbox_interval', type=int, default=-1, help='W&B: Set bounding-box image logging interval')
    parser.add_argument('--artifact_alias', type=str, default='latest', help='W&B: Version of dataset artifact to use')

    opt = parser.parse_known_args()[0] if known else parser.parse_args()
    return opt

epochs：指的就是训练过程中整个数据集将被迭代多少次,显卡不行你就调小点。
batch-size：一次看完多少张图片才进行权重更新，梯度下降的mini-batch,显卡不行你就调小点。
cfg：存储模型结构的配置文件
data：存储训练、测试数据的文件
img-size：输入图片宽高,显卡不行你就调小点。
rect：进行矩形训练
resume：恢复最近保存的模型开始训练
nosave：仅保存最终checkpoint
notest：仅测试最后的epoch
evolve：进化超参数
bucket：gsutil bucket
cache-images：缓存图像以加快训练速度
weights：权重文件路径
name：重命名results.txt to results_name.txt
device：cuda device, i.e. 0 or 0,1,2,3 or cpu
adam：使用adam优化
multi-scale：多尺度训练，img-size +/- 50%
single-cls：单类别的训练集

6.可视化

待更新。