最近在做图像分割任务的任务,使用的是DeepLabV3Plus框架,所以需要制作VOC格式的数据集。在制作的时候发现VOC数据集中的Mask尽管是彩色的,但是竟然是8位深的。
而一般情况下,图像的每个通道都有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))
完成后的图像标签如下:
要特别注意的是,训练模型时,一定要用PIL来操作图像,不要使用OpenCV,否则打开的图像还是24位,切记!!!