Fix affine balanxced scale 1

.pre-commit-config.yaml

@@ -81,7 +81,7 @@ repos:
     hooks:
       - id: pyproject-fmt
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.12.0
+    rev: v1.12.1
     hooks:
       - id: mypy
         files: ^albumentations/

albumentations/augmentations/geometric/transforms.py

@@ -756,28 +756,36 @@ def apply_to_keypoints(
 
     @staticmethod
     def get_scale(
-        scale: dict[str, tuple[float, float]],
+        scale: dict[str, float | tuple[float, float]],
         keep_ratio: bool,
         balanced_scale: bool,
     ) -> fgeometric.ScaleDict:
         result_scale = {}
-        if balanced_scale:
-            for key, value in scale.items():
-                lower_interval = (value[0], 1.0) if value[0] < 1 else None
-                upper_interval = (1.0, value[1]) if value[1] > 1 else None
-
-                if lower_interval is not None and upper_interval is not None:
-                    selected_interval = random.choice([lower_interval, upper_interval])
-                elif lower_interval is not None:
-                    selected_interval = lower_interval
-                elif upper_interval is not None:
-                    selected_interval = upper_interval
+        for key, value in scale.items():
+            if isinstance(value, (int, float)):
+                result_scale[key] = float(value)
+            elif isinstance(value, tuple):
+                if balanced_scale:
+                    lower_interval = (value[0], 1.0) if value[0] < 1 else None
+                    upper_interval = (1.0, value[1]) if value[1] > 1 else None
+
+                    if lower_interval is not None and upper_interval is not None:
+                        selected_interval = random.choice([lower_interval, upper_interval])
+                    elif lower_interval is not None:
+                        selected_interval = lower_interval
+                    elif upper_interval is not None:
+                        selected_interval = upper_interval
+                    else:
+                        result_scale[key] = 1.0
+                        continue
+
+                    result_scale[key] = random.uniform(*selected_interval)
                 else:
-                    raise ValueError(f"Both lower_interval and upper_interval are None for key: {key}")
-
-                result_scale[key] = random.uniform(*selected_interval)
-        else:
-            result_scale = {key: random.uniform(*value) for key, value in scale.items()}
+                    result_scale[key] = random.uniform(*value)
+            else:
+                raise TypeError(
+                    f"Invalid scale value for key {key}: {value}. Expected a float or a tuple of two floats.",
+                )
 
         if keep_ratio:
             result_scale["y"] = result_scale["x"]

tests/functional/test_affine.py

@@ -846,3 +846,37 @@ def test_rotate_by_90_bboxes_symmetric_bbox(transform_class, transform_params, f
     bbox_out1 = transform(image=img, bboxes=[bbox], cl=[0])["bboxes"][0]
 
     np.testing.assert_allclose(bbox_out1, bbox, atol=1e-6)
+
+
+
+@pytest.mark.parametrize("scale, keep_ratio, balanced_scale, expected", [
+    ({"x": 1, "y": 1}, False, False, {"x": 1, "y": 1}),
+    ({"x": 1, "y": 1}, True, False, {"x": 1, "y": 1}),
+    ({"x": (0.5, 1.5), "y": (0.5, 1.5)}, False, False, {"x": pytest.approx(1, abs=0.5), "y": pytest.approx(1, abs=0.5)}),
+    ({"x": (0.5, 1.5), "y": (0.5, 1.5)}, True, False, lambda x: x["x"] == x["y"] and 0.5 <= x["x"] <= 1.5),
+    ({"x": (0.5, 1.5), "y": (0.5, 1.5)}, False, True, lambda x: all(0.5 <= v <= 1.5 for v in x.values())),
+    ({"x": (0.5, 2.0), "y": 1}, False, False, lambda x: 0.5 <= x["x"] <= 2.0 and x["y"] == 1),
+    ({"x": 0.5, "y": (0.5, 1.5)}, True, False, lambda x: x["x"] == 0.5 and x["y"] == 0.5),
+    ({"x": (0.8, 1.2), "y": (0.8, 1.2)}, False, True, lambda x: all(0.8 <= v <= 1.2 for v in x.values())),
+])
+def test_get_scale(scale, keep_ratio, balanced_scale, expected):
+    result = A.Affine.get_scale(scale, keep_ratio, balanced_scale)
+
+    if callable(expected):
+        assert expected(result)
+    else:
+        assert result == expected
+
+def test_get_scale_balanced_scale_behavior():
+    scale = {"x": (0.5, 2.0), "y": (0.5, 2.0)}
+    result = A.Affine.get_scale(scale, keep_ratio=False, balanced_scale=True)
+
+    assert all(0.5 <= v <= 2.0 for v in result.values())
+    assert any(v < 1.0 for v in result.values()) or any(v > 1.0 for v in result.values())
+
+def test_get_scale_keep_ratio():
+    scale = {"x": (0.5, 1.5), "y": (0.8, 1.2)}
+    result = A.Affine.get_scale(scale, keep_ratio=True, balanced_scale=False)
+
+    assert result["x"] == result["y"]
+    assert 0.5 <= result["x"] <= 1.5