shift_hsv optimization (#332)

* shift_hsv optimization

* Split shift_hsv to _shift_hsv_uint8 and _shift_hsv_non_uint8

albumentations/augmentations/functional.py

@@ -291,12 +291,34 @@ def clamping_crop(img, x_min, y_min, x_max, y_max):
     return img[int(y_min):int(y_max), int(x_min):int(x_max)]
 
 
-def shift_hsv(img, hue_shift, sat_shift, val_shift):
+def _shift_hsv_uint8(img, hue_shift, sat_shift, val_shift):
+    dtype = img.dtype
+    img = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
+    hue, sat, val = cv2.split(img)
+
+    lut_hue = np.arange(0, 256, dtype=np.int16)
+    lut_hue = np.mod(lut_hue + hue_shift, 180).astype(dtype)
+
+    lut_sat = np.arange(0, 256, dtype=np.int16)
+    lut_sat = np.clip(lut_sat + sat_shift, 0, 255).astype(dtype)
+
+    lut_val = np.arange(0, 256, dtype=np.int16)
+    lut_val = np.clip(lut_val + val_shift, 0, 255).astype(dtype)
+
+    hue = cv2.LUT(hue, lut_hue)
+    sat = cv2.LUT(sat, lut_sat)
+    val = cv2.LUT(val, lut_val)
+
+    img = cv2.merge((hue, sat, val)).astype(dtype)
+    img = cv2.cvtColor(img, cv2.COLOR_HSV2RGB)
+    return img
+
+
+def _shift_hsv_non_uint8(img, hue_shift, sat_shift, val_shift):
     dtype = img.dtype
     img = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
-    if dtype == np.uint8:
-        img = img.astype(np.int32)
     hue, sat, val = cv2.split(img)
+
     hue = cv2.add(hue, hue_shift)
     hue = np.where(hue < 0, hue + 180, hue)
     hue = np.where(hue > 180, hue - 180, hue)
@@ -308,6 +330,13 @@ def shift_hsv(img, hue_shift, sat_shift, val_shift):
     return img
 
 
+def shift_hsv(img, hue_shift, sat_shift, val_shift):
+    if img.dtype == np.uint8:
+        return _shift_hsv_uint8(img, hue_shift, sat_shift, val_shift)
+
+    return _shift_hsv_non_uint8(img, hue_shift, sat_shift, val_shift)
+
+
 def solarize(img, threshold=128):
     """Invert all pixel values above a threshold.