Add ToRGB (#1323)

* Add ToRGB * Update README.md Co-authored-by: Mikhail Druzhinin <dipetm@gmail.com>
albumentations-team · Oct 25, 2022 · 0c6f13b · 0c6f13b
1 parent a51a641
commit 0c6f13b
Show file tree

Hide file tree

Showing 5 changed files with 42 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -166,6 +166,7 @@ Pixel-level transforms will change just an input image and will leave any additi
 - [TemplateTransform](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.TemplateTransform)
 - [ToFloat](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.ToFloat)
 - [ToGray](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.ToGray)
+- [ToRGB](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.ToRGB)
 - [ToSepia](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.ToSepia)
 - [UnsharpMask](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.UnsharpMask)
 - [ZoomBlur](https://albumentations.ai/docs/api_reference/augmentations/blur/transforms/#albumentations.augmentations.blur.transforms.ZoomBlur)

diff --git a/albumentations/augmentations/functional.py b/albumentations/augmentations/functional.py
@@ -58,6 +58,7 @@
     "swap_tiles_on_image",
     "to_float",
     "to_gray",
+    "gray_to_rgb",
     "unsharp_mask",
 ]
 
@@ -883,6 +884,10 @@ def to_gray(img):
     return cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
 
 
+def gray_to_rgb(img):
+    return cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
+
+
 @preserve_shape
 def downscale(img, scale, down_interpolation=cv2.INTER_AREA, up_interpolation=cv2.INTER_LINEAR):
     h, w = img.shape[:2]

diff --git a/albumentations/augmentations/transforms.py b/albumentations/augmentations/transforms.py
@@ -44,6 +44,7 @@
     "ChannelShuffle",
     "InvertImg",
     "ToGray",
+    "ToRGB",
     "ToSepia",
     "JpegCompression",
     "ImageCompression",
@@ -1422,6 +1423,35 @@ def get_transform_init_args_names(self):
         return ()
 
 
+class ToRGB(ImageOnlyTransform):
+    """Convert the input grayscale image to RGB.
+
+    Args:
+        p (float): probability of applying the transform. Default: 1.
+
+    Targets:
+        image
+
+    Image types:
+        uint8, float32
+    """
+
+    def __init__(self, always_apply=True, p=1.0):
+        super(ToRGB, self).__init__(always_apply=always_apply, p=p)
+
+    def apply(self, img, **params):
+        if is_rgb_image(img):
+            warnings.warn("The image is already an RGB.")
+            return img
+        if not is_grayscale_image(img):
+            raise TypeError("ToRGB transformation expects 2-dim images or 3-dim with the last dimension equal to 1.")
+
+        return F.gray_to_rgb(img)
+
+    def get_transform_init_args_names(self):
+        return ()
+
+
 class ToSepia(ImageOnlyTransform):
     """Applies sepia filter to the input RGB image
 

diff --git a/tests/test_augmentations.py b/tests/test_augmentations.py
@@ -266,6 +266,7 @@ def test_augmentations_wont_change_float_input(augmentation_cls, params, float_i
             A.RandomShadow,
             A.RandomSnow,
             A.RandomSunFlare,
+            A.ToRGB,
             A.ToSepia,
             A.PixelDistributionAdaptation,
             A.UnsharpMask,
@@ -424,6 +425,7 @@ def test_mask_fill_value(augmentation_cls, params):
             A.RandomSunFlare,
             A.ToFloat,
             A.ToGray,
+            A.ToRGB,
             A.ToSepia,
             A.FancyPCA,
             A.PixelDistributionAdaptation,
@@ -483,6 +485,7 @@ def test_multichannel_image_augmentations(augmentation_cls, params):
             A.RandomSnow,
             A.RandomSunFlare,
             A.ToGray,
+            A.ToRGB,
             A.ToSepia,
             A.Equalize,
             A.FancyPCA,
@@ -537,6 +540,7 @@ def test_float_multichannel_image_augmentations(augmentation_cls, params):
             A.RandomSunFlare,
             A.ToFloat,
             A.ToGray,
+            A.ToRGB,
             A.ToSepia,
             A.FancyPCA,
             A.FDA,
@@ -591,6 +595,7 @@ def test_multichannel_image_augmentations_diff_channels(augmentation_cls, params
             A.RandomSnow,
             A.RandomSunFlare,
             A.ToGray,
+            A.ToRGB,
             A.ToSepia,
             A.Equalize,
             A.FancyPCA,

diff --git a/tests/test_serialization.py b/tests/test_serialization.py
@@ -312,6 +312,7 @@ def test_augmentations_serialization(augmentation_cls, params, p, seed, image, m
     [A.FancyPCA, dict(alpha=0.3)],
     [A.RandomRotate90, {}],
     [A.ToGray, {}],
+    [A.ToRGB, {}],
     [A.ToSepia, {}],
     [A.Transpose, {}],
     [A.VerticalFlip, {}],