fix linearwrap + tf function

examples/GAN/Image2Image.py

@@ -155,6 +155,8 @@ def split_input(img):
 
 def colorization_input(img):
     assert img.ndim == 3
+    if min(img.shape[:2]) < SHAPE:
+        return None   # skip the image
     # create gray + RGB pairs
     gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)[:, :, np.newaxis]
     return [gray, img]
@@ -162,7 +164,6 @@ def colorization_input(img):
 
 def get_data():
     datadir = args.data
-    # assume each image is 512x256 split to left and right
     imgs = glob.glob(os.path.join(datadir, '*.jpg'))
     ds = ImageFromFile(imgs, channel=3, shuffle=True)
 

tensorpack/dataflow/image.py

@@ -7,6 +7,7 @@ import cv2
 from .base import RNGDataFlow
 from .common import MapDataComponent, MapData
 from .imgaug import AugmentorList
+from ..utils import logger
 
 __all__ = ['ImageFromFile', 'AugmentImageComponent', 'AugmentImageComponents']
 
@@ -81,14 +82,23 @@ class AugmentImageComponents(MapData):
         """
         self.augs = AugmentorList(augmentors)
         self.ds = ds
+        self._nr_error = 0
 
         def func(dp):
-            im = dp[index[0]]
-            im, prms = self.augs._augment_return_params(im)
-            dp[index[0]] = im
-            for idx in index[1:]:
-                dp[idx] = self.augs._augment(dp[idx], prms)
-            return dp
+            try:
+                im = dp[index[0]]
+                im, prms = self.augs._augment_return_params(im)
+                dp[index[0]] = im
+                for idx in index[1:]:
+                    dp[idx] = self.augs._augment(dp[idx], prms)
+                return dp
+            except KeyboardInterrupt:
+                raise
+            except Exception:
+                self._nr_error += 1
+                if self._nr_error % 1000 == 0:
+                    logger.warn("Got {} augmentation errors.".format(self._nr_error))
+                return None
 
         super(AugmentImageComponents, self).__init__(ds, func)
 

tensorpack/models/linearwrap.py

@@ -60,7 +60,9 @@ class LinearWrap(object):
                 logger.warn(
                     "You're calling LinearWrap.__getattr__ with {}:"
                     " neither a layer nor 'tf'!".format(layer_name))
-            assert isinstance(layer, ModuleType)
+            import tensorflow as tf  # noqa
+            layer = eval(layer_name)
+            assert isinstance(layer, ModuleType), layer
             return LinearWrap._TFModuleFunc(layer, self._t)
 
     def apply(self, func, *args, **kwargs):