OpenMMLab Detection Toolbox and Benchmark
https://mmdetection.readthedocs.io/
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
268 lines
8.2 KiB
268 lines
8.2 KiB
# 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
|
|
|