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