BatchBoundingBoxes

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

torch.Tensor subclass for batch of bounding boxes.

Note

There should be only one BatchBoundingBoxes instance per sample e.g. {"img": img, "bbox": BatchBoundingBoxes(...)}.

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 batch of images or videos.

• idx_sample – Each element is the index of the first bounding box of the corresponding sample in the batch of N samples. Contains N+1 elements whose last value is the number of bounding boxes.

• 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 cat(bounding_boxes_batches)

Concatenates the given sequence of BatchBoundingBoxes along the first dimension.

Parameters:

bounding_boxes_batches (Sequence[BatchBoundingBoxes]) – The sequence of BatchBoundingBoxes to concatenate.

Returns:

The concatenated batch of bounding boxes.

Return type:

BatchBoundingBoxes

get_chunk(chunk_indices)

Get a chunk of the batch of bounding boxes.

Parameters:

chunk_indices (torch.Tensor) – The indices of the chunk to get.

Returns:

The chunk of the batch bounding boxes.

Return type:

BatchBoundingBoxes

get_sample(idx)

Get the bounding boxes for a sample in the batch.

Parameters:

idx (int) – The index of the sample to get.

Return type:

BoundingBoxes

Returns:

The bounding boxes for the sample.

classmethod masked_remove(bboxes, mask)

Remove boxes from the batch of bounding boxes.

Parameters:

• bboxes (BatchBoundingBoxes) – The batch of bounding boxes to remove boxes from.

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

Returns:

The updated batch of bounding boxes.

Return type:

BatchBoundingBoxes

update_chunk_(chunk, chunk_indices)

Update a chunk of the batch of bounding boxes.

Parameters:

• chunk (BatchBoundingBoxes) – The chunk update.

• chunk_indices (torch.Tensor) – The indices of the chunk to update.

Returns:

The updated batch of bounding boxes.

Return type:

BatchBoundingBoxes