Updating target assigner utils to allow for weights applied in box regions.

PiperOrigin-RevId: 377208217

research/object_detection/core/target_assigner_test.py

@@ -1734,9 +1734,9 @@ class CenterNetKeypointTargetAssignerTest(test_case.TestCase):
     # the region of the 1st instance boxing box should be blacked out
     # (0.0, 0.0, 0.5, 0.5), transfering to (0, 0, 15, 10) in absolute output
     # space.
-    self.assertAlmostEqual(np.sum(valid_mask[:, 0:16, 0:11]), 0.0)
-    # All other values are 1.0 so the sum is: 30 * 20 - 16 * 11 = 424.
-    self.assertAlmostEqual(np.sum(valid_mask), 424.0)
+    self.assertAlmostEqual(np.sum(valid_mask[:, 0:15, 0:10]), 0.0)
+    # All other values are 1.0 so the sum is: 30 * 20 - 15 * 10 = 450.
+    self.assertAlmostEqual(np.sum(valid_mask), 450.0)
 
   def test_assign_keypoints_offset_targets(self):
     def graph_fn():

research/object_detection/utils/target_assigner_utils.py

@@ -289,12 +289,38 @@ def get_valid_keypoint_mask_for_class(keypoint_coordinates,
   return mask, keypoints_nan_to_zeros
 
 
-def blackout_pixel_weights_by_box_regions(height, width, boxes, blackout):
-  """Blackout the pixel weights in the target box regions.
+def blackout_pixel_weights_by_box_regions(height, width, boxes, blackout,
+                                          weights=None):
+  """Apply weights at pixel locations.
 
   This function is used to generate the pixel weight mask (usually in the output
   image dimension). The mask is to ignore some regions when computing loss.
 
+  Weights are applied as follows:
+  - Any region outside of a box gets the default weight 1.0
+  - Any box for which an explicit weight is specifed gets that weight. If
+    multiple boxes overlap, the maximum of the weights is applied.
+  - Any box for which blackout=True is specified will get a weight of 0.0,
+    regardless of whether an equivalent non-zero weight is specified. Also, the
+    blackout region takes precedence over other boxes which may overlap with
+    non-zero weight.
+
+    Example:
+    height = 4
+    width = 4
+    boxes = [[0., 0., 2., 2.],
+             [0., 0., 4., 2.],
+             [3., 0., 4., 4.]]
+    blackout = [False, False, True]
+    weights = [4.0, 3.0, 2.0]
+    blackout_pixel_weights_by_box_regions(height, width, boxes, blackout,
+                                          weights)
+    >> [[4.0, 4.0, 1.0, 1.0],
+        [4.0, 4.0, 1.0, 1.0],
+        [3.0, 3.0, 1.0, 1.0],
+        [0.0, 0.0, 0.0, 0.0]]
+
+
   Args:
     height: int, height of the (output) image.
     width: int, width of the (output) image.
@@ -302,10 +328,15 @@ def blackout_pixel_weights_by_box_regions(height, width, boxes, blackout):
       coordinates of the four corners of the boxes.
     blackout: A boolean tensor with shape [num_instances] indicating whether to
       blackout (zero-out) the weights within the box regions.
+    weights: An optional float32 tensor with shape [num_instances] indicating
+      a value to apply in each box region. Note that if blackout=True for a
+      given box, the weight will be zero. If None, all weights are assumed to be
+      1.
 
   Returns:
     A float tensor with shape [height, width] where all values within the
-    regions of the blackout boxes are 0.0 and 1.0 else where.
+    regions of the blackout boxes are 0.0 and 1.0 (or weights if supplied)
+    elsewhere.
   """
   num_instances, _ = shape_utils.combined_static_and_dynamic_shape(boxes)
   # If no annotation instance is provided, return all ones (instead of
@@ -323,22 +354,36 @@ def blackout_pixel_weights_by_box_regions(height, width, boxes, blackout):
 
   # Make the mask with all 1.0 in the box regions.
   # Shape: [num_instances, height, width]
-  in_boxes = tf.cast(
-      tf.logical_and(
-          tf.logical_and(y_grid >= y_min, y_grid <= y_max),
-          tf.logical_and(x_grid >= x_min, x_grid <= x_max)),
-      dtype=tf.float32)
-
-  # Shape: [num_instances, height, width]
-  blackout = tf.tile(
-      tf.expand_dims(tf.expand_dims(blackout, axis=-1), axis=-1),
-      [1, height, width])
-
-  # Select only the boxes specified by blackout.
-  selected_in_boxes = tf.where(blackout, in_boxes, tf.zeros_like(in_boxes))
-  out_boxes = tf.reduce_max(selected_in_boxes, axis=0)
-  out_boxes = tf.ones_like(out_boxes) - out_boxes
-  return out_boxes
+  in_boxes = tf.math.logical_and(
+      tf.math.logical_and(y_grid >= y_min, y_grid < y_max),
+      tf.math.logical_and(x_grid >= x_min, x_grid < x_max))
+
+  if weights is None:
+    weights = tf.ones_like(blackout, dtype=tf.float32)
+
+  # Compute a [height, width] tensor with the maximum weight in each box, and
+  # 0.0 elsewhere.
+  weights_tiled = tf.tile(
+      weights[:, tf.newaxis, tf.newaxis], [1, height, width])
+  weights_3d = tf.where(in_boxes, weights_tiled,
+                        tf.zeros_like(weights_tiled))
+  weights_2d = tf.math.maximum(
+      tf.math.reduce_max(weights_3d, axis=0), 0.0)
+
+  # Add 1.0 to all regions outside a box.
+  weights_2d = tf.where(
+      tf.math.reduce_any(in_boxes, axis=0),
+      weights_2d,
+      tf.ones_like(weights_2d))
+
+  # Now enforce that blackout regions all have zero weights.
+  keep_region = tf.cast(tf.math.logical_not(blackout), tf.float32)
+  keep_region_tiled = tf.tile(
+      keep_region[:, tf.newaxis, tf.newaxis], [1, height, width])
+  keep_region_3d = tf.where(in_boxes, keep_region_tiled,
+                            tf.ones_like(keep_region_tiled))
+  keep_region_2d = tf.math.reduce_min(keep_region_3d, axis=0)
+  return weights_2d * keep_region_2d
 
 
 def _get_yx_indices_offset_by_radius(radius):

research/object_detection/utils/target_assigner_utils_test.py

@@ -196,13 +196,36 @@ class TargetUtilTest(parameterized.TestCase, test_case.TestCase):
       return output
 
     output = self.execute(graph_fn, [])
-    # All zeros in region [0:6, 0:6].
-    self.assertAlmostEqual(np.sum(output[0:6, 0:6]), 0.0)
-    # All zeros in region [12:19, 6:9].
-    self.assertAlmostEqual(np.sum(output[6:9, 12:19]), 0.0)
+    # All zeros in region [0:5, 0:5].
+    self.assertAlmostEqual(np.sum(output[0:5, 0:5]), 0.0)
+    # All zeros in region [12:18, 6:8].
+    self.assertAlmostEqual(np.sum(output[6:8, 12:18]), 0.0)
     # All other pixel weights should be 1.0.
-    # 20 * 10 - 6 * 6 - 3 * 7 = 143.0
-    self.assertAlmostEqual(np.sum(output), 143.0)
+    # 20 * 10 - 5 * 5 - 2 * 6 = 163.0
+    self.assertAlmostEqual(np.sum(output), 163.0)
+
+  def test_blackout_pixel_weights_by_box_regions_with_weights(self):
+    def graph_fn():
+      boxes = tf.constant(
+          [[0.0, 0.0, 2.0, 2.0],
+           [0.0, 0.0, 4.0, 2.0],
+           [3.0, 0.0, 4.0, 4.0]],
+          dtype=tf.float32)
+      blackout = tf.constant([False, False, True], dtype=tf.bool)
+      weights = tf.constant([0.4, 0.3, 0.2], tf.float32)
+      blackout_pixel_weights_by_box_regions = tf.function(
+          ta_utils.blackout_pixel_weights_by_box_regions)
+      output = blackout_pixel_weights_by_box_regions(
+          4, 4, boxes, blackout, weights)
+      return output
+
+    output = self.execute(graph_fn, [])
+    expected_weights = [
+        [0.4, 0.4, 1.0, 1.0],
+        [0.4, 0.4, 1.0, 1.0],
+        [0.3, 0.3, 1.0, 1.0],
+        [0.0, 0.0, 0.0, 0.0]]
+    np.testing.assert_array_almost_equal(expected_weights, output)
 
   def test_blackout_pixel_weights_by_box_regions_zero_instance(self):
     def graph_fn():