自制VOC格式图像分割数据集:使用python+PIL生成8位深的RGB图像

最近在做图像分割任务的任务,使用的是DeepLabV3Plus框架,所以需要制作VOC格式的数据集。在制作的时候发现VOC数据集中的Mask尽管是彩色的,但是竟然是8位深的。
自制VOC格式图像分割数据集:使用python+PIL生成8位深的RGB图像_第1张图片
而一般情况下,图像的每个通道都有8位,因此RGB图像的位深度应该为24位。8位深的图像一般都是灰度图像。起初,我并没有发现VOC2012中mask的特殊性,直接用24位的Mask作为自制VOC数据集的Mask,在pytorch上跑程序报了张量维度错误:

1only batches of spatial targets supported (non-empty 3D tensors) but got targets of size: : [16, 256, 256, 3]

于是利用代码将自己数据集的24位深的彩色Mask转换成8位深的RGB图像,利用的是PIL的“P”模式和python的imgviz包,具体什么是PIL的“P”模式可参考这篇文章。
代码如下(可实现多分类):

from random import randint
import os
import numpy as np
import imgviz
from PIL import Image

def get_gray_cls(van_lbl, array_lbl):
    cls = [0]  # 用来存储灰度图像中每种类别所对应的像素,默认背景色为0
    for x in range(van_lbl.size[0]):
        for y in range(van_lbl.size[1]):
            if array_lbl[x, y] not in cls:
                cls.append(array_lbl[x, y])
    return cls

def get_P_cls(cls_gray):
    cls_P = []  # 将灰度图像中的每类像素用0~N表示
    for i in range(len(cls_gray)):
        cls_P.append(i)
    return cls_P

def array_gray_to_P(cls_gray, cls_P, array):
    for i in range(len(cls_gray)):
        array[array == cls_gray[i]] = cls_P[i]
    return array

if __name__ == '__main__':
    label_from_PATH = "E:\label\label24"  # 存放原始label的路径
    label_to_PATH = "E:\label\label8"  # 存放转换后图像当路径
    van_file = 'E:\label\\van_label.png'  # 必须是一张包含所有类别的图像,称之为先锋图像
    van_lbl = Image.open(van_file).convert('L')  # 将先锋图像转换为灰度图像

    array_lbl = np.array(van_lbl)  # 获得灰度图像的numpy矩阵

    cls_gray = get_gray_cls(van_lbl, array_lbl)  # 获取灰度图像中每种类别所对应的像素值
    cls_P = get_P_cls(cls_gray)  # 将灰度图像中的每种类别所对应的像素值映射为0~N

    # print(cls_gray)
    # print(cls_P)

    file_list = os.listdir(label_from_PATH)
    if not os.path.isdir(label_to_PATH):
        os.mkdir(label_to_PATH)
    # 遍历每一张原始图像
    for file_name in file_list:
        file_path = os.path.join(label_from_PATH, file_name)
        orig_lbl = Image.open(file_path).convert('L')  # 将图像转换为灰度图像
        array_gray = np.array(orig_lbl)  # 获得灰度图像的numpy矩阵
        array_P = array_gray_to_P(cls_gray, cls_P, array_gray)  # 将灰度图像的numpy矩阵值映射为0~N
        label = Image.fromarray(array_P.astype(np.uint8), mode='P')  # 转换为PIL的P模式
        # 转换成VOC格式的P模式图像
        colormap = imgviz.label_colormap()
        label.putpalette(colormap.flatten())
        label.save(os.path.join(label_to_PATH, file_name))

完成后的图像标签如下:
自制VOC格式图像分割数据集:使用python+PIL生成8位深的RGB图像_第2张图片

要特别注意的是,训练模型时,一定要用PIL来操作图像,不要使用OpenCV,否则打开的图像还是24位,切记!!!

你可能感兴趣的:(计算机视觉,python,计算机视觉,深度学习)