在多维度的 numpy 中广播

Question

我有一张黑色背景的图像，其中包含不同颜色的不同形状。我想为每个形状生成一个图像，其中形状为白色，背景为黑色。我已经能够用 numpy 做到这一点，但我想使用矢量化来优化我的代码。这是我目前所拥有的：

import numpy as np
import cv2

image = cv2.imread('mask.png')

image.shape
# (720, 1280, 3)

# Get all colors that are not black
colors = np.unique(image.reshape(-1,3), axis=0)
colors = np.delete(colors, [0,0,0], axis=0)

colors.shape
# (5, 3)

# Example for one color. I could do a for-loop, but I want to vectorize instead
c = colors[0]
query = (image == c).all(axis=2)

# Make the image all black, except for the pixels that match the shape
image[query] = [255,255,255]
image[np.logical_not(query)] = [0,0,0]

Answer 1

我能够通过以下方式解决它：

import numpy as np
import cv2

# Read the image
image = cv2.imread('0-mask.png')

# Get unique colors (remove black)
colors = np.unique(image.reshape(-1,3), axis=0)
colors = np.delete(colors, [0,0,0], axis=0)

# Get number of unique colors
instances = colors.shape[0]

# Reshape colors and image for broadcasting
colors = colors.reshape(instances,1,1,3)
image = image[np.newaxis]

# Generate multiple images, one per instance
mask = np.ones((instances, 1, 1, 1))
images = (image * mask)

# Run query with the original image
query = (image == colors).all(axis=3)

# For every image, color the shape white, everything else black
images[query] = [255,255,255]
images[np.logical_not(query)] = [0,0,0]

Answer 2

方法 #1

你可以节省很多中间数组数据，将唯一的 colors 扩展到更高的 dim，然后与原始数据数组进行比较，然后直接使用掩码得到最终输出 -

# Get unique colors (remove black)
colors = np.unique(image.reshape(-1,3), axis=0)
colors = np.delete(colors, [0,0,0], axis=0)

mask = (colors[:,None,None,:]==image).all(-1)
out = mask[...,None]*np.array([255,255,255])

方法 #2

获得 mask 的 better/memory-efficient 方法是这样的 -

u,ids = np.unique(image.reshape(-1,3), axis=0, return_inverse=1)
m,n = image.shape[:-1]
ids = ids.reshape(m,n)-1
mask = np.zeros((ids.max()+1,m,n),dtype=bool)
mask[ids,np.arange(m)[:,None],np.arange(n)] = ids>=0

因此，这是获得最终输出的更好方法，就像这样 -

out = np.zeros(mask.shape + (3,), dtype=np.uint8)
out[mask] = [255,255,255]

获得 ids 的更好方法可能是使用 matrix-multiplication。因此：

u,ids = np.unique(image.reshape(-1,3), axis=0, return_inverse=1)

可以替换为：

image2D = np.tensordot(image,256**np.arange(3),axes=(-1,-1))
ids = np.unique(image2D,return_inverse=1)[1]

在多维度的 numpy 中广播

Broadcasting in numpy with multiple dimensions

python

opencv

numpy

numpy-ndarray

array-broadcasting