Merge pull request #393 from google/sync

Sync github with internal code.

Author: mingxingtan

efficientdet/README.md

@@ -81,7 +81,7 @@ Then you will get:
  - frozen graph with name savedmodeldir/efficientdet-d0_frozen.pb
  - tflite file with name efficientdet-d0.tflite
 
-Notably, --tflite_path is optional, and it only works with tensorflow >= 2.2.0-rc4.
+Notably, --tflite_path does not work for now. It requires some extra fixes in future TensorFlow rerlease (> 2.2.0-rc4).
 
 
 ## 4. Benchmark model latency.
@@ -304,7 +304,7 @@ If you want to do inference for custom data, you can run
       --model_name=efficientdet-d0   --ckpt_path=efficientdet-d0 \
       --hparams=voc_config.yaml  \
       --input_image=img.png --output_image_dir=/tmp/
-  
+
 You should check more details of runmode which is written in caption-4.
 
 ## 10. Training EfficientDets on TPUs.

efficientdet/anchors.py

@@ -291,6 +291,7 @@ def _generate_detections_tf(cls_outputs,
     image_scale: a float tensor representing the scale between original image
       and input image for the detector. It is used to rescale detections for
       evaluating with the original groundtruth annotations.
+    image_size: a tuple (height, width) or an integer for image size.
     min_score_thresh: A float representing the threshold for deciding when to
       remove boxes based on score.
     max_boxes_to_draw: Max number of boxes to draw.
@@ -307,7 +308,7 @@ def _generate_detections_tf(cls_outputs,
       [image_id, ymin, xmin, ymax, xmax, score, class]
   """
   if not image_size:
-    raise ValueError('image_size cannot be empty.')
+    raise ValueError('tf version generate_detection needs non-empty image_size')
 
   logging.info('Using tf version of post-processing.')
   anchor_boxes = tf.gather(anchor_boxes, indices)

efficientdet/aug/autoaugment.py

@@ -20,15 +20,17 @@
 
 from __future__ import absolute_import
 from __future__ import division
+# gtype import
 from __future__ import print_function
 
 import inspect
 import math
 from absl import logging
+import numpy as np
 import tensorflow.compat.v1 as tf
 import tensorflow_probability as tfp
+
 import hparams_config
-import numpy as np
 
 try:
   # addon image_ops are simpler, but they have some issues on GPU and TPU.
@@ -1547,7 +1549,13 @@ def select_and_apply_random_policy(policies, image, bboxes):
         lambda: (image, bboxes))
   return (image, bboxes)
 
-def select_and_apply_random_policy_augmix(policies, image, bboxes, mixture_width=3, mixture_depth=-1, alpha=1):
+
+def select_and_apply_random_policy_augmix(policies,
+                                          image,
+                                          bboxes,
+                                          mixture_width=3,
+                                          mixture_depth=-1,
+                                          alpha=1):
   """Select a random policy from `policies` and apply it to `image`."""
   policy_to_select = tf.random_uniform([], maxval=len(policies), dtype=tf.int32)
   # Note that using tf.case instead of tf.conds would result in significantly
@@ -1562,12 +1570,13 @@ def select_and_apply_random_policy_augmix(policies, image, bboxes, mixture_width
       for (i, policy) in enumerate(policies):
         aug_image, bboxes = tf.cond(
             tf.equal(i, policy_to_select),
-            lambda selected_policy=policy: selected_policy(aug_image, bboxes),
-            lambda: (aug_image, bboxes))
+            lambda policy_fn=policy, img=aug_image: policy_fn(img, bboxes),
+            lambda img=aug_image: (img, bboxes))
     mix += ws[j] * tf.cast(aug_image, tf.float32)
   mixed = tf.cast((1 - m) * tf.cast(image, tf.float32) + m * mix, tf.uint8)
   return (mixed, bboxes)
 
+
 def build_and_apply_nas_policy(policies, image, bboxes,
                                augmentation_hparams, use_augmix=False,
                                mixture_width=3, mixture_depth=-1, alpha=1):
@@ -1583,9 +1592,9 @@ def build_and_apply_nas_policy(policies, image, bboxes,
       normalized between [0, 1].
     augmentation_hparams: Hparams associated with the NAS learned policy.
     use_augmix: whether use augmix[https://arxiv.org/pdf/1912.02781.pdf]
-    width: Width of augmentation chain
-    depth: Depth of augmentation chain. -1 enables stochastic depth uniformly
-      from [1, 3]
+    mixture_width: Width of augmentation chain
+    mixture_depth: Depth of augmentation chain. -1 enables stochastic depth
+      uniformly from [1, 3].
     alpha: Probability coefficient for Beta and Dirichlet distributions.
 
   Returns:
@@ -1631,8 +1640,13 @@ def final_policy(image_, bboxes_):
   return (augmented_images, augmented_bboxes)
 
 
-def distort_image_with_autoaugment(image, bboxes, augmentation_name, use_augmix=False,
-                                   mixture_width=3, mixture_depth=-1, alpha=1):
+def distort_image_with_autoaugment(image,
+                                   bboxes,
+                                   augmentation_name,
+                                   use_augmix=False,
+                                   mixture_width=3,
+                                   mixture_depth=-1,
+                                   alpha=1):
   """Applies the AutoAugment policy to `image` and `bboxes`.
 
   Args:
@@ -1646,6 +1660,11 @@ def distort_image_with_autoaugment(image, bboxes, augmentation_name, use_augmix=
       found on the COCO dataset that have slight variation in what operations
       were used during the search procedure along with how many operations are
       applied in parallel to a single image (2 vs 3).
+    use_augmix: whether use augmix[https://arxiv.org/pdf/1912.02781.pdf]
+    mixture_width: Width of augmentation chain
+    mixture_depth: Depth of augmentation chain. -1 enables stochastic depth
+      uniformly from [1, 3].
+    alpha: Probability coefficient for Beta and Dirichlet distributions.
 
   Returns:
     A tuple containing the augmented versions of `image` and `bboxes`.
@@ -1668,4 +1687,5 @@ def distort_image_with_autoaugment(image, bboxes, augmentation_name, use_augmix=
 
   with tf.device('/cpu:0'):
     return build_and_apply_nas_policy(policy, image, bboxes,
-        augmentation_hparams, use_augmix, mixture_width, mixture_depth, alpha)
+                                      augmentation_hparams, use_augmix,
+                                      mixture_width, mixture_depth, alpha)

efficientdet/aug/autoaugment_test.py

@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from absl import logging
 import tensorflow.compat.v1 as tf
 
 from aug import autoaugment
@@ -30,10 +31,12 @@ def test_autoaugment_policy(self):
     image = tf.placeholder(tf.uint8, shape=[640, 640, 3])
     bboxes = tf.placeholder(tf.float32, shape=[4, 4])
     autoaugment.distort_image_with_autoaugment(image, bboxes, 'test')
-    autoaugment.distort_image_with_autoaugment(image, bboxes, 'test', use_augmix=True)
+    autoaugment.distort_image_with_autoaugment(
+        image, bboxes, 'test', use_augmix=True)
 
 
 if __name__ == '__main__':
+  logging.set_verbosity(logging.WARNING)
   tf.disable_eager_execution()
   tf.test.main()
 

efficientdet/backbone/efficientnet_builder_test.py

@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from absl import logging
 import numpy as np
 import tensorflow.compat.v1 as tf
 
@@ -109,6 +110,7 @@ def test_efficientnet_b0_base(self):
 
 
 if __name__ == '__main__':
+  logging.set_verbosity(logging.WARNING)
   # Disable eager to allow tf.profile works for #params/#flops.
   tf.disable_eager_execution()
   tf.test.main()

efficientdet/backbone/efficientnet_lite_builder_test.py

@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from absl import logging
 import numpy as np
 import tensorflow.compat.v1 as tf
 
@@ -67,6 +68,7 @@ def test_efficientnet_b4(self):
 
 
 if __name__ == '__main__':
+  logging.set_verbosity(logging.WARNING)
   # Disable eager to allow tf.profile works for #params/#flops.
   tf.disable_eager_execution()
   tf.test.main()

efficientdet/backbone/efficientnet_model_test.py

@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from absl import logging
 import tensorflow.compat.v1 as tf
 
 import utils
@@ -252,4 +253,5 @@ def test_reduction_endpoint_with_single_block_without_sp(self):
     self.assertNotIn('reduction_2', model.endpoints)
 
 if __name__ == '__main__':
+  logging.set_verbosity(logging.WARNING)
   tf.test.main()

efficientdet/dataloader.py

@@ -277,7 +277,8 @@ def _dataset_parser(value):
         if params.get('autoaugment_policy', None) and self._is_training:
           from aug import autoaugment  # pylint: disable=g-import-not-at-top
           image, boxes = autoaugment.distort_image_with_autoaugment(
-              image, boxes, params['autoaugment_policy'], params['use_augmix'], *params['augmix_params'])
+              image, boxes, params['autoaugment_policy'], params['use_augmix'],
+              *params['augmix_params'])
 
         input_processor = DetectionInputProcessor(
             image, params['image_size'], boxes, classes)

efficientdet/dataset/create_coco_tfrecord.py

@@ -31,10 +31,12 @@
 import hashlib
 import io
 import json
-import logging
 import multiprocessing
 import os
+
+from absl import app
 from absl import flags
+from absl import logging
 import numpy as np
 import PIL.Image
 
@@ -43,7 +45,6 @@
 from dataset import label_map_util
 from dataset import tfrecord_util
 
-
 flags.DEFINE_boolean(
     'include_masks', False, 'Whether to include instance segmentations masks '
     '(PNG encoded) in the result. default: False.')
@@ -61,6 +62,7 @@
                     'captions.')
 flags.DEFINE_string('output_file_prefix', '/tmp/train', 'Path to output file')
 flags.DEFINE_integer('num_shards', 32, 'Number of shards for output file.')
+flags.DEFINE_integer('num_threads', None, 'Number of threads to run.')
 FLAGS = flags.FLAGS
 
 
@@ -285,9 +287,8 @@ def _create_tf_record_from_coco_annotations(images_info_file,
 
   logging.info('writing to output path: %s', output_path)
   writers = [
-      tf.python_io.TFRecordWriter(
-          output_path + '-%05d-of-%05d.tfrecord' % (i, num_shards))
-      for i in range(num_shards)
+      tf.python_io.TFRecordWriter(output_path + '-%05d-of-%05d.tfrecord' %
+                                  (i, num_shards)) for i in range(num_shards)
   ]
   images = _load_images_info(images_info_file)
 
@@ -313,13 +314,14 @@ def _get_caption_annotation(image_id):
     else:
       return None
 
-  pool = multiprocessing.Pool()
+  pool = multiprocessing.Pool(FLAGS.num_threads)
   total_num_annotations_skipped = 0
   for idx, (_, tf_example, num_annotations_skipped) in enumerate(
-      pool.imap(_pool_create_tf_example,
-                [(image, image_dir, _get_object_annotation(image['id']),
-                  category_index, _get_caption_annotation(image['id']),
-                  include_masks) for image in images])):
+      pool.imap(
+          _pool_create_tf_example,
+          [(image, image_dir, _get_object_annotation(image['id']),
+            category_index, _get_caption_annotation(image['id']), include_masks)
+           for image in images])):
     if idx % 100 == 0:
       logging.info('On image %d of %d', idx, len(images))
 
@@ -361,4 +363,4 @@ def main(_):
 
 
 if __name__ == '__main__':
-  tf.app.run(main)
+  app.run(main)

efficientdet/dataset/create_coco_tfrecord_test.py

@@ -18,6 +18,8 @@
 import json
 import os
 
+from absl import flags
+from absl import logging
 import numpy as np
 import PIL.Image
 import six
@@ -28,6 +30,10 @@
 
 class CreateCocoTFRecordTest(tf.test.TestCase):
 
+  def setUp(self):
+    super(CreateCocoTFRecordTest, self).setUp()
+    flags.FLAGS.num_threads = 1
+
   def _assertProtoEqual(self, proto_field, expectation):
     """Helper function to assert if a proto field equals some value.
 
@@ -251,4 +257,5 @@ def test_create_sharded_tf_record(self):
 
 
 if __name__ == '__main__':
+  logging.set_verbosity(logging.WARNING)
   tf.test.main()