【语义分割】——labelme标注的分割数据处理

原创

怡宝2号 2021-09-07 14:10:35 ©著作权

©著作权归作者所有：来自51CTO博客作者怡宝2号的原创作品，请联系作者获取转载授权，否则将追究法律责任

可以参考:labelme标注语义数据

code

import argparse
import base64
import json
import glob
import os
import os.path as osp
import numpy as np

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."
    )
    oAnnotationDir  = "./annotations/"
    oImageDir       = "./images"
    oVisualDir      = "./visual"
    iAnnotationPath = "./labelme/*.json"
    jsonFiles = glob.glob(iAnnotationPath)
    jsonFiles = sorted(jsonFiles)

    for path in [oAnnotationDir, oImageDir, oVisualDir]:
        if not osp.exists(path):
            os.makedirs(path)

    parser = argparse.ArgumentParser()
    parser.add_argument("-o", "--out", default='./output')
    args = parser.parse_args()

    for idx, json_file in enumerate(jsonFiles):

        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)

        # labelme的标签转成具体的类别数字
        label_name_to_value = {"_background_": 0,               
                            "floor": 1,
                            "person": 4,
                            "door": 5,
                            "table": 6,
                            "chair": 7,
                            "refrigerator": 11,
                            "wall": 14, 
                            "plant": 16}
        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

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

        # 保存图片
        name = "{}/cd_{:05d}.jpg".format(oImageDir, idx)
        PIL.Image.fromarray(img).save( name )

        # 保存标签图片
        name = "{}/cd_{:05d}.png".format(oAnnotationDir, idx)
        PIL.Image.fromarray(lbl.astype(np.uint8)).save( name )

        # 保存可视化图片
        name = "{}/cd_{:05d}.png".format(oVisualDir, idx)
        PIL.Image.fromarray(lbl_viz).save( name )

        PIL.Image.fromarray(img).save(osp.join(out_dir, "img.png"))
        utils.lblsave(osp.join(out_dir, "label.png"), lbl)
        lbl_pil = PIL.Image.fromarray(lbl.astype(np.uint8))             # 添加语义分割label图片的类别
        lbl_pil.save(osp.join(out_dir, "label1.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")

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


if __name__ == "__main__":
    main()