Wrap aleju/imgaug (#399)

README.md

@@ -30,7 +30,7 @@ It's Yet Another TF high-level API, with __speed__, __readability__ and __flexib
 	+ There are too many symbolic function wrappers in the world. Tensorpack includes only a few common models.
 	  But you can use any symbolic function library inside Tensorpack, including tf.layers/Keras/slim/tflearn/tensorlayer/....
 
-See [tutorials](http://tensorpack.readthedocs.io/tutorial/index.html#user-tutorials) to know more about these features.
+See [tutorials and documentations](http://tensorpack.readthedocs.io/tutorial/index.html#user-tutorials) to know more about these features.
 
 ## [Examples](examples):
 

docs/conf.py

@@ -30,7 +30,21 @@ MOCK_MODULES = ['tabulate', 'h5py',
                 'scipy', 'scipy.misc', 'scipy.io',
                 'tornado', 'tornado.concurrent',
                 'horovod', 'horovod.tensorflow',
-                'subprocess32', 'functools32']
+                'subprocess32', 'functools32',
+                'imgaug']
+
+# it's better to have tensorflow installed (for some docs to show)
+# but it's OK to mock it as well
+try:
+    import tensorflow
+except ImportError:
+    mod = sys.modules['tensorflow'] = mock.Mock(name='tensorflow')
+    mod.__version__ = mod.VERSION = '1.12'
+    MOCK_MODULES.extend(['tensorflow.python.training.monitored_session'])
+    MOCK_MODULES.extend(['tensorflow.python.training'])
+    MOCK_MODULES.extend(['tensorflow.python.client'])
+    MOCK_MODULES.extend(['tensorflow.contrib.graph_editor'])
+
 for mod_name in MOCK_MODULES:
     sys.modules[mod_name] = mock.Mock(name=mod_name)
 sys.modules['cv2'].__version__ = '3.2.1'    # fake version

docs/modules/dataflow.imgaug.rst

 tensorpack.dataflow.imgaug package
 ==================================
 
+This package contains Tensorpack's augmentors.
+Note that other image augmentation libraries can be wrapped into Tensorpack's interface as well.
+For example, see `imgaug.IAAugmentor <#tensorpack.dataflow.imgaug.IAAugmentor>`_.
+
 .. container:: custom-index
 
     .. raw:: html

docs/tutorial/index.rst

@@ -22,7 +22,6 @@ User Tutorials
   save-load
   summary
   inference
-  export
   faq
 
 

tensorpack/dataflow/imgaug/external.py

+#!/usr/bin/env python
+
+import numpy as np
+
+from .base import ImageAugmentor
+
+
+__all__ = ['IAAugmentor']
+
+
+class IAAugmentor(ImageAugmentor):
+    """
+    Wrap an augmentor form the IAA library: https://github.com/aleju/imgaug
+    Both images and coordinates are supported.
+
+    Note:
+        1. It's NOT RECOMMENDED
+           to use coordinates because the IAA library does not handle coordinates accurately.
+
+        2. Only uint8 images are supported by the IAA library.
+
+        3. The IAA library can only produces images of the same shape.
+    """
+
+    def __init__(self, augmentor):
+        """
+        Args:
+            augmentor (iaa.Augmenter):
+        """
+        super(IAAugmentor, self).__init__()
+        self._aug = augmentor
+
+    def _get_augment_params(self, img):
+        return (self._aug.to_deterministic(), img.shape)
+
+    def _augment(self, img, p):
+        aug, _ = p
+        return aug.augment_image(img)
+
+    def _augment_coords(self, coords, p):
+        aug, shape = p
+        points = [IA.Keypoint(x=x, y=y) for x, y in coords]
+        points = IA.KeypointsOnImage(points, shape=shape)
+        augmented = aug.augment_keypoints([points])[0].keypoints
+        return np.asarray([[p.x, p.y] for p in augmented])
+
+
+from ...utils.develop import create_dummy_class   # noqa
+try:
+    import imgaug as IA
+except ImportError:
+    IAAugmentor = create_dummy_class('IAAugmentor', 'imgaug')  # noqa