Marco DeepLabV3Plus de fabricación propia para entrenar conjuntos de datos en formato VOC de segmentación semántica

El conjunto de datos utilizado por el marco DeepLabV3Plus para entrenar la segmentación semántica es una imagen en color de un solo canal en formato VOC (un modo de color en Pillow).

Poco sentido común: OpenCV abre imágenes en un formato de tres canales de forma predeterminada, y muestra y guarda solo imágenes de tres canales. Es decir, cv::imshow() solo puede mostrar imágenes de tres canales. Si cv::imwrite() y cv::VideoWriter quieren guardar imágenes en escala de grises, primero deben llamar a cv::cvtColor() para convertir la imagen en una imagen de tres canales. .

Estoy haciendo dos clasificaciones y solo uso los dos primeros colores de colours_map en el código siguiente. La máscara del conjunto de datos existente es de tres canales de 24 bits y debe convertirse al formato VOC. El código se muestra a continuación:

# 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)