Added 1 channel support to CLAHE (albumentations-team#190)

* Added 1 channel support to CLAHE

* PEP8

albumentations/augmentations/functional.py

@@ -1,8 +1,8 @@
 from __future__ import division
 
 import math
-from functools import wraps
 import random
+from functools import wraps
 from warnings import warn
 
 import cv2
@@ -297,10 +297,16 @@ def shift_rgb(img, r_shift, g_shift, b_shift):
 def clahe(img, clip_limit=2.0, tile_grid_size=(8, 8)):
     if img.dtype != np.uint8:
         raise TypeError('clahe supports only uint8 inputs')
-    img = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
+
     clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
-    img[:, :, 0] = clahe.apply(img[:, :, 0])
-    img = cv2.cvtColor(img, cv2.COLOR_LAB2RGB)
+
+    if len(img.shape) == 2:
+        img = clahe.apply(img)
+    else:
+        img = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
+        img[:, :, 0] = clahe.apply(img[:, :, 0])
+        img = cv2.cvtColor(img, cv2.COLOR_LAB2RGB)
+
     return img