【语义分割数据集制作】使用labelme制作自己的U-net语义分割数据集

  最近有一个工作需要使用Unet的语义分割结果作为上游任务，语义分割数据集需要自己制作，特此记录一下。

1、下载labelme

  第一步当然是下载labelme工具，主要有两种方式：
  1. 系统下可执行程序，如Win下的exe
  2. python脚本，通过执行python脚本运行程序，好在labelme已经被打包成为pip包
  这里就不展开了，下载、使用方法，官方讲的比我详细：labelme下载传送门
  

2、标注完成之后

  目录结构大概这样：
  
  其中标注完之后的原图放在before里，即每张图片和对应的json文件。
  
  class_name.txt主要是语义信息，例如我要分割出桥墩，那么就是。这个语义信息和labme里的是对应的，后面会用到
  

3、数据格式转换

  执行两个python脚本转换数据格式，代码如下：
  1、先执行 json_to_dataset.py：

import argparse
import json
import os
import os.path as osp
import warnings
import imgviz
import PIL.Image
import yaml
 
from labelme import utils
import base64
 
def main():
    count = os.listdir("./before/") 
    for i in range(0, len(count)):
        path = os.path.join("./before", count[i])
 
        # 找到before文件中以json为结尾的文件并且打开
        if os.path.isfile(path) and path.endswith('json'):
            data = json.load(open(path))
            
            if data['imageData']:
                imageData = data['imageData']
            else:
                imagePath = os.path.join(os.path.dirname(path), data['imagePath'])
                with open(imagePath, 'rb') as f:
                    imageData = f.read()
                    imageData = base64.b64encode(imageData).decode('utf-8')
            img = utils.img_b64_to_arr(imageData)
            label_name_to_value = {'_background_': 0}
            for shape in data['shapes']:
                label_name = shape['label']
                if label_name in label_name_to_value:
                    label_value = label_name_to_value[label_name]
                else:
                    label_value = len(label_name_to_value)
                    label_name_to_value[label_name] = label_value
            
            # label_values must be dense
            label_values, label_names = [], []
            for ln, lv in sorted(label_name_to_value.items(), key=lambda x: x[1]):
                label_values.append(lv)
                label_names.append(ln)
            assert label_values == list(range(len(label_values)))
            
            lbl, _ = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)

            # captions = ['{}: {}'.format(lv, ln)
            #     for ln, lv in label_name_to_value.items()]
            # lbl_viz = utils.labelme_shapes_to_label(lbl, img, captions)
            label_names = [None] * (max(label_name_to_value.values()) + 1)
            for name, value in label_name_to_value.items():
                label_names[value] = name
            # 4
            # lbl_viz = imgviz.label2rgb(
            #     lbl, imgviz.asgray(img), label_names=label_names, loc="rb"
            # )

            out_dir = osp.basename(count[i]).replace('.', '_')
            out_dir = osp.join(osp.dirname(count[i]), out_dir)
            out_dir = osp.join("output",out_dir)
    
            if not osp.exists(out_dir):
                os.mkdir(out_dir)
            # image
            PIL.Image.fromarray(img).save(osp.join(out_dir, 'img.png'))
            # label.png
            utils.lblsave(osp.join(out_dir, 'label.png'), lbl)
            # PIL.Image.fromarray(lbl_viz).save(osp.join(out_dir, 'label_viz.png'))
            # label_viz.png
            # PIL.Image.fromarray(lbl_viz).save(osp.join(out_dir, "label_viz.png"))

            with open(osp.join(out_dir, 'label_names.txt'), 'w') as f:
                for lbl_name in label_names:
                    f.write(lbl_name + '\n')
 
            # warnings.warn('info.yaml is being replaced by label_names.txt')
            # info = dict(label_names=label_names)
            # with open(osp.join(out_dir, 'info.yaml'), 'w') as f:
            #     yaml.safe_dump(info, f, default_flow_style=False)
 
            print('Saved to: %s' % out_dir)
 
if __name__ == '__main__':
    main()
 

  2、再执行draw_label_png.py

import os
from PIL import Image
import numpy as np

# 文件路径处理
root = os.getcwd()
before = os.path.join(root, "before")
output = os.path.join(root, "output")
assert(os.path.exists(before)), "please check before folder"
assert(os.path.exists(output)), "please check output folder"
jpg = os.path.join(root, "jpg")
png = os.path.join(root, "png")
if not os.path.exists(jpg):
    os.mkdir(jpg)
if not os.path.exists(png):
    os.mkdir(png)

def main():
    # 读取原文件夹
    count = os.listdir("./before/") 
    for i in range(0, len(count)):
        # 如果里的文件以jpg结尾
        # 则寻找它对应的png
        if count[i].endswith("png"):
            path = os.path.join("./before", count[i])
            img = Image.open(path)

            # 对应原图像未标注 则遍历下一张图片
            if not os.path.exists("./output/" + count[i].split(".")[0] + "_json/label.png"):
                continue

            img.save(os.path.join("./jpg", count[i]))
 
 
            # 找到对应的png
            path = "./output/" + count[i].split(".")[0] + "_json/label.png"
            img = Image.open(path)
 
 
            # 找到全局的类
            class_txt = open("./before/class_name.txt","r")
            class_name = class_txt.read().splitlines()
            # ["bk","cat","dog"] 全局的类
            # 打开x_json文件里面存在的类，称其为局部类
            with open("./output/" + count[i].split(".")[0] + "_json/label_names.txt","r") as f:
                names = f.read().splitlines()
                # ["bk","dog"] 局部的类
                # 新建一张空白图片, 单通道
                new = Image.new("P", (img.width, img.height))
 
 
                # 找到局部的类在全局中的类的序号
                for name in names:
                    # index_json是x_json文件里存在的类label_names.txt，局部类
                    index_json = names.index(name)
                    # index_all是全局的类,
                    index_all = class_name.index(name)
 
 
                    # 将局部类转换成为全局类
                    # 将原图img中像素点的值为index_json的像素点乘以其在全局中的像素点的所对应的类的序号 得到 其实际在数据集中像素点的值
                    # 比如dog,在局部类（output/x_json/label_names）中它的序号为1,dog在原图中的像素点的值也为1.
                    # 但是在全局的类（before/classes.txt）中其对应的序号为2，所以在新的图片中要将局部类的像素点的值*全局类的序号，从而得到标签文件
                    new = new + (index_all*(np.array(img) == index_json))
 
 
            new = Image.fromarray(np.uint8(new))
            # 将转变后的得到的新的最终的标签图片保存到make_dataset/png文件夹下
            new.save(os.path.join("./png", count[i].replace("jpg","png")))
            # 找到新的标签文件中像素点值的最大值和最小值，最大值为像素点对应的类在class_name.txt中的序号，最小值为背景，即0
            print(np.max(new),np.min(new))
 
 
if __name__ == '__main__':
    main()

  如果有不理解的，可以对照这两段代码看。很多细节都在代码里。
  注意更改自己的路径。

4、移动到U-Net目录下

  这一步具体看U-net的项目说明就可以，传送门