Last commit not found
| from typing import Dict, List | |
| import torch | |
| import colorsys | |
| import random | |
| import numpy as np | |
| from skimage.draw import line_aa, circle_perimeter_aa | |
| import cv2 | |
| from .util import select_data | |
| def _gen_random_colors(N, bright=True): | |
| brightness = 1.0 if bright else 0.7 | |
| hsv = [(i / N, 1, brightness) for i in range(N)] | |
| colors = list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv)) | |
| random.shuffle(colors) | |
| return colors | |
| _static_label_colors = [ | |
| np.array((1.0, 1.0, 1.0), np.float32), | |
| np.array((255, 250, 79), np.float32) / 255.0, # face | |
| np.array([255, 125, 138], np.float32) / 255.0, # lb | |
| np.array([213, 32, 29], np.float32) / 255.0, # rb | |
| np.array([0, 144, 187], np.float32) / 255.0, # le | |
| np.array([0, 196, 253], np.float32) / 255.0, # re | |
| np.array([255, 129, 54], np.float32) / 255.0, # nose | |
| np.array([88, 233, 135], np.float32) / 255.0, # ulip | |
| np.array([0, 117, 27], np.float32) / 255.0, # llip | |
| np.array([255, 76, 249], np.float32) / 255.0, # imouth | |
| np.array((1.0, 0.0, 0.0), np.float32), # hair | |
| np.array((255, 250, 100), np.float32) / 255.0, # lr | |
| np.array((255, 250, 100), np.float32) / 255.0, # rr | |
| np.array((250, 245, 50), np.float32) / 255.0, # neck | |
| np.array((0.0, 1.0, 0.5), np.float32), # cloth | |
| np.array((1.0, 0.0, 0.5), np.float32), | |
| ] + _gen_random_colors(256) | |
| _names_in_static_label_colors = [ | |
| 'background', 'face', 'lb', 'rb', 'le', 're', 'nose', | |
| 'ulip', 'llip', 'imouth', 'hair', 'lr', 'rr', 'neck', | |
| 'cloth', 'eyeg', 'hat', 'earr' | |
| ] | |
| def _blend_labels(image, labels, label_names_dict=None, | |
| default_alpha=0.6, color_offset=None): | |
| assert labels.ndim == 2 | |
| bg_mask = labels == 0 | |
| if label_names_dict is None: | |
| colors = _static_label_colors | |
| else: | |
| colors = [np.array((1.0, 1.0, 1.0), np.float32)] | |
| for i in range(1, labels.max() + 1): | |
| if isinstance(label_names_dict, dict) and i not in label_names_dict: | |
| bg_mask = np.logical_or(bg_mask, labels == i) | |
| colors.append(np.zeros((3))) | |
| continue | |
| label_name = label_names_dict[i] | |
| if label_name in _names_in_static_label_colors: | |
| color = _static_label_colors[ | |
| _names_in_static_label_colors.index( | |
| label_name)] | |
| else: | |
| color = np.array((1.0, 1.0, 1.0), np.float32) | |
| colors.append(color) | |
| if color_offset is not None: | |
| ncolors = [] | |
| for c in colors: | |
| nc = np.array(c) | |
| if (nc != np.zeros(3)).any(): | |
| nc += color_offset | |
| ncolors.append(nc) | |
| colors = ncolors | |
| if image is None: | |
| image = orig_image = np.zeros( | |
| [labels.shape[0], labels.shape[1], 3], np.float32) | |
| alpha = 1.0 | |
| else: | |
| orig_image = image / np.max(image) | |
| image = orig_image * (1.0 - default_alpha) | |
| alpha = default_alpha | |
| for i in range(1, np.max(labels) + 1): | |
| image += alpha * \ | |
| np.tile( | |
| np.expand_dims( | |
| (labels == i).astype(np.float32), -1), | |
| [1, 1, 3]) * colors[(i) % len(colors)] | |
| image[np.where(image > 1.0)] = 1.0 | |
| image[np.where(image < 0)] = 0.0 | |
| image[np.where(bg_mask)] = orig_image[np.where(bg_mask)] | |
| return image | |
| def _draw_hwc(image: torch.Tensor, data: Dict[str, torch.Tensor]): | |
| device = image.device | |
| image = np.array(image.cpu().numpy(), copy=True) | |
| dtype = image.dtype | |
| h, w, _ = image.shape | |
| draw_score_error = False | |
| for tag, batch_content in data.items(): | |
| if tag == 'rects': | |
| for cid, content in enumerate(batch_content): | |
| x1, y1, x2, y2 = [int(v) for v in content] | |
| y1, y2 = [max(min(v, h-1), 0) for v in [y1, y2]] | |
| x1, x2 = [max(min(v, w-1), 0) for v in [x1, x2]] | |
| for xx1, yy1, xx2, yy2 in [ | |
| [x1, y1, x2, y1], | |
| [x1, y2, x2, y2], | |
| [x1, y1, x1, y2], | |
| [x2, y1, x2, y2] | |
| ]: | |
| rr, cc, val = line_aa(yy1, xx1, yy2, xx2) | |
| val = val[:, None][:, [0, 0, 0]] | |
| image[rr, cc] = image[rr, cc] * (1.0-val) + val * 255 | |
| if 'scores' in data: | |
| try: | |
| import cv2 | |
| score = data['scores'][cid].item() | |
| score_str = f'{score:0.3f}' | |
| image_c = np.array(image).copy() | |
| cv2.putText(image_c, score_str, org=(int(x1), int(y2)), | |
| fontFace=cv2.FONT_HERSHEY_TRIPLEX, | |
| fontScale=0.6, color=(255, 255, 255), thickness=1) | |
| image[:, :, :] = image_c | |
| except Exception as e: | |
| if not draw_score_error: | |
| print(f'Failed to draw scores on image.') | |
| print(e) | |
| draw_score_error = True | |
| if tag == 'points': | |
| for content in batch_content: | |
| # content: npoints x 2 | |
| for x, y in content: | |
| x = max(min(int(x), w-1), 0) | |
| y = max(min(int(y), h-1), 0) | |
| rr, cc, val = circle_perimeter_aa(y, x, 1) | |
| valid = np.all([rr >= 0, rr < h, cc >= 0, cc < w], axis=0) | |
| rr = rr[valid] | |
| cc = cc[valid] | |
| val = val[valid] | |
| val = val[:, None][:, [0, 0, 0]] | |
| image[rr, cc] = image[rr, cc] * (1.0-val) + val * 255 | |
| if tag == 'seg': | |
| label_names = batch_content['label_names'] | |
| for seg_logits in batch_content['logits']: | |
| # content: nclasses x h x w | |
| seg_probs = seg_logits.softmax(dim=0) | |
| seg_labels = seg_probs.argmax(dim=0).cpu().numpy() | |
| image = (_blend_labels(image.astype(np.float32) / | |
| 255, seg_labels, | |
| label_names_dict=label_names) * 255).astype(dtype) | |
| return torch.from_numpy(image).to(device=device) | |
| def draw_bchw(images: torch.Tensor, data: Dict[str, torch.Tensor]) -> torch.Tensor: | |
| images2 = [] | |
| for image_id, image_chw in enumerate(images): | |
| selected_data = select_data(image_id == data['image_ids'], data) | |
| images2.append( | |
| _draw_hwc(image_chw.permute(1, 2, 0), selected_data).permute(2, 0, 1)) | |
| return torch.stack(images2, dim=0) | |
| def draw_landmarks(img, bbox=None, landmark=None, color=(0, 255, 0)): | |
| """ | |
| Input: | |
| - img: gray or RGB | |
| - bbox: type of BBox | |
| - landmark: reproject landmark of (5L, 2L) | |
| Output: | |
| - img marked with landmark and bbox | |
| """ | |
| img = cv2.UMat(img).get() | |
| if bbox is not None: | |
| x1, y1, x2, y2 = np.array(bbox)[:4].astype(np.int32) | |
| cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2) | |
| if landmark is not None: | |
| for x, y in np.array(landmark).astype(np.int32): | |
| cv2.circle(img, (int(x), int(y)), 2, color, -1) | |
| return img |