O arquivo json marcado por labelme se torna uma imagem de máscara em preto e branco para cada categoria

        Use labelme para rotulagem. Após a rotulagem, há apenas um arquivo json. Na tarefa de segmentação, ao ler os dados, você obtém basicamente a imagem e a máscara correspondente à imagem. Em seguida, converta esses dois em tensores e coloque-os no modelo para treinamento.

        Para obter os dados de que precisamos, o arquivo json precisa ser processado. Aqui está uma demonstração de um script processado:

Ao converter o mapa de máscara de numpy para tensor, não há necessidade de realizar a normalização, caso contrário o modelo treinado será muito ineficaz.

O efeito é o seguinte: 


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)

Acho que você gosta

Origin blog.csdn.net/m0_48095841/article/details/128397443
Recomendado
Clasificación