create multi-colored thresholded image from grayscale heatmap in python

Question

I have a gray-scale images and want to threshold them in 4 range of values (0.035, 0.7, 0.75), displayed by 4 different colors. I need the result saved as images in UINT8 format. The gray-scale image information is as follows:

print(type(grads))
print(grads.shape)
print(grads.min())
print(grads.max())
cv2.imshow('1_grads', grads)
cv2.waitKey()

### OUTPUT
<class 'numpy.ndarray'>
(512, 512)
0.0
1.0

I have tried the following:

thresh_map = Image.new('RGB', grads.shape, color='white')
thresh_map = np.where(grads < 0.035, (0, 0, 0), (0, 0, 0))
thresh_map = np.where(0.035 < grads < 0.7, (0, 0, 255), (0, 0, 0))
thresh_map = np.where(0.7 < grads < 0.75, (0, 255, 0), (0, 0, 0))
thresh_map = np.where(0.75 < grads, (0, 255, 0), (0, 0, 0))

This returns this Error: ValueError: operands could not be broadcast together with shapes (512,512) (3,) (3,)

I have somehow solved the problem by using for loops and pasting the pixel values one by one. But it is not nice and takes forever, considering the fact that I am going to apply this to ~4000 images.

thresh_map = Image.new('RGB', grads.shape, color='white')
blac = Image.new('RGB', (1, 1), color='black')
blue = Image.new('RGB', (1, 1), color='blue')
redd = Image.new('RGB', (1, 1), color='red')
gree = Image.new('RGB', (1, 1), color='green')
for i in range(grads.shape[0]):
    for j in range(grads.shape[1]):
        print(i, j)
        if grads[i, j] < 0.035:
            thresh_map.paste(blac, (i, j))

        elif .035 < grads[i, j] < 0.7:
            thresh_map.paste(redd, (i, j))

        elif 0.7 < grads[i, j] < 0.75:
            thresh_map.paste(gree, (i, j))

        elif 0.75 < grads[i, j]:
            thresh_map.paste(blue, (i, j))

np_thresh_map = np.asarray(thresh_map)
cv2.imshow('1_thresh', np_thresh_map)
cv2.waitKey()

Is there a more sophisticated and efficient way of doing this?

Answer 1

Here's my solution using NumPy's boolean array indexing. The code should be straightforward. If not, please ask. I'll then provide some more explanations.

import cv2
import numpy as np

# Set up some random grayscale image; and sort for better visualization
image = np.sort(np.random.rand(300, 400))

# Given thresholds
thresholds = [0.035, 0.7, 0.75]

# Expand threshold with boundaries
thresholds = np.concatenate(([0], thresholds, [1]))

# Initialize output, and map colors
map = np.zeros((image.shape[0], image.shape[1], 3), np.uint8)
colors = np.array([[255, 0, 0], [128, 128, 0], [0, 0, 255], [0, 255, 0]])

# Iterate ranges; find proper pixel indices; set proper color in map at these indices
for i in range(len(thresholds)-1):
    idx = (thresholds[i] <= image) & (image < thresholds[i+1])
    map[idx, :] = colors[i, :]

# Show image and map
cv2.imshow('image', image)
cv2.imshow('map', map)
cv2.waitKey(0)
cv2.destroyAllWindows()

The image might look like this:

And, the corresponding map looks like this:

Hope that helps!