Add Solarize transform from PIL

albumentations/augmentations/functional.py

@@ -308,6 +308,24 @@ def shift_hsv(img, hue_shift, sat_shift, val_shift):
     return img
 
 
+def solarize(img: np.ndarray, threshold=128):
+    """Invert all pixel values above a threshold.
+
+    Args:
+        img: The image to solarize.
+        threshold: All pixels above this greyscale level are inverted.
+
+    Returns:
+        Solarized image.
+
+    """
+    dtype = img.dtype
+    max_val = MAX_VALUES_BY_DTYPE[dtype]
+
+    lut = [(i if i < threshold else max_val - i) for i in range(max_val + 1)]
+    return cv2.LUT(img, np.array(lut, dtype=dtype))
+
+
 @clipped
 def shift_rgb(img, r_shift, g_shift, b_shift):
     if img.dtype == np.uint8:

albumentations/augmentations/transforms.py

@@ -25,7 +25,7 @@
     'Resize', 'RandomSizedCrop', 'RandomBrightnessContrast',
     'RandomCropNearBBox', 'RandomSizedBBoxSafeCrop', 'RandomSnow',
     'RandomRain', 'RandomFog', 'RandomSunFlare', 'RandomShadow', 'Lambda',
-    'ChannelDropout', 'ISONoise'
+    'ChannelDropout', 'ISONoise', 'Solarize'
 ]
 
 
@@ -1554,6 +1554,35 @@ def get_transform_init_args_names(self):
         return ('hue_shift_limit', 'sat_shift_limit', 'val_shift_limit')
 
 
+class Solarize(ImageOnlyTransform):
+    """Invert all pixel values above a threshold.
+
+        Args:
+            threshold ((int, int) or int): range for solarizing threshold. If threshold is a single int, the range
+                will be (-threshold, threshold). Default: 128.
+            p (float): probability of applying the transform. Default: 0.5.
+
+        Targets:
+        image
+
+        Image types:
+            uint8, uint16, uint32
+
+        """
+    def __init__(self, threshold=128, always_apply=False, p=0.5):
+        super(Solarize, self).__init__(always_apply, p)
+        self.threshold = to_tuple(threshold)
+
+    def apply(self, image, threshold=0, **params):
+        return F.solarize(image, threshold)
+
+    def get_params(self):
+        return {'threshold': random.uniform(self.threshold[0], self.threshold[1])}
+
+    def get_transform_init_args_names(self):
+        return ('threshold', )
+
+
 class RGBShift(ImageOnlyTransform):
     """Randomly shift values for each channel of the input RGB image.