Compute the bounding boxes around the provided masks The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions. Returns a [N, 4] tensors, with the boxes in xyxy format
(masks)
| 62 | |
| 63 | |
| 64 | def masks_to_boxes(masks): |
| 65 | """Compute the bounding boxes around the provided masks |
| 66 | |
| 67 | The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions. |
| 68 | |
| 69 | Returns a [N, 4] tensors, with the boxes in xyxy format |
| 70 | """ |
| 71 | if masks.numel() == 0: |
| 72 | return torch.zeros((0, 4), device=masks.device) |
| 73 | |
| 74 | h, w = masks.shape[-2:] |
| 75 | |
| 76 | y = torch.arange(0, h, dtype=torch.float) |
| 77 | x = torch.arange(0, w, dtype=torch.float) |
| 78 | y, x = torch.meshgrid(y, x) |
| 79 | |
| 80 | x_mask = (masks * x.unsqueeze(0)) |
| 81 | x_max = x_mask.flatten(1).max(-1)[0] |
| 82 | x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] |
| 83 | |
| 84 | y_mask = (masks * y.unsqueeze(0)) |
| 85 | y_max = y_mask.flatten(1).max(-1)[0] |
| 86 | y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] |
| 87 | |
| 88 | return torch.stack([x_min, y_min, x_max, y_max], 1) |