Как преобразовать изображение 3D-метки RGB (в семантической сегментации) в 2D-изображение серого цвета, а индексы классов начинаются с 0?

Question

У меня есть метка семантической сегментации rgb, если в ней есть 3 класса, и каждое значение RGB является одним из:

[255, 255, 0], [0, 255, 255], [255, 255, 255]

соответственно, тогда я хочу отобразить все значения в файле rgb в новое 2-мерное изображение метки в соответствии с dict:

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

после этого все значения в новом файле серой метки равны 0 или 2. Есть ли эффективный способ решить эту проблему? Например трансляция в numpy.

Answer 1

Как насчет этого:

mask_mapping = {
    (255, 255, 0):   0,
    (0, 255, 255):   1,
    (255, 255, 255): 2,
}
for k in mask_mapping:
    label[(label == k).all(axis=2)] = mask_mapping[k]

Я думаю, что он основан на той же идее, что и принятый метод, но выглядит более четко.

Answer 2

Я попробовал это ...

Во-первых, я заметил, что в следующей таблице значений RGB значения Грина одинаковы, поэтому нет смысла проверять их.

Во-вторых, если вы разделите значения в массиве на 255, вы получите нули и те, которые очень близки к нужной вам маркировке. Итак, если вы сделаете немного математики:

t = R/255 + 2B/255 -1

тогда вы получите это для значений в словаре:

  R   G   B    t
==================
255 255   0    0
 0  255 255    1
255 255 255    2

Код для сравнения с парой других ответов выглядит следующим образом:

#!/usr/bin/env python3

import numpy as np

def me(img): 
    """Return R + 2B - 1 as label"""
    return np.uint8((img[:,:,0]/255) + 2*(img[:,:,2]/255) - 1) 

def deepak(img):
    r = np.array([255, 255, 0])
    g = np.array([0, 255, 255])
    b = np.array([255, 255, 255])

    label_seg = np.zeros((img.shape[:2]), dtype=np.uint8)
    label_seg[(img==r).all(axis=2)] = 0
    label_seg[(img==g).all(axis=2)] = 1
    label_seg[(img==b).all(axis=2)] = 2
    return label_seg

def marios(label):
    mask_mapping = {
       (255, 255, 0):   0,
       (0, 255, 255):   1,
       (255, 255, 255): 2,
    }
    for k in mask_mapping:
        label[(label == k).all(axis=2)] = mask_mapping[k]

    return label

# Generate a sample image
img = np.zeros((480,640,3), dtype=np.uint8)
img[:160,:,:]    = [255,255,0]
img[160:320,:,:] = [0,255,255]
img[320:,:,:]    = [255,255,255]

Время получается так:

In [134]: %timeit deepak(img)
15.4 ms ± 181 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [135]: %timeit marios(img)
15.4 ms ± 166 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [172]: %timeit me(img)                                                                           
869 µs ± 8.93 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)each)

Стоит ли ускорение в 18 раз меньше читаемого кода - это еще один аргумент, хотя комментарии могут помочь: -)

---

Обратите внимание, что, честно говоря, Дипаку его время можно уменьшить до 0 10,3 мс, удалив ненужную строку, ниже которой обнуляются некоторые элементы в массиве нулей:

label_seg[(img==r).all(axis=2)] = 0

Answer 3

Вы можете сделать это:

# the three channels
r = np.array([255, 255, 0])
g = np.array([0, 255, 255])
b = np.array([255, 255, 255])

label_seg = np.zeros((img.shape[:2]), dtype=np.int)
label_seg[(img==r).all(axis=2)] = 0
label_seg[(img==g).all(axis=2)] = 1
label_seg[(img==b).all(axis=2)] = 2

Так что, если

img = np.array([[r,g,b],[r,r,r],[b,g,r],[b,g,r]])

тогда

label_seg = array([[0, 1, 2],
                   [0, 0, 0],
                   [2, 1, 0],
                   [2, 1, 0]])