绘制numpy数组的argmax

Question

我有一个形状为(20,20,6(的numpy数组。

import numpy as np
import matplotlib.cm as cm
num_channels = 6
a = np.random.choice(range(100),(20,20,num_channels))

我想得到一个形状为(20,20,4(的数组，它是RGBA颜色的20x20倍，指示a中的哪个值是沿最后一个轴的argmax。因此，我生成了6种RGBA颜色，并在a上使用argmax。

color_list = cm.rainbow(np.linspace(0, 1, num_channels))
a = a.argmax(axis=-1)

现在CCD_ 3具有尺寸CCD_。如何使用颜色列表将其发送到20x20x4？

这就是颜色列表的样子：

[[0.5  0.   1.   1.  ]
[0.1  0.59 0.95 1.  ]
[0.3  0.95 0.81 1.  ]
[0.7  0.95 0.59 1.  ]
[1.   0.59 0.31 1.  ]
[1.   0.   0.   1.  ]]

Answer 1

这应该是您想要的。

import random
import numpy as np
import matplotlib.cm as cm
# get colormap
num_channels = 6
color_list = cm.rainbow(np.linspace(0, 1, num_channels))
# generate an array
a = np.random.rand(20,20,6)
a = np.argmax(a, axis=-1)
print(a)
# create new array
new_data = np.zeros((20,20,4))
for i in range(20):
for j in range(20):
# put RGBA value in new_data
color = color_list[a[i][j]]
new_data[i][j][:] = color
print(new_data)

如果这个问题对你有帮助，请将其标记为已解决！