针对某一特定类别，JSON文件转PNG

在写论文的时候，有的朋友需要将某一特定类别的标签展现出来，与分割结果做对比，以下代码可以实现这一功能：

import os
import os.path as osp
import json
import base64
import numpy as np
from PIL import Image
from labelme import utils
from labelme.logger import logger

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

    json_file = r"E:\heatmap\trunk"
    out_dir1 = r"E:\heatmap\trunk"

    if not osp.exists(out_dir1):
        os.mkdir(out_dir1)

    os.environ['OPENAI_GPT_UTF8_MODE'] = '1'

    target_class = "Diseased_trunk"  # 在自己的类别中，指定要转换的类别
    background_color = [68,1,84]  # 背景颜色，这里是白色

    for file_name in os.listdir(json_file):
        if file_name.endswith(".json"):
            path = os.path.join(json_file, file_name)
            if os.path.isfile(path):
                data = json.load(open(path, 'r', encoding='utf-8'))

                imageData = data.get("imageData")
                if not imageData:
                    imagePath = os.path.join(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)

                # 创建一个数组来存储标签
                lbl = np.zeros((img.shape[0], img.shape[1], 3), dtype=np.uint8)
                lbl[:, :] = background_color  # 将背景设置为指定的颜色

                for shape in data["shapes"]:
                    label_name = shape["label"]
                    if label_name == target_class:
                        polygon = shape["points"]
                        mask = utils.shape_to_mask(img.shape[:2], polygon)
                        lbl[mask] = [230,240,244]  # 将目标类别的区域设置为红色

                lbl_image = Image.fromarray(lbl)
                lbl_image.save(osp.join(out_dir1, "%s_label.png" % file_name.split(".")[0]))

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

if __name__ == "__main__":
    main()