Internal change

PiperOrigin-RevId: 332331988

official/vision/beta/ops/anchor.py

@@ -18,11 +18,9 @@ import collections
 # Import libraries
 import tensorflow as tf
 from official.vision import keras_cv
-from official.vision.detection.utils.object_detection import argmax_matcher
 from official.vision.detection.utils.object_detection import balanced_positive_negative_sampler
 from official.vision.detection.utils.object_detection import box_list
 from official.vision.detection.utils.object_detection import faster_rcnn_box_coder
-from official.vision.detection.utils.object_detection import target_assigner
 
 
 class Anchor(object):
@@ -134,18 +132,13 @@ class AnchorLabeler(object):
         upper-bound threshold to assign negative labels for anchors. An anchor
         with a score below the threshold is labeled negative.
     """
-    similarity_calc = keras_cv.ops.IouSimilarity()
-    matcher = argmax_matcher.ArgMaxMatcher(
-        match_threshold,
-        unmatched_threshold=unmatched_threshold,
-        negatives_lower_than_unmatched=True,
+    self.similarity_calc = keras_cv.ops.IouSimilarity()
+    self.anchor_labeler = keras_cv.ops.AnchorLabeler()
+    self.matcher = keras_cv.ops.BoxMatcher(
+        positive_threshold=match_threshold,
+        negative_threshold=unmatched_threshold,
         force_match_for_each_row=True)
-    box_coder = faster_rcnn_box_coder.FasterRcnnBoxCoder()
-
-    self._target_assigner = target_assigner.TargetAssigner(
-        similarity_calc, matcher, box_coder)
-    self._match_threshold = match_threshold
-    self._unmatched_threshold = unmatched_threshold
+    self.box_coder = faster_rcnn_box_coder.FasterRcnnBoxCoder()
 
   def label_anchors(self, anchor_boxes, gt_boxes, gt_labels):
     """Labels anchors with ground truth inputs.
@@ -176,30 +169,17 @@ class AnchorLabeler(object):
         1.0 for positive matched anchors, and 0.0 for negative and ignored
         anchors.
     """
-    gt_box_list = box_list.BoxList(gt_boxes)
     flattened_anchor_boxes = []
     for anchors in anchor_boxes.values():
       flattened_anchor_boxes.append(tf.reshape(anchors, [-1, 4]))
     flattened_anchor_boxes = tf.concat(flattened_anchor_boxes, axis=0)
+    similarity_matrix = self.similarity_calc(gt_boxes, flattened_anchor_boxes)
+    match_results = self.matcher(similarity_matrix)
+    cls_targets, box_targets, cls_weights, box_weights = self.anchor_labeler(
+        gt_boxes, gt_labels, match_results)
+    box_targets_list = box_list.BoxList(box_targets)
     anchor_box_list = box_list.BoxList(flattened_anchor_boxes)
-
-    # The cls_weights, box_weights are not used.
-    (cls_targets, cls_weights, box_targets, box_weights,
-     matches) = self._target_assigner.assign(anchor_box_list, gt_box_list,
-                                             gt_labels)
-
-    # Labels definition in matches.match_results:
-    # (1) match_results[i]>=0, meaning that column i is matched with row
-    #     match_results[i].
-    # (2) match_results[i]=-1, meaning that column i is not matched.
-    # (3) match_results[i]=-2, meaning that column i is ignored.
-    match_results = tf.expand_dims(matches.match_results, axis=1)
-    cls_targets = tf.cast(cls_targets, tf.int32)
-    cls_targets = tf.where(
-        tf.equal(match_results, -1), -tf.ones_like(cls_targets), cls_targets)
-    cls_targets = tf.where(
-        tf.equal(match_results, -2), -2 * tf.ones_like(cls_targets),
-        cls_targets)
+    box_targets = self.box_coder.encode(box_targets_list, anchor_box_list)
 
     # Unpacks labels into multi-level representations.
     cls_targets_dict = unpack_targets(cls_targets, anchor_boxes)
@@ -284,19 +264,33 @@ class RpnAnchorLabeler(AnchorLabeler):
         width_l represent the dimension of bounding box regression output at
         l-th level.
     """
-    gt_box_list = box_list.BoxList(gt_boxes)
     flattened_anchor_boxes = []
     for anchors in anchor_boxes.values():
       flattened_anchor_boxes.append(tf.reshape(anchors, [-1, 4]))
     flattened_anchor_boxes = tf.concat(flattened_anchor_boxes, axis=0)
-    anchor_box_list = box_list.BoxList(flattened_anchor_boxes)
-
+    similarity_matrix = self.similarity_calc(gt_boxes, flattened_anchor_boxes)
+    match_results = self.matcher(similarity_matrix)
     # cls_targets, cls_weights, box_weights are not used.
-    _, _, box_targets, _, matches = self._target_assigner.assign(
-        anchor_box_list, gt_box_list, gt_labels)
+    _, box_targets, _, _ = self.anchor_labeler(
+        gt_boxes, gt_labels, match_results)
+    box_targets_list = box_list.BoxList(box_targets)
+    anchor_box_list = box_list.BoxList(flattened_anchor_boxes)
+    box_targets = self.box_coder.encode(box_targets_list, anchor_box_list)
+
+    # Zero out the unmatched and ignored regression targets.
+    num_matches = match_results.shape.as_list()[0] or tf.shape(match_results)[0]
+    unmatched_ignored_box_targets = tf.zeros([num_matches, 4], dtype=tf.float32)
+    matched_anchors_mask = tf.greater_equal(match_results, 0)
+    # To broadcast matched_anchors_mask to the same shape as
+    # matched_reg_targets.
+    matched_anchors_mask = tf.tile(
+        tf.expand_dims(matched_anchors_mask, 1),
+        [1, tf.shape(box_targets)[1]])
+    box_targets = tf.where(matched_anchors_mask, box_targets,
+                           unmatched_ignored_box_targets)
 
     # score_targets contains the subsampled positive and negative anchors.
-    score_targets, _, _ = self._get_rpn_samples(matches.match_results)
+    score_targets, _, _ = self._get_rpn_samples(match_results)
 
     # Unpacks labels.
     score_targets_dict = unpack_targets(score_targets, anchor_boxes)

official/vision/keras_cv/ops/__init__.py

@@ -14,4 +14,6 @@
 # ==============================================================================
 """Keras-CV layers package definition."""
 from official.vision.keras_cv.ops.anchor_generator import AnchorGenerator
+from official.vision.keras_cv.ops.anchor_labeler import AnchorLabeler
+from official.vision.keras_cv.ops.box_matcher import BoxMatcher
 from official.vision.keras_cv.ops.iou_similarity import IouSimilarity

official/vision/keras_cv/ops/anchor_labeler.py

+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Definition of anchor labeler, which assigns ground truth boxes to anchors."""
+
+import tensorflow as tf
+
+
+class AnchorLabeler:
+  """Labeler for dense object detector."""
+
+  def __init__(
+      self,
+      positive_class_weight=1.0,
+      positive_regression_weight=1.0,
+      negative_class_weight=1.0,
+      negative_regression_weight=0.0,
+      negative_class_label=-1,
+      ignore_class_label=-2,
+      negative_regression_label=0.,
+      ignore_regression_label=0.):
+    """Constructs Anchor Labeler.
+
+    Args:
+      positive_class_weight: classification weight to be associated to positive
+        matched anchor. Defaults to 1.0.
+      positive_regression_weight: regression weight to be associated to positive
+        matched anchor. Defaults to 1.0.
+      negative_class_weight: classification weight to be associated to negative
+        matched anchor. Default to 1.0
+      negative_regression_weight: classification weight to be associated to
+        negative matched anchor. Default to 0.0.
+      negative_class_label: An integer for classification label to be associated
+        for negative matched anchor. Defaults to -1.
+      ignore_class_label: An integer for classification label to be associated
+        for ignored anchor. Defaults to -2.
+      negative_regression_label: A float for regression label to be associated
+        for negative matched anchor. Defaults to 0.
+      ignore_regression_label: A float for regression label to be associated
+        for ignored anchor. Defaults to 0.
+    """
+    self.positive_class_weight = positive_class_weight
+    self.positive_regression_weight = positive_regression_weight
+    self.negative_class_weight = negative_class_weight
+    self.negative_regression_weight = negative_regression_weight
+    self.negative_class_label = negative_class_label
+    self.ignore_class_label = ignore_class_label
+    self.negative_regression_label = negative_regression_label
+    self.ignore_regression_label = ignore_regression_label
+
+  def __call__(self, boxes, labels, matches):
+    """Labels anchors with ground truth inputs.
+
+    Args:
+      boxes: A float tensor with shape [N, 4] representing groundtruth boxes.
+        For each row, it stores [y0, x0, y1, x1] for four corners of a box.
+      labels: An integer tensor with shape [N, 1] representing groundtruth
+        classes.
+      matches: An integer tensor with shape [N] representing match results, must
+        be -1 for negative matched anchor, and -2 for ignored anchor.
+
+    Returns:
+      class_targets: A integer Tensor with shape [num_anchors].
+      box_targets: A float Tensor with shape [num_anchors, 4].
+      class_weights: A float Tensor with shape [num_anchors], that
+        serves as masking / sample weight for classification loss. Its value
+        is 1.0 for positive and negative matched anchors, and 0.0 for ignored
+        anchors.
+      box_weights: A float Tensor with shape [num_anchors], that
+        serves as masking / sample weight for regression loss. Its value is
+        1.0 for positive matched anchors, and 0.0 for negative and ignored
+        anchors.
+    """
+
+    class_targets = self._gather_based_on_match(
+        matches, tf.cast(labels, tf.int32),
+        negative_value=tf.constant([self.negative_class_label], tf.int32),
+        ignored_value=tf.constant([self.ignore_class_label], tf.int32))
+
+    negative_reg_value = tf.constant(
+        [self.negative_regression_label] * 4, dtype=tf.float32)
+    ignore_reg_value = tf.constant(
+        [self.ignore_regression_label] * 4, dtype=tf.float32)
+    reg_targets = self._gather_based_on_match(
+        matches, boxes, negative_reg_value, ignore_reg_value)
+
+    num_gt_boxes = boxes.shape.as_list()[0] or tf.shape(boxes)[0]
+
+    groundtruth_class_weights = self.positive_class_weight * tf.ones(
+        [num_gt_boxes], dtype=tf.float32)
+    class_weights = self._gather_based_on_match(
+        matches, groundtruth_class_weights,
+        negative_value=self.negative_class_weight,
+        ignored_value=0.)
+
+    groundtruth_reg_weights = self.positive_regression_weight * tf.ones(
+        [num_gt_boxes], dtype=tf.float32)
+    reg_weights = self._gather_based_on_match(
+        matches, groundtruth_reg_weights,
+        negative_value=self.negative_regression_weight, ignored_value=0.)
+
+    return class_targets, reg_targets, class_weights, reg_weights
+
+  def _gather_based_on_match(
+      self, matches, inputs, negative_value, ignored_value):
+    """Gathers elements from `input_tensor` based on match results.
+
+    For columns that are matched to a row, gathered_tensor[col] is set to
+    input_tensor[match[col]]. For columns that are unmatched,
+    gathered_tensor[col] is set to negative_value. Finally, for columns that
+    are ignored gathered_tensor[col] is set to ignored_value.
+
+    Note that the input_tensor.shape[1:] must match with unmatched_value.shape
+    and ignored_value.shape
+
+    Args:
+      matches: A integer tensor with shape [N] representing the
+        matching results of anchors. (1) match_results[i]>=0,
+        meaning that column i is matched with row match_results[i].
+        (2) match_results[i]=-1, meaning that column i is not matched.
+        (3) match_results[i]=-2, meaning that column i is ignored.
+      inputs: Tensor to gather values from.
+      negative_value: Constant tensor value for unmatched columns.
+      ignored_value: Constant tensor value for ignored columns.
+
+    Returns:
+      gathered_tensor: A tensor containing values gathered from input_tensor.
+        The shape of the gathered tensor is [match.shape[0]] +
+        input_tensor.shape[1:].
+    """
+    inputs = tf.concat(
+        [tf.stack([ignored_value, negative_value]), inputs], axis=0)
+    gather_indices = tf.maximum(matches + 2, 0)
+    gathered_tensor = tf.gather(inputs, gather_indices)
+    return gathered_tensor

official/vision/keras_cv/ops/box_matcher.py

+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+
+"""Box matcher implementation."""
+
+
+import tensorflow as tf
+
+
+class BoxMatcher:
+  """Matcher based on highest value.
+
+  This class computes matches from a similarity matrix. Each column is matched
+  to a single row.
+
+  To support object detection target assignment this class enables setting both
+  positive_threshold (upper threshold) and negative_threshold (lower thresholds)
+  defining three categories of similarity which define whether examples are
+  positive, negative, or ignored:
+  (1) similarity >= positive_threshold: Highest similarity. Matched/Positive!
+  (2) positive_threshold > similarity >= negative_threshold: Medium similarity.
+        This is Ignored.
+  (3) negative_threshold > similarity: Lowest similarity for Negative Match.
+  For ignored matches this class sets the values in the Match object to -2.
+  """
+
+  def __init__(
+      self,
+      positive_threshold,
+      negative_threshold=None,
+      force_match_for_each_row=False,
+      negative_value=-1,
+      ignore_value=-2):
+    """Construct BoxMatcher.
+
+    Args:
+      positive_threshold: Threshold for positive matches. Positive if
+        sim >= positive_threshold, where sim is the maximum value of the
+        similarity matrix for a given column. Set to None for no threshold.
+      negative_threshold: Threshold for negative matches. Negative if
+        sim < negative_threshold or
+        positive_threshold > sim >= negative_threshold.
+        Defaults to positive_threshold when set to None.
+      force_match_for_each_row: If True, ensures that each row is matched to
+        at least one column (which is not guaranteed otherwise if the
+        positive_threshold is high). Defaults to False.
+      negative_value: An integer to fill for negative matches.
+      ignore_value: An integer to fill for ignored matches.
+
+    Raises:
+      ValueError: If negative_threshold > positive_threshold.
+    """
+    self._positive_threshold = positive_threshold
+    if negative_threshold is None:
+      self._negative_threshold = positive_threshold
+    else:
+      if negative_threshold > positive_threshold:
+        raise ValueError('negative_threshold needs to be smaller or equal'
+                         'to positive_threshold')
+      self._negative_threshold = negative_threshold
+
+    self._negative_value = negative_value
+    self._ignore_value = ignore_value
+    self._force_match_for_each_row = force_match_for_each_row
+
+  def __call__(self, similarity_matrix):
+    """Tries to match each column of the similarity matrix to a row.
+
+    Args:
+      similarity_matrix: A float tensor of shape [N, M] representing any
+        similarity metric.
+
+    Returns:
+      A integer tensor with corresponding match indices for each of M columns,
+      for positive match, the match result will be the corresponding row index,
+      for negative match, the match will be `negative_value`, for ignored match,
+      the match result will be `ignore_value`.
+    """
+
+    def _match_when_rows_are_empty():
+      """Performs matching when the rows of similarity matrix are empty.
+
+      When the rows are empty, all detections are false positives. So we return
+      a tensor of -1's to indicate that the columns do not match to any rows.
+
+      Returns:
+        matches:  int32 tensor indicating the row each column matches to.
+      """
+      static_shape = similarity_matrix.shape.as_list()
+      num_cols = static_shape[1] or tf.shape(similarity_matrix)[1]
+      return -1 * tf.ones([num_cols], dtype=tf.int32)
+
+    def _match_when_rows_are_non_empty():
+      """Performs matching when the rows of similarity matrix are non empty.
+
+      Returns:
+        matches:  int32 tensor indicating the row each column matches to.
+      """
+      # Matches for each column
+      matches = tf.argmax(input=similarity_matrix, axis=0, output_type=tf.int32)
+
+      # Deal with matched and unmatched threshold
+      if self._positive_threshold is not None:
+        # Get logical indices of ignored and unmatched columns as tf.int64
+        matched_vals = tf.reduce_max(similarity_matrix, axis=0)
+        below_negative_threshold = tf.greater(self._negative_threshold,
+                                              matched_vals)
+        between_thresholds = tf.logical_and(
+            tf.greater_equal(matched_vals, self._negative_threshold),
+            tf.greater(self._positive_threshold, matched_vals))
+
+        matches = self._set_values_using_indicator(matches,
+                                                   below_negative_threshold,
+                                                   self._negative_value)
+        matches = self._set_values_using_indicator(matches,
+                                                   between_thresholds,
+                                                   self._ignore_value)
+
+      if self._force_match_for_each_row:
+        num_gt_boxes = similarity_matrix.shape.as_list()[1] or tf.shape(
+            similarity_matrix)[1]
+        force_match_column_ids = tf.argmax(
+            input=similarity_matrix, axis=1, output_type=tf.int32)
+        force_match_column_indicators = tf.one_hot(
+            force_match_column_ids, depth=num_gt_boxes)
+        force_match_row_ids = tf.argmax(
+            input=force_match_column_indicators, axis=0, output_type=tf.int32)
+        force_match_column_mask = tf.cast(
+            tf.reduce_max(force_match_column_indicators, axis=0),
+            tf.bool)
+        final_matches = tf.where(force_match_column_mask, force_match_row_ids,
+                                 matches)
+        return final_matches
+      else:
+        return matches
+
+    if similarity_matrix.shape.is_fully_defined():
+      if similarity_matrix.shape.dims[0].value == 0:
+        return _match_when_rows_are_empty()
+      else:
+        return _match_when_rows_are_non_empty()
+    else:
+      return tf.cond(
+          pred=tf.greater(tf.shape(similarity_matrix)[0], 0),
+          true_fn=_match_when_rows_are_non_empty,
+          false_fn=_match_when_rows_are_empty)
+
+  def _set_values_using_indicator(self, x, indicator, val):
+    """Set the indicated fields of x to val.
+
+    Args:
+      x: tensor.
+      indicator: boolean with same shape as x.
+      val: scalar with value to set.
+
+    Returns:
+      modified tensor.
+    """
+    indicator = tf.cast(indicator, x.dtype)
+    return tf.add(tf.multiply(x, 1 - indicator), val * indicator)