セマンティック セグメンテーション VOC 形式のデータ セットをトレーニングするための自作 DeepLabV3Plus フレームワーク

セマンティック セグメンテーションをトレーニングするために DeepLabV3Plus フレームワークによって使用されるデータセットは、VOC 形式のシングル チャネル カラー画像 (Pillow のカラー モード) です。

ちょっとした常識: OpenCV はデフォルトで 3 チャネル形式で画像を開き、表示と保存は 3 チャネル画像のみをサポートします。つまり、 cv::imshow() は 3 チャンネルのイメージのみを表示できます。 cv::imwrite() と cv::VideoWriter がグレースケール イメージを保存したい場合は、最初に cv::cvtColor() を呼び出してイメージを次の形式に変換する必要があります。 3チャンネル画像です。

2 つの分類を行っており、以下のコードでは color_map の最初の 2 色のみを使用しています。既存のデータ セットのマスクは 24 ビット 3 チャネルであり、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)