Internal change

PiperOrigin-RevId: 471911078

official/vision/ops/box_ops.py

@@ -783,3 +783,66 @@ def box_matching(boxes, gt_boxes, gt_classes):
 
   return (matched_gt_boxes, matched_gt_classes, matched_gt_indices,
           matched_iou, iou)
+
+
+def bbox2mask(bbox: tf.Tensor,
+              *,
+              image_height: int,
+              image_width: int,
+              dtype: tf.DType = tf.bool) -> tf.Tensor:
+  """Converts bounding boxes to bitmasks.
+
+  Args:
+    bbox: A tensor in shape (..., 4) with arbitrary numbers of batch dimensions,
+      representing the absolute coordinates (ymin, xmin, ymax, xmax) for each
+      bounding box.
+    image_height: an integer representing the height of the image.
+    image_width: an integer representing the width of the image.
+    dtype: DType of the output bitmasks.
+
+  Returns:
+    A tensor in shape (..., height, width) which stores the bitmasks created
+    from the bounding boxes. For example:
+
+    >>> bbox2mask(tf.constant([[1,2,4,4]]),
+                  image_height=5,
+                  image_width=5,
+                  dtype=tf.int32)
+    <tf.Tensor: shape=(1, 5, 5), dtype=int32, numpy=
+    array([[[0, 0, 0, 0, 0],
+            [0, 0, 1, 1, 0],
+            [0, 0, 1, 1, 0],
+            [0, 0, 1, 1, 0],
+            [0, 0, 0, 0, 0]]], dtype=int32)>
+  """
+  bbox_shape = bbox.get_shape().as_list()
+  if bbox_shape[-1] != 4:
+    raise ValueError(
+        'Expected the last dimension of `bbox` has size == 4, but the shape '
+        'of `bbox` was: %s' % bbox_shape)
+
+  # (..., 1)
+  ymin = bbox[..., 0:1]
+  xmin = bbox[..., 1:2]
+  ymax = bbox[..., 2:3]
+  xmax = bbox[..., 3:4]
+  # (..., 1, width)
+  ymin = tf.expand_dims(tf.repeat(ymin, repeats=image_width, axis=-1), axis=-2)
+  # (..., height, 1)
+  xmin = tf.expand_dims(tf.repeat(xmin, repeats=image_height, axis=-1), axis=-1)
+  # (..., 1, width)
+  ymax = tf.expand_dims(tf.repeat(ymax, repeats=image_width, axis=-1), axis=-2)
+  # (..., height, 1)
+  xmax = tf.expand_dims(tf.repeat(xmax, repeats=image_height, axis=-1), axis=-1)
+
+  # (height, 1)
+  y_grid = tf.expand_dims(tf.range(image_height, dtype=bbox.dtype), axis=-1)
+  # (1, width)
+  x_grid = tf.expand_dims(tf.range(image_width, dtype=bbox.dtype), axis=-2)
+
+  # (..., height, width)
+  ymin_mask = y_grid >= ymin
+  xmin_mask = x_grid >= xmin
+  ymax_mask = y_grid < ymax
+  xmax_mask = x_grid < xmax
+  return tf.cast(ymin_mask & xmin_mask & ymax_mask & xmax_mask, dtype)

official/vision/ops/mask_ops.py

@@ -18,7 +18,6 @@ import math
 # Import libraries
 import cv2
 import numpy as np
-import tensorflow as tf
 
 
 def paste_instance_masks(masks: np.ndarray, detected_boxes: np.ndarray,
@@ -184,66 +183,3 @@ def paste_instance_masks_v2(masks: np.ndarray, detected_boxes: np.ndarray,
 
   segms = np.array(segms)
   return segms
-
-
-def bbox2mask(bbox: tf.Tensor,
-              *,
-              image_height: int,
-              image_width: int,
-              dtype: tf.DType = tf.bool) -> tf.Tensor:
-  """Converts bounding boxes to bitmasks.
-
-  Args:
-    bbox: A tensor in shape (..., 4) with arbitrary numbers of batch dimensions,
-      representing the absolute coordinates (ymin, xmin, ymax, xmax) for each
-      bounding box.
-    image_height: an integer representing the height of the image.
-    image_width: an integer representing the width of the image.
-    dtype: DType of the output bitmasks.
-
-  Returns:
-    A tensor in shape (..., height, width) which stores the bitmasks created
-    from the bounding boxes. For example:
-
-    >>> bbox2mask(tf.constant([[1,2,4,4]]),
-                  image_height=5,
-                  image_width=5,
-                  dtype=tf.int32)
-    <tf.Tensor: shape=(1, 5, 5), dtype=int32, numpy=
-    array([[[0, 0, 0, 0, 0],
-            [0, 0, 1, 1, 0],
-            [0, 0, 1, 1, 0],
-            [0, 0, 1, 1, 0],
-            [0, 0, 0, 0, 0]]], dtype=int32)>
-  """
-  bbox_shape = bbox.get_shape().as_list()
-  if bbox_shape[-1] != 4:
-    raise ValueError(
-        'Expected the last dimension of `bbox` has size == 4, but the shape '
-        'of `bbox` was: %s' % bbox_shape)
-
-  # (..., 1)
-  ymin = bbox[..., 0:1]
-  xmin = bbox[..., 1:2]
-  ymax = bbox[..., 2:3]
-  xmax = bbox[..., 3:4]
-  # (..., 1, width)
-  ymin = tf.expand_dims(tf.repeat(ymin, repeats=image_width, axis=-1), axis=-2)
-  # (..., height, 1)
-  xmin = tf.expand_dims(tf.repeat(xmin, repeats=image_height, axis=-1), axis=-1)
-  # (..., 1, width)
-  ymax = tf.expand_dims(tf.repeat(ymax, repeats=image_width, axis=-1), axis=-2)
-  # (..., height, 1)
-  xmax = tf.expand_dims(tf.repeat(xmax, repeats=image_height, axis=-1), axis=-1)
-
-  # (height, 1)
-  y_grid = tf.expand_dims(tf.range(image_height, dtype=bbox.dtype), axis=-1)
-  # (1, width)
-  x_grid = tf.expand_dims(tf.range(image_width, dtype=bbox.dtype), axis=-2)
-
-  # (..., height, width)
-  ymin_mask = y_grid >= ymin
-  xmin_mask = x_grid >= xmin
-  ymax_mask = y_grid < ymax
-  xmax_mask = x_grid < xmax
-  return tf.cast(ymin_mask & xmin_mask & ymax_mask & xmax_mask, dtype)

official/vision/ops/mask_ops_test.py

@@ -49,57 +49,6 @@ class MaskUtilsTest(tf.test.TestCase):
         np.array(masks > 0.5, dtype=np.uint8),
         1e-5)
 
-  def testBbox2mask(self):
-    bboxes = tf.constant([[1, 2, 4, 4], [-1, -1, 3, 3], [2, 3, 6, 8],
-                          [1, 1, 2, 2], [1, 1, 1, 4]])
-    masks = mask_ops.bbox2mask(
-        bboxes, image_height=5, image_width=6, dtype=tf.int32)
-    expected_masks = tf.constant(
-        [
-            [  # bbox = [1, 2, 4, 4]
-                [0, 0, 0, 0, 0, 0],
-                [0, 0, 1, 1, 0, 0],
-                [0, 0, 1, 1, 0, 0],
-                [0, 0, 1, 1, 0, 0],
-                [0, 0, 0, 0, 0, 0],
-            ],
-            [  # bbox = [-1, -1, 3, 3]
-                [1, 1, 1, 0, 0, 0],
-                [1, 1, 1, 0, 0, 0],
-                [1, 1, 1, 0, 0, 0],
-                [0, 0, 0, 0, 0, 0],
-                [0, 0, 0, 0, 0, 0],
-            ],
-            [  # bbox = [2, 3, 6, 8]
-                [0, 0, 0, 0, 0, 0],
-                [0, 0, 0, 0, 0, 0],
-                [0, 0, 0, 1, 1, 1],
-                [0, 0, 0, 1, 1, 1],
-                [0, 0, 0, 1, 1, 1],
-            ],
-            [  # bbox =  [1, 1, 2, 2]
-                [0, 0, 0, 0, 0, 0],
-                [0, 1, 0, 0, 0, 0],
-                [0, 0, 0, 0, 0, 0],
-                [0, 0, 0, 0, 0, 0],
-                [0, 0, 0, 0, 0, 0],
-            ],
-            [  # bbox = [1, 1, 1, 4]
-                [0, 0, 0, 0, 0, 0],
-                [0, 0, 0, 0, 0, 0],
-                [0, 0, 0, 0, 0, 0],
-                [0, 0, 0, 0, 0, 0],
-                [0, 0, 0, 0, 0, 0],
-            ]
-        ],
-        dtype=tf.int32)
-    self.assertAllEqual(expected_masks, masks)
-
-  def testBbox2maskInvalidInput(self):
-    bboxes = tf.constant([[1, 2, 4, 4, 4], [-1, -1, 3, 3, 3]])
-    with self.assertRaisesRegex(ValueError, 'bbox.*size == 4'):
-      mask_ops.bbox2mask(bboxes, image_height=5, image_width=6, dtype=tf.int32)
-
 
 if __name__ == '__main__':
   tf.test.main()