target_assigner_test.py

# Copyright 2017 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.
# ==============================================================================

"""Tests for object_detection.core.target_assigner."""
import numpy as np
import tensorflow as tf

from box_coders import mean_stddev_box_coder
from box_coders import faster_rcnn_rbox_coder
from core import box_list
from core import rbox_list
from core import region_similarity_calculator
from core import target_assigner as targetassigner
from matchers import argmax_matcher
from matchers import bipartite_matcher


class TargetAssignerTest(tf.test.TestCase):

    def test_assign_agnostic(self):
        similarity_calc = region_similarity_calculator.NegSqDistSimilarity()
        matcher = bipartite_matcher.GreedyBipartiteMatcher()
        box_coder = mean_stddev_box_coder.MeanStddevBoxCoder()
        target_assigner = targetassigner.TargetAssigner(
            similarity_calc, matcher, box_coder, unmatched_cls_target=None)

        prior_means = tf.constant([[0.0, 0.0, 0.5, 0.5],
                                   [0.5, 0.5, 1.0, 0.8],
                                   [0, 0.5, .5, 1.0]])
        prior_stddevs = tf.constant(3 * [4 * [.1]])
        priors = box_list.BoxList(prior_means)
        priors.add_field('stddev', prior_stddevs)

        box_corners = [[0.0, 0.0, 0.5, 0.5], [0.5, 0.5, 0.9, 0.9]]
        boxes = box_list.BoxList(tf.constant(box_corners))
        exp_cls_targets = [[1], [1], [0]]
        exp_cls_weights = [1, 1, 1]
        exp_reg_targets = [[0, 0, 0, 0],
                           [0, 0, -1, 1],
                           [0, 0, 0, 0]]
        exp_reg_weights = [1, 1, 0]
        exp_matching_anchors = [0, 1]

        result = target_assigner.assign(priors, boxes, num_valid_rows=2)
        (cls_targets, cls_weights, reg_targets, reg_weights, match) = result

        with self.test_session() as sess:
            (cls_targets_out, cls_weights_out,
             reg_targets_out, reg_weights_out, matching_anchors_out) = sess.run(
                [cls_targets, cls_weights, reg_targets, reg_weights,
                 match.matched_column_indices()])

            self.assertAllClose(cls_targets_out, exp_cls_targets)
            self.assertAllClose(cls_weights_out, exp_cls_weights)
            self.assertAllClose(reg_targets_out, exp_reg_targets)
            self.assertAllClose(reg_weights_out, exp_reg_weights)
            self.assertAllClose(matching_anchors_out, exp_matching_anchors)
            self.assertEquals(cls_targets_out.dtype, np.float32)
            self.assertEquals(cls_weights_out.dtype, np.float32)
            self.assertEquals(reg_targets_out.dtype, np.float32)
            self.assertEquals(reg_weights_out.dtype, np.float32)
            self.assertEquals(matching_anchors_out.dtype, np.int32)

    def test_assign_with_ignored_matches(self):
        # Note: test is very similar to above. The third box matched with an IOU
        # of 0.35, which is between the matched and unmatched threshold. This means
        # That like above the expected classification targets are [1, 1, 0].
        # Unlike above, the third target is ignored and therefore expected
        # classification weights are [1, 1, 0].
        similarity_calc = region_similarity_calculator.IouSimilarity()
        matcher = argmax_matcher.ArgMaxMatcher(matched_threshold=0.5,
                                               unmatched_threshold=0.3)
        box_coder = mean_stddev_box_coder.MeanStddevBoxCoder()
        target_assigner = targetassigner.TargetAssigner(
            similarity_calc, matcher, box_coder)

        prior_means = tf.constant([[0.0, 0.0, 0.5, 0.5],
                                   [0.5, 0.5, 1.0, 0.8],
                                   [0.0, 0.5, .9, 1.0]])
        prior_stddevs = tf.constant(3 * [4 * [.1]])
        priors = box_list.BoxList(prior_means)
        priors.add_field('stddev', prior_stddevs)

        box_corners = [[0.0, 0.0, 0.5, 0.5],
                       [0.5, 0.5, 0.9, 0.9]]
        boxes = box_list.BoxList(tf.constant(box_corners))
        exp_cls_targets = [[1], [1], [0]]
        exp_cls_weights = [1, 1, 0]
        exp_reg_targets = [[0, 0, 0, 0],
                           [0, 0, -1, 1],
                           [0, 0, 0, 0]]
        exp_reg_weights = [1, 1, 0]
        exp_matching_anchors = [0, 1]

        result = target_assigner.assign(priors, boxes)
        (cls_targets, cls_weights, reg_targets, reg_weights, match) = result
        with self.test_session() as sess:
            (cls_targets_out, cls_weights_out,
             reg_targets_out, reg_weights_out, matching_anchors_out) = sess.run(
                [cls_targets, cls_weights, reg_targets, reg_weights,
                 match.matched_column_indices()])

            self.assertAllClose(cls_targets_out, exp_cls_targets)
            self.assertAllClose(cls_weights_out, exp_cls_weights)
            self.assertAllClose(reg_targets_out, exp_reg_targets)
            self.assertAllClose(reg_weights_out, exp_reg_weights)
            self.assertAllClose(matching_anchors_out, exp_matching_anchors)
            self.assertEquals(cls_targets_out.dtype, np.float32)
            self.assertEquals(cls_weights_out.dtype, np.float32)
            self.assertEquals(reg_targets_out.dtype, np.float32)
            self.assertEquals(reg_weights_out.dtype, np.float32)
            self.assertEquals(matching_anchors_out.dtype, np.int32)

    def test_assign_multiclass(self):
        similarity_calc = region_similarity_calculator.NegSqDistSimilarity()
        matcher = bipartite_matcher.GreedyBipartiteMatcher()
        box_coder = mean_stddev_box_coder.MeanStddevBoxCoder()
        unmatched_cls_target = tf.constant([1, 0, 0, 0, 0, 0, 0], tf.float32)
        target_assigner = targetassigner.TargetAssigner(
            similarity_calc, matcher, box_coder,
            unmatched_cls_target=unmatched_cls_target)

        prior_means = tf.constant([[0.0, 0.0, 0.5, 0.5],
                                   [0.5, 0.5, 1.0, 0.8],
                                   [0, 0.5, .5, 1.0],
                                   [.75, 0, 1.0, .25]])
        prior_stddevs = tf.constant(4 * [4 * [.1]])
        priors = box_list.BoxList(prior_means)
        priors.add_field('stddev', prior_stddevs)

        box_corners = [[0.0, 0.0, 0.5, 0.5],
                       [0.5, 0.5, 0.9, 0.9],
                       [.75, 0, .95, .27]]
        boxes = box_list.BoxList(tf.constant(box_corners))

        groundtruth_labels = tf.constant([[0, 1, 0, 0, 0, 0, 0],
                                          [0, 0, 0, 0, 0, 1, 0],
                                          [0, 0, 0, 1, 0, 0, 0]], tf.float32)

        exp_cls_targets = [[0, 1, 0, 0, 0, 0, 0],
                           [0, 0, 0, 0, 0, 1, 0],
                           [1, 0, 0, 0, 0, 0, 0],
                           [0, 0, 0, 1, 0, 0, 0]]
        exp_cls_weights = [1, 1, 1, 1]
        exp_reg_targets = [[0, 0, 0, 0],
                           [0, 0, -1, 1],
                           [0, 0, 0, 0],
                           [0, 0, -.5, .2]]
        exp_reg_weights = [1, 1, 0, 1]
        exp_matching_anchors = [0, 1, 3]

        result = target_assigner.assign(priors, boxes, groundtruth_labels,
                                        num_valid_rows=3)
        (cls_targets, cls_weights, reg_targets, reg_weights, match) = result
        with self.test_session() as sess:
            (cls_targets_out, cls_weights_out,
             reg_targets_out, reg_weights_out, matching_anchors_out) = sess.run(
                [cls_targets, cls_weights, reg_targets, reg_weights,
                 match.matched_column_indices()])

            self.assertAllClose(cls_targets_out, exp_cls_targets)
            self.assertAllClose(cls_weights_out, exp_cls_weights)
            self.assertAllClose(reg_targets_out, exp_reg_targets)
            self.assertAllClose(reg_weights_out, exp_reg_weights)
            self.assertAllClose(matching_anchors_out, exp_matching_anchors)
            self.assertEquals(cls_targets_out.dtype, np.float32)
            self.assertEquals(cls_weights_out.dtype, np.float32)
            self.assertEquals(reg_targets_out.dtype, np.float32)
            self.assertEquals(reg_weights_out.dtype, np.float32)
            self.assertEquals(matching_anchors_out.dtype, np.int32)

    def test_assign_multiclass_unequal_class_weights(self):
        similarity_calc = region_similarity_calculator.NegSqDistSimilarity()
        matcher = bipartite_matcher.GreedyBipartiteMatcher()
        box_coder = mean_stddev_box_coder.MeanStddevBoxCoder()
        unmatched_cls_target = tf.constant([1, 0, 0, 0, 0, 0, 0], tf.float32)
        target_assigner = targetassigner.TargetAssigner(
            similarity_calc, matcher, box_coder,
            positive_class_weight=1.0, negative_class_weight=0.5,
            unmatched_cls_target=unmatched_cls_target)

        prior_means = tf.constant([[0.0, 0.0, 0.5, 0.5],
                                   [0.5, 0.5, 1.0, 0.8],
                                   [0, 0.5, .5, 1.0],
                                   [.75, 0, 1.0, .25]])
        prior_stddevs = tf.constant(4 * [4 * [.1]])
        priors = box_list.BoxList(prior_means)
        priors.add_field('stddev', prior_stddevs)

        box_corners = [[0.0, 0.0, 0.5, 0.5],
                       [0.5, 0.5, 0.9, 0.9],
                       [.75, 0, .95, .27]]
        boxes = box_list.BoxList(tf.constant(box_corners))

        groundtruth_labels = tf.constant([[0, 1, 0, 0, 0, 0, 0],
                                          [0, 0, 0, 0, 0, 1, 0],
                                          [0, 0, 0, 1, 0, 0, 0]], tf.float32)

        exp_cls_weights = [1, 1, .5, 1]
        result = target_assigner.assign(priors, boxes, groundtruth_labels,
                                        num_valid_rows=3)
        (_, cls_weights, _, _, _) = result
        with self.test_session() as sess:
            cls_weights_out = sess.run(cls_weights)
            self.assertAllClose(cls_weights_out, exp_cls_weights)

    def test_assign_multidimensional_class_targets(self):
        similarity_calc = region_similarity_calculator.NegSqDistSimilarity()
        matcher = bipartite_matcher.GreedyBipartiteMatcher()
        box_coder = mean_stddev_box_coder.MeanStddevBoxCoder()
        unmatched_cls_target = tf.constant([[0, 0], [0, 0]], tf.float32)
        target_assigner = targetassigner.TargetAssigner(
            similarity_calc, matcher, box_coder,
            unmatched_cls_target=unmatched_cls_target)

        prior_means = tf.constant([[0.0, 0.0, 0.5, 0.5],
                                   [0.5, 0.5, 1.0, 0.8],
                                   [0, 0.5, .5, 1.0],
                                   [.75, 0, 1.0, .25]])
        prior_stddevs = tf.constant(4 * [4 * [.1]])
        priors = box_list.BoxList(prior_means)
        priors.add_field('stddev', prior_stddevs)

        box_corners = [[0.0, 0.0, 0.5, 0.5],
                       [0.5, 0.5, 0.9, 0.9],
                       [.75, 0, .95, .27]]
        boxes = box_list.BoxList(tf.constant(box_corners))

        groundtruth_labels = tf.constant([[[0, 1], [1, 0]],
                                          [[1, 0], [0, 1]],
                                          [[0, 1], [1, .5]]], tf.float32)

        exp_cls_targets = [[[0, 1], [1, 0]],
                           [[1, 0], [0, 1]],
                           [[0, 0], [0, 0]],
                           [[0, 1], [1, .5]]]
        exp_cls_weights = [1, 1, 1, 1]
        exp_reg_targets = [[0, 0, 0, 0],
                           [0, 0, -1, 1],
                           [0, 0, 0, 0],
                           [0, 0, -.5, .2]]
        exp_reg_weights = [1, 1, 0, 1]
        exp_matching_anchors = [0, 1, 3]

        result = target_assigner.assign(priors, boxes, groundtruth_labels,
                                        num_valid_rows=3)
        (cls_targets, cls_weights, reg_targets, reg_weights, match) = result
        with self.test_session() as sess:
            (cls_targets_out, cls_weights_out,
             reg_targets_out, reg_weights_out, matching_anchors_out) = sess.run(
                [cls_targets, cls_weights, reg_targets, reg_weights,
                 match.matched_column_indices()])

            self.assertAllClose(cls_targets_out, exp_cls_targets)
            self.assertAllClose(cls_weights_out, exp_cls_weights)
            self.assertAllClose(reg_targets_out, exp_reg_targets)
            self.assertAllClose(reg_weights_out, exp_reg_weights)
            self.assertAllClose(matching_anchors_out, exp_matching_anchors)
            self.assertEquals(cls_targets_out.dtype, np.float32)
            self.assertEquals(cls_weights_out.dtype, np.float32)
            self.assertEquals(reg_targets_out.dtype, np.float32)
            self.assertEquals(reg_weights_out.dtype, np.float32)
            self.assertEquals(matching_anchors_out.dtype, np.int32)

    def test_assign_empty_groundtruth(self):
        similarity_calc = region_similarity_calculator.NegSqDistSimilarity()
        matcher = bipartite_matcher.GreedyBipartiteMatcher()
        box_coder = mean_stddev_box_coder.MeanStddevBoxCoder()
        unmatched_cls_target = tf.constant([0, 0, 0], tf.float32)
        target_assigner = targetassigner.TargetAssigner(
            similarity_calc, matcher, box_coder,
            unmatched_cls_target=unmatched_cls_target)

        prior_means = tf.constant([[0.0, 0.0, 0.5, 0.5],
                                   [0.5, 0.5, 1.0, 0.8],
                                   [0, 0.5, .5, 1.0],
                                   [.75, 0, 1.0, .25]])
        prior_stddevs = tf.constant(4 * [4 * [.1]])
        priors = box_list.BoxList(prior_means)
        priors.add_field('stddev', prior_stddevs)

        box_corners_expanded = tf.constant([[0.0, 0.0, 0.0, 0.0]])
        box_corners = tf.slice(box_corners_expanded, [0, 0], [0, 4])
        boxes = box_list.BoxList(box_corners)

        groundtruth_labels_expanded = tf.constant([[0, 0, 0]], tf.float32)
        groundtruth_labels = tf.slice(groundtruth_labels_expanded, [0, 0], [0, 3])

        exp_cls_targets = [[0, 0, 0],
                           [0, 0, 0],
                           [0, 0, 0],
                           [0, 0, 0]]
        exp_cls_weights = [1, 1, 1, 1]
        exp_reg_targets = [[0, 0, 0, 0],
                           [0, 0, 0, 0],
                           [0, 0, 0, 0],
                           [0, 0, 0, 0]]
        exp_reg_weights = [0, 0, 0, 0]
        exp_matching_anchors = []

        result = target_assigner.assign(priors, boxes, groundtruth_labels)
        (cls_targets, cls_weights, reg_targets, reg_weights, match) = result
        with self.test_session() as sess:
            (cls_targets_out, cls_weights_out,
             reg_targets_out, reg_weights_out, matching_anchors_out) = sess.run(
                [cls_targets, cls_weights, reg_targets, reg_weights,
                 match.matched_column_indices()])

            self.assertAllClose(cls_targets_out, exp_cls_targets)
            self.assertAllClose(cls_weights_out, exp_cls_weights)
            self.assertAllClose(reg_targets_out, exp_reg_targets)
            self.assertAllClose(reg_weights_out, exp_reg_weights)
            self.assertAllClose(matching_anchors_out, exp_matching_anchors)
            self.assertEquals(cls_targets_out.dtype, np.float32)
            self.assertEquals(cls_weights_out.dtype, np.float32)
            self.assertEquals(reg_targets_out.dtype, np.float32)
            self.assertEquals(reg_weights_out.dtype, np.float32)
            self.assertEquals(matching_anchors_out.dtype, np.int32)

    def test_raises_error_on_invalid_groundtruth_labels(self):
        similarity_calc = region_similarity_calculator.NegSqDistSimilarity()
        matcher = bipartite_matcher.GreedyBipartiteMatcher()
        box_coder = mean_stddev_box_coder.MeanStddevBoxCoder()
        unmatched_cls_target = tf.constant([[0, 0], [0, 0], [0, 0]], tf.float32)
        target_assigner = targetassigner.TargetAssigner(
            similarity_calc, matcher, box_coder,
            unmatched_cls_target=unmatched_cls_target)

        prior_means = tf.constant([[0.0, 0.0, 0.5, 0.5]])
        prior_stddevs = tf.constant([[1.0, 1.0, 1.0, 1.0]])
        priors = box_list.BoxList(prior_means)
        priors.add_field('stddev', prior_stddevs)

        box_corners = [[0.0, 0.0, 0.5, 0.5],
                       [0.5, 0.5, 0.9, 0.9],
                       [.75, 0, .95, .27]]
        boxes = box_list.BoxList(tf.constant(box_corners))

        groundtruth_labels = tf.constant([[[0, 1], [1, 0]]], tf.float32)

        with self.assertRaises(ValueError):
            target_assigner.assign(priors, boxes, groundtruth_labels,
                                   num_valid_rows=3)

    def test_assign_multidimensional_class_targets_rbox(self):
        similarity_calc = region_similarity_calculator.IouSimilarity()
        matcher = argmax_matcher.ArgMaxMatcher(matched_threshold=0.5, unmatched_threshold=0.3)
        box_coder = faster_rcnn_rbox_coder.FasterRcnnRBoxCoder()
        unmatched_cls_target = tf.constant([[0, 0], [0, 0]], tf.float32)
        target_assigner = targetassigner.TargetAssigner(similarity_calc, matcher, box_coder,
                                                        unmatched_cls_target=unmatched_cls_target)

        anchors = tf.constant([[0.0, 0.0, 0.5, 0.5, 0.0],
                               [0.5, 0.5, 1.0, 0.8, 0.0],
                               [0, 0.5, .5, 1.0, 0.0],
                               [.75, 0, 1.0, .25, 0.0]])
        anchors_rbox = rbox_list.RBoxList(anchors)

        box_corners = [[0.0, 0.0, 0.5, 0.5, 0.0],
                       [0.5, 0.5, 0.9, 0.9, 0.0],
                       [.75, 0, .95, .27, 0.0]]
        rboxes = rbox_list.RBoxList(tf.constant(box_corners))

        groundtruth_labels = tf.constant([[[0, 1], [1, 0]],
                                          [[1, 0], [0, 1]],
                                          [[0, 1], [1, .5]]], tf.float32)

        exp_cls_targets = [[[0, 1], [1, 0]],
                           [[1, 0], [0, 1]],
                           [[0, 0], [0, 0]],
                           [[0, 1], [1, .5]]]
        exp_cls_weights = [1, 1, 1, 1]
        exp_reg_targets = [[0, 0, 0, 0, 0],
                           [0, 0, -0.105361, 0.117783, 0],
                           [0, 0, 0, 0, 0],
                           [0, 0, -0.0512933, .0769611, 0]]
        exp_reg_weights = [1, 1, 0, 1]
        exp_matching_anchors = [0, 1, 3]

        result = target_assigner.assign(anchors_rbox, rboxes, groundtruth_labels)
        (cls_targets, cls_weights, reg_targets, reg_weights, match) = result
        with self.test_session() as sess:
            (cls_targets_out, cls_weights_out, reg_targets_out, reg_weights_out, matching_anchors_out) = \
                sess.run([cls_targets, cls_weights, reg_targets, reg_weights, match.matched_column_indices()])

            self.assertAllClose(cls_targets_out, exp_cls_targets)
            self.assertAllClose(cls_weights_out, exp_cls_weights)
            self.assertAllClose(reg_targets_out, exp_reg_targets)
            self.assertAllClose(reg_weights_out, exp_reg_weights)
            self.assertAllClose(matching_anchors_out, exp_matching_anchors)
            self.assertEquals(cls_targets_out.dtype, np.float32)
            self.assertEquals(cls_weights_out.dtype, np.float32)
            self.assertEquals(reg_targets_out.dtype, np.float32)
            self.assertEquals(reg_weights_out.dtype, np.float32)
            self.assertEquals(matching_anchors_out.dtype, np.int32)


class BatchTargetAssignerTest(tf.test.TestCase):

    def _get_agnostic_target_assigner(self):
        similarity_calc = region_similarity_calculator.NegSqDistSimilarity()
        matcher = bipartite_matcher.GreedyBipartiteMatcher()
        box_coder = mean_stddev_box_coder.MeanStddevBoxCoder()
        return targetassigner.TargetAssigner(
            similarity_calc, matcher, box_coder,
            positive_class_weight=1.0,
            negative_class_weight=1.0,
            unmatched_cls_target=None)

    def _get_multi_class_target_assigner(self, num_classes):
        similarity_calc = region_similarity_calculator.NegSqDistSimilarity()
        matcher = bipartite_matcher.GreedyBipartiteMatcher()
        box_coder = mean_stddev_box_coder.MeanStddevBoxCoder()
        unmatched_cls_target = tf.constant([1] + num_classes * [0], tf.float32)
        return targetassigner.TargetAssigner(
            similarity_calc, matcher, box_coder,
            positive_class_weight=1.0,
            negative_class_weight=1.0,
            unmatched_cls_target=unmatched_cls_target)

    def _get_multi_dimensional_target_assigner(self, target_dimensions):
        similarity_calc = region_similarity_calculator.NegSqDistSimilarity()
        matcher = bipartite_matcher.GreedyBipartiteMatcher()
        box_coder = mean_stddev_box_coder.MeanStddevBoxCoder()
        unmatched_cls_target = tf.constant(np.zeros(target_dimensions),
                                           tf.float32)
        return targetassigner.TargetAssigner(
            similarity_calc, matcher, box_coder,
            positive_class_weight=1.0,
            negative_class_weight=1.0,
            unmatched_cls_target=unmatched_cls_target)

    def _get_multi_dimensional_target_assigner_rbox(self, target_dimensions):
        similarity_calc = region_similarity_calculator.IouSimilarity()
        matcher = argmax_matcher.ArgMaxMatcher(matched_threshold=0.5, unmatched_threshold=0.3)
        box_coder = faster_rcnn_rbox_coder.FasterRcnnRBoxCoder()
        unmatched_cls_target = tf.constant(np.zeros(target_dimensions), tf.float32)
        return targetassigner.TargetAssigner(
            similarity_calc, matcher, box_coder,
            positive_class_weight=1.0,
            negative_class_weight=1.0,
            unmatched_cls_target=unmatched_cls_target)

    def test_batch_assign_targets(self):
        box_list1 = box_list.BoxList(tf.constant([[0., 0., 0.2, 0.2]]))
        box_list2 = box_list.BoxList(tf.constant(
            [[0, 0.25123152, 1, 1],
             [0.015789, 0.0985, 0.55789, 0.3842]]
        ))

        gt_box_batch = [box_list1, box_list2]
        gt_class_targets = [None, None]

        prior_means = tf.constant([[0, 0, .25, .25],
                                   [0, .25, 1, 1],
                                   [0, .1, .5, .5],
                                   [.75, .75, 1, 1]])
        prior_stddevs = tf.constant([[.1, .1, .1, .1],
                                     [.1, .1, .1, .1],
                                     [.1, .1, .1, .1],
                                     [.1, .1, .1, .1]])
        priors = box_list.BoxList(prior_means)
        priors.add_field('stddev', prior_stddevs)

        exp_reg_targets = [[[0, 0, -0.5, -0.5],
                            [0, 0, 0, 0],
                            [0, 0, 0, 0, ],
                            [0, 0, 0, 0, ], ],
                           [[0, 0, 0, 0, ],
                            [0, 0.01231521, 0, 0],
                            [0.15789001, -0.01500003, 0.57889998, -1.15799987],
                            [0, 0, 0, 0]]]
        exp_cls_weights = [[1, 1, 1, 1],
                           [1, 1, 1, 1]]
        exp_cls_targets = [[[1], [0], [0], [0]],
                           [[0], [1], [1], [0]]]
        exp_reg_weights = [[1, 0, 0, 0],
                           [0, 1, 1, 0]]
        exp_match_0 = [0]
        exp_match_1 = [1, 2]

        agnostic_target_assigner = self._get_agnostic_target_assigner()
        (cls_targets, cls_weights, reg_targets, reg_weights,
         match_list) = targetassigner.batch_assign_targets(
            agnostic_target_assigner, priors, gt_box_batch, gt_class_targets)
        self.assertTrue(isinstance(match_list, list) and len(match_list) == 2)
        with self.test_session() as sess:
            (cls_targets_out, cls_weights_out, reg_targets_out, reg_weights_out,
             match_out_0, match_out_1) = sess.run([
                                                      cls_targets, cls_weights, reg_targets, reg_weights] + [
                                                      match.matched_column_indices() for match in match_list])
            self.assertAllClose(cls_targets_out, exp_cls_targets)
            self.assertAllClose(cls_weights_out, exp_cls_weights)
            self.assertAllClose(reg_targets_out, exp_reg_targets)
            self.assertAllClose(reg_weights_out, exp_reg_weights)
            self.assertAllClose(match_out_0, exp_match_0)
            self.assertAllClose(match_out_1, exp_match_1)

    def test_batch_assign_multiclass_targets(self):
        box_list1 = box_list.BoxList(tf.constant([[0., 0., 0.2, 0.2]]))

        box_list2 = box_list.BoxList(tf.constant(
            [[0, 0.25123152, 1, 1],
             [0.015789, 0.0985, 0.55789, 0.3842]]
        ))

        gt_box_batch = [box_list1, box_list2]

        class_targets1 = tf.constant([[0, 1, 0, 0]], tf.float32)
        class_targets2 = tf.constant([[0, 0, 0, 1],
                                      [0, 0, 1, 0]], tf.float32)

        gt_class_targets = [class_targets1, class_targets2]

        prior_means = tf.constant([[0, 0, .25, .25],
                                   [0, .25, 1, 1],
                                   [0, .1, .5, .5],
                                   [.75, .75, 1, 1]])
        prior_stddevs = tf.constant([[.1, .1, .1, .1],
                                     [.1, .1, .1, .1],
                                     [.1, .1, .1, .1],
                                     [.1, .1, .1, .1]])
        priors = box_list.BoxList(prior_means)
        priors.add_field('stddev', prior_stddevs)

        exp_reg_targets = [[[0, 0, -0.5, -0.5],
                            [0, 0, 0, 0],
                            [0, 0, 0, 0],
                            [0, 0, 0, 0]],
                           [[0, 0, 0, 0],
                            [0, 0.01231521, 0, 0],
                            [0.15789001, -0.01500003, 0.57889998, -1.15799987],
                            [0, 0, 0, 0]]]
        exp_cls_weights = [[1, 1, 1, 1],
                           [1, 1, 1, 1]]
        exp_cls_targets = [[[0, 1, 0, 0],
                            [1, 0, 0, 0],
                            [1, 0, 0, 0],
                            [1, 0, 0, 0]],
                           [[1, 0, 0, 0],
                            [0, 0, 0, 1],
                            [0, 0, 1, 0],
                            [1, 0, 0, 0]]]
        exp_reg_weights = [[1, 0, 0, 0],
                           [0, 1, 1, 0]]
        exp_match_0 = [0]
        exp_match_1 = [1, 2]

        multiclass_target_assigner = self._get_multi_class_target_assigner(
            num_classes=3)

        (cls_targets, cls_weights, reg_targets, reg_weights,
         match_list) = targetassigner.batch_assign_targets(
            multiclass_target_assigner, priors, gt_box_batch, gt_class_targets)
        self.assertTrue(isinstance(match_list, list) and len(match_list) == 2)
        with self.test_session() as sess:
            (cls_targets_out, cls_weights_out, reg_targets_out, reg_weights_out,
             match_out_0, match_out_1) = sess.run([cls_targets, cls_weights, reg_targets, reg_weights] +
                                                  [match.matched_column_indices() for match in match_list])
            self.assertAllClose(cls_targets_out, exp_cls_targets)
            self.assertAllClose(cls_weights_out, exp_cls_weights)
            self.assertAllClose(reg_targets_out, exp_reg_targets)
            self.assertAllClose(reg_weights_out, exp_reg_weights)
            self.assertAllClose(match_out_0, exp_match_0)
            self.assertAllClose(match_out_1, exp_match_1)

    def test_batch_assign_multidimensional_targets(self):
        box_list1 = box_list.BoxList(tf.constant([[0., 0., 0.2, 0.2]]))

        box_list2 = box_list.BoxList(tf.constant(
            [[0, 0.25123152, 1, 1],
             [0.015789, 0.0985, 0.55789, 0.3842]]
        ))

        gt_box_batch = [box_list1, box_list2]
        class_targets1 = tf.constant([[[0, 1, 1],
                                       [1, 1, 0]]], tf.float32)
        class_targets2 = tf.constant([[[0, 1, 1],
                                       [1, 1, 0]],
                                      [[0, 0, 1],
                                       [0, 0, 1]]], tf.float32)

        gt_class_targets = [class_targets1, class_targets2]

        prior_means = tf.constant([[0, 0, .25, .25],
                                   [0, .25, 1, 1],
                                   [0, .1, .5, .5],
                                   [.75, .75, 1, 1]])
        prior_stddevs = tf.constant([[.1, .1, .1, .1],
                                     [.1, .1, .1, .1],
                                     [.1, .1, .1, .1],
                                     [.1, .1, .1, .1]])
        priors = box_list.BoxList(prior_means)
        priors.add_field('stddev', prior_stddevs)

        exp_reg_targets = [[[0, 0, -0.5, -0.5],
                            [0, 0, 0, 0],
                            [0, 0, 0, 0],
                            [0, 0, 0, 0]],
                           [[0, 0, 0, 0],
                            [0, 0.01231521, 0, 0],
                            [0.15789001, -0.01500003, 0.57889998, -1.15799987],
                            [0, 0, 0, 0]]]
        exp_cls_weights = [[1, 1, 1, 1],
                           [1, 1, 1, 1]]

        exp_cls_targets = [[[[0., 1., 1.],
                             [1., 1., 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., 1., 1.],
                             [1., 1., 0.]],
                            [[0., 0., 1.],
                             [0., 0., 1.]],
                            [[0., 0., 0.],
                             [0., 0., 0.]]]]
        exp_reg_weights = [[1, 0, 0, 0],
                           [0, 1, 1, 0]]
        exp_match_0 = [0]
        exp_match_1 = [1, 2]

        multiclass_target_assigner = self._get_multi_dimensional_target_assigner(
            target_dimensions=(2, 3))

        (cls_targets, cls_weights, reg_targets, reg_weights,
         match_list) = targetassigner.batch_assign_targets(
            multiclass_target_assigner, priors, gt_box_batch, gt_class_targets)
        self.assertTrue(isinstance(match_list, list) and len(match_list) == 2)
        with self.test_session() as sess:
            (cls_targets_out, cls_weights_out, reg_targets_out, reg_weights_out,
             match_out_0, match_out_1) = sess.run([cls_targets, cls_weights, reg_targets, reg_weights] +
                                                  [match.matched_column_indices() for match in match_list])
            self.assertAllClose(cls_targets_out, exp_cls_targets)
            self.assertAllClose(cls_weights_out, exp_cls_weights)
            self.assertAllClose(reg_targets_out, exp_reg_targets)
            self.assertAllClose(reg_weights_out, exp_reg_weights)
            self.assertAllClose(match_out_0, exp_match_0)
            self.assertAllClose(match_out_1, exp_match_1)

    def test_batch_assign_empty_groundtruth(self):
        box_coords_expanded = tf.zeros((1, 4), tf.float32)
        box_coords = tf.slice(box_coords_expanded, [0, 0], [0, 4])
        box_list1 = box_list.BoxList(box_coords)
        gt_box_batch = [box_list1]

        prior_means = tf.constant([[0, 0, .25, .25],
                                   [0, .25, 1, 1]])
        prior_stddevs = tf.constant([[.1, .1, .1, .1],
                                     [.1, .1, .1, .1]])
        priors = box_list.BoxList(prior_means)
        priors.add_field('stddev', prior_stddevs)

        exp_reg_targets = [[[0, 0, 0, 0],
                            [0, 0, 0, 0]]]
        exp_cls_weights = [[1, 1]]
        exp_cls_targets = [[[1, 0, 0, 0],
                            [1, 0, 0, 0]]]
        exp_reg_weights = [[0, 0]]
        exp_match_0 = []

        num_classes = 3
        pad = 1
        gt_class_targets = tf.zeros((0, num_classes + pad))
        gt_class_targets_batch = [gt_class_targets]

        multiclass_target_assigner = self._get_multi_class_target_assigner(
            num_classes=3)

        (cls_targets, cls_weights, reg_targets, reg_weights,
         match_list) = targetassigner.batch_assign_targets(
            multiclass_target_assigner, priors,
            gt_box_batch, gt_class_targets_batch)
        self.assertTrue(isinstance(match_list, list) and len(match_list) == 1)
        with self.test_session() as sess:
            (cls_targets_out, cls_weights_out, reg_targets_out, reg_weights_out,
             match_out_0) = sess.run([cls_targets, cls_weights, reg_targets, reg_weights] +
                                     [match.matched_column_indices() for match in match_list])
            self.assertAllClose(cls_targets_out, exp_cls_targets)
            self.assertAllClose(cls_weights_out, exp_cls_weights)
            self.assertAllClose(reg_targets_out, exp_reg_targets)
            self.assertAllClose(reg_weights_out, exp_reg_weights)
            self.assertAllClose(match_out_0, exp_match_0)

    def test_batch_assign_multidimensional_targets_rbox(self):
        rbox_list1 = rbox_list.RBoxList(tf.constant([[0., 0., 0.2, 0.2, 0.0]]))
        rbox_list2 = rbox_list.RBoxList(tf.constant([[0, 0.25123152, 1, 1, 0.0],
                                                     [0.015789, 0.0985, 0.55789, 0.3842, 0.0]]))

        gt_box_batch = [rbox_list1, rbox_list2]
        class_targets1 = tf.constant([[[0, 1, 1],
                                       [1, 1, 0]]], tf.float32)
        class_targets2 = tf.constant([[[0, 1, 1],
                                       [1, 1, 0]],
                                      [[0, 0, 1],
                                       [0, 0, 1]]], tf.float32)

        gt_class_targets = [class_targets1, class_targets2]

        anchors = tf.constant([[0, 0, .25, .25, 0.0],
                               [0, .25, 1, 1, 0.0],
                               [0, .1, .5, .5, 0.0],
                               [.75, .75, 1, 1, 0.0]])
        anchors_rbox = rbox_list.RBoxList(anchors)

        exp_reg_targets = [[[0, 0, -0.223143, -0.223143, 0.0],
                            [0, 0, 0, 0, 0],
                            [0, 0, 0, 0, 0],
                            [0, 0, 0, 0, 0]],
                           [[0, 0, 0, 0, 0],
                            [0, 0.00123152, 0, 0, 0.0],
                            [0.031578, -0.00300001, 0.109554, -0.263445, 0.0],
                            [0, 0, 0, 0, 0]]]
        exp_cls_weights = [[1, 1, 1, 1],
                           [1, 1, 1, 1]]

        exp_cls_targets = [[[[0., 1., 1.],
                             [1., 1., 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., 1., 1.],
                             [1., 1., 0.]],
                            [[0., 0., 1.],
                             [0., 0., 1.]],
                            [[0., 0., 0.],
                             [0., 0., 0.]]]]
        exp_reg_weights = [[1, 0, 0, 0],
                           [0, 1, 1, 0]]
        exp_match_0 = [0]
        exp_match_1 = [1, 2]

        multiclass_target_assigner = self._get_multi_dimensional_target_assigner_rbox(target_dimensions=(2, 3))

        (cls_targets, cls_weights, reg_targets, reg_weights, match_list) = \
            targetassigner.batch_assign_targets(multiclass_target_assigner,
                                                anchors_rbox,
                                                gt_box_batch,
                                                gt_class_targets)
        self.assertTrue(isinstance(match_list, list) and len(match_list) == 2)
        with self.test_session() as sess:
            (cls_targets_out, cls_weights_out, reg_targets_out, reg_weights_out,
             match_out_0, match_out_1) = sess.run([cls_targets, cls_weights, reg_targets, reg_weights] +
                                                  [match.matched_column_indices() for match in match_list])
            self.assertAllClose(cls_targets_out, exp_cls_targets)
            self.assertAllClose(cls_weights_out, exp_cls_weights)
            self.assertAllClose(reg_targets_out, exp_reg_targets)
            self.assertAllClose(reg_weights_out, exp_reg_weights)
            self.assertAllClose(match_out_0, exp_match_0)
            self.assertAllClose(match_out_1, exp_match_1)


class CreateTargetAssignerTest(tf.test.TestCase):

    def test_create_target_assigner(self):
        """Tests that named constructor gives working target assigners.

        TODO: Make this test more general.
        """
        corners = [[0.0, 0.0, 1.0, 1.0]]
        groundtruth = box_list.BoxList(tf.constant(corners))

        priors = box_list.BoxList(tf.constant(corners))
        prior_stddevs = tf.constant([[1.0, 1.0, 1.0, 1.0]])
        priors.add_field('stddev', prior_stddevs)
        multibox_ta = (targetassigner.create_target_assigner('Multibox', stage='proposal'))
        multibox_ta.assign(priors, groundtruth)
        # No tests on output, as that may vary arbitrarily as new target assigners
        # are added. As long as it is constructed correctly and runs without errors,
        # tests on the individual assigners cover correctness of the assignments.

        anchors = box_list.BoxList(tf.constant(corners))
        faster_rcnn_proposals_ta = (targetassigner.create_target_assigner('FasterRCNN',stage='proposal'))
        faster_rcnn_proposals_ta.assign(anchors, groundtruth)

        fast_rcnn_ta = (targetassigner.create_target_assigner('FastRCNN'))
        fast_rcnn_ta.assign(anchors, groundtruth)

        faster_rcnn_detection_ta = (targetassigner.create_target_assigner('FasterRCNN', stage='detection'))
        faster_rcnn_detection_ta.assign(anchors, groundtruth)

        corners = [[0.0, 0.0, 1.0, 1.0, 0.0]]
        groundtruth = rbox_list.RBoxList(tf.constant(corners))
        anchors = rbox_list.RBoxList(tf.constant(corners))
        rfaster_rcnn_proposals_ta = (targetassigner.create_target_assigner('RFasterRCNN', stage='proposal'))
        rfaster_rcnn_proposals_ta.assign(anchors, groundtruth)

        rfaster_rcnn_detection_ta = (targetassigner.create_target_assigner('RFasterRCNN', stage='detection'))
        rfaster_rcnn_detection_ta.assign(anchors, groundtruth)

        with self.assertRaises(ValueError):
            targetassigner.create_target_assigner('InvalidDetector', stage='invalid_stage')


if __name__ == '__main__':
    tf.test.main()