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