Вещание в 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

Подход # 1

Вы можете сэкономить много на данных промежуточного массива с расширением уникального colors до более высокого затемнения, а затем сравнивать с исходным массивом данных и затем использовать маску напрямуючтобы получить окончательный результат -

# 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 будет с чем-то вроде этого -

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]

Answer 2

Я смог решить это следующим образом:

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]