lighting augmentation

tensorpack/dataflow/imgaug/imgproc.py

@@ -118,8 +118,9 @@ class Clip(ImageAugmentor):
         return img
 
 class Saturation(ImageAugmentor):
-    """ Saturation, see 'fb.resnet.torch' https://github.com/facebook/fb.resnet.torch/blob/master/datasets/transforms.lua#L218"""
     def __init__(self, alpha=0.4):
+        """ Saturation, see 'fb.resnet.torch' https://github.com/facebook/fb.resnet.torch/blob/master/datasets/transforms.lua#L218
+        """
         super(Saturation, self).__init__()
         assert alpha < 1
         self._init(locals())
@@ -130,3 +131,29 @@ class Saturation(ImageAugmentor):
     def _augment(self, img, v):
         grey = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
         return img * v + (grey * (1 - v))[:,:,np.newaxis]
+
+class Lighting(ImageAugmentor):
+    def __init__(self, std, eigval, eigvec):
+    """ Lighting noise.
+        See `ImageNet Classification with Deep Convolutional Neural Networks - Alex`
+        The implementation follows 'fb.resnet.torch': https://github.com/facebook/fb.resnet.torch/blob/master/datasets/transforms.lua#L184
+
+        :param eigvec: each column is one eigen vector
+    """
+        eigval = np.asarray(eigval)
+        eigvec = np.asarray(eigvec)
+        assert eigval.shape == (3,)
+        assert eigvec.shape == (3,3)
+        self._init(locals())
+
+    def _get_augment_params(self, img):
+        assert img.shape[2] == 3
+        return self.rng.randn(3) * self.std
+
+    def _augment(self, img, v):
+        v = v * self.eigval
+        v = v.reshape((3, 1))
+        inc = np.dot(self.eigvec, v).reshape((3,))
+        img += inc
+        return img
+