【实例分割】Labelme数据集转换YOLO格式

YOLOv5的7代版本以及YOLOv8都开始支持了实例分割的操作,其效果十分优秀,但是在自己的数据集上进行训练其格式是YOLO自己的格式(如coco128-seg),而通常的自己标注的数据集为labelme的格式(图片+json),为了方便使用YOLO来进行实例分割,可以使用python脚本来进行转换

原始的Labelme格式的实例分割数据集通常如下组织,图片存放在images文件夹下,json中存放同名的标记数据。

【实例分割】Labelme数据集转换YOLO格式_第1张图片

我们在original_data文件夹下创建目录labels,同时将如下脚本放在和original_data一级的目录下并运行

【实例分割】Labelme数据集转换YOLO格式_第2张图片

import os
import json

def findAllFile(path):
    for root, ds, fs in os.walk(path):
        for f in fs:
            if f.endswith('.json'):
                yield f[:-5]

# 使用字典对于标签进行存储
label_dict=dict()
label_num = 0

base = "./original_data"

for name in findAllFile(base+"/json"):
    # 获取json数据
    with open(f"{base}/json/{name}.json", 'r') as jf:
        data = json.load(jf)
    height = data["imageHeight"]
    width = data["imageWidth"]

    s="" # 用来储藏txt中的内容
    for item in data["shapes"]: # 遍历数据集中每一个分割子类
        # 对于标签进行收集和处理
        if (label:=item["label"]) not in label_dict.keys():
            label_dict[label]=label_num
            label_num+=1

        s = s+str(label_dict[item["label"]])+" "

        points = item["points"]
        for point in points:
            s=s+str(point[0]/width)+" "
            s=s+str(point[1]/height)+" "
        s = s[:-1]+"\n"

    # 将数据集进行写入labels文件夹下
    with open(f"{base}/labels/{name}.txt", 'w') as tf:
        tf.write(s)

with open("./dataset.yaml","w") as f:
    f.write(f"path: {base}\n")
    f.write("train: images\n")
    f.write("val: images\n")
    f.write("test: \n\n")
    f.write("names:\n")
    for key,num in label_dict.items():
        f.write(f"  {num}:{key}\n")

运行之后得到如下的组织结构

【实例分割】Labelme数据集转换YOLO格式_第3张图片

其中images,labels和dataset.yaml组成了我们训练yolo实例分割模型所需要的数据集。(其中的json没用,可以直接进行删除)

你可能感兴趣的:(YOLO,json,python)