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)