Repository: albumentations-team/albumentations

Issue:
Failed to serialize ImageCompression class with albumentations.core.serialization.save
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

Steps to reproduce the behavior:

1. Wrap ImageCompression class with Compose
2. Save the compose in yaml format using albumentations.core.serialization.save
3. An error happens

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->
Traceback (most recent call last):
  File "augmentation_yaml_creator.py", line 31, in <module>
    save(train_compose, '..' + os.sep + 'configs' + os.sep + 'train_augmentation.yaml', data_format='yaml')
  File "lib\site-packages\albumentations\core\serialization.py", line 128, in save
    dump_fn(transform_dict, f)
  File "lib\site-packages\yaml\__init__.py", line 306, in safe_dump
    return dump_all([data], stream, Dumper=SafeDumper, **kwds)
  File "lib\site-packages\yaml\__init__.py", line 278, in dump_all
    dumper.represent(data)
  File "lib\site-packages\yaml\representer.py", line 27, in represent
    node = self.represent_data(data)
  File "lib\site-packages\yaml\representer.py", line 48, in represent_data
    node = self.yaml_representers[data_types[0]](self, data)
  File "lib\site-packages\yaml\representer.py", line 207, in represent_dict
    return self.represent_mapping('tag:yaml.org,2002:map', data)
  File "lib\site-packages\yaml\representer.py", line 118, in represent_mapping
    node_value = self.represent_data(item_value)
  File "lib\site-packages\yaml\representer.py", line 48, in represent_data
    node = self.yaml_representers[data_types[0]](self, data)
  File "lib\site-packages\yaml\representer.py", line 207, in represent_dict
    return self.represent_mapping('tag:yaml.org,2002:map', data)
  File "lib\site-packages\yaml\representer.py", line 118, in represent_mapping
    node_value = self.represent_data(item_value)
  File "lib\site-packages\yaml\representer.py", line 48, in represent_data
    node = self.yaml_representers[data_types[0]](self, data)
  File "lib\site-packages\yaml\representer.py", line 199, in represent_list
    return self.represent_sequence('tag:yaml.org,2002:seq', data)
  File "lib\site-packages\yaml\representer.py", line 92, in represent_sequence
    node_item = self.represent_data(item)
  File "lib\site-packages\yaml\representer.py", line 48, in represent_data
    node = self.yaml_representers[data_types[0]](self, data)
  File "lib\site-packages\yaml\representer.py", line 207, in represent_dict
    return self.represent_mapping('tag:yaml.org,2002:map', data)
  File "lib\site-packages\yaml\representer.py", line 118, in represent_mapping
    node_value = self.represent_data(item_value)
  File "lib\site-packages\yaml\representer.py", line 58, in represent_data
    node = self.yaml_representers[None](self, data)
  File "lib\site-packages\yaml\representer.py", line 231, in represent_undefined
    raise RepresenterError("cannot represent an object", data)
yaml.representer.RepresenterError: ('cannot represent an object', <ImageCompressionType.JPEG: 0>)

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
yaml file which contains augmentation infi is created.
This error does not happen if I use JPEGCompression instead of ImageCompression.

## Environment

 - Albumentations version (e.g., 0.1.8): 0.4.3
 - Python version (e.g., 3.7): 3.7
 - OS (e.g., Linux): Windows10
 - How you installed albumentations (`conda`, `pip`, source): conda

## Code to reproduce the error
```
from albumentations import Compose
from albumentations.core.serialization import save
from albumentations.augmentations.transforms import ImageCompression

train_compose = Compose([
        ImageCompression(quality_lower=50, quality_upper=50, p=1.0),
    ])
save(train_compose, 'test.yaml', data_format='yaml')
```


Discussion:

Patch:
diff --git a/albumentations/augmentations/transforms.py b/albumentations/augmentations/transforms.py
index cf69dc0..869d725 100644
--- a/albumentations/augmentations/transforms.py
+++ b/albumentations/augmentations/transforms.py
@@ -3,12 +3,12 @@ from __future__ import absolute_import, division
 import math
 import random
 import warnings
-from enum import Enum
+from enum import IntEnum
 from types import LambdaType
-from skimage.measure import label
 
 import cv2
 import numpy as np
+from skimage.measure import label
 
 from . import functional as F
 from .bbox_utils import denormalize_bbox, normalize_bbox, union_of_bboxes
@@ -1587,7 +1587,7 @@ class ImageCompression(ImageOnlyTransform):
         quality_upper (float): upper bound on the image quality.
                                Should be in [0, 100] range for jpeg and [1, 100] for webp.
         compression_type (ImageCompressionType): should be ImageCompressionType.JPEG or ImageCompressionType.WEBP.
-            Defaul: ImageCompressionType.JPEG
+            Default: ImageCompressionType.JPEG
 
     Targets:
         image
@@ -1596,7 +1596,7 @@ class ImageCompression(ImageOnlyTransform):
         uint8, float32
     """
 
-    class ImageCompressionType(Enum):
+    class ImageCompressionType(IntEnum):
         JPEG = 0
         WEBP = 1
 
@@ -1610,7 +1610,7 @@ class ImageCompression(ImageOnlyTransform):
     ):
         super(ImageCompression, self).__init__(always_apply, p)
 
-        self.compression_type = compression_type
+        self.compression_type = ImageCompression.ImageCompressionType(compression_type)
         low_thresh_quality_assert = 0
 
         if self.compression_type == ImageCompression.ImageCompressionType.WEBP:
@@ -1633,8 +1633,12 @@ class ImageCompression(ImageOnlyTransform):
 
         return {"quality": random.randint(self.quality_lower, self.quality_upper), "image_type": image_type}
 
-    def get_transform_init_args_names(self):
-        return ("quality_lower", "quality_upper", "compression_type")
+    def get_transform_init_args(self):
+        return {
+            "quality_lower": self.quality_lower,
+            "quality_upper": self.quality_upper,
+            "compression_type": self.compression_type.value,
+        }
 
 
 class JpegCompression(ImageCompression):
@@ -2496,7 +2500,7 @@ class GaussNoise(ImageOnlyTransform):
 
     def __init__(self, var_limit=(10.0, 50.0), mean=0, always_apply=False, p=0.5):
         super(GaussNoise, self).__init__(always_apply, p)
-        if isinstance(var_limit, tuple):
+        if isinstance(var_limit, (tuple, list)):
             if var_limit[0] < 0:
                 raise ValueError("Lower var_limit should be non negative.")
             if var_limit[1] < 0:
@@ -2504,9 +2508,13 @@ class GaussNoise(ImageOnlyTransform):
             self.var_limit = var_limit
         elif isinstance(var_limit, (int, float)):
             if var_limit < 0:
-                raise ValueError(" var_limit should be non negative.")
+                raise ValueError("var_limit should be non negative.")
 
             self.var_limit = (0, var_limit)
+        else:
+            raise TypeError(
+                "Expected var_limit type to be one of (int, float, tuple, list), got {}".format(type(var_limit))
+            )
 
         self.mean = mean
 


Test patch:
diff --git a/tests/test_serialization.py b/tests/test_serialization.py
index 697172c..6440630 100644
--- a/tests/test_serialization.py
+++ b/tests/test_serialization.py
@@ -1,4 +1,5 @@
 import random
+from unittest.mock import patch
 
 import cv2
 import pytest
@@ -7,7 +8,7 @@ import imgaug as ia
 
 import albumentations as A
 import albumentations.augmentations.functional as F
-
+from .utils import OpenMock
 
 TEST_SEEDS = (0, 1, 42, 111, 9999)
 
@@ -83,8 +84,7 @@ def test_augmentations_serialization(augmentation_cls, params, p, seed, image, m
     assert np.array_equal(aug_data["mask"], deserialized_aug_data["mask"])
 
 
-@pytest.mark.parametrize(
-    ["augmentation_cls", "params"],
+AUGMENTATION_CLS_PARAMS = (
     [
         [
             A.ImageCompression,
@@ -218,6 +218,9 @@ def test_augmentations_serialization(augmentation_cls, params, p, seed, image, m
         [A.MultiplicativeNoise, {"multiplier": (0.7, 2.3), "per_channel": True, "elementwise": True}],
     ],
 )
+
+
+@pytest.mark.parametrize(["augmentation_cls", "params"], *AUGMENTATION_CLS_PARAMS)
 @pytest.mark.parametrize("p", [0.5, 1])
 @pytest.mark.parametrize("seed", TEST_SEEDS)
 @pytest.mark.parametrize("always_apply", (False, True))
@@ -235,6 +238,27 @@ def test_augmentations_serialization_with_custom_parameters(
     assert np.array_equal(aug_data["mask"], deserialized_aug_data["mask"])
 
 
+@pytest.mark.parametrize(["augmentation_cls", "params"], *AUGMENTATION_CLS_PARAMS)
+@pytest.mark.parametrize("p", [0.5, 1])
+@pytest.mark.parametrize("seed", TEST_SEEDS)
+@pytest.mark.parametrize("always_apply", (False, True))
+@pytest.mark.parametrize("data_format", ("yaml",))
+def test_augmentations_serialization_to_file_with_custom_parameters(
+    augmentation_cls, params, p, seed, image, mask, always_apply, data_format
+):
+    with patch("builtins.open", OpenMock()):
+        aug = augmentation_cls(p=p, always_apply=always_apply, **params)
+        filepath = "serialized.{}".format(data_format)
+        A.save(aug, filepath, data_format=data_format)
+        deserialized_aug = A.load(filepath, data_format=data_format)
+        set_seed(seed)
+        aug_data = aug(image=image, mask=mask)
+        set_seed(seed)
+        deserialized_aug_data = deserialized_aug(image=image, mask=mask)
+        assert np.array_equal(aug_data["image"], deserialized_aug_data["image"])
+        assert np.array_equal(aug_data["mask"], deserialized_aug_data["mask"])
+
+
 @pytest.mark.parametrize(
     ["augmentation_cls", "params"],
     [
diff --git a/tests/test_transforms.py b/tests/test_transforms.py
index 340cc3d..f4f24c4 100644
--- a/tests/test_transforms.py
+++ b/tests/test_transforms.py
@@ -645,3 +645,10 @@ def test_grid_dropout_params(ratio, holes_number_x, holes_number_y, unit_size_mi
     elif holes_number_x and holes_number_y:
         assert (holes[0][2] - holes[0][0]) == max(1, int(ratio * 320 // holes_number_x))
         assert (holes[0][3] - holes[0][1]) == max(1, int(ratio * 256 // holes_number_y))
+
+
+def test_gauss_noise_incorrect_var_limit_type():
+    with pytest.raises(TypeError) as exc_info:
+        A.GaussNoise(var_limit={"low": 70, "high": 90})
+    message = "Expected var_limit type to be one of (int, float, tuple, list), got <class 'dict'>"
+    assert str(exc_info.value) == message
diff --git a/tests/utils.py b/tests/utils.py
index 3030c45..868cbc1 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,3 +1,5 @@
+from io import StringIO
+
 import numpy as np
 
 
@@ -19,3 +21,32 @@ def convert_2d_to_target_format(arrays, target):
         return convert_2d_to_3d(arrays, num_channels=4)
     else:
         raise ValueError("Unknown target {}".format(target))
+
+
+class InMemoryFile(StringIO):
+    def __init__(self, value, save_value, file):
+        super().__init__(value)
+        self.save_value = save_value
+        self.file = file
+
+    def close(self):
+        self.save_value(self.getvalue(), self.file)
+        super().close()
+
+
+class OpenMock:
+    """
+    Mocks the `open` built-in function. A call to the instance of OpenMock returns an in-memory file which is
+    readable and writable. The actual in-memory file implementation should call the passed `save_value` method
+    to save the file content in the cache when the file is being closed to preserve the file content.
+    """
+
+    def __init__(self):
+        self.values = {}
+
+    def __call__(self, file, *args, **kwargs):
+        value = self.values.get(file)
+        return InMemoryFile(value, self.save_value, file)
+
+    def save_value(self, value, file):
+        self.values[file] = value