PyTorch图像分割模型——多类别图像分割数据集制作

上一篇文章中，主要介绍了使用segmentation_models_pytorch库进行UNet++模型训练，我们使用的数据集是已经标注好的CamVid数据集，但在实际应用中，我们需要标注自己的多分类图像分割数据集，这篇文章中，就重点介绍下怎么创建自己的图像分割数据集。

首先需强调的是在这里我们用的数据集都是png格式的，生成的标注图像也都是png格式的，因为png图像可以做到无损压缩，能在保证最不失真的情况下尽可能压缩图像文件的大小，而且png用来存储灰度图像时，灰度图像的深度可多到16位，存储彩色图像时，彩色图像的深度可多到48位，并且还可存储多到16位的α通道数据，而jpg不适用于所含颜色很少、具有大块颜色相近的区域或亮度差异十分明显的较简单的图片。简单来说就是jpg格式会在存储时对图像进行压缩，导致mask图像不准确，而png格式不存在这个问题，后面不再对此进行说明。

一、统一图像分辨率

segmentation_models_pytorch库要求我们训练集的图片分辨率的宽和高必须是32的倍数，由于我还未更改训练的代码使其适应多分辨率的数据集，所以我们先提前把数据值的分辨率转换为统一的尺寸，这里我将数据集统一调整为480X320的分辨率。原始图像放在data_mask/orgimages文件夹中，调整分辨率后的图像放在data_mask/images文件夹中。

import cv2
import os


imgpath_list = os.listdir('data_mask/orgimages')

for i in imgpath_list:
    pic_org = cv2.imread(os.path.join('data_mask/orgimages', i))
    pic_new = cv2.resize(pic_org, (480, 320))
    pic_new_name = 'data_mask/images/' + i.split('.')[0] + '.png'
    cv2.imwrite(pic_new_name, pic_new)

二、图像标注

统一分辨率后，我们就可以开始进行图像标注了，这里采用labelme，labelme的使用教程可参考：数据标注软件labelme详解。

使用labelme打开一张图片，开始标注。这里标注了两只猫，两只猫的名字分别为’Gingerbread’和’Coconutmilk’，目标分类也为这两类。

标注完成后进行保存，会生成与图片名相同的json文件，将其存放在data_mask/jsons文件夹中。

{
  "version": "4.5.9",
  "flags": {},
  "shapes": [
    {
      "label": "Gingerbread",
      "points": [
        [
          888.2439024390243,
          409.80487804878044
        ],
        [
          823.6097560975609,
          404.92682926829264
        ],
		...
      ],
      "group_id": null,
      "shape_type": "polygon",
      "flags": {}
    },
    {
      "label": "Coconutmilk",
      "points": [
        [
          882.1463414634146,
          198.8292682926829
        ],
        [
          887.0243902439024,
          263.4634146341463
        ],
		...
      ],
      "group_id": null,
      "shape_type": "polygon",
      "flags": {}
    }
  ],
  "imagePath": "20210801150111.jpg",
  "imageData": "...",
  "imageHeight": 320,
  "imageWidth": 480
}

三、将json转化为标注图像

全部数据集标注完成后，将json文件批量转化为训练所需的标注图像。标注图像所存放的文件夹为data_mask/masks。代码如下。

import os
import cv2
import numpy as np
import json


'''
制作一个只包含分类标注的标签图像，假如我们分类的标签为cat和dog，那么该标签图像中，Background为0，cat为1，dog为2。
我们首先要创建一个和原图大小一致的空白图像，该图像所有像素都是0，这表示在该图像中所有的内容都是Background。
然后根据标签对应的区域使用与之对应的类别索引来填充该图像，也就是说，将cat对应的区域用1填充，dog对应的区域用2填充。
特别注意的是，一定要有Background这一项且一定要放在index为0的位置。
'''

# 分类标签，一定要包含'Background'且必须放在最前面
category_types = ['Background', 'Gingerbread', 'Coconutmilk']
# 将图片标注json文件批量生成训练所需的标签图像png
imgpath_list = os.listdir('data_mask/images')
for img_path in imgpath_list:
    img_name = img_path.split('.')[0]
    img = cv2.imread(os.path.join('data_mask/images', img_path))
    h, w = img.shape[:2]
    # 创建一个大小和原图相同的空白图像
    mask = np.zeros([h, w, 1], np.uint8)

    with open('data_mask/jsons/'+img_name+'.json', encoding='utf-8') as f:
        label = json.load(f)

    shapes = label['shapes']
    for shape in shapes:
        category = shape['label']
        points = shape['points']
        # 将图像标记填充至空白图像
        points_array = np.array(points, dtype=np.int32)
        mask = cv2.fillPoly(mask, [points_array], category_types.index(category))

    # 生成的标注图像必须为png格式
    cv2.imwrite('data_mask/masks/'+img_name+'.png', mask)

此时，我们打开查看生成的标注图像，纳尼，怎么是全黑的图像？

此时，千万不要怀疑，因为这就是正确的标注图像。该标注图像包含背景和我们设定的两个标签，背景位置对应的像素值应该为0，标签图像对应的像素值应该为1或2，该图像只包含0、1和2，因为1和2这个像素值太小了，无法看清，如果打开图像查看的话，就是一个全黑的图像。

如果想查看某个类别的标注情况，可以在mask = cv2.fillPoly(mask, [points_array], category_types.index(category))中将该颜色改为一个清晰可见的颜色，例如将Gingerbread的像素值改成125，将Coconutmilk的像素值改为255，就可以清晰的看到我们标注的图像了。但是最终制作标签的时候要记得将其改回对应的索引值。

for shape in shapes:
    category = shape['label']
    points = shape['points']
    # 将图像标记填充至空白图像
    points_array = np.array(points, dtype=np.int32)
    # mask = cv2.fillPoly(mask, [points_array], category_types.index(category))

    if category == 'Gingerbread':
        # 调试时将某种标注的填充颜色改为255，便于查看用，实际时不需进行该操作
        mask = cv2.fillPoly(mask, [points_array], 125)
    elif category == 'Coconutmilk':
        mask = cv2.fillPoly(mask, [points_array], 255)
    else:
        mask = cv2.fillPoly(mask, [points_array], category_types.index(category))

cv2.imshow('mask', mask)
cv2.waitKey(0)

四、数据集分类

以上步骤完成后，我们需将数据集分为训练集、验证集和测试集进行训练，在这里划分数据集的代码我也写好了，按照7:2:1进行划分。代码如下。

import os
import random
import shutil


'''
├── data(按照7:2:1比例划分)
│   ├── train 存放用于训练的图片
│   ├── trainannot 存放用于训练的图片标注
│   ├── val 存放用于验证的图片
│   ├── valannot 存放用于验证的图片标注
│   ├── test 存放用于测试的图片
│   ├── testannot 存放用于测试的图片标注
'''

# 创建数据集文件夹
dirpath_list = ['data/train', 'data/trainannot', 'data/val', 'data/valannot', 'data/test', 'data/testannot']
for dirpath in dirpath_list:
    if os.path.exists(dirpath):
        shutil.rmtree(dirpath)   # 删除原有的文件夹
        os.makedirs(dirpath)   # 创建文件夹
    elif not os.path.exists(dirpath):
        os.makedirs(dirpath)

# 训练集、验证集、测试集所占比例
train_percent = 0.7
val_percent = 0.2
test_percent = 0.1

# 数据集原始图片所存放的文件夹，必须为png文件
imagefilepath = 'data_mask/images'
total_img = os.listdir(imagefilepath)
# 所有数据集的图片名列表
total_name_list = [row.split('.')[0] for row in total_img]
num = len(total_name_list)
num_list = range(num)
# 训练集、验证集、测试集所包含的图片数目
train_tol = int(num * train_percent)
val_tol = int(num * val_percent)
test_tol = int(num * test_percent)

# 训练集在total_name_list中的index
train_numlist = random.sample(num_list, train_tol)
# 验证集在total_name_list中的index
val_test_numlist = list(set(num_list) - set(train_numlist))
val_numlist = random.sample(val_test_numlist, val_tol)
# 测试集在total_name_list中的index
test_numlist = list(set(val_test_numlist) - set(val_numlist))

# 将数据集和标签图片安装分类情况依次复制到对应的文件夹
for i in train_numlist:
    img_path = 'data_mask/images/'+total_name_list[i]+'.png'
    new_path = 'data/train/'+total_name_list[i]+'.png'
    shutil.copy(img_path, new_path)
    img_path = 'data_mask/masks/' + total_name_list[i] + '.png'
    new_path = 'data/trainannot/' + total_name_list[i] + '.png'
    shutil.copy(img_path, new_path)
for i in val_numlist:
    img_path = 'data_mask/images/'+total_name_list[i]+'.png'
    new_path = 'data/val/'+total_name_list[i]+'.png'
    shutil.copy(img_path, new_path)
    img_path = 'data_mask/masks/' + total_name_list[i] + '.png'
    new_path = 'data/valannot/' + total_name_list[i] + '.png'
    shutil.copy(img_path, new_path)
for i in test_numlist:
    img_path = 'data_mask/images/'+total_name_list[i]+'.png'
    new_path = 'data/test/'+total_name_list[i]+'.png'
    shutil.copy(img_path, new_path)
    img_path = 'data_mask/masks/' + total_name_list[i] + '.png'
    new_path = 'data/testannot/' + total_name_list[i] + '.png'
    shutil.copy(img_path, new_path)

全部完成后，我们用于多类别图像分割的数据集已经制作完毕，接下来的训练就不用多说了，可以参考我的上一篇文章。
PyTorch图像分割模型——segmentation_models_pytorch库的使用