Update the bounding box utility to be more readable (keras-team#51)

* Update the bounding box utility to be more readable

1. Added docstring header to contain more details about whats the formats we usually use for bounding box, and what's the math conversion between them.

2. Slight change the implementation of conversion function to be more readable. The current index slice based approach is very hard to understand.

3. Added unit test for bbox util.

* Fix the format issue.

* Fix more format issue

Author: (Ian Stenbit)

keras_cv/utils/bbox.py

@@ -12,42 +12,86 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Shared utility functions for working with bounding boxes."""
+"""Shared utility functions for working with bounding boxes.
+
+Usually bounding boxes is a 2D Tensor with shape [batch, 4]. The second dimension
+will contain 4 numbers based on 2 different formats:
+
+1. LEFT, TOP, RIGHT, BOTTOM, where LEFT, TOP represent the top-left corner
+   coordinates, and RIGHT, BOTTOM represent the bottom-right corner coordinates.
+2. X, Y, WIDTH, HEIGHT, where X and Y are the coordinates for the center of the box.
+
+Math wise:
+LEFT = X - WIDTH / 2
+TOP = Y - HEIGHT / 2
+RIGHT = X + WIDTH / 2
+BOTTOM = Y + HEIGHT / 2
+
+X = (LEFT + RIGHT) / 2
+Y = (TOP + BOTTOM) / 2
+WIDTH = RIGHT - LEFT
+HEIGHT = BOTTOM - TOP
+
+Note that these two formats are both commonly used. Corners format are mostly used
+for IOU computation, whereas XYWH are easy for bounding box generation with different
+center and width/height ratio.
+"""
 
 import tensorflow as tf
 
-# These are the dimensions used in Tensors to represent each corresponding side.
+# These are the indexes used in Tensors to represent each corresponding side.
 LEFT, TOP, RIGHT, BOTTOM = 0, 1, 2, 3
 
-# These are the dimensions that you can use for bboxes in corners format.
+# These are the indexes that you can use for bounding box in XYWH format.
 X, Y, WIDTH, HEIGHT = 0, 1, 2, 3
 
 # Regardless of format these constants are consistent.
-# Class is held in the 4th index
+# Class is held in the 5th index
 CLASS = 4
-# Confidence exists only on y_pred, and is in the 5th index.
+# Confidence exists only on y_pred, and is in the 6th index.
 CONFIDENCE = 5
 
 
-def convert_corners_to_xywh(bboxes):
-    """Converts bboxes in corners format to xywh format."""
+def corners_to_xywh(bboxes):
+    """Converts bboxes in corners format to XYWH format.
+
+    Args:
+        bboxes: a Tensor which has at least 2D rank, with shape [..., 4]
+
+    Returns:
+        converted bboxes with same shape, but in XYWH format.
+    """
+    left, top, right, bottom, rest = tf.split(bboxes, [1, 1, 1, 1, -1], axis=-1)
     return tf.concat(
         [
-            (bboxes[..., :2] + bboxes[..., 2:4]) / 2.0,
-            bboxes[..., 2:4] - bboxes[..., :2],
-            bboxes[..., 4:],
+            # We use ... here in case user has higher rank of inputs.
+            (left + right) / 2.0,  # X
+            (top + bottom) / 2.0,  # Y
+            right - left,  # WIDTH
+            bottom - top,  # HEIGHT
+            rest,  # In case there is any more index after the BOTTOM.
         ],
         axis=-1,
     )
 
 
 def xywh_to_corners(bboxes):
-    """Converts bboxes in xywh format to corners format."""
+    """Converts bboxes in XYWH format to corners format.
+
+    Args:
+        bboxes: a Tensor which has at least 2D rank, with shape [..., 4]
+
+    Returns:
+        converted bboxes with same shape, but in corners format.
+    """
+    x, y, width, height, rest = tf.split(bboxes, [1, 1, 1, 1, -1], axis=-1)
     return tf.concat(
         [
-            bboxes[..., :2] - bboxes[..., 2:4] / 2.0,
-            bboxes[..., :2] + bboxes[..., 2:4] / 2.0,
-            bboxes[..., 4:],
+            x - width / 2.0,
+            y - height / 2.0,
+            x + width / 2.0,
+            y + height / 2.0,
+            rest,  # In case there is any more index after the HEIGHT.
         ],
         axis=-1,
     )

keras_cv/utils/bbox_test.py

@@ -0,0 +1,76 @@
+# Copyright 2022 The KerasCV 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 tensorflow as tf
+
+from keras_cv.utils import bbox
+
+
+class BBOXTestCase(tf.test.TestCase):
+    def setUp(self):
+        super().setUp()
+        self.corner_bbox = tf.constant(
+            [[10, 10, 110, 110], [20, 20, 120, 120]], dtype=tf.float32
+        )
+        self.xywh_bbox = tf.constant(
+            [[60, 60, 100, 100], [70, 70, 100, 100]], dtype=tf.float32
+        )
+
+    def test_corner_to_xywh(self):
+        self.assertAllClose(bbox.corners_to_xywh(self.corner_bbox), self.xywh_bbox)
+
+        # Make sure it also accept higher rank than 2
+        corner_bbox_3d = tf.expand_dims(self.corner_bbox, 0)
+        xywh_bbox_3d = tf.expand_dims(self.xywh_bbox, 0)
+        self.assertAllClose(bbox.corners_to_xywh(corner_bbox_3d), xywh_bbox_3d)
+
+        # Make sure it also accept more value after last index.
+        padded_corner_bbox = tf.pad(
+            self.corner_bbox, [[0, 0], [0, 2]]
+        )  # Right pad 2 more value
+        padded_xywh_bbox = tf.pad(self.xywh_bbox, [[0, 0], [0, 2]])
+        self.assertAllClose(
+            bbox.corners_to_xywh(padded_corner_bbox), padded_xywh_bbox
+        )
+
+        # Same for higher rank
+        padded_corner_bbox_3d = tf.expand_dims(padded_corner_bbox, 0)
+        padded_xywh_bbox_3d = tf.expand_dims(padded_xywh_bbox, 0)
+        self.assertAllClose(
+            bbox.corners_to_xywh(padded_corner_bbox_3d), padded_xywh_bbox_3d
+        )
+
+    def test_xywh_to_corner(self):
+        self.assertAllClose(bbox.xywh_to_corners(self.xywh_bbox), self.corner_bbox)
+
+        # Make sure it also accept higher rank than 2
+        corner_bbox_3d = tf.expand_dims(self.corner_bbox, 0)
+        xywh_bbox_3d = tf.expand_dims(self.xywh_bbox, 0)
+        self.assertAllClose(bbox.xywh_to_corners(xywh_bbox_3d), corner_bbox_3d)
+
+        # Make sure it also accept more value after last index.
+        padded_corner_bbox = tf.pad(
+            self.corner_bbox, [[0, 0], [0, 2]]
+        )  # Right pad 2 more value
+        padded_xywh_bbox = tf.pad(self.xywh_bbox, [[0, 0], [0, 2]])
+        self.assertAllClose(
+            bbox.xywh_to_corners(padded_xywh_bbox), padded_corner_bbox
+        )
+
+        # Same for higher rank
+        padded_corner_bbox_3d = tf.expand_dims(padded_corner_bbox, 0)
+        padded_xywh_bbox_3d = tf.expand_dims(padded_xywh_bbox, 0)
+        self.assertAllClose(
+            bbox.xywh_to_corners(padded_xywh_bbox_3d), padded_corner_bbox_3d
+        )