Change to_colorspace default in UniformColorQuant. to None

imgaug/augmenters/color.py

@@ -1824,7 +1824,9 @@ class KMeansColorQuantization(_AbstractColorQuantization):
     >>> aug = iaa.KMeansColorQuantization(
     >>>     from_colorspace=iaa.ChangeColorspace.BGR)
 
-    Creates an augmenter that quantizes images that are in ``BGR`` colorspace.
+    Creates an augmenter that quantizes input images that are in
+    ``BGR`` colorspace. The quantization happens in ``RGB`` or ``Lab``
+    colorspace into which the images are temporarily converted.
 
     >>> aug = iaa.KMeansColorQuantization(
     >>>     to_colorspace=[iaa.ChangeColorspace.RGB, iaa.ChangeColorspace.HSV])
@@ -2041,8 +2043,6 @@ class UniformColorQuantization(_AbstractColorQuantization):
 
     Creates an augmenter to apply uniform color quantization to images using a
     random amount of colors, sampled uniformly from the interval ``[2..16]``.
-    It assumes the input image colorspace to be ``RGB`` and clusters colors
-    randomly in ``RGB`` or ``Lab`` colorspace.
 
     >>> aug = iaa.UniformColorQuantization(n_colors=8)
 
@@ -2055,22 +2055,21 @@ class UniformColorQuantization(_AbstractColorQuantization):
     ``[4, 32]``.
 
     >>> aug = iaa.UniformColorQuantization(
-    >>>     from_colorspace=iaa.ChangeColorspace.BGR)
-
-    Creates an augmenter that quantizes images that are in ``BGR`` colorspace.
-
-    >>> aug = iaa.UniformColorQuantization(
+    >>>     from_colorspace=iaa.ChangeColorspace.BGR,
     >>>     to_colorspace=[iaa.ChangeColorspace.RGB, iaa.ChangeColorspace.HSV])
 
-    Creates an augmenter that quantizes images by clustering colors randomly
-    in either ``RGB`` or ``HSV`` colorspace. The assumed input colorspace
-    of images is ``RGB``.
+    Creates an augmenter that uniformly quantizes images in either ``RGB``
+    or ``HSV`` colorspace (randomly picked per image). The input colorspace
+    of all images has to be ``BGR``.
 
     """
 
-    def __init__(self, n_colors=(2, 16), from_colorspace=ChangeColorspace.RGB,
-                 to_colorspace=[ChangeColorspace.RGB, ChangeColorspace.Lab],
-                 max_size=None, interpolation="linear",
+    def __init__(self,
+                 n_colors=(2, 16),
+                 from_colorspace=ChangeColorspace.RGB,
+                 to_colorspace=None,
+                 max_size=None,
+                 interpolation="linear",
                  name=None, deterministic=False, random_state=None):
         # pylint: disable=dangerous-default-value
         super(UniformColorQuantization, self).__init__(

test/augmenters/test_color.py

@@ -1263,9 +1263,7 @@ def test___init___defaults(self):
         assert aug.n_colors.a.value == 2
         assert aug.n_colors.b.value == 16
         assert aug.from_colorspace == iaa.ChangeColorspace.RGB
-        assert isinstance(aug.to_colorspace, list)
-        assert aug.to_colorspace == [iaa.ChangeColorspace.RGB,
-                                     iaa.ChangeColorspace.Lab]
+        assert aug.to_colorspace is None
         assert aug.max_size is None
         assert aug.interpolation == "linear"
 
@@ -1356,6 +1354,30 @@ def test_images_with_4_channels_integrationtest(self):
 
         assert np.array_equal(observed, expected)
 
+    def test_from_colorspace(self):
+        # Actual to_colorspace doesn't matter here as it is overwritten
+        # via return_value. Important is just to set it to a non-None value
+        # so that a colorspace conversion actually happens.
+        aug = self.augmenter(from_colorspace="BGR",
+                             to_colorspace="Lab")
+        mock_change_colorspace = mock.MagicMock()
+        mock_change_colorspace.return_value = mock_change_colorspace
+        mock_change_colorspace.augment_image.side_effect = lambda img: img
+        mock_change_colorspace._draw_samples.return_value = (None, ["foo"])
+
+        fname = "imgaug.augmenters.color.ChangeColorspace"
+        with mock.patch(fname, mock_change_colorspace):
+            _ = aug.augment_image(np.zeros((4, 4, 3), dtype=np.uint8))
+
+        # call 0, kwargs, argument 'from_colorspace'
+        assert (
+            mock_change_colorspace.call_args_list[0][1]["from_colorspace"]
+            == "BGR")
+        # call 1, kwargs, argument 'from_colorspace' (inverse transform)
+        assert (
+            mock_change_colorspace.call_args_list[1][1]["from_colorspace"]
+            == "foo")
+
 
 class Test_quantize_colors_uniform(unittest.TestCase):
     def setUp(self):