Convert multi-dimensional Numpy array to 2-dimensional array based on color values

Question

I have an image which is read as a uint8 array with the shape (512,512,3). Now I would like to convert this array to a uint8 array of shape (512,512,1)

Answer 1

Here's one based on views -

# https://stackoverflow.com/a/45313353/ @Divakar
def view1D(a, b): # a, b are arrays
    # This function gets 1D view into 2D input arrays
    a = np.ascontiguousarray(a)
    b = np.ascontiguousarray(b)
    void_dt = np.dtype((np.void, a.dtype.itemsize * a.shape[-1]))
    return a.view(void_dt).ravel(),  b.view(void_dt).ravel()

def img2label(a, maps):
    # Get one-dimension reduced view into input image and map arrays.
    # We need to reshape image to 2D, then feed it to view1D to get 1D
    # outputs and then reshape 1D image to 2D 
    A,B = view1D(a.reshape(-1,a.shape[-1]),maps)
    A = A.reshape(a.shape[:2])

    # Trace back positions of A in B using searchsorted. This gives us
    # original order, which is the final output.
    sidx = B.argsort()
    return sidx[np.searchsorted(B,A,sorter=sidx)]

Given that your labels start from 1, you might want to add 1 to the output.

Sample run -

In [100]: # Mapping array
     ...: maps = np.array([[0, 0, 0],[0, 0, 255],\
     ...:                  [255, 0, 0],[150, 30, 150]],dtype=np.uint8)
     ...: 
     ...: # Setup random image array
     ...: idx = np.array([[0,2,1,3],[1,3,2,0]])
     ...: img = maps[idx]

In [101]: img2label(img, maps) # should retrieve back idx
Out[101]: 
array([[0, 2, 1, 3],
       [1, 3, 2, 0]])