labelme语义分割数据集制作【附数据划分代码和注释】

目录

1. 标注数据转换
2. 数据集划分
3. 使用说明（需要自己修改的地方）

1、标注数据转换

一般语义分割模型需要三种数据（以lanenet举例）：原图像、标注图像、实例分割图像，能产生这三种数据的标注工具有很多，本文以最常用的labelme标注工具举例（labelme的安装此处略过，网上也有很多教程，大家可以根据自己的实际情况选择合适的安装教程）。

labelme标注好的数据般被保存为json类型文件，json类型文件是不便于被神经网络使用的，因此我们需要将其转为图像文件。将json类型文件转为图像文件不需要我们自己写，在labelme安装的文件夹下就有提供。

linux用户可以使用以下命令找到自己的labelme安装路径：

whereis labelme

windows用户可以通过以下命令找到自己的labelme安装位置：

pip show labelme

找到自己的labelme安装位置后，我们进入labelme安装的文件夹，在该文件夹下有一个名为cli的文件夹，在该文件夹中有一份名为json_to_dataset.py的文件，这份文件可以帮助我们将标注好的json类型文件转为图像文件，将其复制到你需要转换json类型文件的文件夹。

例如此处，我的原图像都放在了data文件夹下，所有标注好的json类型文件都放到了result文件夹下，那么我们就将labelme提供的json_to_dataset.py文件放到同级目录下。

2、数据集划分

话不多说，直接看代码，详细注释也在代码里：

import os
import cv2
import shutil
import random
from glob import glob
from tqdm import tqdm

# 获取result文件夹下所有的json类型文件
# 这里的glob函数获取result文件夹下所有的json类型文件后会返回一个列表
# 列表中的每个元素都是字符串，代表了每个json文件的位置
label_path = glob('./result/*.json')

image_path = './image/'	# 原图像——存放位置
binary_path = './gt_image_binary/'	# 标注图像——存放位置
instance_path = './gt_image_instance/'	# 实例分割图像（或原图像叠加标注图像）——存放位置

# 以下三句os。makedirs是检查当前文件夹是否存在这三个文件夹
# 如果存在，则不创建对应文件夹
# 如果不存在，则创建对应文件夹
os.makedirs(image_path, exist_ok=True)
os.makedirs(binary_path, exist_ok=True)
os.makedirs(instance_path, exist_ok=True)

# 这里tqdm的使用方法仅能返回int类型数据
for i in tqdm(range(len(label_path)), desc='Json to dataset'):
    path_name = label_path[i]	# 获取文件列表的元素
    
    # 编辑好命令行的命令，该命令会自动创建一个临时文件夹temporary，用以存放临时的转换数据
    cmd = 'python ./json_to_dataset.py ' + path_name + ' -o ' + 'temporary/'
    
    # 运行上述命令行的命令
    os.system(cmd)
    
    # 29-32行完成后，程序已经将一份json文件转为了图像数据
    # 此时temporary文件夹里已经存放了四份文件
    # 分别为：img.png、label.png、label_viz.png、label_names.txt
    # img.png：原图像
    # label.png：标注图像
    # label_viz.png：原图像叠加标注图像
    # label_names.txt：标签
    
    # 修改img.png的文件名字
    os.rename('./temporary/img.png', './temporary/' + path_name[9:-5] + '.png')
    # 移动该文件，至image文件夹
    shutil.move('./temporary/' + path_name[9:-5] + '.png', image_path + path_name[9:-5] + '.png')

    # 利用opencv读取label.png文件，第二个参数设置为0，使得读入的图像为灰度图像
    img_binary = cv2.imread('./temporary/label.png', 0)
    # 对图像进行二值化
    ret, threshold_binary = cv2.threshold(img_binary, 20, 255, cv2.THRESH_BINARY)
    # 将处理好的图像保存至gt_image_binary文件夹
    cv2.imwrite(binary_path + path_name[9:-5] + '.png', threshold_binary)

    # 利用opencv读取label.png文件，第二个参数设置为0，使得读入的图像为灰度图像
    img_temporary = cv2.imread('./temporary/label.png', 0)
    # 对图像进行二值化
    ret, threshold_temporary = cv2.threshold(img_binary, 20, 200, cv2.THRESH_BINARY)
    # 将处理好的图像保存至gt_image_instance文件夹
    cv2.imwrite(instance_path + path_name[9:-5] + '.png', threshold_temporary)

# 清除命令行中的数据，仅是为了方便查看之后的信息
os.system('clear')

# 获取以下三个文件夹中的文件列表
image_path = glob(image_path + '*.png')
binary_path = glob(binary_path + '*.png')
instance_path = glob(instance_path + '*png')


for i in tqdm(range(len(image_path)), desc='Write txt'):
    # 编辑好要写入txt文件的信息
    text = '../data/training_data' + image_path[i][1:] + ' ' + binary_path[i] + ' ' + instance_path[i] + '\n'
    
    # 定义一个随机数，便于分开训练集和测试集
    # 这里的随机数可以写成按比例生成的，此处只是使用了一个最简单的随机数生成方式
    random_int = random.randint(1, 10)
    if random_int > 2:
        with open('./train.txt', 'a') as f:
            f.write(text)
    else:
        with open('./val.txt', 'a') as f:
            f.write(text)

运行截图：

运行完成截图：

运行完成后文件夹截图：

train.txt截图：

val.txt截图：

3、使用说明

1. 上述代码中，第48～52，55～59两段其实是一样的，这是因为我所使用的模型需要这样的二值化数据，这两段可以根据自己的需求修改一下。
2. 72行中写入txt的内容，前面的路径部分需要改为你自己的数据路径，否则无法用至模型训练。
3. 76行的随机数生成可以写得更严谨一些，比如按比例生成随机数，使得区分出来的数据集和训练集成一定的比例。
4. 15行的注释中提到的：原图像叠加标注图像，其实长这样：