在 numpy 中比较数组的最有效方法？

Question

背景：

我正在研究 Google 对 Python 的 body 分段 API。由于它最初是为 js (tensorflow.js) 编写的库，因此 python equivalent 似乎非常有限。因此，获得 body 分割的唯一方法似乎是将生成的掩码的 RGB 颜色与该颜色应代表的 body 部分进行比较。即：

躯干应该是绿色的，所以我知道:

torso = np.array([175, 240, 91])

头部的右侧部分应该是紫色的，所以它是 [110, 64, 170]，依此类推...

我的做法：

# get prediction result
img = "front_pic"
img_filename = img + ".png"
image = tf.keras.preprocessing.image.load_img(img_filename)
image_array = tf.keras.preprocessing.image.img_to_array(image)
result = bodypix_model.predict_single(image_array)

# simple mask
mask = result.get_mask(threshold=0.75)

# colored mask (separate colour for each body part)
colored_mask = result.get_colored_part_mask(mask)
tf.keras.preprocessing.image.save_img(img+'_cmask'+'.jpg',
    colored_mask
)

# color codes
right_head = np.array([110, 64, 170])
left_head = np.array([143, 61, 178])
torso = np.array([175, 240, 91])
left_feet = np.array([84, 101, 214])
right_feet = np.array([99, 81, 195])
left_arm_shoulder = np.array([210, 62, 167])
right_arm_shoulder = np.array([255, 78, 125])

# (x,y) coordinates
coordinate_x = 0
coordinate_y = 0
for vertical_pixels in colored_mask:
    coordinate_y = coordinate_y + 1
    #if coordinate_y > height:
    #    coordinate_y = 0
    for pixels in vertical_pixels:
        coordinate_x = coordinate_x + 1
        if coordinate_x > width:
            coordinate_x = 1
        # Current Pixel
        np_pixels = np.array(pixels)
        current_coordinate = np.array([[coordinate_x,coordinate_y]])
        #print(current_coordinate)
        if np.array_equal(np_pixels,right_head) or np.array_equal(np_pixels,left_head): # right head or left head
            pixels_head = pixels_head + 1
            head_coordinates = np.concatenate((head_coordinates,current_coordinate),axis=0) # Save coordinates
        if np.array_equal(np_pixels,torso): # Torso
            torso_pixels = torso_pixels + 1
            torso_coordinates = np.concatenate((torso_coordinates,current_coordinate),axis=0) # Save coordinates
        if np.array_equal(np_pixels,left_feet) or np.array_equal(np_pixels,right_feet): # feet_pixels
            feet_pixels = feet_pixels + 1
            feet_coordinates = np.concatenate((feet_coordinates,current_coordinate),axis=0) # Save coordinates
        if np.array_equal(np_pixels,left_arm_shoulder): # left_arm_shoulder
            left_arm_shoulder_pixels = left_arm_shoulder_pixels + 1
            left_arm_shoulder_coordinates = np.concatenate((left_arm_shoulder_coordinates,current_coordinate),axis=0) # Save coordinates
        if np.array_equal(np_pixels,right_arm_shoulder): # right_arm_shoulder
            right_arm_shoulder_pixels = right_arm_shoulder_pixels + 1
            right_arm_shoulder_coordinates = np.concatenate((right_arm_shoulder_coordinates,current_coordinate),axis=0) # Save coordinates

问题：

我的方法的问题是它太慢了！例如，这些代码行：

if np.array_equal(np_pixels,torso): # Torso

占用大量执行时间。必须将每个像素与其等效的 RGB 进行比较太慢了。

我的问题

最好的解决方案是什么？所以要么：

在 python-tf-bodypix library/API 中有更好的方法来获取分段的 body 部分的像素坐标。（有人知道 bodypix 库中是否存在这种方法吗？）

或...

任何 better/faster 比较两个 numpy 数组的方法？
您在我的方法中看到了我应该更改的任何其他低效代码吗？

Answer 1

来自回答：Finding the (x,y) indexes of specific (R,G,B) color values from images stored in NumPy ndarrays

您的问题的解决方案是：

cords = list(zip(*np.where(np.all(np_pixels == torso, axis=-1))))

Answer 2

您可以利用每个 RGB 三元组求和为不同值的事实。

right_head = np.array([110, 64, 170]) # SUM = 344
left_head  = np.array([143, 61, 178]) # SUM = 382
...

因此您可以沿 RGB 维度对像素值求和：

x = np.sum(colored_mask,axis=0)

并创建一个包含所有不同可能总和的向量，对应于 body 部分：

val = np.array([344,382,506,399,375,439,458]) # [right_head_sum, left_head_sum...]

然后使用广播比较这些值：

compare = x == val[:,None,None]

计算您在 7 个不同类别中有多少像素：

count = np.sum(compare,axis=(1,2))

您可以使用 np.where 和 np.split 检索坐标：

coord3d = np.where(compare)
split = np.where(np.diff(coord3d[0]))[0]+1
coord_x = np.split(coord3d[1],split)
coord_y = np.split(coord3d[2],split)

3x5x5 图像示例：

x = array([[375, 399, 458, 382, 506],
          [375, 382, 506, 458, 382],
          [439, 344, 382, 344, 375],
          [439, 439, 344, 382, 344],
          [382, 399, 506, 399, 382]])

val = array([344, 382, 506, 399, 375, 439, 458])

count = array([4, 7, 3, 3, 3, 3, 2])    # [right_head, left_head,...

coord_x = [array([2, 2, 3, 3]),         # right_head
           array([0, 1, 1, 2, 3, 4, 4]),# left_head
           array([0, 1, 4]),            # ...
           array([0, 4, 4]),
           array([0, 1, 2]),
           array([2, 3, 3]),
           array([0, 1]]

coord_y = [array([1, 3, 2, 4],          # right_head 
           array([3, 1, 4, 2, 3, 0, 4], # left_head
           array([4, 2, 2],             # ...
           array([1, 1, 3],
           array([0, 0, 4],
           array([0, 0, 1]),
           array([2, 3]]

它应该比你的 for 循环快得多。

在 numpy 中比较数组的最有效方法？

Most efficient ways to compare arrays in numpy?

python

numpy

tensorflow

bodypix