|
使用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)
|