自制多分类cityscapes格式数据集用于HRNet网络进行语义分割

！说在前面（原创小文，如有需要请标明转载出处）

语义分割对图像进行分割是基于像素点对其进行分类的，所以我今天要说的分割，用的除了公共数据集之外，我是按照cityscapes数据集的格式进行数据集的制作用于训练。所以，要制作cityscapes格式数据集的话，务必对cityscapes公共数据集进行了解，并且熟知用于训练的图像数据格式，主要强调的是原图像对应的灰度图，这个也是制作数据集的关键。

1. 查看cityscapes格式

import numpy as np
from PIL import Image
img = Image.open（‘***.png’） #图像所在位置
img = np.array(img)
np.unique(img)

标注类别读取结果如下：

2. 数据集制作

首先，标注用的是labele工具，具体可参考：https://github.com/wkentaro/labelme

2.1 数据集的标注

简单标注如下图所示，左边是标注的边缘框架，右边是命名的类别，标注自己想要的类别进行保存即可。

2.2 json文件的转换

json文件进行转换的时候，可以直接在cmd用单个命令直接转，也可以自己写个批量转换的脚本，这个网上也一大堆，在此不做重复。下面这个代码里面对labelme所在文件位置的json_to_dataset.py进行了改进，对标注的类进行了初始化。（记得在代码所在目录建一个文件夹，我的命名是before）

import argparse
import json
import os
import os.path as osp
import warnings
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])
        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.draw_label(lbl, img, captions)
            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)
            PIL.Image.fromarray(img).save(osp.join(out_dir, 'img.png'))
            utils.lblsave(osp.join(out_dir, 'label.png'), lbl)
            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()

转换后的_json文件夹输出到output文件夹，其中包含五个文件，如下：

这里需要解释的是，我们后期用于训练的是原图（img.png）和标签图（label.png），其中，label.png是8bit图，但并不能直接用于训练，还得继续作转换，要将其转为8bit的灰度图(下一步为具体实现)

2.3 从转换后的_json文件夹中获取mask图

这里是最关键也是最重要的一环，我们需要将图像里面标注的类别进行一个局部到全局的映射.因为cityscapes数据集是为35类，其中用于训练的是19类，其他不参与训练的类别标签将其忽略掉即可（ignore=255）。我自己的类为13+1（background）,在before文件夹中新建一个class_name用于存放我们的类，代码如下：（get_png.py）

import os
from PIL import Image
import numpy as np
def main():
    # 读取原文件夹
    count = os.listdir("./before/")
    for i in range(0, len(count)):
        # 如果里的文件以jpg结尾
        # 则寻找它对应的png
        if count[i].endswith("jpg"):
            path = os.path.join("./before", count[i])
            img = Image.open(path)
            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", "r")
            class_name = class_txt.read().splitlines()
            # ["_background_","a","b"]
            # 打开json文件里面存在的类，称其为局部类
            with open("./output/" + count[i].split(".")[0] + "_json/label_names.txt", "r") as f:
                names = f.read().splitlines()
                # ["_background_","b"]
                new = Image.new("RGB", [np.shape(img)[1], np.shape(img)[0]])
                # print('new:',new)
                for name in names:
                    index_json = names.index(name)          
                    index_all = class_name.index(name)
                    # 将局部类转换成为全局类
                    new = new + np.expand_dims(index_all * (np.array(img) == index_json), -1)
            new = Image.fromarray(np.uint8(new))
            print('new:',new)
            new.save(os.path.join("./png", count[i].replace("jpg", "png")))
            print(np.max(new), np.min(new))
if __name__ == '__main__':
    main()

需要解释的是为了区分用于训练的图，我将原图置为a.jpg格式，mask图为a.png ，然后在before文件夹里传入对应的json文件（上一步已传）和.jpg图像
这里我们需要将cityscapes数据集的标签和我们自己的标签进行类比做出来，这里我们直接上图来说明，下面两幅图中，上面的image_4.png 是cityscapes数据集类的对应标签的列表，下面的image_5.png是自己类的标签列表

图image_4.png

图image_5.png

对应好后，在转训练的png的时候，按照制作好的标签顺序，进行mask的png的制作，运行代码get_png.py,输出图保存到png文件夹中。
这里补充一下大家很多都踩得坑，那就是如下这幅图中涉及的内容：实际参与训练的是trainid，并非id

2.4 转灰度图

得到的图是24bit的图，需要将其转换为8bit的灰度图用于训练，转换代码如下：（get_gray.py）

import cv2
import os
input_dir = './png'      #上一步保存.png图像文件夹
out_dir = 'F:/regular walks/2048X1024_2'
a = os.listdir(input_dir)
for i in a:
    img = cv2.imread(input_dir+'/'+i)
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    cv2.imencode('.png', gray)[1].tofile(out_dir+'/'+i)

2.5 查证灰度图的类别标签

最后，为了证实我们的制作结果，进行类别的验证，还是用最开始的简单方式进行测试，结果如下：

3.总结

以上训练的灰度图就算是制作成功了，然后制作好参与训练需要的文本列表（train.lst,val.lst,testval.lst,test.lst）。制作图像数据集需要大把时间，战线较长，需要耐心。然后hrnet网络相关的内容网上很多，就不多做介绍了，可以直接github找star最多的跑就可以了，暂时就先写这些，以上内容，仅供大家参考

参考

[1]https://www.bilibili.com/video/BV1qJ411S7Pn?p=10