Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,314 Bytes
a249588 |
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 |
# Copyright (c) OpenMMLab. All rights reserved.
import torch
def fp16_clamp(x, min_val=None, max_val=None):
if not x.is_cuda and x.dtype == torch.float16:
return x.float().clamp(min_val, max_val).half()
return x.clamp(min_val, max_val)
def bbox_overlaps(bboxes1,
bboxes2,
mode='iou',
is_aligned=False,
eps=1e-6) -> torch.Tensor:
"""Calculate overlap between two sets of bounding boxes.
Args:
bboxes1 (torch.Tensor): Bounding boxes of shape (..., m, 4) or empty.
bboxes2 (torch.Tensor): Bounding boxes of shape (..., n, 4) or empty.
mode (str): "iou" (intersection over union),
"iof" (intersection over foreground),
or "giou" (generalized intersection over union).
Defaults to "iou".
is_aligned (bool, optional): If True, then m and n must be equal.
Default False.
eps (float, optional): A small constant added to the denominator for
numerical stability. Default 1e-6.
Returns:
torch.Tensor: Overlap values of shape (..., m, n) if is_aligned is
False, else shape (..., m).
Example:
>>> bboxes1 = torch.FloatTensor([
>>> [0, 0, 10, 10],
>>> [10, 10, 20, 20],
>>> [32, 32, 38, 42],
>>> ])
>>> bboxes2 = torch.FloatTensor([
>>> [0, 0, 10, 20],
>>> [0, 10, 10, 19],
>>> [10, 10, 20, 20],
>>> ])
>>> overlaps = bbox_overlaps(bboxes1, bboxes2)
>>> assert overlaps.shape == (3, 3)
>>> overlaps = bbox_overlaps(bboxes1, bboxes2, is_aligned=True)
>>> assert overlaps.shape == (3, )
"""
assert mode in ['iou', 'iof', 'giou'], f'Unsupported mode {mode}'
assert (bboxes1.size(-1) == 4 or bboxes1.size(0) == 0)
assert (bboxes2.size(-1) == 4 or bboxes2.size(0) == 0)
if bboxes1.ndim == 1:
bboxes1 = bboxes1.unsqueeze(0)
if bboxes2.ndim == 1:
bboxes2 = bboxes2.unsqueeze(0)
assert bboxes1.shape[:-2] == bboxes2.shape[:-2]
batch_shape = bboxes1.shape[:-2]
rows = bboxes1.size(-2)
cols = bboxes2.size(-2)
if is_aligned:
assert rows == cols
if rows * cols == 0:
if is_aligned:
return bboxes1.new(batch_shape + (rows, ))
else:
return bboxes1.new(batch_shape + (rows, cols))
area1 = (bboxes1[..., 2] - bboxes1[..., 0]) * (
bboxes1[..., 3] - bboxes1[..., 1])
area2 = (bboxes2[..., 2] - bboxes2[..., 0]) * (
bboxes2[..., 3] - bboxes2[..., 1])
if is_aligned:
lt = torch.max(bboxes1[..., :2], bboxes2[..., :2])
rb = torch.min(bboxes1[..., 2:], bboxes2[..., 2:])
wh = fp16_clamp(rb - lt, min_val=0)
overlap = wh[..., 0] * wh[..., 1]
if mode in ['iou', 'giou']:
union = area1 + area2 - overlap
else:
union = area1
if mode == 'giou':
enclosed_lt = torch.min(bboxes1[..., :2], bboxes2[..., :2])
enclosed_rb = torch.max(bboxes1[..., 2:], bboxes2[..., 2:])
else:
lt = torch.max(bboxes1[..., :, None, :2], bboxes2[..., None, :, :2])
rb = torch.min(bboxes1[..., :, None, 2:], bboxes2[..., None, :, 2:])
wh = fp16_clamp(rb - lt, min_val=0)
overlap = wh[..., 0] * wh[..., 1]
if mode in ['iou', 'giou']:
union = area1[..., None] + area2[..., None, :] - overlap
else:
union = area1[..., None]
if mode == 'giou':
enclosed_lt = torch.min(bboxes1[..., :, None, :2],
bboxes2[..., None, :, :2])
enclosed_rb = torch.max(bboxes1[..., :, None, 2:],
bboxes2[..., None, :, 2:])
eps_tensor = union.new_tensor([eps])
union = torch.max(union, eps_tensor)
ious = overlap / union
if mode in ['iou', 'iof']:
return ious
elif mode == 'giou':
enclose_wh = fp16_clamp(enclosed_rb - enclosed_lt, min_val=0)
enclose_area = enclose_wh[..., 0] * enclose_wh[..., 1]
enclose_area = torch.max(enclose_area, eps_tensor)
gious = ious - (enclose_area - union) / enclose_area
return gious
|