语义分割数据集图像映射、ValueError: Target size (torch.Size([4, 20, 320, 320])) must be the same as input size (

目录

在处理语义分割数据集时需要将RGB的mask图像按照类别进一步处理成对应的灰度图。

1.结果图

2.做法（RGBmask得到label灰度图）：

2.1 修改自己的Clss字典

2.2 修改问价夹

3. 灰度图label恢复RGBmask

---

在处理语义分割数据集时需要将RGB的mask图像按照类别进一步处理成对应的灰度图。

1.结果图

以下图为例。

原图：

使用labelme打完标签之后的RGBmask图： 

把上图按像素点类别映射成灰度图，得到label图如下（由于灰度值是0-4，灰度值较小,所以肉眼看起来看不到区别）：

2.做法（RGBmask得到label灰度图）：

2.1 修改自己的Clss字典

多类别的时候，指定自己的mask图每一种颜色（RGB值）代表什么类别。

Cls = namedtuple('cls', ['name', 'id', 'color'])  # 姓名元组
Clss = [
    Cls('_background_', 0, (0, 0, 0)),
    Cls('wall', 1, (0, 0, 128)),
    Cls('car', 2, (128, 0, 0)),
    Cls('ground', 3, (128, 128, 0)),
    Cls('door', 4, (0, 128, 0)),
]  # 把每一个类别使用id与像素值对应上，并且颜色通道值按照r,g,b的顺序排列


# 其中(0,0,0）代表黑色
# （0,0,128）代表蓝色
# （128,0，0）代表红色
# （128,128,0)表示红绿混合色黄色
# （0,128,0）代表绿色

2.2 修改问价夹

修改放置灰度图的文件夹名字和放置RGB图的文件夹名字（我这里是把RGB图放到了F盘的mask目录里），之后通过

color_gray(color_dict)

函数run一下后就在gray文件夹里面得到了label图。

gts_gray_path = 'F:\gray'  # 放置灰度图的文件夹名字
gts_color_path = 'F:\mask'  # 放置RGB图的文件夹名字

完整代码如下： 

import os, sys, time, cv2
import numpy as np
from collections import namedtuple

# 类别信息
gts_gray_path = 'F:\gray'  # 放置灰度图的文件夹名字
gts_color_path = 'F:\mask'  # 放置RGB图的文件夹名字
# gts_color_path = 'color'  # 放置RGB图的文件夹名字
Cls = namedtuple('cls', ['name', 'id', 'color'])  # 姓名元组
Clss = [
    Cls('_background_', 0, (0, 0, 0)),
    Cls('wall', 1, (0, 0, 128)),
    Cls('car', 2, (128, 0, 0)),
    Cls('ground', 3, (128, 128, 0)),
    Cls('door', 4, (0, 128, 0)),
]  # 把每一个类别使用id与像素值对应上，并且颜色通道值按照r,g,b的顺序排列


# 其中(0,0,0）代表黑色
# （0,0,128）代表蓝色
# （128,0，0）代表红色
# （128,128,0)表示红绿混合色黄色
# （0,128,0）代表绿色
def gray_color(color_dict, gray_path=gts_gray_path, color_path=gts_color_path):
    '''
    swift gray image to color, by color mapping relationship
    :param color_dict:color mapping relationship, dict format
    :param gray_path:gray imgs path
    :param color_path:color imgs path
    :return:
    '''
    pass
    t1 = time.time()
    gt_list = os.listdir(gray_path)
    for index, gt_name in enumerate(gt_list):
        gt_gray_path = os.path.join(gray_path, gt_name)
        gt_color_path = os.path.join(color_path, gt_name)
        gt_gray = cv2.imread(gt_gray_path, cv2.IMREAD_GRAYSCALE)
        assert len(gt_gray.shape) == 2  # make sure gt_gray is 1band

        # # region method 1: swift by pix, slow
        # gt_color = np.zeros((gt_gray.shape[0],gt_gray.shape[1],3),np.uint8)
        # for i in range(gt_gray.shape[0]):
        #     for j in range(gt_gray.shape[1]):
        #         gt_color[i][j] = color_dict[gt_gray[i][j]]      # gray to color
        # # endregion

        # region method 2: swift by array
        # gt_color = np.array(np.vectorize(color_dict.get)(gt_gray),np.uint8).transpose(1,2,0)
        # endregion

        # region method 3: swift by matrix, fast
        gt_color = matrix_mapping(color_dict, gt_gray)
        # endregion

        gt_color = cv2.cvtColor(gt_color,
                                cv2.COLOR_RGB2BGR)  # cv2.cvtColor(p1,p2) 是颜色空间转换函数，p1是需要转换的图片，p2是转换成何种格式。,得到的图矩阵是按R,G,B排列的，CV2是按B，G，R排列的
        cv2.imwrite(gt_color_path, gt_color, )
        cv2.imshow('color image:', gt_color)
        cv2.waitKey(0)
        process_show(index + 1, len(gt_list))
    print(time.time() - t1)


def color_gray(color_dict, color_path=gts_color_path, gray_path=gts_gray_path, ):
    '''
    swift color image to gray, by color mapping relationship
    :param color_dict:color mapping relationship, dict format
    :param gray_path:gray imgs path
    :param color_path:color imgs path
    :return:
    '''
    gray_dict = {}
    for k, v in color_dict.items():
        gray_dict[v] = k
    t1 = time.time()
    gt_list = os.listdir(color_path)
    for index, gt_name in enumerate(gt_list):
        gt_gray_path = os.path.join(gray_path, gt_name)
        gt_color_path = os.path.join(color_path, gt_name)
        color_array = cv2.imread(gt_color_path, cv2.IMREAD_COLOR)
        assert len(color_array.shape) == 3
        print(color_array.shape)
        gt_gray = np.zeros((color_array.shape[0], color_array.shape[1]), np.uint8)
        b, g, r = cv2.split(color_array)
        # zeros = np.zeros(color_array.shape[:2], dtype="uint8")
        # merged_b= cv2.merge([zeros, g, zeros])
        # cv2.imshow('b',merged_b)
        # cv2.waitKey(0)
        color_array = np.array([r, g, b])
        for cls_color, cls_index in gray_dict.items():  # 将图像的像素点的三个通道值与约定的类别代表的像素值逐个逐个地比较，直到可以将整个图像遍历
            cls_pos = arrays_jd(color_array, cls_color)  # 将图像的像素点的三个通道的像素与约定类别对应的三通道像素值进行比较
            # print(cls_pos.shape,gt_gray.shape,cls_index)
            gt_gray[cls_pos] = cls_index  # 如果比较发现是这个类别的三通道像素值，那么就把这个地方的灰度值换成对应的类别值
        # 其中cls_pos的形状必须要与gt_gray一致，cls_pos是有False和True构成的数组，True代表这个位置的数字替换，False代表这个位置的数字不替换。
        # print(gt_gray)
        cv2.imwrite(gt_gray_path, gt_gray)
        # cv2.imshow('gray image:', gt_gray)
        # cv2.waitKey(0)
        process_show(index + 1, len(gt_list))
    print(time.time() - t1)


def arrays_jd(arrays, cond_nums):  # 用于比较像素通道值是否属于某一个类别
    r = arrays[0] == cond_nums[0]  # 比较r通道的像素值
    g = arrays[1] == cond_nums[1]  # 比较g通道的像素值
    b = arrays[2] == cond_nums[2]  # 比较b通道的像素值
    #  print(r & g & b)
    return r & g & b  # 只有三个通道都与约定的id对应的三个通道颜色像素值都一致时才认为这个像素点是属于这个id对应的类别


def matrix_mapping(color_dict, gt):  # 用于将灰度图按照类别转化为三通道图
    colorize = np.zeros([len(color_dict), 3], 'uint8')
    # print(colorize.shape)
    for cls, color in color_dict.items():
        colorize[cls, :] = list(color)  # 直接替换行，得到一个代表像素点类别的带三通道值的矩阵。
    # print(colorize)
    # print(gt)
    ims = colorize[gt,
          :]  # 将gt里面的元素作为索引逐个放入colorize中，当成访问colorize中的每一行，从而按照colorize的每一行按照灰度值代表的id逐个填入，从而将colorize从二维变成三维，按照灰度图的id索引将图像还原为BGR图像
    # print(ims)
    # ims = ims.reshape([gt.shape[0], gt.shape[1], 3])
    return ims


def nt_dic(nt=Clss):
    '''
    swift nametuple to color dict
    :param nt: nametuple
    :return:
    '''
    pass
    color_dict = {}
    for cls in nt:
        color_dict[cls.id] = cls.color
    return color_dict


def process_show(num, nums, pre_fix='', suf_fix=''):
    '''
    auxiliary function, print work progress
    :param num:
    :param nums:
    :param pre_fix:
    :param suf_fix:
    :return:
    '''
    rate = num / nums
    ratenum = round(rate, 3) * 100
    bar = '\r%s %g/%g [%s%s]%.1f%% %s' % \
          (pre_fix, num, nums, '#' * (int(ratenum) // 5), '_' * (20 - (int(ratenum) // 5)), ratenum, suf_fix)
    sys.stdout.write(bar)
    sys.stdout.flush()


if __name__ == '__main__':
    pass
    color_dict = nt_dic()
    # gray_color(color_dict)
    color_gray(color_dict)
gt_gray = np.zeros((2, 2), np.uint8)
cls_pos = np.array([[True, False], [True, True]])
gt_gray[cls_pos] = 9
print(gt_gray)

3. 灰度图label恢复RGBmask

代码中使用函数：

gray_color(color_dict)

将color_gray(color_dict)注释掉

并指定好映射和文件夹位置

pycharm中点击运行if __name__ == '__main__':就可以将灰度label恢复为RGBmask图了。