Add symmetric jitter option.

PiperOrigin-RevId: 374301186

research/object_detection/core/preprocessor.py

@@ -1307,8 +1307,8 @@ def random_distort_color(image, color_ordering=0, preprocess_vars_cache=None):
     return image
 
 
-def random_jitter_boxes(boxes, ratio=0.05, jitter_mode='random', seed=None):
-  """Randomly jitter boxes in image.
+def random_jitter_boxes(boxes, ratio=0.05, jitter_mode='default', seed=None):
+  """Randomly jitters boxes in image.
 
   Args:
     boxes: rank 2 float32 tensor containing the bounding boxes -> [N, 4].
@@ -1321,40 +1321,87 @@ def random_jitter_boxes(boxes, ratio=0.05, jitter_mode='random', seed=None):
     jitter_mode: One of
       shrink - Only shrinks boxes.
       expand - Only expands boxes.
+      expand_symmetric - Expands the boxes symmetrically along height and width
+        dimensions without changing the box center. The ratios of expansion
+        along X, Y dimensions are independent
+      shrink_symmetric - Shrinks the boxes symmetrically along height and width
+        dimensions without changing the box center. The ratios of shrinking
+        along X, Y dimensions are independent
+      expand_symmetric_xy - Expands the boxes symetrically along height and
+        width dimensions and the ratio of expansion is same for both.
+      shrink_symmetric_xy - Shrinks the boxes symetrically along height and
+        width dimensions and the ratio of shrinking is same for both.
       default - Randomly and independently perturbs each box boundary.
     seed: random seed.
 
   Returns:
     boxes: boxes which is the same shape as input boxes.
   """
-  with tf.name_scope('RandomJitterBoxes', values=[boxes]):
+  with tf.name_scope('RandomJitterBoxes'):
     ymin, xmin, ymax, xmax = (boxes[:, i] for i in range(4))
 
-    height, width = ymax - ymin, xmax - xmin
-    ycenter, xcenter = (ymin + ymax) / 2.0, (xmin + xmax) / 2.0
+    blist = box_list.BoxList(boxes)
+    ycenter, xcenter, height, width = blist.get_center_coordinates_and_sizes()
 
-    height = tf.abs(height)
-    width = tf.abs(width)
+    height = tf.maximum(tf.abs(height), 1e-6)
+    width = tf.maximum(tf.abs(width), 1e-6)
 
-    if jitter_mode == 'shrink':
+    if jitter_mode in ['shrink', 'shrink_symmetric', 'shrink_symmetric_xy']:
       min_ratio, max_ratio = -ratio, 0
-    elif jitter_mode == 'expand':
+    elif jitter_mode in ['expand', 'expand_symmetric', 'expand_symmetric_xy']:
       min_ratio, max_ratio = 0, ratio
-    else:
+    elif jitter_mode == 'default':
       min_ratio, max_ratio = -ratio, ratio
+    else:
+      raise ValueError('Unknown jitter mode - %s' % jitter_mode)
 
     num_boxes = tf.shape(boxes)[0]
-    distortion = 1.0 + tf.random_uniform(
-        [num_boxes, 4], minval=min_ratio, maxval=max_ratio, dtype=tf.float32,
-        seed=seed)
 
-    ymin_jitter = height * distortion[:, 0]
-    xmin_jitter = width * distortion[:, 1]
-    ymax_jitter = height * distortion[:, 2]
-    xmax_jitter = width * distortion[:, 3]
+    if jitter_mode in ['expand_symmetric', 'shrink_symmetric',
+                       'expand_symmetric_xy', 'shrink_symmetric_xy']:
+      distortion = 1.0 + tf.random.uniform(
+          [num_boxes, 2], minval=min_ratio, maxval=max_ratio, dtype=tf.float32,
+          seed=seed)
+      height_distortion = distortion[:, 0]
+      width_distortion = distortion[:, 1]
+
+      # This is to ensure that all boxes are augmented symmetrically. We clip
+      # each boundary to lie within the image, and when doing so, we also
+      # adjust its symmetric counterpart.
+      max_height_distortion = tf.abs(tf.minimum(
+          (2.0 * ycenter) / height, 2.0 * (1 - ycenter) / height))
+      max_width_distortion = tf.abs(tf.minimum(
+          (2.0 * xcenter) / width, 2.0 * (1 - xcenter) / width))
+
+      if jitter_mode in ['expand_symmetric_xy', 'shrink_symmetric_xy']:
+        height_distortion = width_distortion = distortion[:, 0]
+        max_height_distortion = max_width_distortion = (
+            tf.minimum(max_width_distortion, max_height_distortion))
+
+      height_distortion = tf.clip_by_value(
+          height_distortion, -max_height_distortion, max_height_distortion)
+      width_distortion = tf.clip_by_value(
+          width_distortion, -max_width_distortion, max_width_distortion)
+
+      ymin = ycenter - (height * height_distortion / 2.0)
+      ymax = ycenter + (height * height_distortion / 2.0)
+      xmin = xcenter - (width * width_distortion / 2.0)
+      xmax = xcenter + (width * width_distortion / 2.0)
+
+    elif jitter_mode in ['expand', 'shrink', 'default']:
+      distortion = 1.0 + tf.random.uniform(
+          [num_boxes, 4], minval=min_ratio, maxval=max_ratio, dtype=tf.float32,
+          seed=seed)
+      ymin_jitter = height * distortion[:, 0]
+      xmin_jitter = width * distortion[:, 1]
+      ymax_jitter = height * distortion[:, 2]
+      xmax_jitter = width * distortion[:, 3]
+
+      ymin, ymax = ycenter - (ymin_jitter / 2.0), ycenter + (ymax_jitter / 2.0)
+      xmin, xmax = xcenter - (xmin_jitter / 2.0), xcenter + (xmax_jitter / 2.0)
 
-    ymin, ymax = ycenter - (ymin_jitter / 2.0), ycenter + (ymax_jitter / 2.0)
-    xmin, xmax = xcenter - (xmin_jitter / 2.0), xcenter + (xmax_jitter / 2.0)
+    else:
+      raise ValueError('Unknown jitter mode - %s' % jitter_mode)
 
     boxes = tf.stack([ymin, xmin, ymax, xmax], axis=1)
     return tf.clip_by_value(boxes, 0.0, 1.0)

research/object_detection/core/preprocessor_test.py

@@ -79,9 +79,24 @@ class PreprocessorTest(test_case.TestCase, parameterized.TestCase):
         [[0.0, 0.25, 0.75, 1.0], [0.25, 0.5, 0.75, 1.0]], dtype=tf.float32)
     return boxes
 
+  def createRandomTextBoxes(self):
+    random_boxes = tf.concat([tf.random.uniform([100, 2], 0.0, 0.5, seed=1),
+                              tf.random.uniform([100, 2], 0.5, 1.0, seed=2)],
+                             axis=1)
+    fixed_boxes = tf.constant(
+        [[0.0, 0.25, 0.75, 1.0],
+         [0.25, 0.5, 0.75, 1.0],
+         [0.0, 0.0, 1.0, 1.0],
+         [0.1, 0.2, 0.3, 0.4]], dtype=tf.float32)
+    zero_boxes = tf.zeros((50, 4))
+    return tf.concat([random_boxes, fixed_boxes, zero_boxes], axis=0)
+
   def createTestGroundtruthWeights(self):
     return tf.constant([1.0, 0.5], dtype=tf.float32)
 
+  def createZeroBoxes(self):
+    return tf.zeros((100, 4))
+
   def createTestMasks(self):
     mask = np.array([
         [[255.0, 0.0, 0.0],
@@ -1253,7 +1268,7 @@ class PreprocessorTest(test_case.TestCase, parameterized.TestCase):
     def graph_fn():
       preprocessing_options = []
       preprocessing_options.append((preprocessor.random_jitter_boxes, {}))
-      boxes = self.createTestBoxes()
+      boxes = self.createRandomTextBoxes()
       boxes_shape = tf.shape(boxes)
       tensor_dict = {fields.InputDataFields.groundtruth_boxes: boxes}
       tensor_dict = preprocessor.preprocess(tensor_dict, preprocessing_options)
@@ -1264,20 +1279,24 @@ class PreprocessorTest(test_case.TestCase, parameterized.TestCase):
     (boxes_shape_, distorted_boxes_shape_) = self.execute_cpu(graph_fn, [])
     self.assertAllEqual(boxes_shape_, distorted_boxes_shape_)
 
-  def testRandomJitterBoxesZeroRatio(self):
+  @parameterized.parameters(
+      ['expand', 'shrink', 'expand_symmetric', 'shrink_symmetric',
+       'expand_symmetric_xy', 'shrink_symmetric_xy']
+  )
+  def testRandomJitterBoxesZeroRatio(self, jitter_mode):
 
     def graph_fn():
       preprocessing_options = []
       preprocessing_options.append((preprocessor.random_jitter_boxes,
-                                    {'ratio': 0.0}))
-      boxes = self.createTestBoxes()
+                                    {'ratio': .0, 'jitter_mode': jitter_mode}))
+      boxes = self.createRandomTextBoxes()
       tensor_dict = {fields.InputDataFields.groundtruth_boxes: boxes}
       tensor_dict = preprocessor.preprocess(tensor_dict, preprocessing_options)
       distorted_boxes = tensor_dict[fields.InputDataFields.groundtruth_boxes]
       return [boxes, distorted_boxes]
 
     (boxes, distorted_boxes) = self.execute_cpu(graph_fn, [])
-    self.assertAllEqual(boxes, distorted_boxes)
+    self.assertAllClose(boxes, distorted_boxes)
 
   def testRandomJitterBoxesExpand(self):
 
@@ -1285,7 +1304,56 @@ class PreprocessorTest(test_case.TestCase, parameterized.TestCase):
       preprocessing_options = []
       preprocessing_options.append((preprocessor.random_jitter_boxes,
                                     {'jitter_mode': 'expand'}))
-      boxes = self.createTestBoxes()
+      boxes = self.createRandomTextBoxes()
+      tensor_dict = {fields.InputDataFields.groundtruth_boxes: boxes}
+      tensor_dict = preprocessor.preprocess(tensor_dict, preprocessing_options)
+      distorted_boxes = tensor_dict[fields.InputDataFields.groundtruth_boxes]
+      return [boxes, distorted_boxes]
+
+    boxes, distorted_boxes = self.execute_cpu(graph_fn, [])
+    ymin, xmin, ymax, xmax = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
+    distorted_ymin, distorted_xmin, distorted_ymax, distorted_xmax = (
+        distorted_boxes[:, 0], distorted_boxes[:, 1], distorted_boxes[:, 2],
+        distorted_boxes[:, 3])
+
+    self.assertTrue(np.all(distorted_ymin <= ymin))
+    self.assertTrue(np.all(distorted_xmin <= xmin))
+    self.assertTrue(np.all(distorted_ymax >= ymax))
+    self.assertTrue(np.all(distorted_xmax >= xmax))
+
+  def testRandomJitterBoxesExpandSymmetric(self):
+
+    def graph_fn():
+      preprocessing_options = []
+      preprocessing_options.append((preprocessor.random_jitter_boxes,
+                                    {'jitter_mode': 'expand_symmetric'}))
+      boxes = self.createRandomTextBoxes()
+      tensor_dict = {fields.InputDataFields.groundtruth_boxes: boxes}
+      tensor_dict = preprocessor.preprocess(tensor_dict, preprocessing_options)
+      distorted_boxes = tensor_dict[fields.InputDataFields.groundtruth_boxes]
+      return [boxes, distorted_boxes]
+
+    boxes, distorted_boxes = self.execute_cpu(graph_fn, [])
+    ymin, xmin, ymax, xmax = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
+    distorted_ymin, distorted_xmin, distorted_ymax, distorted_xmax = (
+        distorted_boxes[:, 0], distorted_boxes[:, 1], distorted_boxes[:, 2],
+        distorted_boxes[:, 3])
+
+    self.assertTrue(np.all(distorted_ymin <= ymin))
+    self.assertTrue(np.all(distorted_xmin <= xmin))
+    self.assertTrue(np.all(distorted_ymax >= ymax))
+    self.assertTrue(np.all(distorted_xmax >= xmax))
+
+    self.assertAllClose(ymin - distorted_ymin, distorted_ymax - ymax, rtol=1e-5)
+    self.assertAllClose(xmin - distorted_xmin, distorted_xmax - xmax, rtol=1e-5)
+
+  def testRandomJitterBoxesExpandSymmetricXY(self):
+
+    def graph_fn():
+      preprocessing_options = []
+      preprocessing_options.append((preprocessor.random_jitter_boxes,
+                                    {'jitter_mode': 'expand_symmetric_xy'}))
+      boxes = self.createRandomTextBoxes()
       tensor_dict = {fields.InputDataFields.groundtruth_boxes: boxes}
       tensor_dict = preprocessor.preprocess(tensor_dict, preprocessing_options)
       distorted_boxes = tensor_dict[fields.InputDataFields.groundtruth_boxes]
@@ -1294,7 +1362,7 @@ class PreprocessorTest(test_case.TestCase, parameterized.TestCase):
     boxes, distorted_boxes = self.execute_cpu(graph_fn, [])
     ymin, xmin, ymax, xmax = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
     distorted_ymin, distorted_xmin, distorted_ymax, distorted_xmax = (
-        distorted_boxes[:, 0], distorted_boxes[:, 1], distorted_boxes[0, 2],
+        distorted_boxes[:, 0], distorted_boxes[:, 1], distorted_boxes[:, 2],
         distorted_boxes[:, 3])
 
     self.assertTrue(np.all(distorted_ymin <= ymin))
@@ -1302,6 +1370,16 @@ class PreprocessorTest(test_case.TestCase, parameterized.TestCase):
     self.assertTrue(np.all(distorted_ymax >= ymax))
     self.assertTrue(np.all(distorted_xmax >= xmax))
 
+    self.assertAllClose(ymin - distorted_ymin, distorted_ymax - ymax, rtol=1e-5)
+    self.assertAllClose(xmin - distorted_xmin, distorted_xmax - xmax, rtol=1e-5)
+
+    height, width = tf.maximum(1e-6, ymax - ymin), tf.maximum(1e-6, xmax - xmin)
+
+    self.assertAllClose((distorted_ymax - ymax) / height,
+                        (distorted_xmax - xmax) / width, rtol=1e-5)
+    self.assertAllLessEqual((distorted_ymax - ymax) / height, 0.05)
+    self.assertAllGreaterEqual((distorted_ymax - ymax) / width, 0.00)
+
   def testRandomJitterBoxesShrink(self):
 
     def graph_fn():
@@ -1317,7 +1395,30 @@ class PreprocessorTest(test_case.TestCase, parameterized.TestCase):
     boxes, distorted_boxes = self.execute_cpu(graph_fn, [])
     ymin, xmin, ymax, xmax = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
     distorted_ymin, distorted_xmin, distorted_ymax, distorted_xmax = (
-        distorted_boxes[:, 0], distorted_boxes[:, 1], distorted_boxes[0, 2],
+        distorted_boxes[:, 0], distorted_boxes[:, 1], distorted_boxes[:, 2],
+        distorted_boxes[:, 3])
+
+    self.assertTrue(np.all(distorted_ymin >= ymin))
+    self.assertTrue(np.all(distorted_xmin >= xmin))
+    self.assertTrue(np.all(distorted_ymax <= ymax))
+    self.assertTrue(np.all(distorted_xmax <= xmax))
+
+  def testRandomJitterBoxesShrinkSymmetric(self):
+
+    def graph_fn():
+      preprocessing_options = []
+      preprocessing_options.append((preprocessor.random_jitter_boxes,
+                                    {'jitter_mode': 'shrink_symmetric'}))
+      boxes = self.createTestBoxes()
+      tensor_dict = {fields.InputDataFields.groundtruth_boxes: boxes}
+      tensor_dict = preprocessor.preprocess(tensor_dict, preprocessing_options)
+      distorted_boxes = tensor_dict[fields.InputDataFields.groundtruth_boxes]
+      return [boxes, distorted_boxes]
+
+    boxes, distorted_boxes = self.execute_cpu(graph_fn, [])
+    ymin, xmin, ymax, xmax = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
+    distorted_ymin, distorted_xmin, distorted_ymax, distorted_xmax = (
+        distorted_boxes[:, 0], distorted_boxes[:, 1], distorted_boxes[:, 2],
         distorted_boxes[:, 3])
 
     self.assertTrue(np.all(distorted_ymin >= ymin))
@@ -1325,6 +1426,41 @@ class PreprocessorTest(test_case.TestCase, parameterized.TestCase):
     self.assertTrue(np.all(distorted_ymax <= ymax))
     self.assertTrue(np.all(distorted_xmax <= xmax))
 
+    self.assertAllClose(ymin - distorted_ymin, distorted_ymax - ymax, rtol=1e-5)
+    self.assertAllClose(xmin - distorted_xmin, distorted_xmax - xmax, rtol=1e-5)
+
+  def testRandomJitterBoxesShrinkSymmetricXY(self):
+
+    def graph_fn():
+      preprocessing_options = []
+      preprocessing_options.append((preprocessor.random_jitter_boxes,
+                                    {'jitter_mode': 'shrink_symmetric_xy'}))
+      boxes = self.createTestBoxes()
+      tensor_dict = {fields.InputDataFields.groundtruth_boxes: boxes}
+      tensor_dict = preprocessor.preprocess(tensor_dict, preprocessing_options)
+      distorted_boxes = tensor_dict[fields.InputDataFields.groundtruth_boxes]
+      return [boxes, distorted_boxes]
+
+    boxes, distorted_boxes = self.execute_cpu(graph_fn, [])
+    ymin, xmin, ymax, xmax = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
+    distorted_ymin, distorted_xmin, distorted_ymax, distorted_xmax = (
+        distorted_boxes[:, 0], distorted_boxes[:, 1], distorted_boxes[:, 2],
+        distorted_boxes[:, 3])
+
+    self.assertTrue(np.all(distorted_ymin >= ymin))
+    self.assertTrue(np.all(distorted_xmin >= xmin))
+    self.assertTrue(np.all(distorted_ymax <= ymax))
+    self.assertTrue(np.all(distorted_xmax <= xmax))
+
+    self.assertAllClose(ymin - distorted_ymin, distorted_ymax - ymax, rtol=1e-5)
+    self.assertAllClose(xmin - distorted_xmin, distorted_xmax - xmax, rtol=1e-5)
+
+    height, width = tf.maximum(1e-6, ymax - ymin), tf.maximum(1e-6, xmax - xmin)
+    self.assertAllClose((ymax - distorted_ymax) / height,
+                        (xmax - distorted_xmax) / width, rtol=1e-5)
+    self.assertAllLessEqual((ymax - distorted_ymax) / height, 0.05)
+    self.assertAllGreaterEqual((ymax - distorted_ymax)/ width, 0.00)
+
   def testRandomCropImage(self):
 
     def graph_fn():

research/object_detection/protos/preprocessor.proto

@@ -171,11 +171,25 @@ message RandomJitterBoxes {
     DEFAULT = 0;
     EXPAND = 1;
     SHRINK = 2;
+    EXPAND_SYMMETRIC = 4;
+    SHRINK_SYMMETRIC = 5;
+    EXPAND_SYMMETRIC_XY = 6;
+    SHRINK_SYMMETRIC_XY = 7;
   }
   // The mode of jittering
   // EXPAND - Only expands boxes
   // SHRINK - Only shrinks boxes
-  // DEFAULT - Jitters each box boundary independently
+  // EXPAND_SYMMETRIC - Expands the boxes symmetrically along height and width
+  //   dimensions without changing the box center. The ratios of expansion along
+  //   X, Y dimensions are independent.
+  // SHRINK_SYMMETRIC - Shrinks the boxes symmetrically along height and width
+  //   dimensions without changing the box center. The ratios of shrinking along
+  //   X, Y dimensions are independent.
+  // EXPAND_SYMMETRIC_XY - Expands the boxes symetrically along height and
+  //   width dimensions and the ratio of expansion is same for both.
+  // SHRINK_SYMMETRIC_XY - Shrinks the boxes symetrically along height and
+  //   width dimensions and the ratio of shrinking is same for both.
+  // DEFAULT - Jitters each box boundary independently.
   optional JitterMode jitter_mode = 2 [default = DEFAULT];
 }