yolov8分割随记

打标又打了一天，其实分割是比检测容易标注一点的，segment_anything真的是很强大的模型

先看看文件夹是怎样的

.json文件存储在json_labels文件夹内，images和labels就是后续存储训练集和测试集的对应文件夹了

得到120个.json文件后，第一步就是将.json文件转换为.txt文件，保留每张图片中每个实例对应多边形的坐标点--->创建一个json_to_txt.py文件

ini 复制代码

import json
import os
import argparse
from tqdm import tqdm
import glob
import cv2
import numpy as np


def convert_label_json(json_dir, save_dir, classes):
    json_paths = os.listdir(json_dir)
    classes = classes.split(',')

    for json_path in tqdm(json_paths):
        # for json_path in json_paths:
        path = os.path.join(json_dir, json_path)
        # print(path)
        with open(path, 'r') as load_f:
            print(load_f)
            json_dict = json.load(load_f, )
        h, w = json_dict['imageHeight'], json_dict['imageWidth']

        # save txt path
        txt_path = os.path.join(save_dir, json_path.replace('json', 'txt'))
        txt_file = open(txt_path, 'w')

        for shape_dict in json_dict['shapes']:
            label = shape_dict['label']
            label_index = classes.index(label)
            points = shape_dict['points']

            points_nor_list = []

            for point in points:
                points_nor_list.append(point[0] / w)
                points_nor_list.append(point[1] / h)

            points_nor_list = list(map(lambda x: str(x), points_nor_list))
            points_nor_str = ' '.join(points_nor_list)

            label_str = str(label_index) + ' ' + points_nor_str + '\n'
            txt_file.writelines(label_str)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='json convert to txt params')
    parser.add_argument('--json-dir', type=str, default='json_labels', help='json path dir')
    parser.add_argument('--save-dir', type=str, default='labels', help='txt save dir')
    parser.add_argument('--classes', type=str, default='1', help='classes')
    args = parser.parse_args()
    json_dir = args.json_dir
    save_dir = args.save_dir
    classes = args.classes

这一步执行后，.json文件对应的边界框信息便存入了labels文件夹中，images图片要手动存放对应.json文件的图片（图片名不能重）

接着，验证一下是否存储成功

css 复制代码

def check_labels(txt_labels, images_dir):
    txt_files = glob.glob(txt_labels + "/*.txt")
    print('---', txt_files)
    # 根据文件名找对应的图片进行输出
    for txt_file in txt_files:
        filename = os.path.splitext(os.path.basename(txt_file))[0]
        print('---', filename)
        pic_path = images_dir + '/' + filename + ".bmp"
        print(pic_path)
        img = cv2.imread(pic_path)
        height, width, _ = img.shape

        file_handle = open(txt_file)
        cnt_info = file_handle.readlines()
        new_cnt_info = [line_str.replace("\n", "").split(" ") for line_str in cnt_info]

        color_map = {"0": (0, 255, 255)}
        for new_info in new_cnt_info:
            print(new_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)])
            cv2.polylines(img, [np.array(s, np.int32)], True, color_map.get(new_info[0]))
        cv2.namedWindow('img', 0)
        cv2.imshow('img', img)
        cv2.waitKey()
    print('---检查完毕')

显示边框对应图片，确实没问题

下一步，就是划分测试集和训练集了。我们常用八二开或者九一开划分训练集。这里选择了九一开。函数如下：

ini 复制代码

import os
import random
import os
import shutil
def data_split(full_list, ratio):

    n_total = len(full_list)
    offset = int(n_total * ratio)
    if n_total == 0 or offset < 1:
        return [], full_list
    random.shuffle(full_list)
    sublist_1 = full_list[:offset]
    sublist_2 = full_list[offset:]
    return sublist_1, sublist_2


train_p="D:\PycharmFiles\graduation project\work\seg_work/train"
val_p="D:\PycharmFiles\graduation project\work\seg_work/val"
imgs_p="images"
labels_p="labels"

#创建训练集
if not os.path.exists(train_p):#指定要创建的目录
    os.mkdir(train_p)
tp1=os.path.join(train_p,imgs_p)
tp2=os.path.join(train_p,labels_p)
print(tp1,tp2)
if not os.path.exists(tp1):#指定要创建的目录
    os.mkdir(tp1)
if not os.path.exists(tp2):  # 指定要创建的目录
    os.mkdir(tp2)

#创建测试集文件夹
if not os.path.exists(val_p):#指定要创建的目录
    os.mkdir(val_p)
vp1=os.path.join(val_p,imgs_p)
vp2=os.path.join(val_p,labels_p)
print(vp1,vp2)
if not os.path.exists(vp1):#指定要创建的目录
    os.mkdir(vp1)
if not os.path.exists(vp2):  # 指定要创建的目录
    os.mkdir(vp2)

#数据集路径
images_dir="D:\PycharmFiles\graduation project\work\seg_work/images"
labels_dir="D:\PycharmFiles\graduation project\work\seg_work/labels"
#划分数据集，设置数据集数量占比
proportion_ = 0.9 #训练集占比

total_file = os.listdir(images_dir)

num = len(total_file)  # 统计所有的标注文件
list_=[]
for i in range(0,num):
    list_.append(i)

list1,list2=data_split(list_,proportion_)

for i in range(0,num):
    file=total_file[i]
    print(i,' - ',total_file[i])
    name=file.split('.')[0]
    if i in list1:
        jpg_1 = os.path.join(images_dir, file)
        jpg_2 = os.path.join(train_p, imgs_p, file)
        txt_1 = os.path.join(labels_dir, name + '.txt')
        txt_2 = os.path.join(train_p, labels_p, name + '.txt')
        if os.path.exists(txt_1) and os.path.exists(jpg_1):
            shutil.copyfile(jpg_1, jpg_2)
            shutil.copyfile(txt_1, txt_2)
        elif os.path.exists(txt_1):
            print(txt_1)
        else:
            print(jpg_1)

    elif i in list2:
        jpg_1 = os.path.join(images_dir, file)
        jpg_2 = os.path.join(val_p, imgs_p, file)
        txt_1 = os.path.join(labels_dir, name + '.txt')
        txt_2 = os.path.join(val_p, labels_p, name + '.txt')
        shutil.copyfile(jpg_1, jpg_2)
        shutil.copyfile(txt_1, txt_2)

print("数据集划分完成： 总数量：",num," 训练集数量：",len(list1)," 验证集数量：",len(list2))

划分好数据集后，要改写.yaml文件---和前面的目标检测类似的。

注意：这里使用的模型是yolov8s-seg.pt，区别主要是这个。

老规矩，激活环境，切换路径，直接开跑。

yolo detect data=mydata.yaml model=yolov8s-seg.pt device=0 imgsz=640 epochs=10

模型运行完成后，结果存在当前文件夹下的runs/segment/train_X文件夹下。

对应图标及参数我就不赘述了，很多博客其实都有- -