Fixes

albumentations-team · Mar 4, 2024 · c2f2449 · c2f2449
1 parent bddf9c8
commit c2f2449
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 12 deletions.
diff --git a/albumentations/augmentations/mixing/transforms.py b/albumentations/augmentations/mixing/transforms.py
@@ -1,5 +1,6 @@
 import random
-from typing import Any, Callable, Dict, Generator, Iterable, Iterator, Optional, Sequence, Tuple, Union
+import types
+from typing import Any, Callable, Dict, Generator, Iterable, Iterator, List, Optional, Sequence, Tuple, Union
 from warnings import warn
 
 import numpy as np
@@ -61,8 +62,8 @@ class MixUp(ReferenceBasedTransform):
 
     def __init__(
         self,
-        reference_data: Optional[Union[Generator[ReferenceImage, None, None], Sequence[ReferenceImage]]] = None,
-        read_fn: Callable[[ReferenceImage], Dict[str, Any]] = lambda x: {"image": x, "mask": None, "class_label": None},
+        reference_data: Optional[Union[Generator[ReferenceImage, None, None], Sequence[Any]]] = None,
+        read_fn: Callable[[ReferenceImage], Any] = lambda x: {"image": x, "mask": None, "class_label": None},
         alpha: float = 0.4,
         always_apply: bool = False,
         p: float = 0.5,
@@ -79,8 +80,13 @@ def __init__(
         if reference_data is None:
             warn("No reference data provided for MixUp. This transform will act as a no-op.")
             # Create an empty generator
-        elif isinstance(reference_data, Iterable) and not isinstance(reference_data, str):
-            self.reference_data = reference_data
+            self.reference_data: List[Any] = []
+        elif (
+            isinstance(reference_data, types.GeneratorType)
+            or isinstance(reference_data, Iterable)
+            and not isinstance(reference_data, str)
+        ):
+            self.reference_data = reference_data  # type: ignore[assignment]
         else:
             msg = "reference_data must be a list, tuple, generator, or None."
             raise TypeError(msg)
@@ -132,7 +138,9 @@ def get_params(self) -> Dict[str, Union[None, float, Dict[str, Any]]]:
                     "Further mixing augmentations will not be applied.",
                     RuntimeWarning,
                 )
-                return {"mix_data": None, "mix_coef": 1}
+                return {"mix_data": {}, "mix_coef": 1}
+        else:
+            return {"mix_data": {}, "mix_coef": 1}
         mix_coef = beta(self.alpha, self.alpha) if mix_data else 1
 
         return {"mix_data": self.read_fn(mix_data) if mix_data else None, "mix_coef": mix_coef}
diff --git a/tests/test_mixing.py b/tests/test_mixing.py
@@ -22,14 +22,19 @@ def complex_read_fn_image(x):
              [(A.MixUp, {
                 "reference_data": [{"image": np.random.randint(0, 256, [100, 100, 3], dtype=np.uint8)}],
                 "read_fn": lambda x: x}),
+              (A.MixUp, {
+                  "reference_data": [1],
+                  "read_fn": lambda x: {"image": np.random.randint(0, 256, [100, 100, 3], dtype=np.uint8)}},
+              ),
+              (A.MixUp, {
+                  "reference_data": None,
+              }),
               (A.MixUp, {
             "reference_data": image_generator(),
             "read_fn": lambda x: x}),
               (A.MixUp, {
             "reference_data": complex_image_generator(),
-            "read_fn": complex_read_fn_image})]
-
-)
+            "read_fn": complex_read_fn_image})] )
 def test_image_only(augmentation_cls, params, image):
     aug = augmentation_cls(p=1, **params)
     data = aug(image=image)
@@ -40,7 +45,13 @@ def test_image_only(augmentation_cls, params, image):
              [(A.MixUp, {
                 "reference_data": [{"image": np.random.randint(0, 256, [100, 100, 3], dtype=np.uint8),
                                    "global_label": np.array([0, 0, 1])}],
-                "read_fn": lambda x: x})]
+                "read_fn": lambda x: x}),
+            (A.MixUp, {
+                  "reference_data": [1],
+                  "read_fn": lambda x: {"image": np.ones((100, 100, 3)).astype(np.uint8),
+                                        "global_label": np.array([0, 0, 1])}},
+              ),
+              ]
 )
 def test_image_global_label(augmentation_cls, params, image, global_label):
     aug = augmentation_cls(p=1, **params)
@@ -49,8 +60,13 @@ def test_image_global_label(augmentation_cls, params, image, global_label):
 
     assert data["image"].dtype == np.uint8
 
-    mix_coeff_image = find_mix_coef(data["image"], image, aug.reference_data[0]["image"])
-    mix_coeff_label = find_mix_coef(data["global_label"], global_label, aug.reference_data[0]["global_label"])
+    reference_item = params["read_fn"](aug.reference_data[0])
+
+    reference_image = reference_item["image"]
+    reference_global_label = reference_item["global_label"]
+
+    mix_coeff_image = find_mix_coef(data["image"], image, reference_image)
+    mix_coeff_label = find_mix_coef(data["global_label"], global_label, reference_global_label)
 
     assert math.isclose(mix_coeff_image, mix_coeff_label, abs_tol=0.01)
     assert 0 <= mix_coeff_image <= 1