labelme标注的json文件变成每个类别一张黑白mask图

        使用labelme进行标注，标注完之后只有一个json文件。在分割任务中，读取数据时基本上都是获取图片以及图片对应的mask图。再将这两个转化为tensor然后放入模型进行训练。

        为了能够获取我们需要的数据，所以需要对json文件进行处理。下面是一个处理的脚本的demo：

将mask图从numpy转变成tensor的时候，不需要进行归一化处理，不然训练出来的模型的效果很拉胯。

效果如下： 

import numpy as np
from utils_func import seg_utils as ut
import cv2
import json
import os


def get_img_mask_path(data_folder):
    # 标注的图片和标注完生成的json是在同一个文件夹下面的
    data_lists = []
    for file_name in os.listdir(data_folder):
        if file_name.endswith("json"):
            json_path = os.path.join(data_folder, file_name)
            img_path = os.path.join(data_folder, "%s.jpg" % file_name.split(".")[0])
            if os.path.exists(img_path):
                data_lists.append([img_path, json_path])
    return data_lists


def gen_image_and_label(img_path, labelme_json_path, img_size, category_types:list):
    """
    :param img_path: 图片路径
    :param labelme_json_path: json文件路径
    :param img_size: 训练时需要resize的图片的大小
    :param category_types: 类别名称：是一个列表
    :return:
    """
    img = ut.p2i(img_path)          # ut.p2i:是一个读图的函数，直接用cv2.imread()也行
    h, w = img.shape[:2]
    img = cv2.resize(img, (img_size[0], img_size[1]))
    cv2.imwrite(r"C:\Users\xx\Desktop\test\%s.png" % (img_path.split(".")[0].split("\\")[-1]), img)
    img_tensor = ut.i2t(img)        # ut.i2t：是把图片转化成tensor的函数，自己写也行

    mask_array = np.zeros([len(category_types), h, w, 1], dtype=np.uint8)       # 有几个类别就创建多少个维度的数组，后面一个维度对应一个类别
    with open(labelme_json_path, "r") as f:
        json_data = json.load(f)
    shapes = json_data["shapes"]
    for shape in shapes:
        category = shape["label"]
        category_idx = category_types.index(category)
        points = shape["points"]
        points_array = np.array(points, dtype=np.int32)
        temp = mask_array[category_idx, ...]
        mask_array[category_idx, ...] = cv2.fillPoly(temp, [points_array], 255)
    mask_array = np.transpose(mask_array, (1, 2, 0, 3)).squeeze(axis=-1)
    mask_array = cv2.resize(mask_array, (img_size[0], img_size[1]))
    mask_tensor = ut.i2t(mask_array,False)
    # for循环的内容打开可以看到 mask 图
    for i in range(len(category_types)):
        cv2.imwrite(r"C:\Users\xx\Desktop\test\%s_mask_%s.png" % (img_path.split(".")[0].split("\\")[-1], category_types[i]), mask_array[:, :, i])
    return img_tensor, mask_tensor


if __name__ == '__main__':
    data_folder = r"D:\python代码\my_project2\datasets\data1"
    img_size = [512, 512]
    classes_name = ["cat", "dog"]
    data_list = get_img_mask_path(data_folder)
    for data in data_list:
        img_tensor, mask_tensor = gen_image_and_label(data[0], data[1], img_size, classes_name)