1

I have an image with a black background that contains different shapes in different colors. I want to generate an image per shape, in which the shape is white and the background is black. I have been able to do this with numpy, but I would like to optimize my code using vectorization. This is what I have so far:

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]
kmario23
  • 57,311
  • 13
  • 161
  • 150
Oliver Hoffman Garcia
  • 78
  • 5

2 Answers2

1

Approach #1

You can save a lot on intermediate array data with extension of unique colors into higher dim and then comparing against original data array and then using the mask directly to get the final output -

# 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])

Approach #2

A better/memory-efficient way to get that mask would be with something like this -

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

and hence, a better way to get the final output, like so -

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

and probably a better way to get ids would be with matrix-multiplication. Hence :

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

could be replaced by :

image2D = np.tensordot(image,256**np.arange(3),axes=(-1,-1))
ids = np.unique(image2D,return_inverse=1)[1]
Divakar
  • 218,885
  • 19
  • 262
  • 358
0

I was able to solve it the following way:

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]
Oliver Hoffman Garcia
  • 78
  • 5