BatchMasks

class torchaug.ta_tensors.BatchMasks(data: Any, *, samples_ranges: List[Tuple[int, int]], dtype: dtype | None = None, device: device | str | int | None = None, requires_grad: bool | None = None)

torch.Tensor subclass for batch of segmentation and detection masks.

Parameters:

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

• dtype – Desired data type. If omitted, will be inferred from data.

• samples_ranges – Each element is the range of the indices of the masks for each sample.

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

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

classmethod cat(masks_batches)

Concatenates a sequence of BatchMasks along the first dimension.

Parameters:

masks_batches (Sequence[BatchMasks]) – A sequence of BatchMasks to concatenate.

Returns:

The concatenated BatchMasks.

get_chunk(chunk_indices)

Get a chunk of the batch of masks.

Parameters:

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

Return type:

BatchMasks

Returns:

The chunk of the batch of masks.

get_sample(idx)

Get the masks for a sample in the batch.

Parameters:

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

Return type:

Mask

Returns:

The masks for the sample.

classmethod masked_select(masks, mask)

Remove masks from the batch of masks.

Parameters:

• masks (BatchMasks) – The batch of masks to remove masks from.

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

Return type:

BatchMasks

Returns:

The updated batch of masks.

to_samples()

Get the tensors.

Return type:

list[Mask]

update_chunk_(chunk, chunk_indices)

Update a chunk of the batch of masks.

Parameters:

• chunk (BatchMasks) – The chunk update.

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

Return type:

BatchMasks

Returns:

The updated batch of masks.