File size: 2,979 Bytes
25986db | 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 | import torch
def rect_to_rel(bb, sz_norm=None):
"""Convert standard rectangular parametrization of the bounding box [x, y, w, h]
to relative parametrization [cx/sw, cy/sh, log(w), log(h)], where [cx, cy] is the center coordinate.
args:
bb - N x 4 tensor of boxes.
sz_norm - [N] x 2 tensor of value of [sw, sh] (optional). sw=w and sh=h if not given.
"""
c = bb[...,:2] + 0.5 * bb[...,2:]
if sz_norm is None:
c_rel = c / bb[...,2:]
else:
c_rel = c / sz_norm
sz_rel = torch.log(bb[...,2:])
return torch.cat((c_rel, sz_rel), dim=-1)
def rel_to_rect(bb, sz_norm=None):
"""Inverts the effect of rect_to_rel. See above."""
sz = torch.exp(bb[...,2:])
if sz_norm is None:
c = bb[...,:2] * sz
else:
c = bb[...,:2] * sz_norm
tl = c - 0.5 * sz
return torch.cat((tl, sz), dim=-1)
def masks_to_bboxes(mask, fmt='c'):
""" Convert a mask tensor to one or more bounding boxes.
Note: This function is a bit new, make sure it does what it says. /Andreas
:param mask: Tensor of masks, shape = (..., H, W)
:param fmt: bbox layout. 'c' => "center + size" or (x_center, y_center, width, height)
't' => "top left + size" or (x_left, y_top, width, height)
'v' => "vertices" or (x_left, y_top, x_right, y_bottom)
:return: tensor containing a batch of bounding boxes, shape = (..., 4)
"""
batch_shape = mask.shape[:-2]
mask = mask.reshape((-1, *mask.shape[-2:]))
bboxes = []
for m in mask:
mx = m.sum(dim=-2).nonzero()
my = m.sum(dim=-1).nonzero()
bb = [mx.min(), my.min(), mx.max(), my.max()] if (len(mx) > 0 and len(my) > 0) else [0, 0, 0, 0]
bboxes.append(bb)
bboxes = torch.tensor(bboxes, dtype=torch.float32, device=mask.device)
bboxes = bboxes.reshape(batch_shape + (4,))
if fmt == 'v':
return bboxes
x1 = bboxes[..., :2]
s = bboxes[..., 2:] - x1 + 1
if fmt == 'c':
return torch.cat((x1 + 0.5 * s, s), dim=-1)
elif fmt == 't':
return torch.cat((x1, s), dim=-1)
raise ValueError("Undefined bounding box layout '%s'" % fmt)
def masks_to_bboxes_multi(mask, ids, fmt='c'):
assert mask.dim() == 2
bboxes = []
for id in ids:
mx = (mask == id).sum(dim=-2).nonzero()
my = (mask == id).float().sum(dim=-1).nonzero()
bb = [mx.min(), my.min(), mx.max(), my.max()] if (len(mx) > 0 and len(my) > 0) else [0, 0, 0, 0]
bb = torch.tensor(bb, dtype=torch.float32, device=mask.device)
x1 = bb[:2]
s = bb[2:] - x1 + 1
if fmt == 'v':
pass
elif fmt == 'c':
bb = torch.cat((x1 + 0.5 * s, s), dim=-1)
elif fmt == 't':
bb = torch.cat((x1, s), dim=-1)
else:
raise ValueError("Undefined bounding box layout '%s'" % fmt)
bboxes.append(bb)
return bboxes
|