Spaces:
Runtime error
Runtime error
File size: 5,668 Bytes
cc0dd3c |
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 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional, Union
import cv2 as cv
import numpy as np
import torch
from torchvision.transforms import ToPILImage
class SimCCVisualizer:
def draw_instance_xy_heatmap(self,
heatmap: torch.Tensor,
overlaid_image: Optional[np.ndarray],
n: int = 20,
mix: bool = True,
weight: float = 0.5):
"""Draw heatmaps of GT or prediction.
Args:
heatmap (torch.Tensor): Tensor of heatmap.
overlaid_image (np.ndarray): The image to draw.
n (int): Number of keypoint, up to 20.
mix (bool):Whether to merge heatmap and original image.
weight (float): Weight of original image during fusion.
Returns:
np.ndarray: the drawn image which channel is RGB.
"""
heatmap2d = heatmap.data.max(0, keepdim=True)[0]
xy_heatmap, K = self.split_simcc_xy(heatmap)
K = K if K <= n else n
blank_size = tuple(heatmap.size()[1:])
maps = {'x': [], 'y': []}
for i in xy_heatmap:
x, y = self.draw_1d_heatmaps(i['x']), self.draw_1d_heatmaps(i['y'])
maps['x'].append(x)
maps['y'].append(y)
white = self.creat_blank(blank_size, K)
map2d = self.draw_2d_heatmaps(heatmap2d)
if mix:
map2d = cv.addWeighted(overlaid_image, 1 - weight, map2d, weight,
0)
self.image_cover(white, map2d, int(blank_size[1] * 0.1),
int(blank_size[0] * 0.1))
white = self.add_1d_heatmaps(maps, white, blank_size, K)
return white
def split_simcc_xy(self, heatmap: Union[np.ndarray, torch.Tensor]):
"""Extract one-dimensional heatmap from two-dimensional heatmap and
calculate the number of keypoint."""
size = heatmap.size()
k = size[0] if size[0] <= 20 else 20
maps = []
for _ in range(k):
xy_dict = {}
single_heatmap = heatmap[_]
xy_dict['x'], xy_dict['y'] = self.merge_maps(single_heatmap)
maps.append(xy_dict)
return maps, k
def merge_maps(self, map_2d):
"""Synthesis of one-dimensional heatmap."""
x = map_2d.data.max(0, keepdim=True)[0]
y = map_2d.data.max(1, keepdim=True)[0]
return x, y
def draw_1d_heatmaps(self, heatmap_1d):
"""Draw one-dimensional heatmap."""
size = heatmap_1d.size()
length = max(size)
np_heatmap = ToPILImage()(heatmap_1d).convert('RGB')
cv_img = cv.cvtColor(np.asarray(np_heatmap), cv.COLOR_RGB2BGR)
if size[0] < size[1]:
cv_img = cv.resize(cv_img, (length, 15))
else:
cv_img = cv.resize(cv_img, (15, length))
single_map = cv.applyColorMap(cv_img, cv.COLORMAP_JET)
return single_map
def creat_blank(self,
size: Union[list, tuple],
K: int = 20,
interval: int = 10):
"""Create the background."""
blank_height = int(
max(size[0] * 2, size[0] * 1.1 + (K + 1) * (15 + interval)))
blank_width = int(
max(size[1] * 2, size[1] * 1.1 + (K + 1) * (15 + interval)))
blank = np.zeros((blank_height, blank_width, 3), np.uint8)
blank.fill(255)
return blank
def draw_2d_heatmaps(self, heatmap_2d):
"""Draw a two-dimensional heatmap fused with the original image."""
np_heatmap = ToPILImage()(heatmap_2d).convert('RGB')
cv_img = cv.cvtColor(np.asarray(np_heatmap), cv.COLOR_RGB2BGR)
map_2d = cv.applyColorMap(cv_img, cv.COLORMAP_JET)
return map_2d
def image_cover(self, background: np.ndarray, foreground: np.ndarray,
x: int, y: int):
"""Paste the foreground on the background."""
fore_size = foreground.shape
background[y:y + fore_size[0], x:x + fore_size[1]] = foreground
return background
def add_1d_heatmaps(self,
maps: dict,
background: np.ndarray,
map2d_size: Union[tuple, list],
K: int,
interval: int = 10):
"""Paste one-dimensional heatmaps onto the background in turn."""
y_startpoint, x_startpoint = [int(1.1*map2d_size[1]),
int(0.1*map2d_size[0])],\
[int(0.1*map2d_size[1]),
int(1.1*map2d_size[0])]
x_startpoint[1] += interval * 2
y_startpoint[0] += interval * 2
add = interval + 10
for i in range(K):
self.image_cover(background, maps['x'][i], x_startpoint[0],
x_startpoint[1])
cv.putText(background, str(i),
(x_startpoint[0] - 30, x_startpoint[1] + 10),
cv.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
self.image_cover(background, maps['y'][i], y_startpoint[0],
y_startpoint[1])
cv.putText(background, str(i),
(y_startpoint[0], y_startpoint[1] - 5),
cv.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
x_startpoint[1] += add
y_startpoint[0] += add
return background[:x_startpoint[1] + y_startpoint[1] +
1, :y_startpoint[0] + x_startpoint[0] + 1]
|