How to turn 3d array into 2d array with labels based on the 3rd dimension?

Question

I have a 3D NumPy Array with dimensions (224, 224, 3). The third dimension represents colors. I have only 21 colors in the whole picture so there are only 21 distinct values for each pixel.
These are the color codes:

colors = np.array([
                  [  0,   0,   0], [128,   0,   0], [  0, 128,   0], [128, 128,   0], [  0,   0, 128],
                  [128,   0, 128], [  0, 128, 128], [128, 128, 128], [ 64,   0,   0], [192,   0,   0],
                  [ 64, 128,   0], [192, 128,   0], [ 64,   0, 128], [192,   0, 128], [ 64, 128, 128],
                  [192, 128, 128], [  0,  64,   0], [128,  64,   0], [  0, 192,   0], [128, 192,   0],
                  [ 0, 64, 128]], dtype=np.uint8)

How can I iterate over each pixel and create a 2D array with shape (224, 224) which consists of integers 0 to 20 based on the value of 3rd dimension? What would be the most efficient way to do this?

Answer 1

Let's treat color as 4 based positional system, then we can create two kind-of hash functions for converting from this 4 based to 10 based system and backwards:

import numpy as np
from typing import Sequence

# I have changed the order I guess
colors = np.array(
    [
        [0, 0, 0],
        [64, 0, 0],
        [128, 0, 0],
        [192, 0, 0],
        [0, 64, 0],
        [64, 64, 0],
        [128, 64, 0],
        [192, 64, 0],
        [0, 128, 0],
        [64, 128, 0],
        [128, 128, 0],
        [192, 128, 0],
        [0, 192, 0],
        [64, 192, 0],
        [128, 192, 0],
        [192, 192, 0],
        [0, 0, 64],
        [64, 0, 64],
        [128, 0, 64],
        [192, 0, 64],
        [0, 64, 64],
    ],
    dtype=np.uint8,
)


def color_to_value(cell: Sequence[int]):
    mult = 1
    value = 0
    for sub in cell:
        value += (sub // 64) * mult
        mult *= 4
    return value


_MAP = (
    0,
    64,
    128,
    192,
)


def value_to_color(value: int):
    seq = [0] * 3
    for i in range(3):
        reminder = value % 4
        value //= 4
        seq[i] = _MAP[reminder]
    return seq


for i in range(21):
    print(value_to_color(i))


img = np.random.random((16, 16, 3)) * 255

converted = np.array(
    [[color_to_value(cell) for cell in row] for row in img], dtype=np.int8
)
print(converted)

Answer 2

Use np.unique.

The following is some boilerplate code for testing.

IMAGE_SHAPE = (224, 224)  # Given image shape

# Given colors.
colors = np.array([
                  [  0,   0,   0], [128,   0,   0], [  0, 128,   0], [128, 128,   0], [  0,   0, 128],
                  [128,   0, 128], [  0, 128, 128], [128, 128, 128], [ 64,   0,   0], [192,   0,   0],
                  [ 64, 128,   0], [192, 128,   0], [ 64,   0, 128], [192,   0, 128], [ 64, 128, 128],
                  [192, 128, 128], [  0,  64,   0], [128,  64,   0], [  0, 192,   0], [128, 192,   0],
                  [ 0, 64, 128]], dtype=np.uint8)

# Generate a random image with pixel colors taken from 'colors'.
image = colors[np.random.randint(0, colors.shape[0], IMAGE_SHAPE)]

Running print(image) will produce a (224, 224, 3) array that looks something like

array([[[128, 128,   0],
        [  0, 128,   0],
        [192,   0,   0],
        ...,
        [192,   0,   0],
        [128,   0,   0],
        [ 64,   0, 128]],

        ...,

       [[ 64, 128,   0],
        [128, 128, 128],
        [  0,  64, 128],
        ...,
        [ 64,   0,   0],
        [192, 128, 128],
        [192,   0,   0]]], dtype=uint8)

To convert the image to an array of pixel-color indices, invoke np.unique as

_colors, index_image = np.unique(image.reshape(-1, 3), return_inverse=True)
index_image = index_image.reshape(IMAGE_SHAPE)

The array _colors will contain the same pixel values as colors (though likely in a different order), and index_image will be a (224, 224) array with entries that give the index of each pixel in _colors:

print(_colors)
print(index_image.shape)
print(index_image)

results in

array([[  0,   0,   0],
       [  0,   0, 128],
       [  0,  64,   0],
       [  0,  64, 128],
       [  0, 128,   0],
       [  0, 128, 128],
       [  0, 192,   0],
       [ 64,   0,   0],
       [ 64,   0, 128],
       [ 64, 128,   0],
       [ 64, 128, 128],
       [128,   0,   0],
       [128,   0, 128],
       [128,  64,   0],
       [128, 128,   0],
       [128, 128, 128],
       [128, 192,   0],
       [192,   0,   0],
       [192,   0, 128],
       [192, 128,   0],
       [192, 128, 128]], dtype=uint8)

(224, 224)

array([[14,  4, 17, ..., 17, 11,  8],
       [ 9,  7, 11, ...,  5, 17,  7],
       ...,
       [ 8, 10,  0, ...,  1,  6, 16],
       [ 9, 15,  3, ...,  7, 20, 17]])

The mapping can be reversed simply by indexing _colors:

print(_colors[index_image])
print(np.all(_color[index_image] == image))

produces

array([[[128, 128,   0],
        [  0, 128,   0],
        [192,   0,   0],
        ...,
        [192,   0,   0],
        [128,   0,   0],
        [ 64,   0, 128]],

        ...,

       [[ 64, 128,   0],
        [128, 128, 128],
        [  0,  64, 128],
        ...,
        [ 64,   0,   0],
        [192, 128, 128],
        [192,   0,   0]]], dtype=uint8)

True

Answer 3

Since your color code is unique, you can consider encoding each color bar, and then use np.searchsorted to search the sorted codes:

>>> m = 256 ** np.arange(3)
>>> coded = colors @ m
>>> perm = coded.argsort()
>>> sort = coded[perm]
>>> img = colors[np.random.choice(len(colors), 4)].reshape(2, 2, 3)
>>> result = perm[np.searchsorted(sort, img.reshape(-1, 3) @ m)].reshape(img.shape[:-1])
>>> img
array([[[128,   0, 128],
        [128,  64,   0]],

       [[ 64,   0, 128],
        [128, 128, 128]]], dtype=uint8)
>>> result
array([[ 5, 17],
       [12,  7]], dtype=int64)
>>> np.all(colors[result] == img)
True