分割模型训练-制作图像分割数据集

数据集格式：VOC数据集

voc数据集格式：

• VOC2007
  • ImageSets 
    • Segmentation
      • train.txt：训练集txt
      • val.txt：验证集txt
  • JPEGImages：要训练的原图（RGB）
  • SegmentationClass：对应的标签图片（0-255的8位灰度图）

 【注意】在用labelme标注图像的时候，可以选择是否保存图像数据，如果保存，那么生成的json文件中，就会比较大，其中的imagePath存储了图片名称，imageData存储了图片数据。

在代码中，可以使用下面代码，根据一个json文件生成对应的img图片

import base64
import json
import os

json_file_path = 'data/right/111524_005130.json'
output_image_path = 'img.jpg'

# 读取 JSON 文件
with open(json_file_path, 'r') as file:
    data = json.load(file)

# 获取 imageData 字段
image_data_base64 = data.get('imageData')

if image_data_base64 is None:
    raise ValueError("JSON file does not contain 'imageData' field")

# 解码 Base64 字符串为二进制数据
image_data = base64.b64decode(image_data_base64)

# 保存为图像文件
with open(output_image_path, 'wb') as image_file:
    image_file.write(image_data)

print(f"Image saved to {output_image_path}")

导入需要的包：

import os
import json
import numpy as np
import cv2
import shutil
import shutil
import random
from tqdm import tqdm

 导入图片

从保存图片的文件夹data中导入所有的json和jpg文件，保存到对应的列表中。

# 获取data中图片和json文件路径
for root, dirs, files in os.walk('data'):
    for file in files:
        if file.endswith(".json"):
            json_flies_path.append(os.path.join(root, file))
        elif file.endswith(".jpg"):
            img_flies_path.append(os.path.join(root, file))

print("json_flies_path:", json_flies_path)
print("img_flies_path:", img_flies_path)

检查masks文件夹是否存在，存在的话就清空文件夹中的所有内容，不存在就创建文件夹，避免重复生成掩码图片。masks文件夹将用来保存根据标签文件生成的掩码图片。

# 清空masks文件夹
if not os.path.exists('masks'):
    os.makedirs('masks')
    print(f"Directory created: {'masks'}")
else:
    for root, dirs, files in os.walk('masks'):
        for file in files:
            os.remove(os.path.join(root, file))

掩码图像生成函数

思路：

1. 读取一张原图
2. 由于并不是每个json文件中都有4个标签：场地，标志，车道线和围栏。所以生成和原图4个一样大小的全0二维掩码图片数组：img_mask_1，2，3，4.用来保存json文件中每个标签对应的标签点。
3. 打开一个json文件到labelme
4. json文件的shape中，label是标签号，points是标签的坐标集合.
5. 使用for循环遍历每个shape，再用if语句判断每个标签是否存在
6. 若标签存在，就坐标点保存到numpy数组points中，然后使用cv2.fillPoly()函数，用1来填充这个全零数组img_mask_1。依次对剩下3个标签进行操作。
7. 得到4个标签对应的二维数组之后，由于标图的过程中，车道线会在场地上，标志会在场地和围栏上。所以，最终得到的掩码图像，应该是在场地的掩码图像img_mask_1上进行操作，通过遍历img_mask_1数组，首先判断这个坐标上，车道线数组是否为3，如果是，则让这个位置的值为3 。再判断是否为4，也就是围栏，是的话，就让它为4 。最后判断是否为2，标志，是的话，就让它为2 。最后就能得到一个拥有所有标签的二维掩码图像img_mask_1了。
8. 最后返回这个掩码图像。

def labelme2mask_single_img(img_path, labelme_json_path):
    '''
    输入原始图像路径和labelme标注路径，输出 mask
    '''
    
    img_bgr = cv2.imread(img_path)
    print(img_path)
    img_mask = np.zeros(img_bgr.shape[:2]) # 创建空白图像 0-背景
    
    img_mask_1 = np.zeros(img_bgr.shape[:2]) # 场地
    img_mask_2 = np.zeros(img_bgr.shape[:2]) # 标志
    img_mask_3 = np.zeros(img_bgr.shape[:2]) # 车道线
    img_mask_4 = np.zeros(img_bgr.shape[:2]) # 围栏
    
    with open(labelme_json_path, 'r', encoding='utf-8') as f:
        labelme = json.load(f)
        
    for each in labelme['shapes']: # 遍历所有标注，找到属于当前类别的标注
        if each['label'] == '1':
            # 获取点的坐标
            points = [np.array(each['points'], dtype=np.int32).reshape((-1, 1, 2))]
            # 在空白图上画 mask（闭合区域）
            img_mask_1 = cv2.fillPoly(img_mask_1, points, color= 1)
        if each['label'] == '2':
            # 获取点的坐标
            points = [np.array(each['points'], dtype=np.int32).reshape((-1, 1, 2))]
            # 在空白图上画 mask（闭合区域）
            img_mask_2 = cv2.fillPoly(img_mask_2, points, color=2)
        if each['label'] == '3':
            # 获取点的坐标
            points = [np.array(each['points'], dtype=np.int32).reshape((-1, 1, 2))]
            # 在空白图上画 mask（闭合区域）
            img_mask_3 = cv2.fillPoly(img_mask_3, points, color=3)
        if each['label'] == '4':
            # 获取点的坐标
            points = [np.array(each['points'], dtype=np.int32).reshape((-1, 1, 2))]
            # 在空白图上画 mask（闭合区域）
            img_mask_4 = cv2.fillPoly(img_mask_4, points, color=4) 
            
    for i in range(img_mask_1.shape[0]):
        for j in range(img_mask_1.shape[1]):
            # 增加围栏
            if img_mask_4[i][j] == 4 :
                img_mask_1[i][j] = 4
            # 增加车道线
            if img_mask_3[i][j] == 3 :
                img_mask_1[i][j] = 3
            # 增加车道线
            if img_mask_2[i][j] == 2 :
                img_mask_1[i][j] = 2                        
                    
    return img_mask_1

生成掩码图像

利用for循环，遍历图片数组和json数组。并将掩码图像img_mask保存到masks文件夹中。

# 开始生成掩码图像    
for img_path, labelme_json_path  in zip(img_flies_path, json_flies_path):
    try:
        img_mask = labelme2mask_single_img(img_path, labelme_json_path)
        prefix1, _ = os.path.splitext(img_path)
        prefix1=prefix1.split('\\')[-1]
        mask_path = prefix1 + '.png'
        cv2.imwrite(os.path.join('masks',mask_path), img_mask)
    except Exception as E:
        print(img_path, '转换失败', E)

划分掩码图像的训练集和测试集

设置划分比例为0.2，然后将图片列表随机打乱，根据比例值，划分训练集列表train_files和测试集列表val_files。最后生成对应的txt文件。

test_frac = 0.2  # 测试集比例
random.seed(123) # 随机数种子，便于复现
folder = 'masks'

img_paths = os.listdir(folder)
random.shuffle(img_paths) # 随机打乱

val_number = int(len(img_paths) * test_frac) # 测试集文件个数
train_files = [i.split('.')[0] for i in img_paths[val_number:]]         # 训练集文件名列表
val_files = [i.split('.')[0] for i in img_paths[:val_number]]           # 测试集文件名列表

# 生成训练集和测试集的txt文件
with open(os.path.join('img_dir/', 'train.txt'), 'w') as f:
    for name in train_files:
        f.write(name + '\n')

with open(os.path.join('img_dir/', 'val.txt'), 'w') as f:
    for name in val_files:
        f.write(name + '\n')
        
print('数据集文件总数', len(img_paths))
print('训练集文件个数', len(train_files))
print('测试集文件个数', len(val_files))

验证

读取一张掩码图片，然后统计每个值出现的次数，发现所有标签都出现了。说明掩码图像生成成功。