[Python知识库]从图片到dataframe——语义分割数据集制作全流程

分享一下从原始图片，到标记图片，再到转换为python里的数据结构语义分割数据集制作全流程。

安装labelme

labelme 是一个图形界面的图像标注软件，可以很方便地划分出多边形边界。

下面在win10系统中使用anaconda安装labelme

// python3
conda create --name=labelme python=3.6
activate labelme
pip install labelme

anaconda会自动安装好labelme及其依赖包
通过conda list命令查看安装好的包：

在每次打开prompt或cmd激活labelme虚拟环境后，输入labelme即可打开labelme

标注图片

打开labelme，在labelme中打开文件夹，在Edit中选择create polygons即可添加多边形标记框，标记完后输入标签名即可完成一片区域的标记。标记完成后保存为json文件

json文件转为数据集

labelme标签不一致问题

labelme库原有的labelme_json_to_dataset.py脚本可以将json文件转换为原图(img)、掩码(mask)、原图和掩码的叠加（viz）和一个标签文件(txt)。默认情况下，导出掩码的像素值和在标注时标记的顺序有关，而不是和标签名对应。
这里参考了https://zhuanlan.zhihu.com/p/159837405的方法对labelme的 _json_to_dataset.py代码进行更改，使label可以自行配置，同时保证了标签的统一。
首先使用everything找到虚拟环境中labelme库的labelme_json_to_dataset.py文件

注释掉部分源代码，加上自己的标签名称和对应编号（0,1,2,…），保存即可

import argparse
import base64
import json
import os
import os.path as osp

import imgviz
import PIL.Image

from labelme.logger import logger
from labelme import utils


def main():
    logger.warning(
        "This script is aimed to demonstrate how to convert the "
        "JSON file to a single image dataset."
    )
    logger.warning(
        "It won't handle multiple JSON files to generate a "
        "real-use dataset."
    )

    parser = argparse.ArgumentParser()
    parser.add_argument("json_file")
    parser.add_argument("-o", "--out", default=None)
    args = parser.parse_args()

    json_file = args.json_file

    if args.out is None:
        out_dir = osp.basename(json_file).replace(".", "_")
        out_dir = osp.join(osp.dirname(json_file), out_dir)
    else:
        out_dir = args.out
    if not osp.exists(out_dir):
        os.mkdir(out_dir)

    data = json.load(open(json_file))
    imageData = data.get("imageData")

    if not imageData:
        imagePath = os.path.join(os.path.dirname(json_file), 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 sorted(data["shapes"], key=lambda x: x["label"]):
    #     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
    # lbl, _ = utils.shapes_to_label(
    #     img.shape, data["shapes"], label_name_to_value
    # )

    # label_names = [None] * (max(label_name_to_value.values()) + 1)
    # for name, value in label_name_to_value.items():
    #     label_names[value] = name
    label_name_to_value={'_background_': 0, 'E': 1, 'Esl': 2, 'Esc': 3, 'F1':4, 'F2':5, 'Fspread':6}
    
    lbl, _ = utils.shapes_to_label(
        img.shape, data["shapes"], label_name_to_value
    )
    label_names=['_background_', 'E', 'Esl', 'Esc','F1', 'F2', 'Fspread']
    # ###########################   结束修改    #############################

    lbl_viz = imgviz.label2rgb(
        label=lbl, img=imgviz.asgray(img), label_names=label_names, loc="rb"
    )

    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")

    logger.info("Saved to: {}".format(out_dir))


if __name__ == "__main__":
    main()

批量转换json文件

labelme库原有的labelme_json_to_dataset脚本一次只能转换一个json文件在此参考博主里先森的文章批量转换labelme标注json数据为图片，并将不同类型的图片分文件夹存储。

https://blog.csdn.net/sements/article/details/110135137

创建文件trans.py

# 调用labelme库中原有的 labelme_json_to_dataset 为核心
# 批量将文件夹中的json文件转换，并抽取对应图片至各自文件夹
# python trans.py -i ./ -r img -m mask -v viz
import os
import shutil
import argparse
 
 
def GetArgs():
    parser = argparse.ArgumentParser(description='将labelme标注后的json文件批量转换为图片')
    parser.add_argument('--input', '-i', required=True, help='json文件目录')
    parser.add_argument('--out-mask', '-m', required=True, help='mask图存储目录')
    parser.add_argument('--out-img', '-r', help='json文件中提取出的原图存储目录')
    parser.add_argument('--out-viz', '-v', help='mask与原图合并viz图存储目录')
    return parser.parse_args()
 
 
if __name__ == '__main__':
    _args = GetArgs()
    _jsonFolder = _args.input
 
    input_files = os.listdir(_jsonFolder)
    for sfn in input_files:                                     # single file name
        if (os.path.splitext(sfn)[1] == ".json"):               # 是否为json文件
 
            # 调用labelme_json_to_dataset执行转换,输出到 temp 文件夹
            os.system("labelme_json_to_dataset %s -o temp" % (_jsonFolder + '/' + sfn))
 
            # 复制json文件中提取出的原图到存储目录
            if _args.out_img:
                if not os.path.exists(_args.out_img):           # 文件夹是否存在
                    os.makedirs(_args.out_img)
 
                src_img = "temp\img.png"
                dst_img = _args.out_img + '/' + os.path.splitext(sfn)[0] + ".png"
                shutil.copyfile(src_img, dst_img)
 
            # 复制mask图到存储目录
            if _args.out_mask:
                if not os.path.exists(_args.out_mask):          # 文件夹是否存在
                    os.makedirs(_args.out_mask)
 
                src_mask = "temp\label.png"
                dst_mask = _args.out_mask + '/' + os.path.splitext(sfn)[0] + ".png"
                shutil.copyfile(src_mask, dst_mask)
 
            # 复制viz图到存储目录
            if _args.out_viz:
                if not os.path.exists(_args.out_viz):           # 文件夹是否存在
                    os.makedirs(_args.out_viz)
 
                src_viz = "temp\label_viz.png"
                dst_viz = _args.out_viz + '/' + os.path.splitext(sfn)[0] + ".png"
                shutil.copyfile(src_viz, dst_viz)
 

将其与待转换的json文件放入同一个文件夹中，命令行切换至该目录，执行

python trans.py -i ./ -r img -m mask -v viz

即可完成转换，生成所需的四个文件夹

读取图片并转换为dataframe格式

json文件转换得到的掩码（mask）是RGB三通道彩色图，而语义分割需要的掩码往往是灰度图。

可以使用cv2库中的imread函数直接将三通道彩色图导入为灰度图，导入后查看六种不同标签的到灰度值如下（背景灰度值为0）

为了便于训练，我自己写了一个类似switch的字典用于将灰度值转换为0,1,2,3,4,5,6的标签值。

# labelconvert.py
def case0():                           
    return 0
def case1():                           
    return 1
def case2():                            
    return 2
def case3():                            
    return 3
def case4():                            
    return 4
def case5():                            
    return 5
def case6():                            
    return 6
def default(): 
    print('Error')                        
    return 255
switch = {38: case1,                
          75: case2,
          113: case3,
          14: case4,
          52: case5,
          89: case6,
          0: case0
          }
def labelconvert(img):
    for i in range (360):
        for j in range (400):
            img[i][j] = switch.get(int(img[i][j]), default)() 
    return img

为了方便存储原图、掩码和文件名，使用pandas的dataframe进行存储，dataframe里的数据也可以转换为np数组或tensor张量。

from tensorflow.keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img
import matplotlib.pyplot as plt
import cv2
import numpy as np
from labelconvert import labelconvert  # 自定义函数，用于灰度标签转换为0-6
import os
import pandas as pd
import time

maskdf = pd.DataFrame(columns=('id', 'mask', 'Fspread'))
inputdf = pd.DataFrame(columns=(['input']))
mask_dir = 'D:\\OriginIonograms\\lyf50\\mask\\'
input_dir = 'D:\\OriginIonograms\\lyf50\\img\\'

i = 1
tic1 = time.perf_counter()
# 导入mask数据
for filename in os.listdir(mask_dir):
    img = cv2.imread(mask_dir + "/" + filename, cv2.IMREAD_GRAYSCALE)   #导入灰度图
    img = img_to_array(img)
    img = labelconvert(img)   # 替换比较耗时
    Fspread = (img == 6).any() 
    # Fspread = (img == 89).any() # 不做标签转换
    maskdf.loc[i] = [filename, img, Fspread]
    i = i + 1
tic2 = time.perf_counter()
print('导入', i-1,'组mask数据用时',tic2-tic1,'s' )

# 导入input数据
i = 1
for filename in os.listdir(input_dir):
    img = cv2.imread(input_dir + "/" + filename, cv2.IMREAD_GRAYSCALE)
    img = img_to_array(img)
    inputdf.loc[i] = [img]
    i = i + 1
tic3 = time.perf_counter()
print('导入', i-1,'组input数据用时',tic3-tic2,'s' )


df = pd.concat([maskdf, inputdf], axis=1)
print(df.columns)

df2 = df.drop(df[df.Fspread == 1].index)    
print('所有数据共', df.shape[0], '条')