labelme标注工具json格式标签转化为mask多值图

import os
import os.path as osp
import io
import math
import base64
import json
import uuid

import numpy as np
import imgviz
import PIL.Image
import cv2


def img_data_to_pil(img_data):
    f = io.BytesIO()
    f.write(img_data)
    img_pil = PIL.Image.open(f)
    return img_pil


def img_data_to_arr(img_data):
    img_pil = img_data_to_pil(img_data)
    img_arr = np.array(img_pil)
    return img_arr


def img_b64_to_arr(img_b64):
    img_data = base64.b64decode(img_b64)
    img_arr = img_data_to_arr(img_data)
    return img_arr


def shape_to_mask(img_shape, points, shape_type=None, line_width=10, point_size=5):
    mask = np.zeros(img_shape[:2], dtype=np.uint8)
    mask = PIL.Image.fromarray(mask)
    draw = PIL.ImageDraw.Draw(mask)
    xy = [tuple(point) for point in points]
    if shape_type == "circle":
        assert len(xy) == 2, "Shape of shape_type=circle must have 2 points"
        (cx, cy), (px, py) = xy
        d = math.sqrt((cx - px) ** 2 + (cy - py) ** 2)
        draw.ellipse([cx - d, cy - d, cx + d, cy + d], outline=1, fill=1)
    elif shape_type == "rectangle":
        assert len(xy) == 2, "Shape of shape_type=rectangle must have 2 points"
        draw.rectangle(xy, outline=1, fill=1)
    elif shape_type == "line":
        assert len(xy) == 2, "Shape of shape_type=line must have 2 points"
        draw.line(xy=xy, fill=1, width=line_width)
    elif shape_type == "linestrip":
        draw.line(xy=xy, fill=1, width=line_width)
    elif shape_type == "point":
        assert len(xy) == 1, "Shape of shape_type=point must have 1 points"
        cx, cy = xy[0]
        r = point_size
        draw.ellipse([cx - r, cy - r, cx + r, cy + r], outline=1, fill=1)
    else:
        assert len(xy) > 2, "Polygon must have points more than 2"
        draw.polygon(xy=xy, outline=1, fill=1)
    mask = np.array(mask, dtype=bool)
    return mask


def shapes_to_label(img_shape, shapes, label_name_to_id):
    cls = np.zeros(img_shape[:2], dtype=np.int32)
    ins = np.zeros_like(cls)
    instances = []
    for shape in shapes:
        points = shape["points"]
        label = shape["label"]
        group_id = shape.get("group_id")
        if group_id is None:
            group_id = uuid.uuid1()
        shape_type = shape.get("shape_type", None)

        cls_name = label
        instance = (cls_name, group_id)

        if instance not in instances:
            instances.append(instance)
        ins_id = instances.index(instance) + 1
        cls_id = label_name_to_id[cls_name]

        mask = shape_to_mask(img_shape[:2], points, shape_type)
        cls[mask] = cls_id
        ins[mask] = ins_id

    return cls, ins


def lblsave(filename, lbl, palette=True):
    if osp.splitext(filename)[1] != ".png":
        filename += ".png"
    # Assume label ranses [-1, 254] for int32, and [0, 255] for uint8 as VOC.
    if lbl.min() >= -1 and lbl.max() < 255:
        lbl_pil = PIL.Image.fromarray(lbl.astype(np.uint8), mode="P")
        if palette:
            colormap = imgviz.label_colormap()
            lbl_pil.putpalette(colormap.flatten())
        lbl_pil.save(filename)
    else:
        raise ValueError("[%s] Cannot save the pixel-wise class label as PNG. Please consider using the .npy format." % filename)


def json_to_dataset(json_file, out_dir=None, demo=False):
    state = True
    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 = 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_id:
            label_value = label_name_to_id[label_name]
        else:
            print("标签名错误！{}".format(label_name))
            state = False
            label_value = len(label_name_to_id)
            label_name_to_id[label_name] = label_value
    try:
        lbl, _ = shapes_to_label(img.shape, data["shapes"], label_name_to_id)
    except Exception as e:
        print(e)
        return False, None, None
    if demo:
        if out_dir is None:
            out_dir = osp.basename(json_file).replace(".", "_")
            out_dir = osp.join(osp.dirname(json_file), out_dir)
        if not osp.exists(out_dir):
            os.mkdir(out_dir)
        print("Saved to: {}".format(out_dir))
        PIL.Image.fromarray(img).save(osp.join(out_dir, "img.png"))
        lblsave(osp.join(out_dir, "label.png"), lbl, palette=False)
        lblsave(osp.join(out_dir, "label_color.png"), lbl)
        label_names = [None] * (max(label_name_to_id.values()) + 1)
        for name, value in label_name_to_id.items():
            label_names[value] = name
        lbl_viz = imgviz.label2rgb(lbl, imgviz.asgray(img), label_names=label_names, loc="rb")
        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")
        return state, None, None
    else:
        # lbl_pil = PIL.Image.fromarray(lbl.astype(np.uint8), mode="P")
        # img_pil = PIL.Image.fromarray(img)
        img_cv = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        lbl_cv = lbl.astype(np.uint8)
        return state, img_cv, lbl_cv


label_id_to_name = {
    0: '_background_',
    1: "crossing",
    2: "speedBump",
    3: "parkingLine",
    4: "whiteSolid",
    5: "whiteDotted",
    6: "yellowSolid",
    7: "yellowDotted"
}
label_name_to_id = dict([val,key] for key,val in label_id_to_name.items())


if __name__ == "__main__":
    # 输出详细的label信息到指定目录
    json_to_dataset("./000054.json", "./out", demo=True)
    # 返回图片和对应的标签map，转化是否成功
    state, img_cv, lbl_cv = json_to_dataset("./000054.json")
    print(state)
    cv2.imwrite("img.png",img_cv)
    cv2.imwrite("lbl.png",lbl_cv)
    img_cv = cv2.imread("img.png")
    lbl_cv = cv2.imread("lbl.png", 0)