语义分割标签处理RBG 到gray再转回RGB（通用）

语义分割标签处理RBG 到gray再转回RGB

import numpy as np
import cv2
import os
from PIL import Image

# moon 编码过程  地面, 天空, 大岩石,  小岩石

color_segmentation = np.asarray([[0, 0, 0], [255, 0, 0], [0, 0, 255], [0, 255, 0]],dtype=np.uint8)
# mars
# color_segmentation = np.asarray([[0, 0, 0],[255, 255, 0]],dtype=np.uint8) # n_classes=2

# color_segmentation = np.asarray([  # n_classes=21
#     [0, 0, 0],  # [0]背景
#     [180, 120, 120],
#     [6, 230, 230],
#     [80, 50, 50],
#     [4, 200, 3],
#     [120, 120, 80],
#     [140, 140, 140],
#     [204, 5, 255],
#     [230, 230, 230],
#     [4, 250, 7],
#     [40, 150, 255],
#     [235, 255, 7],
#     [150, 5, 61],
#     [120, 120, 70],
#     [8, 255, 51],
#     [255, 6, 82],
#     [143, 255, 140],
#     [204, 255, 4],
#     [255, 51, 7],
#     [204, 70, 3],
#     [0, 102, 200],
#     [61, 230, 250],
#     [255, 6, 51],
#     [11, 102, 255],
#     [255, 7, 71],
#     [255, 9, 224],
#     [9, 7, 230],
#     [220, 220, 220],
#     [255, 9, 92]
# ], dtype=np.uint8)


def decode_segmap(label_mask, n_classes=2):
    r = label_mask.copy()
    g = label_mask.copy()
    b = label_mask.copy()

    for ll in range(0, n_classes):
        position = label_mask == ll
        b[label_mask == ll] = color_segmentation[ll, 0]
        g[label_mask == ll] = color_segmentation[ll, 1]
        r[label_mask == ll] = color_segmentation[ll, 2]
    rgb = np.zeros((label_mask.shape[0], label_mask.shape[1], 3))
    print("rgb:",rgb.shape)
    
    # 使用Image.PIL 读取
    rgb[:, :, 0] = r
    rgb[:, :, 1] = g
    rgb[:, :, 2] = b

    # 若使用 cv2.imread()读取图片, 使用以下代码,(读取BGR图像 颜色反了。。bgr)
    # rgb[:, :, 0] = b
    # rgb[:, :, 1] = g
    # rgb[:, :, 2] = r

    rgb = rgb.astype(np.uint8)  ##重要！ 要不然opencv显示有问题
    return rgb

def all_change():

    root_dir = "./voc/VOCdevkit/VOC2012/SegmentationClassAug/"
    list_gray_img = os.listdir(root_dir)
    for img_name in list_gray_img:
        path_gray = root_dir + img_name
        laber_mask = cv2.imread(path_gray, 0)  ##灰度 单通道读取
        # print(laber_mask.shape)
        color_img = decode_segmap(laber_mask)

        cv2.imshow("laber_mask", laber_mask * 20)
        cv2.imshow("color", color_img)
        cv2.waitKey()


# 输入灰度标签和解码的关系
def mydecode(img_name, label2color_dict):
    png = Image.open(img_name)
    w, h = png.size
    # ----------------(gray--->gray)---------------------
    png = np.array(png, np.uint8)  # 输入为灰度 h, w
    out_png = np.zeros([h, w, 3])  # 新建的RGB为输出 h, w, c
    # 关系映射
    for i in range(png.shape[0]):  # i for h
        for j in range(png.shape[1]):
            color = label2color_dict[png[i, j]]
            out_png[i, j, 0] = color[0]
            out_png[i, j, 1] = color[1]
            out_png[i, j, 2] = color[2]

    # print("out_png:", out_png.shape)
    out_png = Image.fromarray(np.array(out_png, np.uint8)) # 再次转化为Imag进行保存
    out_png.show()

if __name__ == '__main__':

    # -------------mars decoder-----------------
    # label2color_dict = {
    #         0: [0, 0, 0], # background
    #         1: [255, 255, 0],  # cornsilk
    #     }

    # -------------moon decoder-----------------
    label2color_dict = {
        0: [0, 0, 0],  # 'groud' 黑色
        1: [255, 0, 0],  # 'sky',
        2: [0, 0, 255],  # 'b_rock'
        3: [0, 255, 0]  # 's_rock'
    }
    # img_name = "F:/rock_project/deeplabv3-plus-pytorch-main/VOCdevkit/VOC2007/SegmentationClass/000032.png"
    img_name = "F:/rock_data/Mars_Rover//moon_rock/rock_datasets/images/test_gray/clean0011.png"
    # img_name = "F:/rock_data/Mars_Rover/MarsRocks/MarsData/MarsRocks/SegmentationClass_gray/331_0331MR0013390000301024E01_DXXX_crop134_0.png"
    # img_name =  "F:/rock_data/Mars_Rover/MarsRocks/MarsData/other/rgb-gray-rgb/gray_1/331_0331MR0013390000301024E01_DXXX.png"
    # img_name = "F:/rock_data/Mars_Rover/MarsRocks/MarsData/other/rgb-gray-rgb/gray_1/331_0331MR0013390000301024E01_DXXX.png"


    path_gray = img_name
    laber_mask = cv2.imread(path_gray, 0)  ## 灰度单通道读取
    # print(laber_mask.shape)
    color_img = decode_segmap(laber_mask, n_classes=4)

    # cv2.imshow("laber_mask", laber_mask * 2)
    cv2.imshow("color", color_img)
    cv2.waitKey()

参考：https://blog.csdn.net/yang332233/article/details/116778827