一些做语义分割的朋友在写论文时,需要放上自己的标签图片,其中就产生了一个是否美观的问题,下面这个脚本可以让大家在进行数据集格式转换的时候,自定义颜色,从而使图片与论文整体色彩更搭
import argparse
import base64
import json
import os
import os.path as osp
import numpy as np
from PIL import Image
# utils.py
import numpy as np
import cv2
import argparse
import base64
import json按照固定颜色转png
import os
import os.path as osp
import imgviz
import PIL.Image
import numpy
from labelme.logger import logger
from labelme import utils
import argparse
import base64
import json
import os
import os.path as osp
import numpy as np
from PIL import Image
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:\DeepLearningProjects\LastDataset\Lastjson"
out_dir1 = r"E:\DeepLearningProjects\LastDataset\Lastjson_color"
if not osp.exists(out_dir1):
os.mkdir(out_dir1)
os.environ['OPENAI_GPT_UTF8_MODE'] = '1'
label_name_to_value = {"background": 0}
# 固定颜色映射表
label_to_color = {
0: [68, 1, 84], # Class 0: Purple
1: [48, 103, 141], # Class 1: Blue
2: [255, 69, 0], # Class 2: Magenta
3: [253, 231, 36], # Class 3: Yellow
4: [230, 240, 244], # Class 4: Olive
5: [255, 20, 147], # Class 5: Orange
}
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)
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 = np.zeros((img.shape[0], img.shape[1], 3), dtype=np.uint8)
lbl[:, :] = label_to_color[0]
for shape in data["shapes"]:
label_name = shape["label"]
label_value = label_name_to_value[label_name]
color = label_to_color[label_value]
polygon = shape["points"]
mask = utils.shape_to_mask(img.shape[:2], polygon)
lbl[mask] = color
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()