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