BoundingBoxes

class torchaug.ta_tensors.BoundingBoxes(data: Any, *, format: BoundingBoxFormat | str, canvas_size: Tuple[int, int], dtype: dtype | None = None, device: device | str | int | None = None, requires_grad: bool | None = None)

torch.Tensor subclass for bounding boxes.

Note

There should be only one BoundingBoxes instance per sample e.g. {"img": img, "bbox": BoundingBoxes(...)}, although one BoundingBoxes object can contain multiple bounding boxes.

Parameters:

• data – Any data that can be turned into a tensor with torch.as_tensor().

• format – Format of the bounding box.

• canvas_size – Height and width of the corresponding image or video.

• dtype – Desired data type of the bounding box. If omitted, will be inferred from data.

• device – Desired device of the bounding box. If omitted and data is a torch.Tensor, the device is taken from it. Otherwise, the bounding box is constructed on the CPU.

• requires_grad – Whether autograd should record operations on the bounding box. If omitted and data is a torch.Tensor, the value is taken from it. Otherwise, defaults to False.

classmethod masked_select(bboxes, mask)

Remove boxes from the bounding boxes.

Parameters:

• bboxes (BoundingBoxes) – The bounding boxes to remove boxes from.

• mask (Tensor) – A boolean mask to keep boxes.

Return type:

BoundingBoxes

Returns:

The updated bounding boxes.