7

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)
Armali
  • 18,255
  • 14
  • 57
  • 171
Kun
  • 91
  • 1
  • 1
  • 7
  • So the input will contain just those 6 distinct colors you listed? If so, conversion from RBG to grayscale will give you the following grayscale values: [255,29,179,150,226,76] -- 6 distinct values. Then just run it through `cv2.LUT` to remap this to 0-5. – Dan Mašek Mar 12 '17 at 17:59
  • 1
    Change the order or for loops, even that will speed up your code – smttsp Mar 12 '17 at 22:04
  • 1
    @DanMašek Thanks! Your solution works fine. Before I didn't realise that the weights for RGB are different when converting to grayscale. The volume of the grayscale image is 0-255. It means the maximal number of classes is 256. Nevertheless, it will be fine in most cases. Problem could be some colors may have the same gray value. – Kun Mar 13 '17 at 00:08
  • 1
    @smttsp Thanks for your comment. But it doesn't help – Kun Mar 13 '17 at 00:09
  • @Kun Yes, this restricts the input. We're basically looking at a hashing function that generates no conflicts on our input. Ideally you want to get rid of any lookups and replace them by some simple algebra (and use existing optimized numpy or opencv functions that work on the whole matrix in one call). – Dan Mašek Mar 13 '17 at 01:16
  • 1
    @DanMašek Exactly. That's what I am looking for. I've tried to implement a numpy ufunc. But it takes only one element one time of the array. If OpenCV provides interface for us to customise the `convert` function, It would be fine. – Kun Mar 14 '17 at 03:27
  • I show a fast Numpy method of quantising to a given palette here... https://stackoverflow.com/a/57204807/2836621 – Mark Setchell Oct 09 '19 at 08:30

7 Answers7

2

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)

Mendrika
  • 31
  • 2
1

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.

Roy Shilkrot
  • 3,079
  • 29
  • 25
1

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

peterzhang
  • 11
  • 1
1

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)
Armali
  • 18,255
  • 14
  • 57
  • 171
1

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

Sandro Braun
  • 167
  • 7
0

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

QB_
  • 109
  • 7
  • My problem is not a normal color conversion from one color space to another, but to map a color to an index. But thanks for your answer. – Kun Mar 13 '17 at 00:11
  • I was looking also for the same thing. But could not find yet a good conversion between rgb and index. – desmond13 Dec 02 '21 at 14:32
0

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.

Smajjk
  • 394
  • 3
  • 16