diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 47e66041c..f4ef8fba2 100644 --- a/.pre-commit-config.yaml +++ b/.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/ diff --git a/albumentations/augmentations/geometric/transforms.py b/albumentations/augmentations/geometric/transforms.py index 5a4e12792..035fbae59 100644 --- a/albumentations/augmentations/geometric/transforms.py +++ b/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"] diff --git a/tests/functional/test_affine.py b/tests/functional/test_affine.py index fee559eab..8e1447b7d 100644 --- a/tests/functional/test_affine.py +++ b/tests/functional/test_affine.py @@ -846,3 +846,66 @@ 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 + + +@pytest.mark.parametrize( + "scale, keep_ratio, balanced_scale, expected_x_range, expected_y_range", + [ + ({"x": (0.5, 2), "y": (0.5, 2)}, False, True, (0.5, 2), (0.5, 2)), + ({"x": (1, 2), "y": (1, 2)}, True, True, (1, 2), (1, 2)), + ({"x": (0.5, 1), "y": (0.5, 1)}, True, True, (0.5, 1), (0.5, 1)), + ({"x": (0.5, 2), "y": (0.5, 2)}, False, False, (0.5, 2), (0.5, 2)), + ({"x": (0.5, 2), "y": (0.5, 2)}, True, False, (0.5, 2), (0.5, 2)), + ], +) +def test_get_random_scale(scale, keep_ratio, balanced_scale, expected_x_range, expected_y_range): + result = A.Affine.get_scale(scale, keep_ratio, balanced_scale) + + assert expected_x_range[0] <= result["x"] <= expected_x_range[1], "x is out of range" + + if keep_ratio: + assert result["y"] == result["x"], "y should be equal to x when keep_ratio is True" + else: + assert expected_y_range[0] <= result["y"] <= expected_y_range[1], "y is out of range" + + if balanced_scale: + assert ( + expected_x_range[0] <= result["x"] < 1 or 1 < result["x"] <= expected_x_range[1] + ), "x should be in the balanced range" + assert ( + expected_y_range[0] <= result["y"] < 1 or 1 < result["y"] <= expected_x_range[1] + ), "x should be in the balanced range" diff --git a/tests/test_transforms.py b/tests/test_transforms.py index d862ab03d..76de00dcf 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -1774,35 +1774,6 @@ def test_gauss_noise(mean, image): assert not (result["image"] >= image).all() -@pytest.mark.parametrize( - "scale, keep_ratio, balanced_scale, expected_x_range, expected_y_range", - [ - ({"x": (0.5, 2), "y": (0.5, 2)}, False, True, (0.5, 2), (0.5, 2)), - ({"x": (1, 2), "y": (1, 2)}, True, True, (1, 2), (1, 2)), - ({"x": (0.5, 1), "y": (0.5, 1)}, True, True, (0.5, 1), (0.5, 1)), - ({"x": (0.5, 2), "y": (0.5, 2)}, False, False, (0.5, 2), (0.5, 2)), - ({"x": (0.5, 2), "y": (0.5, 2)}, True, False, (0.5, 2), (0.5, 2)), - ], -) -def test_get_random_scale(scale, keep_ratio, balanced_scale, expected_x_range, expected_y_range): - result = A.Affine.get_scale(scale, keep_ratio, balanced_scale) - - assert expected_x_range[0] <= result["x"] <= expected_x_range[1], "x is out of range" - - if keep_ratio: - assert result["y"] == result["x"], "y should be equal to x when keep_ratio is True" - else: - assert expected_y_range[0] <= result["y"] <= expected_y_range[1], "y is out of range" - - if balanced_scale: - assert ( - expected_x_range[0] <= result["x"] < 1 or 1 < result["x"] <= expected_x_range[1] - ), "x should be in the balanced range" - assert ( - expected_y_range[0] <= result["y"] < 1 or 1 < result["y"] <= expected_x_range[1] - ), "x should be in the balanced range" - - @pytest.mark.parametrize( "params, expected", [