numpy高效的數組乘法

Question

我有一個形狀爲[1200,1600,3]的三維數組img和形狀爲[1200,1600]的二維數組標籤。第一個數組來自圖像，第二個數組來自圖像中的標籤。 img數組中的位置[i，j]對應於圖像像素。 我想要創建一個與img數組尺寸相同的新數組，因此對於標籤爲0的像素，原始數組不變，但所有其他像素都變白（255,255,255）。

我使用的代碼是：

import numpy as np 

newimg=np.zeros((img.shape[0],img.shape[1],img.shape[2])) 
for i in range(0,img.shape[0]): 
    for j in range(0,img.shape[1]): 
      if labels[i][j]==0: 
      newimg[i][j]=img[i][j] 
    else: 
     newimg[i][j]=np.array([255,255,255]) 

是否有這樣做的一個更快的方法？

Answer 1

一般來說，你會做類似的東西：

newimg = img.copy() 
newimg[labels != 0, :] = 255 

或者：

newimg = np.where(labels[..., None] != 0, img, 255)