Replace int values in 2D np.array with list of 3 values to make it 3D-CodePudding

I came along this problem when helping on this question where OP does some image processing. Regardless if there are other ways to do the whole thing, in one part, I have a 2D np.array filles with integers. The integers are just mask values, each standing for a RGB color.

I have a dictionary with integers as keys and arrays of RGB colors as value. This is the mapping and the goal is to replace each int in the array with the colors.

Starting with this array where all RGB-array where already replaced by integers so now it is an array of shape (2,3) (originially it was shape(2,3,3))

import numpy as np
arr = np.array([0,2,4,1,3,5]).reshape(2,3)
print(arr)

array([[0, 2, 4],
       [1, 3, 5]])

Here is the dictionary (chosen numbers are just random for the example):

dic = {0 : [10,20,30], 1 : [12,22,32], 2 : [15,25,35], 3 : [40,50,60], 4 : [100,200,300], 5 : [250,350,450]}

replacing all these values with the arrays makes it an array with shape (2,3,3) like this:

array([[[ 10,  20,  30],
        [ 15,  25,  35],
        [100, 200, 300]],
        
       [[ 12,  22,  32],
        [ 40,  50,  60],
        [250, 350, 450]]])

I looked into np.where because I thought it is the most obvious to me but I always got the error that the shapes are incorrect. I don't know where exactly I'm stuck, when googling, I came across np.dstack, np.concatenate, reading about changing the shape with np.newaxis / None but I just don't get it done. Maybe creating a new array with np.zeros_like and go from there. Do I need to create something like a placeholder before I'm able to insert an array holding these 3 RGB values?

Since every single key is in the array because it is created based on that, I thought about loop through the dict, check for key in array and replace it with the dict.value. Am I at least in the right direction or does that lead to nothing?

Any help much appreciated!!!

CodePudding user response：

In this regard, we can create an array of dictionary values by unpacking that and then order them based on the specified orders in the arr. So:

np.array([*dic.values()])[arr]

If the dictionary keys were not in a sorted order, we can create a mask array for ordering based on them, using np.argsort. So, after sorting the array of dictionary values based on the mask array, we can get the results again e.g.:

dic = {0: [10, 20, 30], 2: [15, 25, 35], 3: [40, 50, 60], 1: [12, 22, 32], 4: [100, 200, 300], 5: [250, 350, 450]}


sort_mask = np.array([*dic.keys()]).argsort()
# [0 3 1 2 4 5]

np.array([*dic.values()])[sort_mask][arr]
# [[[ 10  20  30]
#   [ 15  25  35]
#   [100 200 300]]
# 
#  [[ 12  22  32]
#   [ 40  50  60]
#   [250 350 450]]]