Compare values with a list of values in np.ndarray

Viewed 73

I want to transform RGB labels of a mask of PASCAL to integer labels, how can I do that using numpy? I searched for a long time, but I cannot find a fast way to do it.

Say, I have a ( W = 3, H = 3, C=3) image, the value in each pixel is a RGB label, say [128, 0, 128] means this label belongs to Class 1, and [128, 256, 0] means this is labeled as Class 2, and [1, 1, 1] means this is labeled as Class 3. For example, in this image, pixels are labeled as [[1, 1, 1], [2, 2, 2], [3, 3, 3]] conceptually and encoded physically by such RGB label convention. And I want to exactly convert this ( W = 3, H = 3, C=3) mask to a ( W = 3, H = 3, C=1) mask, which in this case is exactly [[[1, 1, 1], [2, 2, 2], [3, 3, 3]]].

For clearity, the input is [[[128, 0, 128],[128, 0, 128],[128, 0, 128]],[[128, 256, 0],[128, 256, 0],[128, 256, 0]], [[1, 1, 1],[1, 1, 1],[1, 1, 1]]], which is (3,3,3), and the output should be [[[1],[1],[1]],[[2],[2],[2]],[[3],[3],[3]]], which is (3,3,1)

Such conversion means we need to pick the vector of each pixel and compare it with a list of encoding of length 3 in this example and get the index of a match. How can I do that in a fast way using numpy?

A very trivial and very slow implementation will be

labels = np.array([[128, 0, 128], [128, 256, 0], [1, 1, 1]])
def convert(x):
    a = np.zeros((x.shape[0], x.shape[1], 1), x.dtype)
    a = a-1 # defaut label is -1
    for i in range(x.shape[0]):
        for j in range(x.shape[1]):
            # for each pixel, compare its color vector with those in `labels`, and find out the idx of matched one
            for idx, l in enumerate(labels):
                if np.array_equal(x[i][j], l):
                    a[i][j] = idx+1
    return a
3 Answers
import numpy as np


inp = [
    [
        [128, 0, 128],
        [128, 0, 128],
        [128, 0, 128]
    ],
    [
        [128, 256, 0],
        [128, 256, 0],
        [128, 256, 0]
    ],
    [
        [1, 1, 1],
        [1, 1, 1],
        [1, 1, 1]
    ]
]
out_expected = [[[1],[1],[1]],[[2],[2],[2]],[[3],[3],[3]]]

inp_np = np.array(inp)
out_np_expected = np.array(out_expected)

label1 = np.array([128, 0, 128])
label2 = np.array([128, 256, 0])
label3 = np.array([1, 1, 1])

label1_mask = np.equal(inp_np, label1)
label2_mask = np.equal(inp_np, label2)
label3_mask = np.equal(inp_np, label3)

out_np = np.zeros((inp_np.shape[0], inp_np.shape[1]))
label1_reduced = np.all(label1_mask, axis=2)
label2_reduced = np.all(label2_mask, axis=2)
label3_reduced = np.all(label3_mask, axis=2)

out_np[label1_reduced] = 1
out_np[label2_reduced] = 2
out_np[label3_reduced] = 3
out_np = out_np[..., np.newaxis]

np.testing.assert_equal(out_np, out_np_expected)
  1. Check input equality against the template, per coordinate (using np.equal)
  2. Force equality over the entire template using np.all
  3. Assign desired labels (which seem to be arbitrary numbers)
  4. Add an axis to the result, as required in the question
  5. Verify using np.testing.assert_equal

@macguffin, I suggest to combine ==, np.all(), np.sum():

labels = np.array([[128, 0, 128], [128, 256, 0], [1, 1, 1]])
inputs = np.array([[[128, 0, 128],[128, 0, 128],[128, 0, 128]],
    [[128, 256, 0],[128, 256, 0],[128, 256, 0]], 
    [[1, 1, 1],[1, 1, 1],[1, 1, 1]]])

result = np.sum(
    [np.expand_dims(
        np.all(inputs == labels[i], axis=-1), axis=-1) * (i+1) \
        for i in range (labels.shape[0])],
    axis = 1)


array([[[1], [1], [1]],    
       [[2], [2], [2]],    
       [[3], [3], [3]]])

Here is a different take towards the solution, one liner-

image = np.array([[[128, 0, 128],[128, 0, 128],[128, 0, 128]],  # channels are along axis 2
                  [[128, 256, 0],[128, 256, 0],[128, 256, 0]], 
                  [[1, 1, 1],[1, 1, 1],[1, 1, 1]]])

# since first label has minimum at 1st index
# since second label has a minimum at 2nd index
# since third label has a minimum at 0th index, since all are equal

labels = np.argmin(image, axis=2)

# output is
# array([[1, 1, 1],
#        [2, 2, 2],
#        [0, 0, 0]], dtype=int64)

# if the 3rd label should definitely be 3
labels[labels == 0] = 3

# output is
# array([[1, 1, 1],
#        [2, 2, 2],
#        [3, 3, 3]], dtype=int64)
Related