将分割图矢量化为 RGB 颜色

Question

我编写了一个函数，该函数返回分割映射输入的 RGB 颜色图。问题是我在操作中使用了“for”循环，它真的很慢。函数如下：

def segment_map_to_rgb_color_image(color_group, segment_map):
    
         rgb_color_image = np.zeros((len(segment_map[0]), len(segment_map[1]), 3), dtype=np.uint8)
    
         for i in range(np.size(segment_map, axis=0)):
             for j in range(np.size(segment_map, axis=1)):
                 idx = segment_map[i][j]
                 for p in range(len(color_group[0])):  
                     rgb_color_image[i][j][p] = color_group[idx][p]                    
                    
         return rgb_color_image

详情：

color_group 是[256, 3] NumPy 数组，其中每个索引对应一个特定的 RGB 颜色，例如，color_group[0] = [184, 212, 52]。

segment_map 是[height, width] NumPy 数组，其中每个像素是 0 到 255 之间的整数。

rgb_color_image 是 [height, width, 3] NumPy 数组。

我在对函数进行向量化时遇到困难，并且当前的非向量化函数对于实时分析来说太慢了。

有人可以帮我对函数进行矢量化吗？感谢所有输入。

Answer 1

这对numpy 来说很简单：

rgb_color_image = color_group[segment_map]

这只是使用fancy indexing。