import numpy as np
[docs]def add_mask(image, mask, channel='r', intensity=100):
r, g, b = image.transpose(2,0,1)
mask = mask.squeeze() * intensity
if channel == 'r':
r = r.astype(np.uint16) + mask.astype(np.uint16)
r = np.clip(r, 0, 255).astype(np.uint8)
if channel == 'g':
g = g.astype(np.uint16) + mask.astype(np.uint16)
g = np.clip(g, 0, 255).astype(np.uint8)
if channel == 'b':
b = b.astype(np.uint16) + mask.astype(np.uint16)
b = np.clip(b, 0, 255).astype(np.uint8)
image = np.stack((r, g, b)).transpose(1,2,0)
return image