Function spatial_crop_boxes

monai/data/box_utils.py:1011–1065

This function generate the new boxes when the corresponding image is cropped to the given ROI. When ``remove_empty=True``, it makes sure the bounding boxes are within the new cropped image. Args: boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed

(
    boxes: NdarrayTensor,
    roi_start: Sequence[int] | NdarrayOrTensor,
    roi_end: Sequence[int] | NdarrayOrTensor,
    remove_empty: bool = True,
)

def spatial_crop_boxes(
    boxes: NdarrayTensor,
    roi_start: Sequence[int] | NdarrayOrTensor,
    roi_end: Sequence[int] | NdarrayOrTensor,
    remove_empty: bool = True,
) -> tuple[NdarrayTensor, NdarrayOrTensor]:
    """
    This function generate the new boxes when the corresponding image is cropped to the given ROI.
    When ``remove_empty=True``, it makes sure the bounding boxes are within the new cropped image.

    Args:
        boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode``
        roi_start: voxel coordinates for start of the crop ROI, negative values allowed.
        roi_end: voxel coordinates for end of the crop ROI, negative values allowed.
        remove_empty: whether to remove the boxes that are actually empty

    Returns:
        - cropped boxes, boxes[keep], does not share memory with original boxes
        - ``keep``, it indicates whether each box in ``boxes`` are kept when ``remove_empty=True``.
    """

    # convert numpy to tensor if needed
    boxes_t = convert_data_type(boxes, torch.Tensor)[0].clone()

    # convert to float32 since torch.clamp_ does not support float16
    boxes_t = boxes_t.to(dtype=COMPUTE_DTYPE)

    roi_start_t = convert_to_dst_type(src=roi_start, dst=boxes_t, wrap_sequence=True)[0].to(torch.int16)
    roi_end_t = convert_to_dst_type(src=roi_end, dst=boxes_t, wrap_sequence=True)[0].to(torch.int16)
    roi_end_t = torch.maximum(roi_end_t, roi_start_t)

    # makes sure the bounding boxes are within the patch
    spatial_dims = get_spatial_dims(boxes=boxes, spatial_size=roi_end)
    for axis in range(spatial_dims):
        boxes_t[:, axis] = boxes_t[:, axis].clamp(min=roi_start_t[axis], max=roi_end_t[axis] - TO_REMOVE)
        boxes_t[:, axis + spatial_dims] = boxes_t[:, axis + spatial_dims].clamp(
            min=roi_start_t[axis], max=roi_end_t[axis] - TO_REMOVE
        )
        boxes_t[:, axis] -= roi_start_t[axis]
        boxes_t[:, axis + spatial_dims] -= roi_start_t[axis]

    # remove the boxes that are actually empty
    if remove_empty:
        keep_t = boxes_t[:, spatial_dims] >= boxes_t[:, 0] + 1 - TO_REMOVE
        for axis in range(1, spatial_dims):
            keep_t = keep_t & (boxes_t[:, axis + spatial_dims] >= boxes_t[:, axis] + 1 - TO_REMOVE)
        boxes_t = boxes_t[keep_t]
    else:
        keep_t = torch.full_like(boxes_t[:, 0], fill_value=True, dtype=torch.bool)

    # convert tensor back to numpy if needed
    boxes_keep, *_ = convert_to_dst_type(src=boxes_t, dst=boxes)
    keep, *_ = convert_to_dst_type(src=keep_t, dst=boxes, dtype=keep_t.dtype)

    return boxes_keep, keep


def clip_boxes_to_image(