Class ClipBoxToImage

monai/apps/detection/transforms/array.py:351–397

Clip the bounding boxes and the associated labels/scores to make sure they are within the image. There might be multiple arrays of labels/scores associated with one array of boxes. Args: remove_empty: whether to remove the boxes and corresponding labels that are actually empty

class ClipBoxToImage(Transform):
    """
    Clip the bounding boxes and the associated labels/scores to make sure they are within the image.
    There might be multiple arrays of labels/scores associated with one array of boxes.

    Args:
        remove_empty: whether to remove the boxes and corresponding labels that are actually empty
    """

    backend = [TransformBackends.TORCH, TransformBackends.NUMPY]

    def __init__(self, remove_empty: bool = False) -> None:
        self.remove_empty = remove_empty

    def __call__(  # type: ignore
        self,
        boxes: NdarrayOrTensor,
        labels: Sequence[NdarrayOrTensor] | NdarrayOrTensor,
        spatial_size: Sequence[int] | int,
    ) -> tuple[NdarrayOrTensor, tuple | NdarrayOrTensor]:
        """
        Args:
            boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode``
            labels: Sequence of array. Each element represents classification labels or scores
                corresponding to ``boxes``, sized (N,).
            spatial_size: The spatial size of the image where the boxes are attached. len(spatial_size) should be in [2, 3].

        Returns:
            - clipped boxes, does not share memory with original boxes
            - clipped labels, does not share memory with original labels

        Example:
            .. code-block:: python

                box_clipper = ClipBoxToImage(remove_empty=True)
                boxes = torch.ones(2, 6)
                class_labels = torch.Tensor([0, 1])
                pred_scores = torch.Tensor([[0.4,0.3,0.3], [0.5,0.1,0.4]])
                labels = (class_labels, pred_scores)
                spatial_size = [32, 32, 32]
                boxes_clip, labels_clip_tuple = box_clipper(boxes, labels, spatial_size)
        """
        spatial_dims: int = get_spatial_dims(boxes=boxes)
        spatial_size = ensure_tuple_rep(spatial_size, spatial_dims)  # match the spatial image dim

        boxes_clip, keep = clip_boxes_to_image(boxes, spatial_size, self.remove_empty)
        return boxes_clip, select_labels(labels, keep)


class BoxToMask(Transform):