Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| from math import sqrt | |
| import torch | |
| import torch.nn.functional as F | |
| def gaussian2D(radius, sigma=1, dtype=torch.float32, device='cpu'): | |
| """Generate 2D gaussian kernel. | |
| Args: | |
| radius (int): Radius of gaussian kernel. | |
| sigma (int): Sigma of gaussian function. Default: 1. | |
| dtype (torch.dtype): Dtype of gaussian tensor. Default: torch.float32. | |
| device (str): Device of gaussian tensor. Default: 'cpu'. | |
| Returns: | |
| h (Tensor): Gaussian kernel with a | |
| ``(2 * radius + 1) * (2 * radius + 1)`` shape. | |
| """ | |
| x = torch.arange( | |
| -radius, radius + 1, dtype=dtype, device=device).view(1, -1) | |
| y = torch.arange( | |
| -radius, radius + 1, dtype=dtype, device=device).view(-1, 1) | |
| h = (-(x * x + y * y) / (2 * sigma * sigma)).exp() | |
| h[h < torch.finfo(h.dtype).eps * h.max()] = 0 | |
| return h | |
| def gen_gaussian_target(heatmap, center, radius, k=1): | |
| """Generate 2D gaussian heatmap. | |
| Args: | |
| heatmap (Tensor): Input heatmap, the gaussian kernel will cover on | |
| it and maintain the max value. | |
| center (list[int]): Coord of gaussian kernel's center. | |
| radius (int): Radius of gaussian kernel. | |
| k (int): Coefficient of gaussian kernel. Default: 1. | |
| Returns: | |
| out_heatmap (Tensor): Updated heatmap covered by gaussian kernel. | |
| """ | |
| diameter = 2 * radius + 1 | |
| gaussian_kernel = gaussian2D( | |
| radius, sigma=diameter / 6, dtype=heatmap.dtype, device=heatmap.device) | |
| x, y = center | |
| height, width = heatmap.shape[:2] | |
| left, right = min(x, radius), min(width - x, radius + 1) | |
| top, bottom = min(y, radius), min(height - y, radius + 1) | |
| masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right] | |
| masked_gaussian = gaussian_kernel[radius - top:radius + bottom, | |
| radius - left:radius + right] | |
| out_heatmap = heatmap | |
| torch.max( | |
| masked_heatmap, | |
| masked_gaussian * k, | |
| out=out_heatmap[y - top:y + bottom, x - left:x + right]) | |
| return out_heatmap | |
| def gaussian_radius(det_size, min_overlap): | |
| r"""Generate 2D gaussian radius. | |
| This function is modified from the `official github repo | |
| <https://github.com/princeton-vl/CornerNet-Lite/blob/master/core/sample/ | |
| utils.py#L65>`_. | |
| Given ``min_overlap``, radius could computed by a quadratic equation | |
| according to Vieta's formulas. | |
| There are 3 cases for computing gaussian radius, details are following: | |
| - Explanation of figure: ``lt`` and ``br`` indicates the left-top and | |
| bottom-right corner of ground truth box. ``x`` indicates the | |
| generated corner at the limited position when ``radius=r``. | |
| - Case1: one corner is inside the gt box and the other is outside. | |
| .. code:: text | |
| |< width >| | |
| lt-+----------+ - | |
| | | | ^ | |
| +--x----------+--+ | |
| | | | | | |
| | | | | height | |
| | | overlap | | | |
| | | | | | |
| | | | | v | |
| +--+---------br--+ - | |
| | | | | |
| +----------+--x | |
| To ensure IoU of generated box and gt box is larger than ``min_overlap``: | |
| .. math:: | |
| \cfrac{(w-r)*(h-r)}{w*h+(w+h)r-r^2} \ge {iou} \quad\Rightarrow\quad | |
| {r^2-(w+h)r+\cfrac{1-iou}{1+iou}*w*h} \ge 0 \\ | |
| {a} = 1,\quad{b} = {-(w+h)},\quad{c} = {\cfrac{1-iou}{1+iou}*w*h} | |
| {r} \le \cfrac{-b-\sqrt{b^2-4*a*c}}{2*a} | |
| - Case2: both two corners are inside the gt box. | |
| .. code:: text | |
| |< width >| | |
| lt-+----------+ - | |
| | | | ^ | |
| +--x-------+ | | |
| | | | | | |
| | |overlap| | height | |
| | | | | | |
| | +-------x--+ | |
| | | | v | |
| +----------+-br - | |
| To ensure IoU of generated box and gt box is larger than ``min_overlap``: | |
| .. math:: | |
| \cfrac{(w-2*r)*(h-2*r)}{w*h} \ge {iou} \quad\Rightarrow\quad | |
| {4r^2-2(w+h)r+(1-iou)*w*h} \ge 0 \\ | |
| {a} = 4,\quad {b} = {-2(w+h)},\quad {c} = {(1-iou)*w*h} | |
| {r} \le \cfrac{-b-\sqrt{b^2-4*a*c}}{2*a} | |
| - Case3: both two corners are outside the gt box. | |
| .. code:: text | |
| |< width >| | |
| x--+----------------+ | |
| | | | | |
| +-lt-------------+ | - | |
| | | | | ^ | |
| | | | | | |
| | | overlap | | height | |
| | | | | | |
| | | | | v | |
| | +------------br--+ - | |
| | | | | |
| +----------------+--x | |
| To ensure IoU of generated box and gt box is larger than ``min_overlap``: | |
| .. math:: | |
| \cfrac{w*h}{(w+2*r)*(h+2*r)} \ge {iou} \quad\Rightarrow\quad | |
| {4*iou*r^2+2*iou*(w+h)r+(iou-1)*w*h} \le 0 \\ | |
| {a} = {4*iou},\quad {b} = {2*iou*(w+h)},\quad {c} = {(iou-1)*w*h} \\ | |
| {r} \le \cfrac{-b+\sqrt{b^2-4*a*c}}{2*a} | |
| Args: | |
| det_size (list[int]): Shape of object. | |
| min_overlap (float): Min IoU with ground truth for boxes generated by | |
| keypoints inside the gaussian kernel. | |
| Returns: | |
| radius (int): Radius of gaussian kernel. | |
| """ | |
| height, width = det_size | |
| a1 = 1 | |
| b1 = (height + width) | |
| c1 = width * height * (1 - min_overlap) / (1 + min_overlap) | |
| sq1 = sqrt(b1**2 - 4 * a1 * c1) | |
| r1 = (b1 - sq1) / (2 * a1) | |
| a2 = 4 | |
| b2 = 2 * (height + width) | |
| c2 = (1 - min_overlap) * width * height | |
| sq2 = sqrt(b2**2 - 4 * a2 * c2) | |
| r2 = (b2 - sq2) / (2 * a2) | |
| a3 = 4 * min_overlap | |
| b3 = -2 * min_overlap * (height + width) | |
| c3 = (min_overlap - 1) * width * height | |
| sq3 = sqrt(b3**2 - 4 * a3 * c3) | |
| r3 = (b3 + sq3) / (2 * a3) | |
| return min(r1, r2, r3) | |
| def get_local_maximum(heat, kernel=3): | |
| """Extract local maximum pixel with given kernel. | |
| Args: | |
| heat (Tensor): Target heatmap. | |
| kernel (int): Kernel size of max pooling. Default: 3. | |
| Returns: | |
| heat (Tensor): A heatmap where local maximum pixels maintain its | |
| own value and other positions are 0. | |
| """ | |
| pad = (kernel - 1) // 2 | |
| hmax = F.max_pool2d(heat, kernel, stride=1, padding=pad) | |
| keep = (hmax == heat).float() | |
| return heat * keep | |
| def get_topk_from_heatmap(scores, k=20): | |
| """Get top k positions from heatmap. | |
| Args: | |
| scores (Tensor): Target heatmap with shape | |
| [batch, num_classes, height, width]. | |
| k (int): Target number. Default: 20. | |
| Returns: | |
| tuple[torch.Tensor]: Scores, indexes, categories and coords of | |
| topk keypoint. Containing following Tensors: | |
| - topk_scores (Tensor): Max scores of each topk keypoint. | |
| - topk_inds (Tensor): Indexes of each topk keypoint. | |
| - topk_clses (Tensor): Categories of each topk keypoint. | |
| - topk_ys (Tensor): Y-coord of each topk keypoint. | |
| - topk_xs (Tensor): X-coord of each topk keypoint. | |
| """ | |
| batch, _, height, width = scores.size() | |
| topk_scores, topk_inds = torch.topk(scores.view(batch, -1), k) | |
| topk_clses = topk_inds // (height * width) | |
| topk_inds = topk_inds % (height * width) | |
| topk_ys = topk_inds // width | |
| topk_xs = (topk_inds % width).int().float() | |
| return topk_scores, topk_inds, topk_clses, topk_ys, topk_xs | |
| def gather_feat(feat, ind, mask=None): | |
| """Gather feature according to index. | |
| Args: | |
| feat (Tensor): Target feature map. | |
| ind (Tensor): Target coord index. | |
| mask (Tensor | None): Mask of feature map. Default: None. | |
| Returns: | |
| feat (Tensor): Gathered feature. | |
| """ | |
| dim = feat.size(2) | |
| ind = ind.unsqueeze(2).repeat(1, 1, dim) | |
| feat = feat.gather(1, ind) | |
| if mask is not None: | |
| mask = mask.unsqueeze(2).expand_as(feat) | |
| feat = feat[mask] | |
| feat = feat.view(-1, dim) | |
| return feat | |
| def transpose_and_gather_feat(feat, ind): | |
| """Transpose and gather feature according to index. | |
| Args: | |
| feat (Tensor): Target feature map. | |
| ind (Tensor): Target coord index. | |
| Returns: | |
| feat (Tensor): Transposed and gathered feature. | |
| """ | |
| feat = feat.permute(0, 2, 3, 1).contiguous() | |
| feat = feat.view(feat.size(0), -1, feat.size(3)) | |
| feat = gather_feat(feat, ind) | |
| return feat | |