Spaces:
Running
Running
File size: 4,570 Bytes
b53fda4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
import sys
import functools
import numpy as np
import cv2
import cmapy
from PIL import Image
import matplotlib
# BUGFIX in cmapy.py
def cmap(cmap_name, rgb_order=False):
"""
Extract colormap color information as a LUT compatible with cv2.applyColormap().
Default channel order is BGR.
Args:
cmap_name: string, name of the colormap.
rgb_order: boolean, if false or not set, the returned array will be in
BGR order (standard OpenCV format). If true, the order
will be RGB.
Returns:
A numpy array of type uint8 containing the colormap.
"""
c_map = matplotlib.colormaps.get_cmap(cmap_name)
rgba_data = matplotlib.cm.ScalarMappable(cmap=c_map).to_rgba(
np.arange(0, 1.0, 1.0 / 256.0), bytes=True
)
rgba_data = rgba_data[:, 0:-1].reshape((256, 1, 3))
# Convert to BGR (or RGB), uint8, for OpenCV.
cmap = np.zeros((256, 1, 3), np.uint8)
if not rgb_order:
cmap[:, :, :] = rgba_data[:, :, ::-1]
else:
cmap[:, :, :] = rgba_data[:, :, :]
return cmap
# If python 3, redefine cmap() to use lru_cache.
if sys.version_info > (3, 0):
cmap = functools.lru_cache(maxsize=200)(cmap)
def alpha_composite(img, msk, opacity=0.5, colormap=None, alpha_image=None, alpha_mask=None, red_mask=False):
"""Alpha composite an RGBA image (img) and a grayscale mask (msk).
- If alpha_image is None, img's alpha channel is used (or, if not present,
initialized to all 255).
- If alpha_mask is None, msk is overlaid on img only where img's alpha
channel is not 0.
- If alpha_mask is not None, the above behavior is overridden and msk is
overlaid on img only where alpha_mask is not 0."""
# only HWC numpy arrays allowed
assert isinstance(img, np.ndarray), f'Input image must be a numpy array. Got {type(img)}'
assert isinstance(msk, np.ndarray), f'Input mask must be a numpy array. Got {type(msk)}'
if alpha_mask is not None:
assert isinstance(alpha_mask, np.ndarray), f'Alpha mask must be a numpy array. Got {type(alpha_mask)}'
assert alpha_mask.dtype in [np.float32, bool], f'Alpha mask must be of type np.float32 or bool. Got {alpha_mask.dtype}'
assert alpha_mask.shape[2] == 1, f'Alpha mask must be formatted as HWC, with C = 1. Got a shape of {msk.shape}'
assert img.shape[2] in [3,4], f'Input image must be formatted as HWC, with C = 3,4. Got a shape of {img.shape}'
assert msk.shape[2] == 1, f'Input mask must be formatted as HWC, with C = 1. Got a shape of {msk.shape}'
assert (opacity >= 0) and (opacity <= 1), f'Mask opacity must be between 0 and 1. Got {opacity}'
# to avoid modifying the original arrays
img = img.copy()
msk = msk.copy()
if img.shape[2] == 3:
# add alpha channel to img
img = np.concatenate([
img,
np.full((img.shape[0], img.shape[1], 1), 255, dtype=np.uint8)
], axis=-1)
if alpha_image is None:
# initialize alpha_image to all Trues
alpha_image = img[:,:,[3]]
# convert alpha image to bool
alpha_image = alpha_image.astype(bool)
if alpha_mask is None:
# initialize alpha_mask to alpha_image
alpha_mask = alpha_image # so that alpha_mask is AT LEAST as restrictive as alpha_image
# convert alpha mask to bool
alpha_mask = alpha_mask.astype(bool)
if msk.dtype != np.uint8:
# convert mask to a uint8 grayscale image ([0,1] -> [0,255])
# NB: normalize the pixels of the mask we are interested in to [0,1]
# before passing it as input!!!
msk = (msk * 255).astype(np.uint8)
# convert mask from grayscale to RGBA
msk = cv2.cvtColor(msk, cv2.COLOR_GRAY2RGBA)
if colormap is not None:
# apply specified colormap to msk
# NB: values near 0 will be converted to the first colors of the chosen
# colormap, whereas values near 255 will be converted to the last colors
msk[:,:,:3] = cmapy.colorize(msk[:,:,:3], colormap, rgb_order=True)
elif red_mask:
# convert white to red
msk[:,:,[1,2]] = 0
# apply alpha_image to img's alpha channel
img[:,:,[3]] = (alpha_image * img[:,:,[3]]).astype(np.uint8)
# apply alpha_mask and opacity to msk's alpha channel
msk[:,:,[3]] = (alpha_mask * opacity * msk[:,:,[3]]).astype(np.uint8)
# alpha compositing
img_pil = Image.fromarray(img)
msk_pil = Image.fromarray(msk)
img_pil.alpha_composite(msk_pil)
return np.array(img_pil)
|