Spaces:
Sleeping
Sleeping
File size: 6,539 Bytes
0d8d0b4 | 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 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 | import cv2
from PIL import Image
import numpy as np, torchvision
import torch, os, json
from skimage.color import label2rgb
from PIL import ImageColor
from cxas.label_mapper import label_mapper, id2label_dict, colors, colors_alpha, category_colors, category_ids
from cxas.file_io import FileLoader
def get_img(
path: str
) -> np.array:
"""
Load image from path.
Parameters
----------
path: Image file path
Returns
----------
image: Image in numpy format
"""
assert os.path.isfile(path)
loader = FileLoader('')
return loader.load_file(path)['orig_data']
def get_label(
path: str
) -> np.array:
"""
Load label from path with .npy-suffix.
Parameters
----------
path: label file path
Returns
----------
mask: Label mask in numpy format
"""
assert os.path.isfile(path)
return np.load(path)
def visualize_label(label: np.array,
img: np.array,
label_to_visualize: np.array,
concat:bool = False,
axis:int = 1
) -> Image:
"""
visualize certain labels from mask
Parameters
----------
label: Mask in shape [classes (159), width, height]
img: Image in shape [classes (159), width, height]
concat: Whether to display image and visualization side by side
axis: axis at which image and visualization are shown side by side
Returns
----------
visualization: Label visualization as PIL Image
"""
out_mask = np.zeros((img.shape[0],img.shape[1],3)).astype(np.uint8)
for i in label_to_visualize:
imgray = (label[i,:,:]*255).astype(np.uint8)
ret, thresh = cv2.threshold(imgray, 127, 255, 0)
x = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
for contour in x:
if (contour is not None) and (len(contour) > 0) and (len(contour[0])>2) and (type(contour) == type(())):
cv2.fillPoly(
out_mask,
contour,
# [int(j*255) for j in colors[i]],
[colors_alpha[i][0]*255,colors_alpha[i][1]*255,colors_alpha[i][2]*255,colors_alpha[i][3]],
)
out_contour = np.zeros((img.shape[0],img.shape[1],3)).astype(np.uint8)
for i in label_to_visualize:
imgray = (label[i,:,:]*255).astype(np.uint8)
ret, thresh = cv2.threshold(imgray, 127, 255, 0)
x = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
for contour in x:
if (contour is not None) and (len(contour) > 0) and (len(contour[0])>2) and (type(contour) == type(())):
cv2.drawContours(out_contour, contour, -1, [int(j*255) for j in colors[i]], 2)
out = cv2.addWeighted(out_contour, 1, out_mask, 0.7, 0.0)
out = cv2.addWeighted(img, 0.5, out, 0.7, 0.0)
if concat:
return Image.fromarray(np.concatenate([img, out],axis)).convert('RGB')
else:
return Image.fromarray(out).convert('RGB')
def visualize_mask(class_names: list,
mask:np.array,
image:np.array,
img_size: int,
cat:bool=True,
axis:int=1,
) -> Image:
"""
Resize image and label to desired size and visualize certain labels
Parameters
----------
class_names: List of classes of interest
mask: Mask in shape [classes (159), width, height]
image: Image in shape [3, width, height]
img_size: Desired image size
cat: Whether to display image and visualization side by side
axis: axis at which image and visualization are shown side by side
Returns
----------
visualization: Label visualization as PIL Image in desired size
"""
# Resize image and label to desired size
img = torch.nn.functional.interpolate(
torch.tensor(image).float().unsqueeze(0),
img_size,
mode='bilinear'
).byte().numpy()[0]
# reshape to match cv2 shapes
img = np.transpose(img, [1,2,0])
label = torch.nn.functional.interpolate(
torch.tensor(mask).float().unsqueeze(0),
img_size,
mode='nearest'
).bool().numpy()[0]
if type(class_names) == list:
pass
else:
class_names = [class_names]
label_to_visualize = np.concatenate([np.array(label_mapper[n]).flatten() for n in class_names]).flatten()
return visualize_label(label, img, label_to_visualize, cat, axis)
def visualize_from_file(class_names: list,
img_path: str,
label_path: str,
img_size: int,
cat:bool = True,
axis:int = 1,
do_store:bool = False,
out_dir:str = '',
) -> Image:
"""
Load Image and label, resize image and label to desired size, and visualize certain labels
Parameters
----------
class_names: list of class names to visualize
img_path: path of the image file to visualize
label_path: path of the label file to visualize
img_size: size at which the image should be visualized
cat: show image and label side by side
axis: axis at which image and label are shown side by side
do_store: boolean indicating whether to store the visualization in the out_dir
with the associated label path and class_name
out_dir: path at which to store visualization
Returns
----------
visualization: Pillow Image with labels overlaying original image
"""
# Load image and label files
img = get_img(img_path)
label = get_label(label_path)
# visualize desired labels
visualization = visualize_mask(class_names, label, img, img_size, cat, axis)
if do_store:
os.makedirs(os.path.dirname(out_dir), exist_ok=True)
visualization.save(os.path.join(
out_dir,
'{}_{}.png'.format(os.path.basename(label_path).split('.')[0], '_'.join(class_names))))
return visualization |