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)
labelme标注工具json格式标签转化为mask多值图
于 2022-01-05 17:55:52 首次发布