Yolov5-seg 训练自己的数据集_yolov5-seg训练

1.环境配置

我这里使用的版本是

yolov5-7.0：https://github.com/ultralytics/yolov5/tree/v7.0

labelme 3.16.7.

进入yolov5文件，找到 requirements.txt 文件，在pycharm终端运行 pip install -r requirements.txt

下载速度慢的可以配置镜像源。

2.使用labelme制作数据集

制作数据集时，可提前划分好训练和测试集

images存放原图

后面的json文件可以放在原图里面

labels存放转化之后的txt书

打开labelme，以多点的方式进行标注，标注完之后的文件为json文件，可以默认存在图片路径中

(因为这篇博客是这样弄的，当然你也可以分开)

数据转化json-->txt

因为yolov5训练需要txt格式，所以我们将保存好的json文件进行转化，以下代码针对的是json文件和图片文件在同一路径而使用：

import os, cv2, json
import numpy as np
 
classes = ['1'] # 修改成对应的类别
 
base_path = '../point/images' # 指定json和图片的位置
path_list = [i.split('.')[0] for i in os.listdir(base_path)]
for path in path_list:
    image = cv2.imread(f'{base_path}/{path}.jpg')
    h, w, c = image.shape
    with open(f'{base_path}/{path}.json') as f:
        masks = json.load(f)['shapes']
    with open(f'{base_path}/{path}.txt', 'w+') as f:
        for idx, mask_data in enumerate(masks):
            mask_label = mask_data['label']
            if '_' in mask_label:
                mask_label = mask_label.split('_')[0]
            mask = np.array([np.array(i) for i in mask_data['points']], dtype=np.float64)
            mask[:, 0] /= w
            mask[:, 1] /= h
            mask = mask.reshape((-1))
            if idx != 0:
                f.write('\n')
            f.write(f'{classes.index(mask_label)} {" ".join(list(map(lambda x:f"{x:.6f}", mask)))}')

json文件任务结束，可以移除，使用得到的txt数据，如果你划分好了训练集和测试集，接下来就可以进行训练

验证转化是否成功

可以运行以下代码进行检测是否成功转化

import cv2
import numpy as np
import glob
 
# 只需要给定图片文件夹和txt标签文件夹即可
pic_path = r"C:/Users/tensorflow/Desktop/point/images/train/"
txt_path = r"C:/Users/tensorflow/Desktop/point/labels/train/"
 
pic = glob.glob(pic_path + "*.jpg")
 
for pic_file in pic:
    img = cv2.imread(pic_file)
    # print("***:",pic_file)
    substrings = pic_file.split('/')
    substrings = substrings[-1].split('.')
    # print("***:",substrings)
    num = substrings[0].split("\\")[1]
    height, width, _ = img.shape
    txt_file = txt_path + num + ".txt"
    file_handle = open(txt_file)
    cnt_info = file_handle.readlines()
    print("***:", cnt_info)
    new_cnt_info = [line_str.replace("\n", "").split(" ") for line_str in cnt_info]
    # print("***:",new_cnt_info)
    color_map = [(0, 255, 255), (255, 0, 255), (255, 255, 0)]
    for new_info in new_cnt_info:
        s = []
        for i in range(1, len(new_info), 2):
            b = [float(tmp) for tmp in new_info[i:i + 2]]
            s.append([int(b[0] * width), int(b[1] * height)])
        class_ = new_info[0]
        index = int(class_)
        cv2.polylines(img, [np.array(s, np.int32)], True, color_map[index], thickness=3)
 
    save_path = 'labelme/all/' + num + '.jpg'
    cv2.imwrite(save_path, img)
    img = cv2.resize(img, (640, 640))
    cv2.imshow("{}".format(num), img)
    cv2.waitKey(0)

检查轮廓没有问题即为转化成功

3.更改配置文件

打开yolov5源码，找到文件

1.打开data文件下，找到coco128-seg.yaml (可以复制一份再进行更改)

2.打开model文件下segment/yolov5s-seg.yaml (可以复制一份再进行更改）

3.开始训练设置好参数，即可开始训练

至此，训练结束

4.使用训练好的模型进行检测，使用predict.py文件，设置好参数，开始检测