| | |
| | 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 |
| |
|
