box_matcher.py

# 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.

"Argmax-based box matching"

from typing import List
from typing import Tuple

import tensorflow as tf


class ArgmaxBoxMatcher:
    """Box matching logic based on argmax of highest value (e.g., IOU).

    This class computes matches from a similarity matrix. Each row will be
    matched to at least one column, the matched result can either be positive
     / negative, or simply ignored depending on the setting.

    The settings include `thresholds` and `match_values`, for example if:
    1) thresholds=[negative_threshold, positive_threshold], and
       match_values=[negative_value=0, ignore_value=-1, positive_value=1]: the rows will
       be assigned to positive_value if its argmax result >=
       positive_threshold; the rows will be assigned to negative_value if its
       argmax result < negative_threshold, and the rows will be assigned
       to ignore_value if its argmax result is between [negative_threshold, positive_threshold).
    2) thresholds=[negative_threshold, positive_threshold], and
       match_values=[ignore_value=-1, negative_value=0, positive_value=1]: the rows will
       be assigned to positive_value if its argmax result >=
       positive_threshold; the rows will be assigned to ignore_value if its
       argmax result < negative_threshold, and the rows will be assigned
       to negative_value if its argmax result is between [negative_threshold ,positive_threshold).
       This is different from case 1) by swapping first two
       values.
    3) thresholds=[positive_threshold], and
       match_values=[negative_values, positive_value]: the rows will be assigned to
       positive value if its argmax result >= positive_threshold; the rows
       will be assigned to negative_value if its argmax result < negative_threshold.

    Args:
        thresholds: A sorted list of floats to classify the matches into
          different results (e.g. positive or negative or ignored match). The
          list will be prepended with -Inf and and appended with +Inf.
        match_values: A list of integers representing matched results (e.g.
          positive or negative or ignored match). len(`match_values`) must
          equal to len(`thresholds`) + 1.
        force_match_for_each_col: each row will be argmax matched to at
          least one column. This means some columns will be matched to
          multiple rows while some columns will not be matched to any rows.
          Filtering by `thresholds` will make less columns match to positive
          result. Setting this to True guarantees that each column will be
          matched to positive result to at least one row.

    Raises:
        ValueError: if `thresholds` not sorted or
        len(`match_values`) != len(`thresholds`) + 1

    Usage:

    ```python
    box_matcher = keras_cv.ops.ArgmaxBoxMatcher([0.3, 0.7], [-1, 0, 1])
    iou_metric = keras_cv.bounding_box.compute_iou(anchors, gt_boxes)
    matched_columns, matched_match_values = box_matcher(iou_metric)
    cls_mask = tf.less_equal(matched_match_values, 0)
    ```

    TODO(tanzhenyu): document when to use which mode.

    """

    def __init__(
        self,
        thresholds: List[float],
        match_values: List[int],
        force_match_for_each_col: bool = False,
    ):
        if sorted(thresholds) != thresholds:
            raise ValueError(f"`threshold` must be sorted, got {thresholds}")
        self.match_values = match_values
        if len(match_values) != len(thresholds) + 1:
            raise ValueError(
                f"len(`match_values`) must be len(`thresholds`) + 1, got "
                f"match_values {match_values}, thresholds {thresholds}"
            )
        thresholds.insert(0, -float("inf"))
        thresholds.append(float("inf"))
        self.thresholds = thresholds
        self.force_match_for_each_col = force_match_for_each_col

    def __call__(self, similarity_matrix: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        """Matches each row to a column based on argmax

        TODO(tanzhenyu): consider swapping rows and cols.
        Args:
          similarity_matrix: A float Tensor of shape [num_rows, num_cols] or
            [batch_size, num_rows, num_cols] representing any similarity metric.

        Returns:
          matched_columns: An integer tensor of shape [num_rows] or [batch_size,
            num_rows] storing the index of the matched colum for each row.
          matched_values: An integer tensor of shape [num_rows] or [batch_size,
            num_rows] storing the match result (positive match, negative match,
            ignored match).
        """
        squeeze_result = False
        if len(similarity_matrix.shape) == 2:
            squeeze_result = True
            similarity_matrix = tf.expand_dims(similarity_matrix, axis=0)
        static_shape = similarity_matrix.shape.as_list()
        num_rows = static_shape[1] or tf.shape(similarity_matrix)[1]
        batch_size = static_shape[0] or tf.shape(similarity_matrix)[0]

        def _match_when_cols_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 rows do not match to any columns.
            Returns:
                matched_columns: An integer tensor of shape [batch_size, num_rows]
                  storing the index of the matched column for each row.
                matched_values: An integer tensor of shape [batch_size, num_rows]
                  storing the match type indicator (e.g. positive or negative
                  or ignored match).
            """
            with tf.name_scope("empty_gt_boxes"):
                matched_columns = tf.zeros([batch_size, num_rows], dtype=tf.int32)
                matched_values = -tf.ones([batch_size, num_rows], dtype=tf.int32)
                return matched_columns, matched_values

        def _match_when_cols_are_non_empty():
            """Performs matching when the rows of similarity matrix are non empty.
            Returns:
                matched_columns: An integer tensor of shape [batch_size, num_rows]
                  storing the index of the matched column for each row.
                matched_values: An integer tensor of shape [batch_size, num_rows]
                  storing the match type indicator (e.g. positive or negative
                  or ignored match).
            """
            with tf.name_scope("non_empty_gt_boxes"):
                matched_columns = tf.argmax(
                    similarity_matrix, axis=-1, output_type=tf.int32
                )

                # Get logical indices of ignored and unmatched columns as tf.int64
                matched_vals = tf.reduce_max(similarity_matrix, axis=-1)
                matched_values = tf.zeros([batch_size, num_rows], tf.int32)

                match_dtype = matched_vals.dtype
                for (ind, low, high) in zip(
                    self.match_values, self.thresholds[:-1], self.thresholds[1:]
                ):
                    low_threshold = tf.cast(low, match_dtype)
                    high_threshold = tf.cast(high, match_dtype)
                    mask = tf.logical_and(
                        tf.greater_equal(matched_vals, low_threshold),
                        tf.less(matched_vals, high_threshold),
                    )
                    matched_values = self._set_values_using_indicator(
                        matched_values, mask, ind
                    )

                if self.force_match_for_each_col:
                    # [batch_size, num_cols], for each column (groundtruth_box), find the
                    # best matching row (anchor).
                    matching_rows = tf.argmax(
                        input=similarity_matrix, axis=1, output_type=tf.int32
                    )
                    # [batch_size, num_cols, num_rows], a transposed 0-1 mapping matrix M,
                    # where M[j, i] = 1 means column j is matched to row i.
                    column_to_row_match_mapping = tf.one_hot(
                        matching_rows, depth=num_rows
                    )
                    # [batch_size, num_rows], for each row (anchor), find the matched
                    # column (groundtruth_box).
                    force_matched_columns = tf.argmax(
                        input=column_to_row_match_mapping, axis=1, output_type=tf.int32
                    )
                    # [batch_size, num_rows]
                    force_matched_column_mask = tf.cast(
                        tf.reduce_max(column_to_row_match_mapping, axis=1), tf.bool
                    )
                    # [batch_size, num_rows]
                    matched_columns = tf.where(
                        force_matched_column_mask,
                        force_matched_columns,
                        matched_columns,
                    )
                    matched_values = tf.where(
                        force_matched_column_mask,
                        self.match_values[-1]
                        * tf.ones([batch_size, num_rows], dtype=tf.int32),
                        matched_values,
                    )

                return matched_columns, matched_values

        num_gt_boxes = (
            similarity_matrix.shape.as_list()[-1] or tf.shape(similarity_matrix)[-1]
        )
        matched_columns, matched_values = tf.cond(
            pred=tf.greater(num_gt_boxes, 0),
            true_fn=_match_when_cols_are_non_empty,
            false_fn=_match_when_cols_are_empty,
        )

        if squeeze_result:
            matched_columns = tf.squeeze(matched_columns, axis=0)
            matched_values = tf.squeeze(matched_values, axis=0)

        return matched_columns, matched_values

    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)

    def get_config(self):
        config = {
            "thresholds": self.thresholds[1:-1],
            "match_values": self.match_values,
            "force_match_for_each_col": self.force_match_for_each_col,
        }
        return config