Spaces:
Running
Running
File size: 2,034 Bytes
84ce4ca |
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 |
import cv2
import numpy as np
import torchvision.transforms as transforms
# Colors for all 20 parts
part_colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 0, 85], [255, 0, 170],
[0, 255, 0], [85, 255, 0], [170, 255, 0], [0, 255, 85], [0, 255, 170],
[0, 0, 255], [85, 0, 255], [170, 0, 255], [0, 85, 255], [0, 170, 255],
[255, 255, 0], [255, 255, 85], [255, 255, 170], [255, 0, 255], [255, 85, 255],
[255, 170, 255], [0, 255, 255], [85, 255, 255], [170, 255, 255]]
colormap = np.array(part_colors, dtype=np.uint8)
def image_to_tensor(image):
return transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])(image)
def decode_segmentation_masks(mask, n_classes=20):
red = np.zeros_like(mask).astype(np.uint8)
green = np.zeros_like(mask).astype(np.uint8)
blue = np.zeros_like(mask).astype(np.uint8)
for chanel in range(0, n_classes):
idx = mask == chanel
red[idx] = colormap[chanel, 0]
green[idx] = colormap[chanel, 1]
blue[idx] = colormap[chanel, 2]
return np.stack([red, green, blue], axis=2)
def vis_parsing_maps(image: np.array, parsing_anno, stride=1):
image = np.array(image)
vis_im = image.copy().astype(np.uint8)
vis_parsing_anno = parsing_anno.copy().astype(np.uint8)
vis_parsing_anno = cv2.resize(vis_parsing_anno, None, fx=stride, fy=stride, interpolation=cv2.INTER_NEAREST)
vis_parsing_anno_color = np.zeros((vis_parsing_anno.shape[0], vis_parsing_anno.shape[1], 3)) + 255
num_of_class = np.max(vis_parsing_anno)
for pi in range(1, num_of_class + 1):
index = np.where(vis_parsing_anno == pi)
vis_parsing_anno_color[index[0], index[1], :] = part_colors[pi]
vis_parsing_anno_color = vis_parsing_anno_color.astype(np.uint8)
vis_im = cv2.addWeighted(cv2.cvtColor(vis_im, cv2.COLOR_RGB2BGR), 0.4, vis_parsing_anno_color, 0.6, 0)
return vis_parsing_anno, vis_im
|