The json file marked by labelme becomes a black and white mask image for each category

        Use labelme for labeling. After labeling, there is only one json file. In the segmentation task, when reading data, you basically get the picture and the mask corresponding to the picture. Then convert these two into tensors and put them into the model for training.

        In order to obtain the data we need, the json file needs to be processed. Here is a demo of a processed script:

When converting the mask map from numpy to tensor, there is no need to perform normalization, otherwise the trained model will be very ineffective.

The effect is as follows: 

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)