[디버깅] segmentation RGB image -> mask label

[상황]

mmsegmentation은 input으로 0부터 시작하는 2D grayscale image를 받도록 되어있음.

근데 간혹 데이터셋의 mask 이미지가 RGB 컬러로 되어있는 경우가 있음.

코드를 수정하려면 mmseg/datasets/pipelines/loading.py 의 LoadAnnotations의 __call__() 함수를 수정해야함.

여기에 RGB -> mask 맵핑을 추가하여 처리하도록하면 됨.

더 좋은 방법이 있는지는 잘 모르겠음.

 

아래는 RGB image -> mask label 변환 코드

뇌빼고 짜니까 너무 느려서 검색해서 찾음. 이거 쓰니까 빠른 속도로 처리됨

 

출처 https://stackoverflow.com/questions/53059201/how-to-convert-3d-rgb-label-image-in-semantic-segmentation-to-2d-gray-image-a

cmap = {(255, 255, 0): 0, (0, 255, 255): 1, (255, 255, 255): 2}

def rgb2mask(img):

    assert len(img.shape) == 3
    height, width, ch = img.shape
    assert ch == 3

    W = np.power(256, [[0],[1],[2]])

    img_id = img.dot(W).squeeze(-1) 
    values = np.unique(img_id)

    mask = np.zeros(img_id.shape)

    for c in enumerate(values):
        try:
            mask[img_id==c] = cmap[tuple(img[img_id==c][0])] 
        except:
            pass
    return mask