__repr__ for augmentors (fix #388)

tensorpack/dataflow/imgaug/base.py

@@ -2,6 +2,8 @@
 # File: base.py
 # Author: Yuxin Wu <ppwwyyxx@gmail.com>
 
+import inspect
+import pprint
 from abc import abstractmethod, ABCMeta
 from ...utils.utils import get_rng
 import six
@@ -64,6 +66,31 @@ class Augmentor(object):
             size = []
         return self.rng.uniform(low, high, size)
 
+    def __repr__(self):
+        """
+        Produce something like:
+        "imgaug.MyAugmentor(field1={self.field1}, field2={self.field2})"
+        """
+        argspec = inspect.getargspec(self.__init__)
+        assert argspec.varargs is None, "The default __repr__ doesn't work for vaargs!"
+        assert argspec.keywords is None, "The default __repr__ doesn't work for kwargs!"
+        fields = argspec.args[1:]
+        index_field_has_default = len(fields) - (0 if argspec.defaults is None else len(argspec.defaults))
+
+        classname = type(self).__name__
+        argstr = []
+        for idx, f in enumerate(fields):
+            assert hasattr(self, f), \
+                "Attribute {} not found! The default __repr__ only works if attributes match the constructor.".format(f)
+            attr = getattr(self, f)
+            if idx >= index_field_has_default:
+                if attr is argspec.defaults[idx - index_field_has_default]:
+                    continue
+            argstr.append("{}={}".format(f, pprint.pformat(attr)))
+        return "imgaug.{}({})".format(classname, ', '.join(argstr))
+
+    __str__ = __repr__
+
 
 class ImageAugmentor(Augmentor):
     def _augment_coords(self, coords, param):

tensorpack/dataflow/imgaug/convert.py

@@ -12,6 +12,7 @@ __all__ = ['ColorSpace', 'Grayscale', 'ToUint8', 'ToFloat32']
 
 class ColorSpace(ImageAugmentor):
     """ Convert into another colorspace.  """
+
     def __init__(self, mode, keepdims=True):
         """
         Args:

tensorpack/dataflow/imgaug/deform.py

@@ -90,6 +90,7 @@ class GaussianDeform(ImageAugmentor):
             self.randrange = self.shape[0] / 8
         else:
             self.randrange = randrange
+        self.sigma = sigma
 
     def _get_augment_params(self, img):
         v = self.rng.rand(self.K, 2).astype('float32') - 0.5

tensorpack/dataflow/imgaug/meta.py

@@ -148,6 +148,7 @@ class MapImage(ImageAugmentor):
         Args:
             func: a function which takes an image array and return an augmented one
         """
+        super(MapImage, self).__init__()
         self.func = func
         self.coord_func = coord_func
 

tensorpack/dataflow/imgaug/misc.py

@@ -33,8 +33,7 @@ class Flip(ImageAugmentor):
             self.code = 0
         else:
             raise ValueError("At least one of horiz or vert has to be True!")
-        self.prob = prob
-        self._init()
+        self._init(locals())
 
     def _get_augment_params(self, img):
         h, w = img.shape[:2]