Mika Mika - 3 months ago 9
Python Question

How to map pixels (R, G, B) in a collection of images to a distinct pixel-color-value indices?

Lets say one has 600 annotated semantic segmentation mask images, which contain 10 different colors, each representing one entity. These images are in a numpy array of shape (600, 3, 72, 96), where n = 600, 3 = RGB channels, 72 = height, 96 = width.

How to map each RGB-pixel in the numpy array to a color-index-value? For example, a color list would be [(128, 128, 0), (240, 128, 0), ...n], and all (240, 128, 0) pixels in the numpy array would be converted to index value in unique mapping (= 1).

How to do this efficiently and with less code? Here's one solution I came up with, but it's quite slow.

# Input imgs.shape = (N, 3, H, W), where (N = count, W = width, H = height)
def unique_map_pixels(imgs):
original_shape = imgs.shape

# imgs.shape = (N, H, W, 3)
imgs = imgs.transpose(0, 2, 3, 1)

# tupleview.shape = (N, H, W, 1); contains tuples [(R, G, B), (R, G, B)]
tupleview = imgs.reshape(-1, 3).view(imgs.dtype.descr * imgs.shape[3])

# get unique pixel values in images, [(R, G, B), ...]
uniques = list(np.unique(tupleview))

# map uniques into hashed list ({"RXBXG": 0, "RXBXG": 1}, ...)
uniqmap = {}
idx = 0
for x in uniques:
uniqmap["%sX%sX%s" % (x[0], x[1], x[2])] = idx
idx = idx + 1
if idx >= np.iinfo(np.uint16).max:
raise Exception("Can handle only %s distinct colors" % np.iinfo(np.uint16).max)

# imgs1d.shape = (N), contains RGB tuples
imgs1d = tupleview.reshape(np.prod(tupleview.shape))

# imgsmapped.shape = (N), contains uniques-index values
imgsmapped = np.empty((len(imgs1d))).astype(np.uint16)

# map each pixel into unique-pixel-ID
idx = 0
for x in imgs1d:
str = ("%sX%sX%s" % (x[0], x[1] ,x[2]))
imgsmapped[idx] = uniqmap[str]
idx = idx + 1

imgsmapped.shape = (original_shape[0], original_shape[2], original_shape[3]) # (N, H, W)
return (imgsmapped, uniques)


Testing it:

import numpy as np
n = 30
pixelvalues = (np.random.rand(10)*255).astype(np.uint8)
images = np.random.choice(pixelvalues, (n, 3, 72, 96))

(mapped, pixelmap) = unique_map_pixels(images)
assert len(pixelmap) == mapped.max()+1
assert mapped.shape == (len(images), images.shape[2], images.shape[3])
assert pixelmap[mapped[int(n*0.5)][60][81]][0] == images[int(n*0.5)][0][60][81]
print("Done: %s" % list(mapped.shape))

Answer

Here's a compact vectorized approach without those error checks -

def unique_map_pixels_vectorized(imgs):
    N,H,W = len(imgs), imgs.shape[2], imgs.shape[3]
    img2D = imgs.transpose(0, 2, 3, 1).reshape(-1,3)
    ID = np.ravel_multi_index(img2D.T,img2D.max(0)+1)
    _, firstidx, tags = np.unique(ID,return_index=True,return_inverse=True)
    return tags.reshape(N,H,W), img2D[firstidx]

Runtime test and verification -

In [24]: # Setup inputs (3x smaller than original ones)
    ...: N,H,W = 200,24,32
    ...: imgs = np.random.randint(0,10,(N,3,H,W))
    ...: 

In [25]: %timeit unique_map_pixels(imgs)
1 loop, best of 3: 2.21 s per loop

In [26]: %timeit unique_map_pixels_vectorized(imgs)
10 loops, best of 3: 37 ms per loop ## 60x speedup!

In [27]: map1,unq1 = unique_map_pixels(imgs)
    ...: map2,unq2 = unique_map_pixels_vectorized(imgs)
    ...: 

In [28]: np.allclose(map1,map2)
Out[28]: True

In [29]: np.allclose(np.array(map(list,unq1)),unq2)
Out[29]: True