Add IOU heatmap loss for CenterNet.

PiperOrigin-RevId: 421140005

Author: Vighnesh Birodkar

research/object_detection/core/target_assigner.py

@@ -925,7 +925,9 @@ def __init__(self,
                compute_heatmap_sparse=False,
                keypoint_class_id=None,
                keypoint_indices=None,
-               keypoint_weights_for_center=None):
+               keypoint_weights_for_center=None,
+               box_heatmap_type='adaptive_gaussian',
+               heatmap_exponent=1.0):
     """Initializes the target assigner.
 
     Args:
@@ -947,6 +949,17 @@ def __init__(self,
         the number of keypoints. The object center is calculated by the weighted
         mean of the keypoint locations. If not provided, the object center is
         determined by the center of the bounding box (default behavior).
+       box_heatmap_type: str, the algorithm used to compute the box heatmap,
+         used when calling the assign_center_targets_from_boxes method.
+         Options are:
+         'adaptaive_gaussian': A box-size adaptive Gaussian from the original
+           paper[1].
+         'iou': IOU based heatmap target where each point is assigned an IOU
+           based on its location, assuming that it produced a box centered at
+           that point with the correct size.
+       heatmap_exponent: float, The generated heatmap is exponentiated with
+         this number. A number > 1 will result in the heatmap being more peaky
+         and a number < 1 will cause the heatmap to be more spreadout.
     """
 
     self._stride = stride
@@ -955,6 +968,8 @@ def __init__(self,
     self._keypoint_class_id = keypoint_class_id
     self._keypoint_indices = keypoint_indices
     self._keypoint_weights_for_center = keypoint_weights_for_center
+    self._box_heatmap_type = box_heatmap_type
+    self._heatmap_exponent = heatmap_exponent
 
   def assign_center_targets_from_boxes(self,
                                        height,
@@ -1018,19 +1033,29 @@ def assign_center_targets_from_boxes(self,
                                              self._min_overlap)
       # Apply the Gaussian kernel to the center coordinates. Returned heatmap
       # has shape of [out_height, out_width, num_classes]
-      heatmap = ta_utils.coordinates_to_heatmap(
-          y_grid=y_grid,
-          x_grid=x_grid,
-          y_coordinates=y_center,
-          x_coordinates=x_center,
-          sigma=sigma,
-          channel_onehot=class_targets,
-          channel_weights=weights,
-          sparse=self._compute_heatmap_sparse)
+
+      if self._box_heatmap_type == 'adaptive_gaussian':
+        heatmap = ta_utils.coordinates_to_heatmap(
+            y_grid=y_grid,
+            x_grid=x_grid,
+            y_coordinates=y_center,
+            x_coordinates=x_center,
+            sigma=sigma,
+            channel_onehot=class_targets,
+            channel_weights=weights,
+            sparse=self._compute_heatmap_sparse)
+      elif self._box_heatmap_type == 'iou':
+        heatmap = ta_utils.coordinates_to_iou(y_grid, x_grid, boxes,
+                                              class_targets, weights)
+      else:
+        raise ValueError(f'Unknown heatmap type - {self._box_heatmap_type}')
+
       heatmaps.append(heatmap)
 
     # Return the stacked heatmaps over the batch.
-    return tf.stack(heatmaps, axis=0)
+    stacked_heatmaps = tf.stack(heatmaps, axis=0)
+    return (tf.pow(stacked_heatmaps, self._heatmap_exponent) if
+            self._heatmap_exponent != 1.0 else stacked_heatmaps)
 
   def assign_center_targets_from_keypoints(self,
                                            height,

research/object_detection/core/target_assigner_test.py

@@ -1678,6 +1678,66 @@ def graph_fn():
     np.testing.assert_array_equal(preds, [[1, 2], [3, 4], [5, 6], [7, 8]])
 
 
+class CenterNetIOUTargetAssignerTest(test_case.TestCase):
+
+  def setUp(self):
+    super(CenterNetIOUTargetAssignerTest, self).setUp()
+
+    self._box_center = [0.0, 0.0, 1.0, 1.0]
+    self._box_center_small = [0.25, 0.25, 0.75, 0.75]
+    self._box_lower_left = [0.5, 0.0, 1.0, 0.5]
+    self._box_center_offset = [0.1, 0.05, 1.0, 1.0]
+    self._box_odd_coordinates = [0.1625, 0.2125, 0.5625, 0.9625]
+
+  def test_center_location(self):
+    """Test that the centers are at the correct location."""
+    def graph_fn():
+      box_batch = [tf.constant([self._box_center, self._box_lower_left]),
+                   tf.constant([self._box_lower_left, self._box_center])]
+      classes = [
+          tf.one_hot([0, 1], depth=4),
+          tf.one_hot([2, 2], depth=4)
+      ]
+      assigner = targetassigner.CenterNetCenterHeatmapTargetAssigner(
+          4, box_heatmap_type='iou')
+      targets = assigner.assign_center_targets_from_boxes(
+          80, 80, box_batch, classes)
+      return targets
+    targets = self.execute(graph_fn, [])
+    self.assertEqual((10, 10), _array_argmax(targets[0, :, :, 0]))
+    self.assertAlmostEqual(1.0, targets[0, 10, 10, 0])
+    self.assertEqual((15, 5), _array_argmax(targets[0, :, :, 1]))
+    self.assertAlmostEqual(1.0, targets[0, 15, 5, 1])
+
+    self.assertAlmostEqual(1.0, targets[1, 15, 5, 2])
+    self.assertAlmostEqual(1.0, targets[1, 10, 10, 2])
+    self.assertAlmostEqual(0.0, targets[1, 0, 19, 1])
+
+  def test_exponent(self):
+    """Test that the centers are at the correct location."""
+    def graph_fn():
+      box_batch = [tf.constant([self._box_center, self._box_lower_left]),
+                   tf.constant([self._box_lower_left, self._box_center])]
+      classes = [
+          tf.one_hot([0], depth=2),
+      ]
+      assigner = targetassigner.CenterNetCenterHeatmapTargetAssigner(
+          1, box_heatmap_type='iou')
+      targets = assigner.assign_center_targets_from_boxes(
+          4, 4, box_batch, classes)
+
+      assigner = targetassigner.CenterNetCenterHeatmapTargetAssigner(
+          1, box_heatmap_type='iou', heatmap_exponent=0.5)
+      targets_pow = assigner.assign_center_targets_from_boxes(
+          4, 4, box_batch, classes)
+      return targets, targets_pow
+
+    targets, targets_pow = self.execute(graph_fn, [])
+    self.assertLess(targets[0, 2, 3, 0], 1.0)
+    self.assertLess(targets_pow[0, 2, 3, 0], 1.0)
+    self.assertAlmostEqual(targets[0, 2, 3, 0], targets_pow[0, 2, 3, 0] ** 2)
+
+
 class CenterNetKeypointTargetAssignerTest(test_case.TestCase):
 
   def test_keypoint_heatmap_targets(self):

research/object_detection/utils/target_assigner_utils.py

@@ -236,6 +236,77 @@ def compute_floor_offsets_with_indices(y_source,
   return offsets, indices
 
 
+def coordinates_to_iou(y_grid, x_grid, blist,
+                       channels_onehot, weights=None):
+  """Computes a per-pixel IoU with groundtruth boxes.
+
+  At each pixel, we return the IoU assuming that we predicted the
+  ideal height and width for the box at that location.
+
+  Args:
+   y_grid: A 2D tensor with shape [height, width] which contains the grid
+      y-coordinates given in the (output) image dimensions.
+    x_grid: A 2D tensor with shape [height, width] which contains the grid
+      x-coordinates given in the (output) image dimensions.
+    blist: A BoxList object with `num_instances` number of boxes.
+    channels_onehot: A 2D tensor with shape [num_instances, num_channels]
+      representing the one-hot encoded channel labels for each point.
+    weights: A 1D tensor with shape [num_instances] corresponding to the
+      weight of each instance.
+
+  Returns:
+    iou_heatmap: A [height, width, num_channels] shapes float tensor denoting
+      the IoU based heatmap.
+  """
+
+  image_height, image_width = tf.shape(y_grid)[0], tf.shape(y_grid)[1]
+  num_pixels = image_height * image_width
+  _, _, height, width = blist.get_center_coordinates_and_sizes()
+  num_boxes = tf.shape(height)[0]
+
+  per_pixel_ymin = (y_grid[tf.newaxis, :, :] -
+                    (height[:, tf.newaxis, tf.newaxis] / 2.0))
+  per_pixel_xmin = (x_grid[tf.newaxis, :, :] -
+                    (width[:, tf.newaxis, tf.newaxis] / 2.0))
+  per_pixel_ymax = (y_grid[tf.newaxis, :, :] +
+                    (height[:, tf.newaxis, tf.newaxis] / 2.0))
+  per_pixel_xmax = (x_grid[tf.newaxis, :, :] +
+                    (width[:, tf.newaxis, tf.newaxis] / 2.0))
+
+  # [num_boxes, height, width] -> [num_boxes * height * width]
+  per_pixel_ymin = tf.reshape(
+      per_pixel_ymin, [num_pixels * num_boxes])
+  per_pixel_xmin = tf.reshape(
+      per_pixel_xmin, [num_pixels * num_boxes])
+  per_pixel_ymax = tf.reshape(
+      per_pixel_ymax, [num_pixels * num_boxes])
+  per_pixel_xmax = tf.reshape(
+      per_pixel_xmax, [num_pixels * num_boxes])
+  per_pixel_blist = box_list.BoxList(
+      tf.stack([per_pixel_ymin, per_pixel_xmin,
+                per_pixel_ymax, per_pixel_xmax], axis=1))
+
+  target_boxes = tf.tile(
+      blist.get()[:, tf.newaxis, :], [1, num_pixels, 1])
+  # [num_boxes, height * width, 4] -> [num_boxes * height * wdith, 4]
+  target_boxes = tf.reshape(target_boxes,
+                            [num_pixels * num_boxes, 4])
+  target_blist = box_list.BoxList(target_boxes)
+
+  ious = box_list_ops.matched_iou(target_blist, per_pixel_blist)
+  ious = tf.reshape(ious, [num_boxes, image_height, image_width])
+  per_class_iou = (
+      ious[:, :, :, tf.newaxis] *
+      channels_onehot[:, tf.newaxis, tf.newaxis, :])
+
+  if weights is not None:
+    per_class_iou = (
+        per_class_iou * weights[:, tf.newaxis, tf.newaxis, tf.newaxis])
+
+  per_class_iou = tf.maximum(per_class_iou, 0.0)
+  return tf.reduce_max(per_class_iou, axis=0)
+
+
 def get_valid_keypoint_mask_for_class(keypoint_coordinates,
                                       class_id,
                                       class_onehot,

research/object_detection/utils/target_assigner_utils_test.py

@@ -18,6 +18,7 @@
 import numpy as np
 import tensorflow.compat.v1 as tf
 
+from object_detection.core import box_list
 from object_detection.utils import target_assigner_utils as ta_utils
 from object_detection.utils import test_case
 
@@ -265,6 +266,31 @@ def graph_fn():
         np.array([[0.0, 3.0, 4.0, 0.0, 4.0]]))
     self.assertAllEqual(valid, [[False, True, True, False, True]])
 
+  def test_coordinates_to_iou(self):
+
+    def graph_fn():
+      y, x = tf.meshgrid(tf.range(32, dtype=tf.float32),
+                         tf.range(32, dtype=tf.float32))
+      blist = box_list.BoxList(
+          tf.constant([[0., 0., 32., 32.],
+                       [0., 0., 16., 16.],
+                       [0.0, 0.0, 4.0, 4.0]]))
+      classes = tf.constant([[0., 1., 0.],
+                             [1., 0., 0.],
+                             [0., 0., 1.]])
+
+      result = ta_utils.coordinates_to_iou(
+          y, x, blist, classes)
+      return result
+
+    result = self.execute(graph_fn, [])
+    self.assertEqual(result.shape, (32, 32, 3))
+    self.assertAlmostEqual(result[0, 0, 0], 1.0 / 7.0)
+    self.assertAlmostEqual(result[0, 0, 1], 1.0 / 7.0)
+    self.assertAlmostEqual(result[0, 16, 0], 1.0 / 7.0)
+    self.assertAlmostEqual(result[2, 2, 2], 1.0)
+    self.assertAlmostEqual(result[8, 8, 2], 0.0)
+
 
 if __name__ == '__main__':
   tf.test.main()