This structure stores a list of boxes as a Nx4 torch.Tensor. It supports some common methods about boxes (`area`, `clip`, `nonempty`, etc), and also behaves like a Tensor (support indexing, `to(device)`, `.device`, and iteration over all boxes) Attributes: tensor (t
| 131 | |
| 132 | |
| 133 | class Boxes: |
| 134 | """ |
| 135 | This structure stores a list of boxes as a Nx4 torch.Tensor. |
| 136 | It supports some common methods about boxes |
| 137 | (`area`, `clip`, `nonempty`, etc), |
| 138 | and also behaves like a Tensor |
| 139 | (support indexing, `to(device)`, `.device`, and iteration over all boxes) |
| 140 | |
| 141 | Attributes: |
| 142 | tensor (torch.Tensor): float matrix of Nx4. Each row is (x1, y1, x2, y2). |
| 143 | """ |
| 144 | |
| 145 | def __init__(self, tensor: torch.Tensor): |
| 146 | """ |
| 147 | Args: |
| 148 | tensor (Tensor[float]): a Nx4 matrix. Each row is (x1, y1, x2, y2). |
| 149 | """ |
| 150 | device = tensor.device if isinstance(tensor, torch.Tensor) else torch.device("cpu") |
| 151 | tensor = torch.as_tensor(tensor, dtype=torch.float32, device=device) |
| 152 | if tensor.numel() == 0: |
| 153 | # Use reshape, so we don't end up creating a new tensor that does not depend on |
| 154 | # the inputs (and consequently confuses jit) |
| 155 | tensor = tensor.reshape((0, 4)).to(dtype=torch.float32, device=device) |
| 156 | assert tensor.dim() == 2 and tensor.size(-1) == 4, tensor.size() |
| 157 | |
| 158 | self.tensor = tensor |
| 159 | |
| 160 | def clone(self) -> "Boxes": |
| 161 | """ |
| 162 | Clone the Boxes. |
| 163 | |
| 164 | Returns: |
| 165 | Boxes |
| 166 | """ |
| 167 | return Boxes(self.tensor.clone()) |
| 168 | |
| 169 | # https://github.com/pytorch/pytorch/issues/47405 |
| 170 | @torch.jit.unused |
| 171 | def to(self, device: torch.device = None): # noqa |
| 172 | # Boxes are assumed float32 and does not support to(dtype) |
| 173 | return Boxes(self.tensor.to(device=device)) |
| 174 | |
| 175 | def area(self) -> torch.Tensor: |
| 176 | """ |
| 177 | Computes the area of all the boxes. |
| 178 | |
| 179 | Returns: |
| 180 | torch.Tensor: a vector with areas of each box. |
| 181 | """ |
| 182 | box = self.tensor |
| 183 | area = (box[:, 2] - box[:, 0]) * (box[:, 3] - box[:, 1]) |
| 184 | return area |
| 185 | |
| 186 | def clip(self, box_size: Tuple[int, int]) -> None: |
| 187 | """ |
| 188 | Clip (in place) the boxes by limiting x coordinates to the range [0, width] |
| 189 | and y coordinates to the range [0, height]. |
| 190 |
no outgoing calls