DeepLabV3Plus框架训练语义分割的数据集是VOC格式,单通道的彩色图(Pillow中的一种色彩模式)。
小常识:OpenCV打开图片默认是三通道格式,显示和保存也是只支持三通道的图像。即cv::imshow()只能显示三通道的图片,cv::imwrite()、cv::VideoWriter如果要保存灰度图,必须将调用cv::cvtColor()先将图像转换成三通道的。
我是做二分类,只用得上下面代码中colors_map前面两种颜色,已有的数据集的mask是24位三通道的,要转换成VOC格式。代码如下:
# coding=utf-8
import os
import numpy as np
from tqdm import tqdm
from PIL import Image, ImageOps
# 颜色映射表,当把颜色表的颜色种类增加到17种时(17>2^4),图片的位深度就会变成8。
colors_map = [
(0, 0, 0), # 黑色(背景像素) 0
(128, 0, 0), # 暗红色 1
(255, 0, 0), # 红色 2
(0, 255, 0), # 绿色 3
(0, 0, 255), # 蓝色 4
(255, 255, 0), # 黄色 5
(0, 255, 255), # 青色 6
(255, 0, 255), # 深红色 7
(255, 192, 203), # 粉红色 8
(0, 199, 140), # 玉色 9
(128, 42, 42), # 棕色 10
(255, 97, 0), # 橙色 11
(238, 130, 238), # 紫罗兰 12
(128, 0, 128), # 紫色 13
(0, 0, 139), # 深蓝色 14
(0, 0, 128), # 海军蓝 15
(211, 242, 231), # 水绿色 16
(0, 139, 139), # 深青色 17
(0, 100, 0) # 深绿色 18
]
def rgb_to_mask(image_file, out_dir):
image = Image.open(image_file)
image_gray = ImageOps.grayscale(image)
width = image.size[0]
height = image.size[1]
for x in range(0, width):
for y in range(0, height):
pix = image_gray.getpixel((x, y))
if pix != 0:
image_gray.putpixel((x, y), 1)
image_p = image_gray.convert('P') # P代表索引图片,
palette = np.array(colors_map).reshape(-1).tolist()
image_p.putpalette(palette)
(filename, extension) = os.path.splitext(image_file)
filename = os.path.basename(filename)
image_p_name = os.path.join(out_dir, filename + ".png")
image_p.save(image_p_name)
# image_p.show()
if __name__ == '__main__':
input_path = 'D:\\BaiduNetdiskDownload\\CelebAMask-HQ\\CelebA-HQ-mask\\'
output_path = 'D:\\BaiduNetdiskDownload\\CelebAMask-HQ\\CelebA-HQ-mask\\result\\'
if not os.path.exists(output_path):
os.makedirs(output_path)
files = os.listdir(input_path)
image_files = list(filter(lambda x: '.png' in x, files))
image_files = tqdm(image_files)
for image_file in image_files:
image_file = input_path + image_file
rgb_to_mask(image_file, output_path)