一. labelme 标记数据
生成的数据如下:
将数据放入到对应文件夹中：
二. 数据转换代码:
config.py 和 json_2_dataset.py 文件。
#!/usr/bin/env python
# _*_ coding:utf-8 _*_
# ============================================
# @Time : 2020/04/15 22:53
# @Author : WanDaoYi
# @FileName : config.py
# ============================================
import os
from easydict import EasyDict as edict
__C = edict()
# Consumers can get config by: from config import cfg
cfg = __C
# common options 公共配置文件
__C.COMMON = edict()
# 相对路径 当前路径
__C.COMMON.RELATIVE_PATH = "./"
# mask 默认背景 类别, 背景为第一个类别
__C.COMMON.DEFAULT_CLASS_INFO = [{"source": "", "id": 0, "name": "BG"}]
__C.COMMON.DATA_SET_PATH = os.path.join(__C.COMMON.RELATIVE_PATH, "dataset")
# 原始图像 文件 路径
__C.COMMON.IMAGE_PATH = os.path.join(__C.COMMON.RELATIVE_PATH, "dataset/images")
# labelme 生成的 json 注释文件 路径
__C.COMMON.JSON_PATH = os.path.join(__C.COMMON.RELATIVE_PATH, "dataset/ann_json")
# 生成的 info.yaml 文件 路径
__C.COMMON.INFO_YAML_PATH = os.path.join(__C.COMMON.RELATIVE_PATH, "dataset/info_yaml")
# 生成的 label.png 图像 文件 路径
__C.COMMON.LABEL_IMAGE_PATH = os.path.join(__C.COMMON.RELATIVE_PATH, "dataset/label_png")
# 生成的 label_name.txt 文件 路径
__C.COMMON.LABEL_NAME_PATH = os.path.join(__C.COMMON.RELATIVE_PATH, "dataset/label_name")
# 生成的 label_viz.png 图像 文件 路径
__C.COMMON.LABEL_VIZ_IMAGE_PATH = os.path.join(__C.COMMON.RELATIVE_PATH, "dataset/label_viz_png")
# 是否删除已有文件，True 为删除，False 为不删除
__C.COMMON.FILE_EXISTS_FLAG = True
# 数据划分比例
__C.COMMON.TEST_PERCENT = 0.7
__C.COMMON.VAL_PERCENT = 0.2
__C.COMMON.TEST_PERCENT = 0.1
# 文件后缀名
__C.COMMON.JSON_SUFFIX = ".json"
__C.COMMON.PNG_SUFFIX = ".png"
__C.COMMON.JPG_SUFFIX = ".jpg"
__C.COMMON.YAML_SUFFIX = ".yaml"
__C.COMMON.TXT_SUFFIX = ".txt"
# 划分数据的保存路径
__C.COMMON.TRAIN_DATA_PATH = os.path.join(__C.COMMON.RELATIVE_PATH, "infos/train_data.txt")
__C.COMMON.VAL_DATA_PATH = os.path.join(__C.COMMON.RELATIVE_PATH, "infos/val_data.txt")
__C.COMMON.TEST_DATA_PATH = os.path.join(__C.COMMON.RELATIVE_PATH, "infos/test_data.txt")
#!/usr/bin/env python
# _*_ coding:utf-8 _*_
# ============================================
# @Time : 2020/04/25 14:40
# @Author : WanDaoYi
# @FileName : json_2_dataset.py
# ============================================
import math
import json
import os
import io
import random
import shutil
import base64
from datetime import datetime
import matplotlib.pyplot as plt
from PIL import Image
from PIL import ImageDraw
import numpy as np
import yaml
from config import cfg
class Json2Dataset(object):
def __init__(self):
# 输入文件路径
self.json_file = cfg.COMMON.JSON_PATH
self.image_file = cfg.COMMON.IMAGE_PATH
# 数据的百分比
self.test_percent = cfg.COMMON.TEST_PERCENT
self.val_percent = cfg.COMMON.VAL_PERCENT
# 各成分数据保存路径
self.train_data_path = cfg.COMMON.TRAIN_DATA_PATH
self.val_data_path = cfg.COMMON.VAL_DATA_PATH
self.test_data_path = cfg.COMMON.TEST_DATA_PATH
# 输出文件路径
self.info_yaml_file = cfg.COMMON.INFO_YAML_PATH
self.label_image_file = cfg.COMMON.LABEL_IMAGE_PATH
self.label_name_file = cfg.COMMON.LABEL_NAME_PATH
self.label_viz_image_file = cfg.COMMON.LABEL_VIZ_IMAGE_PATH
# 生成文件的文件夹处理方式
self.deal_data_file(file_exists_flag=cfg.COMMON.FILE_EXISTS_FLAG)
# 文件后缀名
self.json_suffix = cfg.COMMON.JSON_SUFFIX
self.png_suffix = cfg.COMMON.PNG_SUFFIX
self.yaml_suffix = cfg.COMMON.YAML_SUFFIX
self.txt_suffix = cfg.COMMON.TXT_SUFFIX
pass
def divide_data(self):
"""
train, val, test 数据划分
:return:
"""
# 图像名字的 list
image_name_list = os.listdir(self.image_file)
# 统计有多少张图像
image_number = len(image_name_list)
# 根据百分比得到各成分 数据量
n_test = int(image_number * self.test_percent)
n_val = int(image_number * self.val_percent)
n_train = image_number - n_test - n_val
if os.path.exists(self.train_data_path):
os.remove(self.train_data_path)
pass
if os.path.exists(self.val_data_path):
os.remove(self.val_data_path)
pass
if os.path.exists(self.test_data_path):
os.remove(self.test_data_path)
pass
# 随机划分数据
n_train_val = n_train + n_val
train_val_list = random.sample(image_name_list, n_train_val)
train_list = random.sample(train_val_list, n_train)
train_file = open(self.train_data_path, "w")
val_file = open(self.val_data_path, "w")
test_file = open(self.test_data_path, "w")
for image_name in image_name_list:
if image_name in train_val_list:
if image_name in train_list:
train_file.write(image_name + "\n")
pass
else:
val_file.write(image_name + "\n")
pass
pass
else:
test_file.write(image_name + "\n")
pass
pass
# 生成文件的文件夹处理方式
def deal_data_file(self, file_exists_flag=True):
# 删除存在的文件夹
if file_exists_flag:
if os.path.exists(self.info_yaml_file):
shutil.rmtree(self.info_yaml_file)
pass
if os.path.exists(self.label_image_file):
shutil.rmtree(self.label_image_file)
pass
if os.path.exists(self.label_name_file):
shutil.rmtree(self.label_name_file)
pass
if os.path.exists(self.label_viz_image_file):
shutil.rmtree(self.label_viz_image_file)
pass
os.mkdir(self.info_yaml_file)
os.mkdir(self.label_image_file)
os.mkdir(self.label_name_file)
os.mkdir(self.label_viz_image_file)
pass
# 不删除存在的
else:
if not os.path.exists(self.info_yaml_file):
os.mkdir(self.info_yaml_file)
if not os.path.exists(self.label_image_file):
os.mkdir(self.label_image_file)
if not os.path.exists(self.label_name_file):
os.mkdir(self.label_name_file)
if not os.path.exists(self.label_viz_image_file):
os.mkdir(self.label_viz_image_file)
pass
pass
def do_data(self):
json_name_list = os.listdir(self.json_file)
for json_name in json_name_list:
json_data_path = os.path.join(self.json_file, json_name)
if os.path.isfile(json_data_path):
# 获取 .json 文件名: 如 000001.json -> 000001
name_info = json_name.split(self.json_suffix)[0]
# 文件的保存路径
info_yaml_path = os.path.join(self.info_yaml_file, name_info + self.yaml_suffix)
label_png_path = os.path.join(self.label_image_file, name_info + self.png_suffix)
label_name_path = os.path.join(self.label_name_file, name_info + self.txt_suffix)
label_viz_png_path = os.path.join(self.label_viz_image_file, name_info + self.png_suffix)
# 加载 .json 文件
data_info = json.load(open(json_data_path))
# print(data_info)
# 获取文件内 图像信息
if data_info["imageData"]:
image_data = data_info["imageData"]
pass
else:
image_path = os.path.join(os.path.dirname(json_data_path), data_info['imagePath'])
with open(image_path, 'rb') as f:
image_data = f.read()
image_data = base64.b64encode(image_data).decode('utf-8')
pass
# 将 base64 str 类型的图像 转为 np.array() 类型
image_arr = self.img_b64_2_arr(image_data)
# 设置 背景 label
label_name_to_value = {'_background_': 0}
for shape in data_info["shapes"]:
label_name = shape["label"]
# 如果 label_name 不在 label_name_to_value 这个 list 里面，则添加进去
if label_name not in label_name_to_value:
label_value = len(label_name_to_value)
label_name_to_value[label_name] = label_value
pass
label_values = []
label_names = []
# ln 为 label_name, lv 为 label_value
for ln, lv in sorted(label_name_to_value.items(), key=lambda x: x[1]):
label_names.append(ln)
label_values.append(lv)
pass
assert label_values == list(range(len(label_values)))
lbl = self.shapes_2_label(image_arr.shape, data_info["shapes"], label_name_to_value)
# 保存 label.png 图像
self.lbl_save(label_png_path, lbl)
captions = ["{}: {}".format(lv, ln) for ln, lv in label_name_to_value.items()]
lbl_viz = self.draw_label(lbl, image_arr, captions)
# 保存 label_viz.png 图像
Image.fromarray(lbl_viz).save(label_viz_png_path)
# 保存 label_name.txt 文件
with open(label_name_path, "w") as file:
for lbl_name in label_names:
file.write(lbl_name + "\n")
pass
info = dict(label_names=label_names)
# print("info: {}".format(info))
# 保存 info.yaml 文件
with open(info_yaml_path, 'w') as f:
yaml.safe_dump(info, f, default_flow_style=False)
pass
print("saved to {}".format(cfg.COMMON.DATA_SET_PATH))
pass
# 将 base64 图像 转为 np.array() list 图像
def img_b64_2_arr(self, img_b64):
file = io.BytesIO()
file.write(base64.b64decode(img_b64))
img_arr = np.array(Image.open(file))
return img_arr
pass
def shapes_2_label(self, image_shape, shapes, label_name_to_value, type='class'):
assert type in ['class', 'instance']
cls = np.zeros(image_shape[: 2], dtype=np.int32)
if type.__eq__("instance"):
ins = np.zeros(image_shape[: 2], dtype=np.int32)
instance_names = ["_background_"]
pass
else:
ins = None
instance_names = None
for shape in shapes:
points = shape["points"]
label = shape["label"]
shape_type = shape.get("shape_type", None)
# 分类处理
if type.__eq__("class"):
cls_name = label
ins_id = None
# 实例分割处理
else:
cls_name = label.split("-")[0]
if label not in instance_names:
instance_names.append(label)
pass
ins_id = instance_names.index(label)
cls_id = label_name_to_value[cls_name]
mask = self.shape_2_mask(image_shape[: 2], points, shape_type)
cls[mask] = cls_id
if type.__eq__("instance"):
ins[mask] = ins_id
pass
pass
if type.__eq__("instance"):
return cls, ins
else:
return cls
pass
def shape_2_mask(self, image_shape, points, shape_type=None,
line_width=10, point_size=5):
mask = np.zeros(image_shape[: 2], dtype=np.uint8)
mask = Image.fromarray(mask)
draw = ImageDraw.Draw(mask)
xy = [tuple(point) for point in points]
if "circle".__eq__(shape_type):
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)
pass
elif "rectangle".__eq__(shape_type):
assert len(xy) == 2, 'Shape of shape_type=rectangle must have 2 points'
draw.rectangle(xy, outline=1, fill=1)
pass
elif "line".__eq__(shape_type):
assert len(xy) == 2, 'Shape of shape_type=line must have 2 points'
draw.line(xy=xy, fill=1, width=line_width)
pass
elif "linestrip".__eq__(shape_type):
draw.line(xy=xy, fill=1, width=line_width)
pass
elif "point".__eq__(shape_type):
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)
pass
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 draw_label(self, label, image=None, label_names=None,
color_map=None, **kwargs):
"""
:param label: ndarray, (H, W), Pixel-wise labels to colorize.
:param image: ndarray, (H, W, 3), optional, Image on which the colorized label will be drawn.
:param label_names: iterable, List of label names.
:param color_map:
:param kwargs:
:return:
"""
backend_org = plt.rcParams["backend"]
plt.switch_backend("agg")
plt.subplots_adjust(left=0, right=1, top=1, bottom=0, wspace=0, hspace=0)
plt.margins(0, 0)
plt.gca().xaxis.set_major_locator(plt.NullLocator())
plt.gca().yaxis.set_major_locator(plt.NullLocator())
if label_names is None:
label_names = [str(l) for l in range(label.max() + 1)]
pass
color_map = self.validate_color_map(color_map, len(label_names))
label_viz = self.label_2_rgb(label, image, n_labels=len(label_names), color_map=color_map)
plt.imshow(label_viz)
plt.axis("off")
plt_handlers = []
plt_titles = []
for label_value, label_name in enumerate(label_names):
if label_value not in label:
continue
pass
fc = color_map[label_value]
p = plt.Rectangle((0, 0), 1, 1, fc=fc)
plt_handlers.append(p)
plt_titles.append("{value}: {name}".format(value=label_value, name=label_name))
file = io.BytesIO()
plt.savefig(file, bbox_inches="tight", pad_inches=0)
plt.cla()
plt.close()
plt.switch_backend(backend_org)
out_size = (label_viz.shape[1], label_viz.shape[0])
out = Image.open(file).resize(out_size, Image.BILINEAR).convert('RGB')
out = np.asarray(out)
return out
pass
pass
def validate_color_map(self, color_map, n_labels):
if color_map is None:
color_map = self.label_color_map(n_labels)
pass
else:
assert color_map.shape == (color_map.shape[0], 3), \
'color_map must be sequence of RGB values'
assert 0 <= color_map.min() and color_map.max() <= 1, \
'color_map must ranges 0 to 1'
pass
return color_map
pass
def label_color_map(self, num=256):
color_map = np.zeros((num, 3))
for i in range(num):
id = i
r, g, b = 0, 0, 0
for j in range(8):
r = np.bitwise_or(r, (self.bit_get(id, 0) << 7 - j))
g = np.bitwise_or(g, (self.bit_get(id, 1) << 7 - j))
b = np.bitwise_or(b, (self.bit_get(id, 2) << 7 - j))
id = (id >> 3)
pass
color_map[i, 0] = r
color_map[i, 1] = g
color_map[i, 2] = b
pass
color_map = color_map.astype(np.float32) / 255
return color_map
pass
def bit_get(self, byte_val, inx):
return ((byte_val & (1 << inx)) != 0)
pass
def label_2_rgb(self, lbl, image=None, n_labels=None, alpha=0.5,
color_map=None):
if n_labels is None:
n_labels = len(np.unique(lbl))
pass
color_map = self.validate_color_map(color_map, n_labels)
color_map = (color_map * 255).astype(np.uint8)
lbl_viz = color_map[lbl]
# 未标记的
lbl_viz[lbl == -1] = (0, 0, 0)
if image is not None:
image_gray = Image.fromarray(image).convert('LA')
image_gray = np.asarray(image_gray.convert("RGB"))
# image_gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
# image_gray = cv2.cvtColor(image_gray, cv2.COLOR_GRAY2RGB)
lbl_viz = alpha * lbl_viz + (1 - alpha) * image_gray
lbl_viz = lbl_viz.astype(np.uint8)
pass
return lbl_viz
pass
def lbl_save(self, file_name, lbl):
if os.path.splitext(file_name)[1] != ".png":
file_name += ".png"
if lbl.min() >= -1 and lbl.max() < 255:
lbl_pil = Image.fromarray(lbl.astype(np.uint8), mode='P')
color_map = self.label_color_map(255)
lbl_pil.putpalette((color_map * 255).astype(np.uint8).flatten())
lbl_pil.save(file_name)
pass
else:
raise ValueError('[%s] Cannot save the pixel-wise class label as PNG. '
'Please consider using the .npy format.' % file_name)
pass
pass
if __name__ == "__main__":
# 代码开始时间
start_time = datetime.now()
print("开始时间: {}".format(start_time))
demo = Json2Dataset()
demo.divide_data()
demo.do_data()
# 代码结束时间
end_time = datetime.now()
print("结束时间: {}, 训练模型耗时: {}".format(end_time, end_time - start_time))
pass
运行代码后，得到
在 实例分割中，有时候会用到这样的数据。
