将 RGB 图像转换为索引图像

Question

I want to convert a 3 channel RGB image to a index image with Python. It's used for handling the labels of training a deep net for semantic segmentation. By index image I mean it has one channel and each pixel is the index, which should starts with zero. And certainly they should have the same size. The conversion is based on the following mapping in Python dict:

color2index = {
        (255, 255, 255) : 0,
        (0,     0, 255) : 1,
        (0,   255, 255) : 2,
        (0,   255,   0) : 3,
        (255, 255,   0) : 4,
        (255,   0,   0) : 5
    }

I've implemented a naive function:

def im2index(im):
    """
    turn a 3 channel RGB image to 1 channel index image
    """
    assert len(im.shape) == 3
    height, width, ch = im.shape
    assert ch == 3
    m_lable = np.zeros((height, width, 1), dtype=np.uint8)
    for w in range(width):
        for h in range(height):
            b, g, r = im[h, w, :]
            m_lable[h, w, :] = color2index[(r, g, b)]
    return m_lable

The input im is a numpy array created by cv2.imread() . However, this code is really slow. Since the im is in numpy array I firstly tried the ufunc of numpy with something like this:

RGB2index = np.frompyfunc(lambda x: color2index(tuple(x)))
indices = RGB2index(im)

But it turns out that the ufunc takes only one element each time. I was unable to give the function three arguments(RGB value) one time.

So is there any other ways to do the optimization? The mapping has not to be that way, if a more efficient data structure exists. I noticed that the access of a Python dict dose not cost much time, but the casting from numpy array to tuple (which is hashable) does.

PS: One idea I got is to implement a kernel in CUDA. But it would be more complicated.

UPDATA1: Dan Mašek's Answer works fine. But first we have to convert the RGB image to grayscale. It could be problematic when two colors have the same grayscale value.

I paste the working code here. Hope it could help others.

lut = np.ones(256, dtype=np.uint8) * 255
lut[[255,29,179,150,226,76]] = np.arange(6, dtype=np.uint8)
im_out = cv2.LUT(cv2.cvtColor(im, cv2.COLOR_BGR2GRAY), lut)

Answer 1

What about this?

color2index = {
    (255, 255, 255) : 0,
    (0,     0, 255) : 1,
    (0,   255, 255) : 2,
    (0,   255,   0) : 3,
    (255, 255,   0) : 4,
    (255,   0,   0) : 5
}

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 i, c in enumerate(values):
        try:
            mask[img_id==c] = color2index[tuple(img[img_id==c][0])] 
        except:
            pass
    return mask

Then just call:

mask = rgb2mask(ing)

Answer 2

actually for-loop takes much time.

binary_mask = (im_array[:,:,0] == 255) & (im_array[:,:,1] == 255) & (im_array[:,:,2] == 0) 

maybe above code can help you

Answer 3

I've implemented a naive function: … I firstly tried the ufunc of numpy with something like this: …

I suggest using an even more naive function which converts just one pixel:

def rgb2index(rgb):
    """
    turn a 3 channel RGB color to 1 channel index color
    """
    return color2index[tuple(rgb)]

Then using a numpy routine is a good idea, but we don't need a ufunc :

np.apply_along_axis(rgb2index, 2, im)

Here numpy.apply_along_axis() is used to apply our rgb2index() function to the RGB slices along the last of the three axes (0, 1, 2) for the whole image im .

We could even do without the function and just write:

np.apply_along_axis(lambda rgb: color2index[tuple(rgb)], 2, im)

Answer 4

Similar to what Armali and Mendrika proposed, I somehow had to tweak it a little bit to get it to work (maybe totally my fault). So I just wanted to share a snippet that works.

COLORS = np.array([
    [0, 0, 0],
    [0, 0, 255],
    [255, 0, 0]
])
W = np.power(255, [0, 1, 2])

HASHES = np.sum(W * COLORS, axis=-1)
HASH2COLOR = {h : c for h, c in zip(HASHES, COLORS)}
HASH2IDX = {h: i for i, h in enumerate(HASHES)}


def rgb2index(segmentation_rgb):
    """
    turn a 3 channel RGB color to 1 channel index color
    """
    s_shape = segmentation_rgb.shape
    s_hashes = np.sum(W * segmentation_rgb, axis=-1)
    func = lambda x: HASH2IDX[int(x)]
    segmentation_idx = np.apply_along_axis(func, 0, s_hashes.reshape((1, -1)))
    segmentation_idx = segmentation_idx.reshape(s_shape[:2])
    return segmentation_idx

segmentation = np.array([[0, 0, 0], [0, 0, 255], [255, 0, 0]] * 3).reshape((3, 3, 3))
rgb2index(segmentation)

Example plot

The code is also available here: https://github.com/theRealSuperMario/supermariopy/blob/dev/scripts/rgb2labels.py

Answer 5

Did you check Pillow library https://python-pillow.org/ ? As I remember, it has some classes and methods to deal with color conversion. See: https://pillow.readthedocs.io/en/4.0.x/reference/Image.html#PIL.Image.Image.convert

Answer 6

Here's a small utility function to convert images (np.array) to per-pixel labels (indices), which can also be a one-hot encoding:

def rgb2label(img, color_codes = None, one_hot_encode=False):
    if color_codes is None:
        color_codes = {val:i for i,val in enumerate(set( tuple(v) for m2d in img for v in m2d ))}
    n_labels = len(color_codes)
    result = np.ndarray(shape=img.shape[:2], dtype=int)
    result[:,:] = -1
    for rgb, idx in color_codes.items():
        result[(img==rgb).all(2)] = idx

    if one_hot_encode:
        one_hot_labels = np.zeros((img.shape[0],img.shape[1],n_labels))
        # one-hot encoding
        for c in range(n_labels):
            one_hot_labels[: , : , c ] = (result == c ).astype(int)
        result = one_hot_labels

    return result, color_codes


img = cv2.imread("input_rgb_for_labels.png")
img_labels, color_codes = rgb2label(img)
print(color_codes) # e.g. to see what the codebook is

img1 = cv2.imread("another_rgb_for_labels.png")
img1_labels, _ = rgb2label(img1, color_codes) # use the same codebook

It calculates (and returns) the color codebook if None is supplied.

Answer 7

If you are happy using MATLAB - maybe saving the result as *.mat and loading with scipy.io.loadmat - there is the rgb2ind function in MATLAB, which does exactly what you are asking for. If not, it could be used as inspiration for a similar implementation in Python.