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)
| 90 | |
| 91 | |
| 92 | def masks_to_boxes(masks): |
| 93 | """Compute the bounding boxes around the provided masks |
| 94 | |
| 95 | The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions. |
| 96 | |
| 97 | Returns a [N, 4] tensors, with the boxes in xyxy format |
| 98 | """ |
| 99 | if masks.numel() == 0: |
| 100 | return torch.zeros((0, 4), device=masks.device) |
| 101 | |
| 102 | h, w = masks.shape[-2:] |
| 103 | |
| 104 | y = torch.arange(0, h, dtype=torch.float) |
| 105 | x = torch.arange(0, w, dtype=torch.float) |
| 106 | y, x = torch.meshgrid(y, x) |
| 107 | |
| 108 | x_mask = (masks * x.unsqueeze(0)) |
| 109 | x_max = x_mask.flatten(1).max(-1)[0] |
| 110 | x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] |
| 111 | |
| 112 | y_mask = (masks * y.unsqueeze(0)) |
| 113 | y_max = y_mask.flatten(1).max(-1)[0] |
| 114 | y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] |
| 115 | |
| 116 | return torch.stack([x_min, y_min, x_max, y_max], 1) |