keras-hub-nightly 0.15.0.dev20240823171555__py3-none-any.whl → 0.16.0.dev2024092017__py3-none-any.whl

keras_hub/src/bounding_box/iou.py

@@ -0,0 +1,263 @@
+# Copyright 2024 The KerasHub Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Contains functions to compute ious of bounding boxes."""
+import math
+
+import keras
+from keras import ops
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.bounding_box.converters import convert_format
+from keras_hub.src.bounding_box.utils import as_relative
+from keras_hub.src.bounding_box.utils import is_relative
+
+
+def _compute_area(box):
+    """Computes area for bounding boxes
+
+    Args:
+      box: [N, 4] or [batch_size, N, 4] float Tensor, either batched
+        or unbatched boxes.
+    Returns:
+      a float Tensor of [N] or [batch_size, N]
+    """
+    y_min, x_min, y_max, x_max = ops.split(box[..., :4], 4, axis=-1)
+    return ops.squeeze((y_max - y_min) * (x_max - x_min), axis=-1)
+
+
+def _compute_intersection(boxes1, boxes2):
+    """Computes intersection area between two sets of boxes.
+
+    Args:
+      boxes1: [N, 4] or [batch_size, N, 4] float Tensor boxes.
+      boxes2: [M, 4] or [batch_size, M, 4] float Tensor boxes.
+    Returns:
+      a [N, M] or [batch_size, N, M] float Tensor.
+    """
+    y_min1, x_min1, y_max1, x_max1 = ops.split(boxes1[..., :4], 4, axis=-1)
+    y_min2, x_min2, y_max2, x_max2 = ops.split(boxes2[..., :4], 4, axis=-1)
+    boxes2_rank = len(boxes2.shape)
+    perm = [1, 0] if boxes2_rank == 2 else [0, 2, 1]
+    # [N, M] or [batch_size, N, M]
+    intersect_ymax = ops.minimum(y_max1, ops.transpose(y_max2, perm))
+    intersect_ymin = ops.maximum(y_min1, ops.transpose(y_min2, perm))
+    intersect_xmax = ops.minimum(x_max1, ops.transpose(x_max2, perm))
+    intersect_xmin = ops.maximum(x_min1, ops.transpose(x_min2, perm))
+
+    intersect_height = intersect_ymax - intersect_ymin
+    intersect_width = intersect_xmax - intersect_xmin
+    zeros_t = ops.cast(0, intersect_height.dtype)
+    intersect_height = ops.maximum(zeros_t, intersect_height)
+    intersect_width = ops.maximum(zeros_t, intersect_width)
+
+    return intersect_height * intersect_width
+
+
+@keras_hub_export("keras_hub.bounding_box.compute_iou")
+def compute_iou(
+    boxes1,
+    boxes2,
+    bounding_box_format,
+    use_masking=False,
+    mask_val=-1,
+    images=None,
+    image_shape=None,
+):
+    """Computes a lookup table vector containing the ious for a given set boxes.
+
+    The lookup vector is to be indexed by [`boxes1_index`,`boxes2_index`] if
+    boxes are unbatched and by [`batch`, `boxes1_index`,`boxes2_index`] if the
+    boxes are batched.
+
+    The users can pass `boxes1` and `boxes2` to be different ranks. For example:
+    1) `boxes1`: [batch_size, M, 4], `boxes2`: [batch_size, N, 4] -> return
+        [batch_size, M, N].
+    2) `boxes1`: [batch_size, M, 4], `boxes2`: [N, 4] -> return
+        [batch_size, M, N]
+    3) `boxes1`: [M, 4], `boxes2`: [batch_size, N, 4] -> return
+        [batch_size, M, N]
+    4) `boxes1`: [M, 4], `boxes2`: [N, 4] -> return [M, N]
+
+    Args:
+      boxes1: a list of bounding boxes in 'corners' format. Can be batched or
+        unbatched.
+      boxes2: a list of bounding boxes in 'corners' format. Can be batched or
+        unbatched.
+      bounding_box_format: a case-insensitive string which is one of `"xyxy"`,
+        `"rel_xyxy"`, `"xyWH"`, `"center_xyWH"`, `"yxyx"`, `"rel_yxyx"`.
+        For detailed information on the supported format, see the
+        [KerasCV bounding box documentation](https://keras.io/api/keras_cv/bounding_box/formats/).
+      use_masking: whether masking will be applied. This will mask all `boxes1`
+        or `boxes2` that have values less than 0 in all its 4 dimensions.
+        Default to `False`.
+      mask_val: int to mask those returned IOUs if the masking is True, defaults
+        to -1.
+
+    Returns:
+      iou_lookup_table: a vector containing the pairwise ious of boxes1 and
+        boxes2.
+    """  # noqa: E501
+
+    boxes1_rank = len(boxes1.shape)
+    boxes2_rank = len(boxes2.shape)
+
+    if boxes1_rank not in [2, 3]:
+        raise ValueError(
+            "compute_iou() expects boxes1 to be batched, or to be unbatched. "
+            f"Received len(boxes1.shape)={boxes1_rank}, "
+            f"len(boxes2.shape)={boxes2_rank}. Expected either "
+            "len(boxes1.shape)=2 AND or len(boxes1.shape)=3."
+        )
+    if boxes2_rank not in [2, 3]:
+        raise ValueError(
+            "compute_iou() expects boxes2 to be batched, or to be unbatched. "
+            f"Received len(boxes1.shape)={boxes1_rank}, "
+            f"len(boxes2.shape)={boxes2_rank}. Expected either "
+            "len(boxes2.shape)=2 AND or len(boxes2.shape)=3."
+        )
+
+    target_format = "yxyx"
+    if is_relative(bounding_box_format):
+        target_format = as_relative(target_format)
+
+    boxes1 = convert_format(
+        boxes1,
+        source=bounding_box_format,
+        target=target_format,
+        images=images,
+        image_shape=image_shape,
+    )
+
+    boxes2 = convert_format(
+        boxes2,
+        source=bounding_box_format,
+        target=target_format,
+        images=images,
+        image_shape=image_shape,
+    )
+
+    intersect_area = _compute_intersection(boxes1, boxes2)
+    boxes1_area = _compute_area(boxes1)
+    boxes2_area = _compute_area(boxes2)
+    boxes2_area_rank = len(boxes2_area.shape)
+    boxes2_axis = 1 if (boxes2_area_rank == 2) else 0
+    boxes1_area = ops.expand_dims(boxes1_area, axis=-1)
+    boxes2_area = ops.expand_dims(boxes2_area, axis=boxes2_axis)
+    union_area = boxes1_area + boxes2_area - intersect_area
+    res = ops.divide(intersect_area, union_area + keras.backend.epsilon())
+
+    if boxes1_rank == 2:
+        perm = [1, 0]
+    else:
+        perm = [0, 2, 1]
+
+    if not use_masking:
+        return res
+
+    mask_val_t = ops.cast(mask_val, res.dtype) * ops.ones_like(res)
+    boxes1_mask = ops.less(ops.max(boxes1, axis=-1, keepdims=True), 0.0)
+    boxes2_mask = ops.less(ops.max(boxes2, axis=-1, keepdims=True), 0.0)
+    background_mask = ops.logical_or(
+        boxes1_mask, ops.transpose(boxes2_mask, perm)
+    )
+    iou_lookup_table = ops.where(background_mask, mask_val_t, res)
+    return iou_lookup_table
+
+
+@keras_hub_export("keras_hub.bounding_box.compute_ciou")
+def compute_ciou(boxes1, boxes2, bounding_box_format):
+    """
+    Computes the Complete IoU (CIoU) between two bounding boxes or between
+    two batches of bounding boxes.
+
+    CIoU loss is an extension of GIoU loss, which further improves the IoU
+    optimization for object detection. CIoU loss not only penalizes the
+    bounding box coordinates but also considers the aspect ratio and center
+    distance of the boxes. The length of the last dimension should be 4 to
+    represent the bounding boxes.
+
+    Args:
+        box1 (tensor): tensor representing the first bounding box with
+            shape (..., 4).
+        box2 (tensor): tensor representing the second bounding box with
+            shape (..., 4).
+        bounding_box_format: a case-insensitive string (for example, "xyxy").
+            Each bounding box is defined by these 4 values. For detailed
+            information on the supported formats, see the [KerasCV bounding box
+            documentation](https://keras.io/api/keras_cv/bounding_box/formats/).
+
+    Returns:
+        tensor: The CIoU distance between the two bounding boxes.
+    """
+    target_format = "xyxy"
+    if is_relative(bounding_box_format):
+        target_format = as_relative(target_format)
+
+    boxes1 = convert_format(
+        boxes1, source=bounding_box_format, target=target_format
+    )
+
+    boxes2 = convert_format(
+        boxes2, source=bounding_box_format, target=target_format
+    )
+
+    x_min1, y_min1, x_max1, y_max1 = ops.split(boxes1[..., :4], 4, axis=-1)
+    x_min2, y_min2, x_max2, y_max2 = ops.split(boxes2[..., :4], 4, axis=-1)
+
+    width_1 = x_max1 - x_min1
+    height_1 = y_max1 - y_min1 + keras.backend.epsilon()
+    width_2 = x_max2 - x_min2
+    height_2 = y_max2 - y_min2 + keras.backend.epsilon()
+
+    intersection_area = ops.maximum(
+        ops.minimum(x_max1, x_max2) - ops.maximum(x_min1, x_min2), 0
+    ) * ops.maximum(
+        ops.minimum(y_max1, y_max2) - ops.maximum(y_min1, y_min2), 0
+    )
+    union_area = (
+        width_1 * height_1
+        + width_2 * height_2
+        - intersection_area
+        + keras.backend.epsilon()
+    )
+    iou = ops.squeeze(
+        ops.divide(intersection_area, union_area + keras.backend.epsilon()),
+        axis=-1,
+    )
+
+    convex_width = ops.maximum(x_max1, x_max2) - ops.minimum(x_min1, x_min2)
+    convex_height = ops.maximum(y_max1, y_max2) - ops.minimum(y_min1, y_min2)
+    convex_diagonal_squared = ops.squeeze(
+        convex_width**2 + convex_height**2 + keras.backend.epsilon(),
+        axis=-1,
+    )
+    centers_distance_squared = ops.squeeze(
+        ((x_min1 + x_max1) / 2 - (x_min2 + x_max2) / 2) ** 2
+        + ((y_min1 + y_max1) / 2 - (y_min2 + y_max2) / 2) ** 2,
+        axis=-1,
+    )
+
+    v = ops.squeeze(
+        ops.power(
+            (4 / math.pi**2)
+            * (ops.arctan(width_2 / height_2) - ops.arctan(width_1 / height_1)),
+            2,
+        ),
+        axis=-1,
+    )
+    alpha = v / (v - iou + (1 + keras.backend.epsilon()))
+
+    return iou - (
+        centers_distance_squared / convex_diagonal_squared + v * alpha
+    )

keras_hub/src/bounding_box/to_dense.py

@@ -0,0 +1,95 @@
+# Copyright 2024 The KerasHub Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import keras_hub.src.bounding_box.validate_format as validate_format
+from keras_hub.src.api_export import keras_hub_export
+
+try:
+    import tensorflow as tf
+except ImportError:
+    tf = None
+
+
+def _box_shape(batched, boxes_shape, max_boxes):
+    # ensure we dont drop the final axis in RaggedTensor mode
+    if max_boxes is None:
+        shape = list(boxes_shape)
+        shape[-1] = 4
+        return shape
+    if batched:
+        return [None, max_boxes, 4]
+    return [max_boxes, 4]
+
+
+def _classes_shape(batched, classes_shape, max_boxes):
+    if max_boxes is None:
+        return None
+    if batched:
+        return [None, max_boxes] + classes_shape[2:]
+    return [max_boxes] + classes_shape[2:]
+
+
+@keras_hub_export("keras_hub.bounding_box.to_dense")
+def to_dense(bounding_boxes, max_boxes=None, default_value=-1):
+    """to_dense converts bounding boxes to Dense tensors
+
+    Args:
+        bounding_boxes: bounding boxes in KerasCV dictionary format.
+        max_boxes: the maximum number of boxes, used to pad tensors to a given
+            shape. This can be used to make object detection pipelines TPU
+            compatible.
+        default_value: the default value to pad bounding boxes with. defaults
+            to -1.
+    """
+    info = validate_format.validate_format(bounding_boxes)
+
+    # guards against errors in metrics regarding modification of inputs.
+    # also guards against unexpected behavior when modifying downstream
+    bounding_boxes = bounding_boxes.copy()
+
+    # Already running in masked mode
+    if not info["ragged"]:
+        # even if already ragged, still copy the dictionary for API consistency
+        return bounding_boxes
+
+    if isinstance(bounding_boxes["classes"], tf.RaggedTensor):
+        bounding_boxes["classes"] = bounding_boxes["classes"].to_tensor(
+            default_value=default_value,
+            shape=_classes_shape(
+                info["is_batched"], bounding_boxes["classes"].shape, max_boxes
+            ),
+        )
+
+    if isinstance(bounding_boxes["boxes"], tf.RaggedTensor):
+        bounding_boxes["boxes"] = bounding_boxes["boxes"].to_tensor(
+            default_value=default_value,
+            shape=_box_shape(
+                info["is_batched"], bounding_boxes["boxes"].shape, max_boxes
+            ),
+        )
+
+    if "confidence" in bounding_boxes:
+        if isinstance(bounding_boxes["confidence"], tf.RaggedTensor):
+            bounding_boxes["confidence"] = bounding_boxes[
+                "confidence"
+            ].to_tensor(
+                default_value=default_value,
+                shape=_classes_shape(
+                    info["is_batched"],
+                    bounding_boxes["confidence"].shape,
+                    max_boxes,
+                ),
+            )
+
+    return bounding_boxes

keras_hub/src/bounding_box/to_ragged.py

@@ -0,0 +1,99 @@
+# Copyright 2024 The KerasHub Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import keras
+
+import keras_hub.src.bounding_box.validate_format as validate_format
+from keras_hub.src.api_export import keras_hub_export
+
+try:
+    import tensorflow as tf
+except ImportError:
+    tf = None
+
+
+@keras_hub_export("keras_hub.bounding_box.to_ragged")
+def to_ragged(bounding_boxes, sentinel=-1, dtype="float32"):
+    """converts a Dense padded bounding box `tf.Tensor` to a `tf.RaggedTensor`.
+
+    Bounding boxes are ragged tensors in most use cases. Converting them to a
+    dense tensor makes it easier to work with Tensorflow ecosystem.
+    This function can be used to filter out the masked out bounding boxes by
+    checking for padded sentinel value of the class_id axis of the
+    bounding_boxes.
+
+    Example:
+    ```python
+    bounding_boxes = {
+        "boxes": tf.constant([[2, 3, 4, 5], [0, 1, 2, 3]]),
+        "classes": tf.constant([[-1, 1]]),
+    }
+    bounding_boxes = bounding_box.to_ragged(bounding_boxes)
+    print(bounding_boxes)
+    # {
+    #     "boxes": [[0, 1, 2, 3]],
+    #     "classes": [[1]]
+    # }
+    ```
+
+    Args:
+        bounding_boxes: a Tensor of bounding boxes. May be batched, or
+            unbatched.
+        sentinel: The value indicating that a bounding box does not exist at the
+            current index, and the corresponding box is padding, defaults to -1.
+        dtype: the data type to use for the underlying Tensors.
+    Returns:
+        dictionary of `tf.RaggedTensor` or 'tf.Tensor' containing the filtered
+        bounding boxes.
+    """
+    if keras.config.backend() != "tensorflow":
+        raise NotImplementedError(
+            "`bounding_box.to_ragged` was called using a backend which does "
+            "not support ragged tensors. "
+            f"Current backend: {keras.backend.backend()}."
+        )
+
+    info = validate_format.validate_format(bounding_boxes)
+
+    if info["ragged"]:
+        return bounding_boxes
+
+    boxes = bounding_boxes.get("boxes")
+    classes = bounding_boxes.get("classes")
+    confidence = bounding_boxes.get("confidence", None)
+
+    mask = classes != sentinel
+
+    boxes = tf.ragged.boolean_mask(boxes, mask)
+    classes = tf.ragged.boolean_mask(classes, mask)
+    if confidence is not None:
+        confidence = tf.ragged.boolean_mask(confidence, mask)
+
+    if isinstance(boxes, tf.Tensor):
+        boxes = tf.RaggedTensor.from_tensor(boxes)
+
+    if isinstance(classes, tf.Tensor) and len(classes.shape) > 1:
+        classes = tf.RaggedTensor.from_tensor(classes)
+
+    if confidence is not None:
+        if isinstance(confidence, tf.Tensor) and len(confidence.shape) > 1:
+            confidence = tf.RaggedTensor.from_tensor(confidence)
+
+    result = bounding_boxes.copy()
+    result["boxes"] = tf.cast(boxes, dtype)
+    result["classes"] = tf.cast(classes, dtype)
+
+    if confidence is not None:
+        result["confidence"] = tf.cast(confidence, dtype)
+
+    return result

keras_hub/src/bounding_box/utils.py

@@ -0,0 +1,194 @@
+# Copyright 2024 The KerasHub Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Utility functions for working with bounding boxes."""
+
+from keras import ops
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.bounding_box import converters
+from keras_hub.src.bounding_box.formats import XYWH
+
+
+@keras_hub_export("keras_hub.bounding_box.is_relative")
+def is_relative(bounding_box_format):
+    """A util to check if a bounding box format uses relative coordinates"""
+    if bounding_box_format.lower() not in converters.TO_XYXY_CONVERTERS:
+        raise ValueError(
+            "`is_relative()` received an unsupported format for the argument "
+            f"`bounding_box_format`. `bounding_box_format` should be one of "
+            f"{converters.TO_XYXY_CONVERTERS.keys()}. "
+            f"Got bounding_box_format={bounding_box_format}"
+        )
+
+    return bounding_box_format.startswith("rel")
+
+
+@keras_hub_export("keras_hub.bounding_box.as_relative")
+def as_relative(bounding_box_format):
+    """A util to get the relative equivalent of a provided bounding box format.
+
+    If the specified format is already a relative format,
+    it will be returned unchanged.
+    """
+
+    if not is_relative(bounding_box_format):
+        return "rel_" + bounding_box_format
+
+    return bounding_box_format
+
+
+def _relative_area(boxes, bounding_box_format):
+    boxes = converters.convert_format(
+        boxes,
+        source=bounding_box_format,
+        target="rel_xywh",
+    )
+    widths = boxes[..., XYWH.WIDTH]
+    heights = boxes[..., XYWH.HEIGHT]
+    # handle corner case where shear performs a full inversion.
+    return ops.where(
+        ops.logical_and(widths > 0, heights > 0), widths * heights, 0.0
+    )
+
+
+@keras_hub_export("keras_hub.bounding_box.clip_to_image")
+def clip_to_image(
+    bounding_boxes, bounding_box_format, images=None, image_shape=None
+):
+    """clips bounding boxes to image boundaries.
+
+    `clip_to_image()` clips bounding boxes that have coordinates out of bounds
+    of an image down to the boundaries of the image. This is done by converting
+    the bounding box to relative formats, then clipping them to the `[0, 1]`
+    range. Additionally, bounding boxes that end up with a zero area have their
+    class ID set to -1, indicating that there is no object present in them.
+
+    Args:
+        bounding_boxes: bounding box tensor to clip.
+        bounding_box_format: the KerasCV bounding box format the bounding boxes
+            are in.
+        images: list of images to clip the bounding boxes to.
+        image_shape: the shape of the images to clip the bounding boxes to.
+    """
+    boxes, classes = bounding_boxes["boxes"], bounding_boxes["classes"]
+
+    boxes = converters.convert_format(
+        boxes,
+        source=bounding_box_format,
+        target="rel_xyxy",
+        images=images,
+        image_shape=image_shape,
+    )
+    boxes, classes, images, squeeze = _format_inputs(boxes, classes, images)
+    x1, y1, x2, y2 = ops.split(boxes, 4, axis=-1)
+    clipped_bounding_boxes = ops.concatenate(
+        [
+            ops.clip(x1, 0, 1),
+            ops.clip(y1, 0, 1),
+            ops.clip(x2, 0, 1),
+            ops.clip(y2, 0, 1),
+        ],
+        axis=-1,
+    )
+    areas = _relative_area(
+        clipped_bounding_boxes, bounding_box_format="rel_xyxy"
+    )
+    clipped_bounding_boxes = converters.convert_format(
+        clipped_bounding_boxes,
+        source="rel_xyxy",
+        target=bounding_box_format,
+        images=images,
+        image_shape=image_shape,
+    )
+    clipped_bounding_boxes = ops.where(
+        ops.expand_dims(areas > 0.0, axis=-1), clipped_bounding_boxes, -1.0
+    )
+    classes = ops.where(areas > 0.0, classes, -1)
+    nan_indices = ops.any(ops.isnan(clipped_bounding_boxes), axis=-1)
+    classes = ops.where(nan_indices, -1, classes)
+
+    # TODO update dict and return
+    clipped_bounding_boxes, classes = _format_outputs(
+        clipped_bounding_boxes, classes, squeeze
+    )
+
+    bounding_boxes.update({"boxes": clipped_bounding_boxes, "classes": classes})
+
+    return bounding_boxes
+
+
+@keras_hub_export("keras_hub.bounding_box.clip_boxes")
+def clip_boxes(boxes, image_shape):
+    """Clip boxes to the boundaries of the image shape"""
+    if boxes.shape[-1] != 4:
+        raise ValueError(
+            "boxes.shape[-1] is {:d}, but must be 4.".format(boxes.shape[-1])
+        )
+
+    if isinstance(image_shape, list) or isinstance(image_shape, tuple):
+        height, width, _ = image_shape
+        max_length = ops.stack([height, width, height, width], axis=-1)
+    else:
+        image_shape = ops.cast(image_shape, dtype=boxes.dtype)
+        height = image_shape[0]
+        width = image_shape[1]
+        max_length = ops.stack([height, width, height, width], axis=-1)
+
+    clipped_boxes = ops.maximum(ops.minimum(boxes, max_length), 0.0)
+    return clipped_boxes
+
+
+def _format_inputs(boxes, classes, images):
+    boxes_rank = len(boxes.shape)
+    if boxes_rank > 3:
+        raise ValueError(
+            "Expected len(boxes.shape)=2, or len(boxes.shape)=3, got "
+            f"len(boxes.shape)={boxes_rank}"
+        )
+    boxes_includes_batch = boxes_rank == 3
+    # Determine if images needs an expand_dims() call
+    if images is not None:
+        images_rank = len(images.shape)
+        if images_rank > 4:
+            raise ValueError(
+                "Expected len(images.shape)=2, or len(images.shape)=3, got "
+                f"len(images.shape)={images_rank}"
+            )
+        images_include_batch = images_rank == 4
+        if boxes_includes_batch != images_include_batch:
+            raise ValueError(
+                "clip_to_image() expects both boxes and images to be batched, "
+                "or both boxes and images to be unbatched. Received "
+                f"len(boxes.shape)={boxes_rank}, "
+                f"len(images.shape)={images_rank}. Expected either "
+                "len(boxes.shape)=2 AND len(images.shape)=3, or "
+                "len(boxes.shape)=3 AND len(images.shape)=4."
+            )
+        if not images_include_batch:
+            images = ops.expand_dims(images, axis=0)
+
+    if not boxes_includes_batch:
+        return (
+            ops.expand_dims(boxes, axis=0),
+            ops.expand_dims(classes, axis=0),
+            images,
+            True,
+        )
+    return boxes, classes, images, False
+
+
+def _format_outputs(boxes, classes, squeeze):
+    if squeeze:
+        return ops.squeeze(boxes, axis=0), ops.squeeze(classes, axis=0)
+    return boxes, classes