Fix test that was flaky due to randomness

PiperOrigin-RevId: 192195514
ktpdpro · Apr 13, 2018 · b0c5c3b · b0c5c3b
1 parent 2624a23
commit b0c5c3b
Showing 1 changed file with 13 additions and 5 deletions.
diff --git a/research/object_detection/core/preprocessor_test.py b/research/object_detection/core/preprocessor_test.py
@@ -1323,16 +1323,24 @@ def testRandomCropImageWithMultiClassScores(self):
     distorted_boxes_rank = tf.rank(distorted_boxes)
     images_rank = tf.rank(images)
     distorted_images_rank = tf.rank(distorted_images)
+    multiclass_scores_rank = tf.rank(multiclass_scores)
+    distorted_multiclass_scores_rank = tf.rank(distorted_multiclass_scores)
 
     with self.test_session() as sess:
-      (boxes_rank_, distorted_boxes_rank_, images_rank_, distorted_images_rank_,
-       multiclass_scores_, distorted_multiclass_scores_) = sess.run([
-           boxes_rank, distorted_boxes_rank, images_rank, distorted_images_rank,
-           multiclass_scores, distorted_multiclass_scores
+      (boxes_rank_, distorted_boxes_, distorted_boxes_rank_, images_rank_,
+       distorted_images_rank_, multiclass_scores_rank_,
+       distorted_multiclass_scores_rank_,
+       distorted_multiclass_scores_) = sess.run([
+           boxes_rank, distorted_boxes, distorted_boxes_rank, images_rank,
+           distorted_images_rank, multiclass_scores_rank,
+           distorted_multiclass_scores_rank, distorted_multiclass_scores
        ])
       self.assertAllEqual(boxes_rank_, distorted_boxes_rank_)
       self.assertAllEqual(images_rank_, distorted_images_rank_)
-      self.assertAllEqual(multiclass_scores_, distorted_multiclass_scores_)
+      self.assertAllEqual(multiclass_scores_rank_,
+                          distorted_multiclass_scores_rank_)
+      self.assertAllEqual(distorted_boxes_.shape[0],
+                          distorted_multiclass_scores_.shape[0])
 
   def testStrictRandomCropImageWithLabelScores(self):
     image = self.createColorfulTestImage()[0]