手把手实战教学！语义分割从0到1：一、数据集制作

本文链接：https://blog.csdn.net/oYeZhou/article/details/115319298

本篇博客，是《手把手实战教学！语义分割从0到1》系列的第一篇实战教学，将重点介绍语义分割相关数据集，以及如何制作自己的数据集。

本系列总的介绍，以及其他章节的汇总，见：手把手实战教学，语义分割从0到1：开篇_语义分割实战-CSDN博客。

1、公开数据集

常用的语义分割数据集有COCO、VOC、Cityscapes、ADE20K等，作为入门，建议从VOC和COCO开始了解，而大多数开源框架也会支持这两种数据集。

由于数据集下载通常比较困难，我就把这两种常用的数据集放到了百度网盘，这里放出VOC和COCO数据集的百度网盘链接：

VOC2012：	链接: https://pan.baidu.com/s/1FRpjEn44bQL0BE95PcpMMw?pwd=k9uj 提取码: k9uj
COCO2017：	链接: https://pan.baidu.com/s/1BKn4_jv2vWIMRr58vVBImw?pwd=wxdi 提取码: wxdi

2、制作自己的语义分割数据集

2.1、准备标注工具

这里需要用到一个标注工具：labelme，该工具的安装可以参考：GitHub - wkentaro/labelme: Image Polygonal Annotation with Python (polygon, rectangle, circle, line, point and image-level flag annotation).。

在Python虚拟环境中，可以用pip直接安装：

pip install labelme

安装成功后，通过一下命令查看帮助：

labelme -h

显示如下内容表示安装成功：

usage: labelme [-h] [--version] [--reset-config]
               [--logger-level {debug,info,warning,fatal,error}]
               [--output OUTPUT] [--config CONFIG] [--nodata] [--autosave]
               [--nosortlabels] [--flags FLAGS] [--labelflags LABEL_FLAGS]
               [--labels LABELS] [--validatelabel {exact}] [--keep-prev]
               [--epsilon EPSILON]
               [filename]

positional arguments:
  filename              image or label filename

optional arguments:
  -h, --help            show this help message and exit
  --version, -V         show version
  --reset-config        reset qt config
  --logger-level {debug,info,warning,fatal,error}
                        logger level
  --output OUTPUT, -O OUTPUT, -o OUTPUT
                        output file or directory (if it ends with .json it is
                        recognized as file, else as directory)
  --config CONFIG       config file or yaml-format string (default:
                        C:\Users\Administrator\.labelmerc)
  --nodata              stop storing image data to JSON file
  --autosave            auto save
  --nosortlabels        stop sorting labels
  --flags FLAGS         comma separated list of flags OR file containing flags
  --labelflags LABEL_FLAGS
                        yaml string of label specific flags OR file containing
                        json string of label specific flags (ex. {person-\d+:
                        [male, tall], dog-\d+: [black, brown, white], .*:
                        [occluded]})
  --labels LABELS       comma separated list of labels OR file containing
                        labels
  --validatelabel {exact}
                        label validation types
  --keep-prev           keep annotation of previous frame
  --epsilon EPSILON     epsilon to find nearest vertex on canvas

2.2、准备待标注的原始数据

原始数据应当是图片格式，建议放在一个统一的文件夹中，如名为“imgs”的文件夹；然后，建议新建一个文件夹用于存放各个图片对应的label，也即json文件；最后，应当新建一个TXT文档（labels.txt）用于指明标注的类别。

最终的目录形如：

其中，labels.txt中的内容形如：

__ignore__
_background_
classA
classB
classC

前两项是必须的，后面几项根据你的类别来设定，有几类就写明几类。labels.txt提前设定好的目的，是为了在标注时让工具已经知道有这些类别，我们只需选择就完了。

2.3、开始标注

先进入到图1所示的目录下：

cd path/to/root_path_of_dataset

然后，启动labelme进行标注（各参数含义参见2.1节展示的帮助）：

labelme imgs --output jsons --nodata --autosave --labels labels.txt

接着就是一张一张的标注了：按Ctrl+N，然后逐点勾勒出目标的轮廓，形成一个闭环即为完成一个实例的标注。

每标完一张图片，就会生成一个对应的json文件，其中存放着我们标注的轮廓上各个点的坐标，以及该图片的一些基本信息；该json文件会存放到我们所制定的名为“json”的文件夹中。

2.4、标注界面

标注界面和labelimg类似，如下图：

2.5、转为VOC格式

我们标注完毕后，得到的是一堆json格式文件，为了可以像VOC数据集那样每张原图对应一个png格式的mask图片，我们需要对标注完的数据集进行转换，从labelme到VOC。

labelme自带了一个脚本（labelme/examples/semantic_segmentation/labelme2voc.py），可以进行转换：

# It generates:
#   - data_dataset_voc/JPEGImages
#   - data_dataset_voc/SegmentationClass
#   - data_dataset_voc/SegmentationClassVisualization
./labelme2voc.py data_annotated data_dataset_voc --labels labels.txt

不过其生成的目录和VOC稍有差别，我在该脚本基础上简单改了一下，可以生成与VOC一致的目录结构，并可以直接在Pycharm等IDE中直接执行，代码如下：

#!/usr/bin/env python

from __future__ import print_function

import argparse
import glob
import os
import os.path as osp
import sys

import imgviz
import numpy as np
import labelme


def main(args):

    if osp.exists(args.output_dir):
        print("Output directory already exists:", args.output_dir)
        sys.exit(1)
    os.makedirs(args.output_dir)
    os.makedirs(osp.join(args.output_dir, "JPEGImages"))
    os.makedirs(osp.join(args.output_dir, "SegmentationClassnpy"))
    os.makedirs(osp.join(args.output_dir, "SegmentationClass"))
    if not args.noviz:
        os.makedirs(
            osp.join(args.output_dir, "SegmentationClassVisualization")
        )
    print("Creating dataset:", args.output_dir)

    class_names = []
    class_name_to_id = {}
    for i, line in enumerate(open(args.labels).readlines()):
        class_id = i - 1  # starts with -1
        class_name = line.strip()
        class_name_to_id[class_name] = class_id
        if class_id == -1:
            assert class_name == "__ignore__"
            continue
        elif class_id == 0:
            assert class_name == "_background_"
        class_names.append(class_name)
    class_names = tuple(class_names)
    print("class_names:", class_names)
    out_class_names_file = osp.join(args.output_dir, "class_names.txt")
    with open(out_class_names_file, "w") as f:
        f.writelines("\n".join(class_names))
    print("Saved class_names:", out_class_names_file)

    for filename in glob.glob(osp.join(args.input_dir, "*.json")):
        print("Generating dataset from:", filename)

        label_file = labelme.LabelFile(filename=filename)

        base = osp.splitext(osp.basename(filename))[0]
        out_img_file = osp.join(args.output_dir, "JPEGImages", base + ".jpg")
        out_lbl_file = osp.join(
            args.output_dir, "SegmentationClassnpy", base + ".npy"
        )
        out_png_file = osp.join(
            args.output_dir, "SegmentationClass", base + ".png"
        )
        if not args.noviz:
            out_viz_file = osp.join(
                args.output_dir,
                "SegmentationClassVisualization",
                base + ".jpg",
            )

        with open(out_img_file, "wb") as f:
            f.write(label_file.imageData)
        img = labelme.utils.img_data_to_arr(label_file.imageData)

        lbl, _ = labelme.utils.shapes_to_label(
            img_shape=img.shape,
            shapes=label_file.shapes,
            label_name_to_value=class_name_to_id,
        )
        labelme.utils.lblsave(out_png_file, lbl)

        np.save(out_lbl_file, lbl)

        if not args.noviz:
            viz = imgviz.label2rgb(
                label=lbl,
#                img=imgviz.rgb2gray(img),
                img=img,
                font_size=15,
                label_names=class_names,
                loc="rb",
            )
            imgviz.io.imsave(out_viz_file, viz)


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_dir", default="../edge_fence_20210203/jsons", type=str, help="input annotated directory")
    parser.add_argument("--output_dir", default="../edge_fence_20210203_voc", type=str, help="output dataset directory")
    parser.add_argument("--labels", default="../edge_fence_20210203/labels.txt", type=str, help="labels file")
    parser.add_argument("--noviz", help="no visualization", action="store_true")
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    args = get_args()
    main(args)

利用labelme/examples/semantic_segmentation/labelme2voc.py或者我上面贴出的代码，可以将标注完毕的数据集转化为VOC格式，其目录结构形如：

其中，JPEGImages中存放的原图，SegmentationClass中存放的是png格式mask，SegmentationClassnpy存放的是numpy格式的mask，SegmentationClassVisualization是将mask叠加到原图得到的可视化结果，class_names.txt中存放的是类别名（包括背景和各个类）。