test_kornia_transforms.py

# ------------------------------------------------------------------------
# RF-DETR
# Copyright (c) 2025 Roboflow. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
"""Tests for Kornia GPU augmentation pipeline builder and bbox utilities.

All tests in this module are CPU-compatible — Kornia operates on CPU tensors identically to GPU tensors, so no
``@pytest.mark.gpu`` is needed.
"""

import pytest
import torch

from rfdetr.datasets.aug_configs import (
    AUG_AERIAL,
    AUG_AGGRESSIVE,
    AUG_CONSERVATIVE,
    AUG_INDUSTRIAL,
)


def _sharpness_sampler_range(transform: torch.nn.Module) -> tuple[float, float]:
    """Read the sampled ``sharpness`` range off a Kornia ``RandomSharpness`` transform.

    This intentionally reads a private Kornia implementation detail (``_param_generator.sampler_dict``) because no
    public equivalent exists: ``RandomSharpness(...).flags`` is empty (verified against the installed Kornia
    version), so the sampled range isn't reachable through any public attribute. If a future Kornia release
    renames or removes ``_param_generator``/``sampler_dict``, this raises a clear, actionable failure instead of a
    raw ``AttributeError``/``KeyError`` deep inside the test body.

    Args:
        transform: A ``kornia.augmentation.RandomSharpness`` instance (or equivalent) built with a sampled range.

    Returns:
        The sampled ``(low, high)`` bounds as floats.
    """
    try:
        sampler = transform._param_generator.sampler_dict["sharpness"]
        return float(sampler.low), float(sampler.high)
    except (AttributeError, KeyError) as exc:
        pytest.fail(
            "Kornia's RandomSharpness no longer exposes the sampled `sharpness` range via the private "
            f"`_param_generator.sampler_dict['sharpness']` path (no public alternative exists): {exc!r}. Update "
            "this helper to match Kornia's new internal parameter-generator shape."
        )


# ---------------------------------------------------------------------------
# TestBuildKorniaPipeline — validates the factory that translates aug_config
# dicts into a Kornia AugmentationSequential pipeline.
# ---------------------------------------------------------------------------


class TestBuildKorniaPipeline:
    """build_kornia_pipeline returns a valid pipeline for every preset and rejects unknown transform keys with a clear
    error."""

    @pytest.fixture(autouse=True)
    def _require_kornia(self):
        pytest.importorskip("kornia")

    @pytest.mark.parametrize(
        "config,config_name",
        [
            pytest.param(AUG_CONSERVATIVE, "AUG_CONSERVATIVE", id="conservative"),
            pytest.param(AUG_AGGRESSIVE, "AUG_AGGRESSIVE", id="aggressive"),
            pytest.param(AUG_AERIAL, "AUG_AERIAL", id="aerial"),
            pytest.param(AUG_INDUSTRIAL, "AUG_INDUSTRIAL", id="industrial"),
        ],
    )
    def test_each_preset_config(self, config, config_name):
        """Each named preset builds a pipeline without errors."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline(config, 560)
        assert pipeline is not None, f"build_kornia_pipeline({config_name}, 560) must return a non-None pipeline"

    def test_unknown_key_raises_value_error(self):
        """An unrecognised transform key raises ValueError immediately."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        with pytest.raises(ValueError, match="FooBarTransform"):
            build_kornia_pipeline({"FooBarTransform": {"p": 0.5}}, 560)

    def test_empty_config_returns_pipeline(self):
        """An empty config dict returns a valid (no-op) pipeline, not None."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({}, 560)
        assert pipeline is not None, "Empty config must still return a pipeline object"

    def test_known_plus_unknown_raises(self):
        """Mixing a valid key with an unknown key still raises ValueError."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        mixed = {"HorizontalFlip": {"p": 0.5}, "BogusTransform": {"p": 0.3}}
        with pytest.raises(ValueError, match="BogusTransform"):
            build_kornia_pipeline(mixed, 560)

    def test_to_gray_builds_random_grayscale(self):
        """``ToGray`` maps onto ``K.RandomGrayscale`` (issue #1227)."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"ToGray": {"p": 0.5}}, 560)
        transform_names = [child.__class__.__name__ for child in pipeline.children()]
        assert "RandomGrayscale" in transform_names

    def test_to_gray_keeps_three_channels_and_greys(self):
        """``ToGray`` matches Albumentations: grayscale content, still three channels."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"ToGray": {"p": 1.0}}, 560)
        image = torch.rand(1, 3, 32, 32)
        boxes = torch.tensor([[[0.0, 0.0, 10.0, 10.0]]])
        out, _ = pipeline(image, boxes)

        assert out.shape == image.shape, "ToGray must preserve the three-channel shape"
        # A greyscale image has identical values across the channel axis.
        assert torch.allclose(out[:, 0], out[:, 1], atol=1e-5)
        assert torch.allclose(out[:, 1], out[:, 2], atol=1e-5)

    def test_to_gray_accepted_by_both_backends(self):
        """The same config is readable by the Albumentations backend too (issue #1227).

        ``ToGray`` resolved on Albumentations via ``getattr`` long before it was a Kornia built-in, so a config that
        worked on one backend raised on the other.
        """
        pytest.importorskip("albumentations")
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline
        from rfdetr.datasets.transforms import AlbumentationsWrapper

        config = {"ToGray": {"p": 0.5}}
        assert build_kornia_pipeline(config, 560) is not None

        wrappers = AlbumentationsWrapper.from_config(config)
        assert len(wrappers) == 1, (
            "from_config(strict=False) silently drops unresolved transforms, so length must be checked"
        )
        built_names = [t.__class__.__name__ for t in wrappers[0].transform.transforms]
        assert "ToGray" in built_names, f"expected a ToGray transform, got {built_names}"

    def test_to_gray_defaults_p_to_point_five_when_omitted(self):
        """Omitting p resolves to 0.5, matching Albumentations (not Kornia's native 0.1 default)."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"ToGray": {}}, 560)
        to_gray = next(child for child in pipeline.children() if child.__class__.__name__ == "RandomGrayscale")
        assert to_gray.p == pytest.approx(0.5)

    def test_to_gray_p_zero_is_a_no_op(self):
        """p=0.0 never applies: forward pass returns the input unchanged."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"ToGray": {"p": 0.0}}, 560)
        image = torch.rand(1, 3, 32, 32)
        boxes = torch.tensor([[[0.0, 0.0, 10.0, 10.0]]])
        out, _ = pipeline(image, boxes)

        assert torch.equal(out, image), "p=0.0 must never apply ToGray"

    def test_to_gray_ignores_method_and_num_output_channels_on_kornia(self):
        """method/num_output_channels have no Kornia equivalent and are currently ignored there.

        Pins the divergence documented on ``_make_to_gray``: Albumentations honors both, Kornia always uses BT.601
        weights and returns 3 channels. If this is ever fixed, this test should be updated to assert the new (parity)
        behavior instead.
        """
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"ToGray": {"method": "max", "num_output_channels": 1, "p": 1.0}}, 560)
        image = torch.rand(1, 3, 32, 32)
        boxes = torch.tensor([[[0.0, 0.0, 10.0, 10.0]]])
        out, _ = pipeline(image, boxes)

        assert out.shape == image.shape, "num_output_channels=1 is currently ignored -- output stays 3-channel"
        assert torch.allclose(out[:, 0], out[:, 1], atol=1e-5), "method='max' is currently ignored on Kornia"

    @pytest.mark.parametrize(
        "config",
        [
            pytest.param({"GaussianBlur": {"blur_limit": (3, 7), "p": 0.5}}, id="blur_limit-pair"),
            pytest.param({"GaussianBlur": {"sigma": 1.5, "p": 0.5}}, id="sigma-scalar"),
            pytest.param({"GaussianBlur": {"sigma": (1.5,), "p": 0.5}}, id="sigma-1elem-seq"),
            pytest.param({"GaussNoise": {"std_range": 0.05, "p": 0.5}}, id="std_range-scalar"),
        ],
    )
    def test_scalar_or_pair_range_params_build(self, config):
        """Range params accept a scalar or a pair, as Albumentations does.

        ``_make_rotate`` already accepts either form for ``limit``; these builders used to raise a bare ``TypeError``
        from inside Kornia on a config that is valid for the CPU path.
        """
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        assert build_kornia_pipeline(config, 560) is not None

    @pytest.mark.parametrize(
        ("blur_limit", "expected"),
        [
            pytest.param((3, 7), (7, 7), id="odd-upper-bound"),
            pytest.param((3, 6), (7, 7), id="even-upper-bound-rounds-up"),
        ],
    )
    def test_blur_limit_pair_uses_upper_bound(self, blur_limit, expected):
        """A ``(min, max)`` ``blur_limit`` resolves to its upper bound, rounded up to an odd kernel size."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"GaussianBlur": {"blur_limit": blur_limit, "p": 1.0}}, 560)
        transform = next(iter(pipeline.children()))
        assert transform.flags["kernel_size"] == expected

    def test_scalar_std_range_is_used_verbatim(self):
        """A scalar ``std_range`` is used as-is, with no per-sample-drift warning emitted."""
        from unittest import mock

        from rfdetr.datasets import kornia_transforms

        with mock.patch.object(kornia_transforms.logger, "warning") as mock_warning:
            pipeline = kornia_transforms.build_kornia_pipeline({"GaussNoise": {"std_range": 0.05, "p": 1.0}}, 560)

        transform = next(iter(pipeline.children()))
        assert transform.flags["std"] == pytest.approx(0.05)
        mock_warning.assert_not_called()

    def test_hflip_disabled_for_keypoint_pipeline(self):
        """Keypoint-mode Kornia augmentation drops hflip transforms with a warning."""
        from unittest import mock

        from rfdetr.datasets import kornia_transforms

        config = {"HorizontalFlip": {"p": 0.5}, "VerticalFlip": {"p": 0.5}}
        mock_warning = mock.patch.object(kornia_transforms.logger, "warning")

        with mock_warning as warning:
            pipeline = kornia_transforms.build_kornia_pipeline(config, 560, include_keypoints=True)

        transform_names = [child.__class__.__name__ for child in pipeline.children()]
        assert "RandomHorizontalFlip" not in transform_names
        assert "RandomVerticalFlip" in transform_names
        assert warning.called
        assert "HorizontalFlip" in str(warning.call_args)

    # --- pixel-level transforms added for issue #1252 -------------------

    @pytest.mark.parametrize(
        ("name", "params", "expected"),
        [
            ("Blur", {"blur_limit": 5}, "RandomBoxBlur"),
            ("Sharpen", {"alpha": (0.2, 0.5)}, "RandomSharpness"),
            ("Equalize", {}, "RandomEqualize"),
            ("CLAHE", {"clip_limit": 4.0}, "RandomClahe"),
        ],
    )
    def test_pixel_transforms_map_to_kornia(self, name, params, expected):
        """Each documented pixel-level name builds its Kornia counterpart (issue #1252)."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({name: params}, 560)
        assert expected in [child.__class__.__name__ for child in pipeline.children()]

    @pytest.mark.parametrize(
        ("name", "params"),
        [
            ("Blur", {"blur_limit": (3, 7)}),
            ("Sharpen", {"alpha": (0.2, 0.5)}),
            ("Equalize", {"p": 0.5}),
            ("CLAHE", {"clip_limit": 4.0, "tile_grid_size": (8, 8)}),
        ],
    )
    def test_pixel_transforms_accepted_by_both_backends(self, name, params):
        """The same config builds on either backend, which is the gap issue #1252 reports."""
        pytest.importorskip("albumentations")
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline
        from rfdetr.datasets.transforms import AlbumentationsWrapper

        config = {name: params}
        assert build_kornia_pipeline(config, 560) is not None

        wrappers = AlbumentationsWrapper.from_config(config)
        assert len(wrappers) == 1, (
            "from_config(strict=False) silently drops unresolved transforms, so length must be checked"
        )
        built = [t.__class__.__name__ for t in wrappers[0].transform.transforms]
        assert name in built, f"expected {name}, got {built}"

    @pytest.mark.parametrize(("blur_limit", "expected"), [(5, 5), (4, 5), ((3, 7), 7), ((3, 6), 7), (2, 3)])
    def test_blur_kernel_is_odd_and_at_least_three(self, blur_limit, expected):
        """Blur resolves its kernel the same way GaussianBlur does: odd, >= 3, upper bound of a pair."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"Blur": {"blur_limit": blur_limit}}, 560)
        transform = next(iter(pipeline.children()))
        assert transform.flags["kernel_size"] == (expected, expected)

    def test_blur_pair_warns_about_the_fixed_kernel(self):
        """A non-degenerate pair the user chose explicitly collapses to one kernel, so it must say so."""
        from unittest import mock

        from rfdetr.datasets import kornia_transforms

        with mock.patch.object(kornia_transforms.logger, "warning") as warning:
            kornia_transforms.build_kornia_pipeline({"Blur": {"blur_limit": (3, 5)}}, 560)
        assert warning.called
        assert "Blur" in str(warning.call_args)

    def test_blur_degenerate_pair_does_not_warn(self):
        """(5, 5) loses nothing, so it should stay quiet."""
        from unittest import mock

        from rfdetr.datasets import kornia_transforms

        with mock.patch.object(kornia_transforms.logger, "warning") as warning:
            kornia_transforms.build_kornia_pipeline({"Blur": {"blur_limit": (5, 5)}}, 560)
        warning.assert_not_called()

    def test_blur_library_default_pair_logs_at_debug_not_warning(self):
        """(3, 7) is Albumentations' own Blur default, an expected divergence, so it must stay off WARNING."""
        from unittest import mock

        from rfdetr.datasets import kornia_transforms

        with (
            mock.patch.object(kornia_transforms.logger, "warning") as warning,
            mock.patch.object(kornia_transforms.logger, "debug") as debug,
        ):
            kornia_transforms.build_kornia_pipeline({"Blur": {"blur_limit": (3, 7)}}, 560)
        warning.assert_not_called()
        assert debug.called
        assert "Blur" in str(debug.call_args)

    @pytest.mark.parametrize(
        "blur_limit",
        [
            pytest.param([], id="empty-sequence"),
            pytest.param((1, 2, 3), id="three-element-sequence"),
        ],
    )
    def test_blur_kernel_rejects_invalid_sequence_length(self, blur_limit):
        """A sequence that is neither a scalar nor a (min, max) pair must raise, not silently misresolve."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        with pytest.raises(ValueError, match="Kernel size parameter must be"):
            build_kornia_pipeline({"Blur": {"blur_limit": blur_limit}}, 560)

    def test_sharpen_shifts_alpha_to_kornias_one_pivoted_sharpness_range(self):
        """Albumentations' alpha (0=no-op) is shifted to Kornia's sharpness (1.0=no-op): sharpness = 1.0 + alpha."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"Sharpen": {"alpha": (0.1, 0.4)}}, 560)
        transform = next(iter(pipeline.children()))
        # Kornia keeps sampled ranges on the parameter generator rather than in `flags`; see
        # `_sharpness_sampler_range` for why this reads a private attribute.
        assert _sharpness_sampler_range(transform) == pytest.approx((1.1, 1.4))

    def test_sharpen_default_alpha_actually_sharpens_not_blurs(self):
        """Regression guard: at the default alpha=(0.2, 0.5), Sharpen must raise edge energy, not lower it.

        Kornia's ``sharpness`` factor is pivoted at 1.0 (0=blur, 1=no-op, >1=sharpen), unlike Albumentations' ``alpha``
        (pivoted at 0). Passing ``alpha`` straight through as ``sharpness`` (the pre-fix bug) resolves to the range
        (0.2, 0.5) — below Kornia's no-op point — which blurs a step edge instead of sharpening it, so this test would
        fail against that code. The fixed mapping resolves to ``sharpness=(1.2, 1.5)``, which sharpens.
        """
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        size = 16
        img = torch.full((1, 3, size, size), 0.3)
        img[:, :, :, size // 2 :] = 0.7  # a single sharp step edge down the middle column
        boxes = torch.tensor([[[0.0, 0.0, float(size), float(size)]]], dtype=torch.float32)

        pipeline = build_kornia_pipeline({"Sharpen": {"p": 1.0}}, 560)
        img_out, _ = pipeline(img, boxes)

        def edge_energy(x: torch.Tensor) -> float:
            # Exclude the outer 2-pixel ring: Kornia's sharpness leaves border pixels unchanged (see
            # kornia.enhance.adjust.sharpness), so including them would dilute the interior sharpening signal.
            interior = x[:, :, 2:-2, 2:-2]
            return (torch.diff(interior, dim=-1).abs().mean() + torch.diff(interior, dim=-2).abs().mean()).item()

        assert edge_energy(img_out) > edge_energy(img), (
            "Sharpen at the default alpha=(0.2, 0.5) must increase edge energy (sharpen); an unchanged or lower "
            "value means the pivot-point bug regressed (sharpness range fell back to (0.2, 0.5), which blurs)."
        )

    def test_sharpen_warns_about_ignored_lightness(self):
        """Lightness has no Kornia equivalent, so dropping it must be announced."""
        from unittest import mock

        from rfdetr.datasets import kornia_transforms

        with mock.patch.object(kornia_transforms.logger, "warning") as warning:
            kornia_transforms.build_kornia_pipeline({"Sharpen": {"lightness": (0.5, 1.0)}}, 560)
        assert warning.called
        assert "lightness" in str(warning.call_args)

    def test_equalize_warns_about_ignored_options(self):
        """mode/by_channels/mask are albumentations-only."""
        from unittest import mock

        from rfdetr.datasets import kornia_transforms

        with mock.patch.object(kornia_transforms.logger, "warning") as warning:
            kornia_transforms.build_kornia_pipeline({"Equalize": {"by_channels": False}}, 560)
        assert warning.called
        assert "by_channels" in str(warning.call_args)

    def test_clahe_maps_both_parameters(self):
        """clip_limit and tile_grid_size map straight onto Kornia's clip_limit and grid_size."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"CLAHE": {"clip_limit": (2.0, 6.0), "tile_grid_size": (4, 4)}}, 560)
        transform = next(iter(pipeline.children()))
        assert tuple(transform.flags["grid_size"]) == (4, 4)
        # Unlike Sharpen's sharpness range, RandomClahe exposes its range on a plain public
        # `clip_limit` attribute (set directly from the constructor arg), so no private access needed.
        assert tuple(transform.clip_limit) == pytest.approx((2.0, 6.0))

    def test_hue_saturation_value_still_unsupported(self):
        """Deliberately out of scope: albumentations shifts additively, Kornia scales multiplicatively."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        with pytest.raises(ValueError, match="HueSaturationValue"):
            build_kornia_pipeline({"HueSaturationValue": {"hue_shift_limit": 20}}, 560)


# ---------------------------------------------------------------------------
# TestCollateBoxes — validates packing of variable-length per-image boxes
# into a zero-padded [B, N_max, 4] tensor with a boolean validity mask.
# ---------------------------------------------------------------------------


class TestCollateBoxes:
    """collate_boxes packs variable-length boxes into [B, N_max, 4] with mask."""

    @pytest.fixture(autouse=True)
    def _require_kornia(self):
        pytest.importorskip("kornia")

    def _make_targets(self, box_counts):
        """Build a list of target dicts with the given per-image box counts.

        Each box is a valid xyxy rectangle within a 100x100 image.
        """
        targets = []
        for n in box_counts:
            boxes = (
                torch.tensor([[10.0, 10.0, 50.0, 50.0]] * n, dtype=torch.float32)
                if n > 0
                else torch.zeros(0, 4, dtype=torch.float32)
            )
            targets.append({"boxes": boxes})
        return targets

    def test_normal_batch(self):
        """Batch of 2 images: output shape is [2, N_max, 4] with valid mask [2, N_max]."""
        from rfdetr.datasets.kornia_transforms import collate_boxes

        targets = self._make_targets([2, 3])
        boxes_padded, valid = collate_boxes(targets, torch.device("cpu"))

        assert boxes_padded.shape == (2, 3, 4), f"Expected shape (2, 3, 4), got {boxes_padded.shape}"
        assert valid.shape == (2, 3), f"Expected valid shape (2, 3), got {valid.shape}"
        assert valid.dtype == torch.bool

    def test_b_zero(self):
        """Empty target list produces shape [0, 0, 4] and valid [0, 0]."""
        from rfdetr.datasets.kornia_transforms import collate_boxes

        boxes_padded, valid = collate_boxes([], torch.device("cpu"))

        assert boxes_padded.shape == (0, 0, 4), f"Expected (0, 0, 4) for empty batch, got {boxes_padded.shape}"
        assert valid.shape == (0, 0), f"Expected valid (0, 0) for empty batch, got {valid.shape}"

    def test_n_zero_per_image(self):
        """One image with 0 boxes: shape [1, 0, 4], valid all-False."""
        from rfdetr.datasets.kornia_transforms import collate_boxes

        targets = self._make_targets([0])
        boxes_padded, valid = collate_boxes(targets, torch.device("cpu"))

        assert boxes_padded.shape == (1, 0, 4), f"Expected (1, 0, 4), got {boxes_padded.shape}"
        assert valid.shape == (1, 0), f"Expected (1, 0), got {valid.shape}"

    def test_single_image(self):
        """B=1 with 3 boxes: output shape is [1, 3, 4]."""
        from rfdetr.datasets.kornia_transforms import collate_boxes

        targets = self._make_targets([3])
        boxes_padded, valid = collate_boxes(targets, torch.device("cpu"))

        assert boxes_padded.shape == (1, 3, 4)
        assert valid.shape == (1, 3)

    def test_valid_mask_matches_box_count(self):
        """The valid mask has True for real boxes and False for padding."""
        from rfdetr.datasets.kornia_transforms import collate_boxes

        targets = self._make_targets([1, 3])
        _, valid = collate_boxes(targets, torch.device("cpu"))

        # Image 0: 1 real box, 2 padding → [True, False, False]
        assert valid[0].tolist() == [True, False, False], f"Image 0 valid mask wrong: {valid[0].tolist()}"
        # Image 1: 3 real boxes, 0 padding → [True, True, True]
        assert valid[1].tolist() == [True, True, True], f"Image 1 valid mask wrong: {valid[1].tolist()}"


# ---------------------------------------------------------------------------
# TestUnpackBoxes — validates the inverse: writing augmented boxes back into
# per-image target dicts with clamping, zero-area removal, and label sync.
# ---------------------------------------------------------------------------


class TestUnpackBoxes:
    """unpack_boxes writes augmented boxes back and removes zero-area entries."""

    @pytest.fixture(autouse=True)
    def _require_kornia(self):
        pytest.importorskip("kornia")

    def _make_inputs(
        self,
        boxes_aug,
        valid_mask,
        original_targets,
        image_height=100,
        image_width=100,
    ):
        """Return tensors suitable for unpack_boxes."""
        boxes_tensor = torch.tensor(boxes_aug, dtype=torch.float32)
        valid_tensor = torch.tensor(valid_mask, dtype=torch.bool)
        return boxes_tensor, valid_tensor, original_targets, image_height, image_width

    def test_all_boxes_removed_after_aug(self):
        """When all augmented boxes are zero-area, output targets have empty boxes."""
        from rfdetr.datasets.kornia_transforms import unpack_boxes

        # B=1, N=2: both boxes are zero-area (x1==x2 or y1==y2)
        boxes_aug = [[[10.0, 10.0, 10.0, 10.0], [20.0, 20.0, 20.0, 20.0]]]
        valid = [[True, True]]
        targets = [
            {
                "boxes": torch.tensor([[10.0, 10.0, 50.0, 50.0], [20.0, 20.0, 60.0, 60.0]]),
                "labels": torch.tensor([1, 2]),
                "area": torch.tensor([1600.0, 1600.0]),
                "iscrowd": torch.tensor([0, 0]),
            }
        ]
        boxes_t, valid_t, tgts, image_height, image_width = self._make_inputs(boxes_aug, valid, targets)
        result = unpack_boxes(boxes_t, valid_t, tgts, image_height, image_width)

        assert result[0]["boxes"].shape[0] == 0, (
            f"Expected 0 boxes after zero-area removal, got {result[0]['boxes'].shape[0]}"
        )
        assert result[0]["labels"].shape[0] == 0

    def test_partial_removal(self):
        """Some boxes survive, some removed; labels/area/iscrowd synced."""
        from rfdetr.datasets.kornia_transforms import unpack_boxes

        # Box 0: valid, non-zero area; Box 1: zero-area
        boxes_aug = [[[10.0, 10.0, 50.0, 50.0], [30.0, 30.0, 30.0, 30.0]]]
        valid = [[True, True]]
        targets = [
            {
                "boxes": torch.tensor([[10.0, 10.0, 50.0, 50.0], [30.0, 30.0, 70.0, 70.0]]),
                "labels": torch.tensor([1, 2]),
                "area": torch.tensor([1600.0, 1600.0]),
                "iscrowd": torch.tensor([0, 1]),
            }
        ]
        boxes_t, valid_t, tgts, image_height, image_width = self._make_inputs(boxes_aug, valid, targets)
        result = unpack_boxes(boxes_t, valid_t, tgts, image_height, image_width)

        assert result[0]["boxes"].shape[0] == 1, f"Expected 1 surviving box, got {result[0]['boxes'].shape[0]}"
        assert result[0]["labels"].tolist() == [1]

    def test_labels_area_iscrowd_sync(self):
        """When boxes are removed, labels/area/iscrowd entries are also removed."""
        from rfdetr.datasets.kornia_transforms import unpack_boxes

        # Box 0: zero-area (removed), Box 1: valid
        boxes_aug = [[[5.0, 5.0, 5.0, 5.0], [10.0, 10.0, 40.0, 40.0]]]
        valid = [[True, True]]
        targets = [
            {
                "boxes": torch.tensor([[5.0, 5.0, 30.0, 30.0], [10.0, 10.0, 40.0, 40.0]]),
                "labels": torch.tensor([7, 9]),
                "area": torch.tensor([625.0, 900.0]),
                "iscrowd": torch.tensor([0, 1]),
            }
        ]
        boxes_t, valid_t, tgts, image_height, image_width = self._make_inputs(boxes_aug, valid, targets)
        result = unpack_boxes(boxes_t, valid_t, tgts, image_height, image_width)

        assert result[0]["labels"].tolist() == [9], (
            f"Expected label [9] after removal of box 0, got {result[0]['labels'].tolist()}"
        )
        assert result[0]["area"].shape[0] == 1
        assert result[0]["iscrowd"].tolist() == [1]

    def test_boxes_clamped_to_image_bounds(self):
        """Boxes outside [0,W]x[0,H] are clamped to image bounds."""
        from rfdetr.datasets.kornia_transforms import unpack_boxes

        # Box extends beyond 100x100 image
        boxes_aug = [[[-10.0, -5.0, 120.0, 110.0]]]
        valid = [[True]]
        targets = [
            {
                "boxes": torch.tensor([[0.0, 0.0, 90.0, 90.0]]),
                "labels": torch.tensor([1]),
                "area": torch.tensor([8100.0]),
                "iscrowd": torch.tensor([0]),
            }
        ]
        image_height, image_width = 100, 100
        boxes_t, valid_t, tgts, image_height, image_width = self._make_inputs(
            boxes_aug,
            valid,
            targets,
            image_height,
            image_width,
        )
        result = unpack_boxes(boxes_t, valid_t, tgts, image_height, image_width)

        result_boxes = result[0]["boxes"]
        assert result_boxes.shape[0] == 1, "Clamped box should survive (non-zero area)"
        # Verify clamping: x1>=0, y1>=0, x2<=W, y2<=H
        assert result_boxes[0, 0].item() >= 0.0, "x1 not clamped to >= 0"
        assert result_boxes[0, 1].item() >= 0.0, "y1 not clamped to >= 0"
        assert result_boxes[0, 2].item() <= image_width, f"x2 not clamped to <= {image_width}"
        assert result_boxes[0, 3].item() <= image_height, f"y2 not clamped to <= {image_height}"


# ---------------------------------------------------------------------------
# TestRotateFactory — validates the Rotate parameter translation from
# Albumentations-style limit (scalar or tuple) to Kornia RandomRotation.
# ---------------------------------------------------------------------------


class TestRotateFactory:
    """Rotate factory translates limit (scalar or tuple) to K.RandomRotation(degrees=...)."""

    @pytest.fixture(autouse=True)
    def _require_kornia(self):
        pytest.importorskip("kornia")

    def test_limit_as_scalar(self):
        """Rotate(limit=45) produces K.RandomRotation(degrees=(-45, 45))."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        # Build a pipeline with just Rotate(limit=45)
        pipeline = build_kornia_pipeline({"Rotate": {"limit": 45, "p": 1.0}}, 560)
        assert pipeline is not None

        # Inspect the pipeline's children to find the RandomRotation and check degrees
        import kornia.augmentation as kornia_augmentation

        rotation_augs = [
            child for child in pipeline.children() if isinstance(child, kornia_augmentation.RandomRotation)
        ]
        assert len(rotation_augs) == 1, f"Expected exactly 1 RandomRotation, found {len(rotation_augs)}"
        degrees = rotation_augs[0].flags["degrees"]
        # degrees should be a tensor representing (-45, 45)
        assert float(degrees[0]) == pytest.approx(-45.0, abs=0.1)
        assert float(degrees[1]) == pytest.approx(45.0, abs=0.1)

    def test_limit_as_tuple(self):
        """Rotate(limit=(90, 90)) produces K.RandomRotation(degrees=(90, 90))."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"Rotate": {"limit": (90, 90), "p": 1.0}}, 560)
        assert pipeline is not None

        import kornia.augmentation as kornia_augmentation

        rotation_augs = [
            child for child in pipeline.children() if isinstance(child, kornia_augmentation.RandomRotation)
        ]
        assert len(rotation_augs) == 1
        degrees = rotation_augs[0].flags["degrees"]
        assert float(degrees[0]) == pytest.approx(90.0, abs=0.1)
        assert float(degrees[1]) == pytest.approx(90.0, abs=0.1)

    def test_flags_include_degrees(self):
        """Rotate factory keeps a legacy degrees entry in Kornia flags for compatibility."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"Rotate": {"limit": 30, "p": 1.0}}, 560)
        assert pipeline is not None

        import kornia.augmentation as kornia_augmentation

        rotation_augs = [
            child for child in pipeline.children() if isinstance(child, kornia_augmentation.RandomRotation)
        ]
        assert len(rotation_augs) == 1
        assert "degrees" in rotation_augs[0].flags
        assert rotation_augs[0].flags["degrees"] == (-30, 30)


# ---------------------------------------------------------------------------
# TestGpuPostprocessFlag — validates that make_coco_transforms respects the
# gpu_postprocess flag to omit augmentation and normalization from CPU path.
# ---------------------------------------------------------------------------


class TestGpuPostprocessFlag:
    """gpu_postprocess flag controls whether aug + normalize appear in CPU pipeline."""

    def test_gpu_postprocess_true_omits_aug_and_normalize_from_train(self):
        """gpu_postprocess=True: train pipeline has no CPU augmentation or Normalize."""
        from rfdetr.datasets._torchvision import RandomHorizontalFlip
        from rfdetr.datasets.coco import make_coco_transforms
        from rfdetr.datasets.transforms import Normalize

        pipeline_gpu = make_coco_transforms("train", 560, gpu_postprocess=True)
        pipeline_cpu = make_coco_transforms("train", 560, gpu_postprocess=False)

        steps_gpu = pipeline_gpu.transforms
        steps_cpu = pipeline_cpu.transforms

        normalize_gpu = [s for s in steps_gpu if isinstance(s, Normalize)]
        assert len(normalize_gpu) == 0, "gpu_postprocess=True must omit Normalize from train pipeline"

        assert not any(isinstance(s, RandomHorizontalFlip) for s in steps_gpu)
        assert any(isinstance(s, RandomHorizontalFlip) for s in steps_cpu)

    def test_gpu_postprocess_false_includes_aug_and_normalize_from_train(self):
        """gpu_postprocess=False (default): train pipeline includes Normalize."""
        from rfdetr.datasets.coco import make_coco_transforms
        from rfdetr.datasets.transforms import Normalize

        pipeline = make_coco_transforms("train", 560, gpu_postprocess=False)
        steps = pipeline.transforms

        normalize_steps = [s for s in steps if isinstance(s, Normalize)]
        assert len(normalize_steps) > 0, "gpu_postprocess=False must include Normalize in train pipeline"

    def test_val_path_unaffected_by_gpu_postprocess(self):
        """Val pipeline is unchanged regardless of gpu_postprocess value."""
        from rfdetr.datasets.coco import make_coco_transforms
        from rfdetr.datasets.transforms import Normalize

        pipeline_default = make_coco_transforms("val", 560, gpu_postprocess=False)
        pipeline_gpu = make_coco_transforms("val", 560, gpu_postprocess=True)

        # Both should have Normalize (val is never stripped)
        norm_default = [s for s in pipeline_default.transforms if isinstance(s, Normalize)]
        norm_gpu = [s for s in pipeline_gpu.transforms if isinstance(s, Normalize)]

        assert len(norm_default) > 0, "Val pipeline (default) must include Normalize"
        assert len(norm_gpu) > 0, "Val pipeline (gpu_postprocess=True) must include Normalize"

        # Same number of pipeline steps
        assert len(pipeline_default.transforms) == len(pipeline_gpu.transforms), (
            "Val pipeline step count must be identical regardless of gpu_postprocess"
        )


# ---------------------------------------------------------------------------
# TestGaussianBlurMinKernel — validates that blur_limit < 3 is clamped so
# Kornia never receives an invalid kernel_size < 3.
# ---------------------------------------------------------------------------


class TestGaussianBlurMinKernel:
    """_make_gaussian_blur enforces kernel_size >= 3 regardless of blur_limit."""

    @pytest.fixture(autouse=True)
    def _require_kornia(self):
        pytest.importorskip("kornia")

    @pytest.mark.parametrize(
        "blur_limit",
        [pytest.param(1, id="blur_limit_1"), pytest.param(2, id="blur_limit_2")],
    )
    def test_small_blur_limit_produces_valid_kernel(self, blur_limit):
        """blur_limit below 3 must be clamped so the resulting kernel_size >= 3."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        # Should not raise; previously blur_limit=1 produced kernel_size=(3,1)
        pipeline = build_kornia_pipeline({"GaussianBlur": {"blur_limit": blur_limit, "p": 1.0}}, 560)
        assert pipeline is not None

        import kornia.augmentation as kornia_augmentation

        blur_augs = [c for c in pipeline.children() if isinstance(c, kornia_augmentation.RandomGaussianBlur)]
        assert len(blur_augs) == 1
        ks = blur_augs[0].flags["kernel_size"]
        assert int(ks[0]) >= 3, f"kernel_size[0]={int(ks[0])} must be >= 3"
        assert int(ks[1]) >= 3, f"kernel_size[1]={int(ks[1])} must be >= 3"

    def test_blur_limit_3_unchanged(self):
        """blur_limit=3 (default) passes through without modification."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"GaussianBlur": {"blur_limit": 3, "p": 1.0}}, 560)
        import kornia.augmentation as kornia_augmentation

        blur_augs = [c for c in pipeline.children() if isinstance(c, kornia_augmentation.RandomGaussianBlur)]
        ks = blur_augs[0].flags["kernel_size"]
        assert int(ks[0]) == 3
        assert int(ks[1]) == 3


# ---------------------------------------------------------------------------
# TestKorniaPipelineForwardPass — validates that a built pipeline produces
# output of the correct shape and dtype on CPU tensors.
# ---------------------------------------------------------------------------


class TestKorniaPipelineForwardPass:
    """build_kornia_pipeline output passes through without shape/dtype errors."""

    @pytest.fixture(autouse=True)
    def _require_kornia(self):
        pytest.importorskip("kornia")

    def test_forward_pass_shape_and_dtype(self):
        """Pipeline output images have same shape as input; boxes shape is [B, N, 4]."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"HorizontalFlip": {"p": 1.0}}, resolution=64)

        batch_size, channels, image_height, image_width = 2, 3, 64, 64
        img = torch.rand(batch_size, channels, image_height, image_width)
        boxes = torch.tensor([[[0.0, 0.0, 32.0, 32.0]], [[10.0, 10.0, 50.0, 50.0]]], dtype=torch.float32)

        img_out, boxes_out = pipeline(img, boxes)

        assert img_out.shape == (batch_size, channels, image_height, image_width), (
            f"Image shape changed: {img_out.shape}"
        )
        assert img_out.dtype == torch.float32
        assert boxes_out.shape == (batch_size, 1, 4), f"Boxes shape wrong: {boxes_out.shape}"

    def test_forward_pass_empty_boxes(self):
        """Pipeline handles a batch where N_max=0 (no boxes) without error."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"HorizontalFlip": {"p": 1.0}}, resolution=32)

        batch_size, channels, image_height, image_width = 2, 3, 32, 32
        img = torch.rand(batch_size, channels, image_height, image_width)
        # [B, 0, 4] — no boxes
        boxes = torch.zeros(batch_size, 0, 4, dtype=torch.float32)

        img_out, boxes_out = pipeline(img, boxes)

        assert img_out.shape == (batch_size, channels, image_height, image_width)
        assert boxes_out.shape == (batch_size, 0, 4)


# ---------------------------------------------------------------------------
# TestCollateMasks — validates packing of variable-length per-image masks
# into a zero-padded [B, N_max, H, W] float32 tensor.
# ---------------------------------------------------------------------------


class TestCollateMasks:
    """collate_masks packs [N_i, H, W] instance masks into [B, N_max, H, W]."""

    def _make_targets_with_masks(self, mask_counts, h=16, w=16):
        """Build target dicts with boolean mask tensors for given instance counts."""
        targets = []
        for n in mask_counts:
            masks = torch.ones(n, h, w, dtype=torch.bool) if n > 0 else torch.zeros(0, h, w, dtype=torch.bool)
            targets.append({"masks": masks, "boxes": torch.zeros(n, 4)})
        return targets

    def test_normal_batch(self):
        """Batch of [2 masks, 3 masks] → shape [2, 3, H, W] float32."""
        from rfdetr.datasets.kornia_transforms import collate_masks

        targets = self._make_targets_with_masks([2, 3])
        masks_padded = collate_masks(targets, torch.device("cpu"), n_max=3, image_height=16, image_width=16)

        assert masks_padded.shape == (2, 3, 16, 16), f"Expected (2, 3, 16, 16), got {masks_padded.shape}"
        assert masks_padded.dtype == torch.float32, f"Expected float32, got {masks_padded.dtype}"

    def test_padding_is_zero(self):
        """Padded slots (beyond real instance count) are filled with zeros."""
        from rfdetr.datasets.kornia_transforms import collate_masks

        targets = self._make_targets_with_masks([1, 3])  # image 0 padded to 3
        masks_padded = collate_masks(targets, torch.device("cpu"), n_max=3, image_height=16, image_width=16)

        # Image 0: slot 0 real (ones), slots 1-2 zero-padded
        assert masks_padded[0, 0].min() == pytest.approx(1.0), "Real mask slot must be all ones"
        assert masks_padded[0, 1].max() == pytest.approx(0.0), "Padded slot 1 must be all zeros"
        assert masks_padded[0, 2].max() == pytest.approx(0.0), "Padded slot 2 must be all zeros"

    def test_n_max_zero_returns_empty(self):
        """n_max=0 → shape [B, 0, H, W]."""
        from rfdetr.datasets.kornia_transforms import collate_masks

        targets = self._make_targets_with_masks([0, 0])
        masks_padded = collate_masks(targets, torch.device("cpu"), n_max=0, image_height=16, image_width=16)

        assert masks_padded.shape == (2, 0, 16, 16), f"Expected (2, 0, 16, 16), got {masks_padded.shape}"

    def test_empty_target_list(self):
        """Empty target list → shape [0, 0, H, W]."""
        from rfdetr.datasets.kornia_transforms import collate_masks

        masks_padded = collate_masks([], torch.device("cpu"), n_max=0, image_height=16, image_width=16)

        assert masks_padded.shape == (0, 0, 16, 16), f"Expected (0, 0, 16, 16), got {masks_padded.shape}"

    def test_targets_without_masks_key(self):
        """Targets without 'masks' key produce all-zero rows."""
        from rfdetr.datasets.kornia_transforms import collate_masks

        targets = [{"boxes": torch.zeros(2, 4)}, {"boxes": torch.zeros(1, 4)}]
        masks_padded = collate_masks(targets, torch.device("cpu"), n_max=2, image_height=8, image_width=8)

        assert masks_padded.shape == (2, 2, 8, 8)
        assert masks_padded.max() == pytest.approx(0.0), "Targets without masks key must produce all-zero output"


# ---------------------------------------------------------------------------
# TestBuildKorniaPipelineWithMasks — validates that with_masks=True produces
# a pipeline with mask data_key included.
# ---------------------------------------------------------------------------


class TestBuildKorniaPipelineWithMasks:
    """build_kornia_pipeline(with_masks=True) includes mask in data_keys."""

    @pytest.fixture(autouse=True)
    def _require_kornia(self):
        """Skip when Kornia is unavailable (optional extra not installed in CPU CI)."""
        pytest.importorskip("kornia")

    def test_with_masks_false_is_default(self):
        """with_masks defaults to False; pipeline returns (img, boxes) on call."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"HorizontalFlip": {"p": 1.0}}, resolution=32)
        img = torch.rand(1, 3, 32, 32)
        boxes = torch.tensor([[[0.0, 0.0, 16.0, 16.0]]])
        result = pipeline(img, boxes)
        assert len(result) == 2, f"Detection pipeline must return 2 values, got {len(result)}"

    def test_with_masks_true_returns_three_values(self):
        """with_masks=True: pipeline(img, boxes, masks) returns (img, boxes, masks)."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"HorizontalFlip": {"p": 1.0}}, resolution=32, with_masks=True)
        img = torch.rand(1, 3, 32, 32)
        boxes = torch.tensor([[[0.0, 0.0, 16.0, 16.0]]])
        masks = torch.ones(1, 1, 32, 32, dtype=torch.float32)
        result = pipeline(img, boxes, masks)
        assert len(result) == 3, f"Segmentation pipeline must return 3 values, got {len(result)}"

    def test_with_masks_true_preserves_mask_shape(self):
        """Mask shape [B, N, H, W] is preserved after pipeline pass."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        pipeline = build_kornia_pipeline({"HorizontalFlip": {"p": 0.0}}, resolution=32, with_masks=True)
        img = torch.rand(2, 3, 32, 32)
        boxes = torch.tensor([[[0.0, 0.0, 16.0, 16.0]], [[8.0, 8.0, 24.0, 24.0]]])
        masks = torch.ones(2, 1, 32, 32, dtype=torch.float32)
        _, _, masks_aug = pipeline(img, boxes, masks)
        assert masks_aug.shape == (2, 1, 32, 32), f"Mask shape must be preserved: {masks_aug.shape}"


# ---------------------------------------------------------------------------
# TestUnpackBoxesWithMasks — validates that unpack_boxes propagates the same
# keep filter to masks when masks_aug is provided.
# ---------------------------------------------------------------------------


class TestUnpackBoxesWithMasks:
    """unpack_boxes with masks_aug keeps/removes masks in sync with boxes."""

    def test_masks_filtered_same_as_boxes(self):
        """Box removed → corresponding mask also removed from output."""
        from rfdetr.datasets.kornia_transforms import unpack_boxes

        # B=1, N=2: box 0 valid, box 1 zero-area (will be removed)
        boxes_aug = torch.tensor([[[5.0, 5.0, 25.0, 25.0], [30.0, 30.0, 30.0, 30.0]]])
        valid = torch.tensor([[True, True]])
        targets = [
            {
                "boxes": torch.tensor([[5.0, 5.0, 25.0, 25.0], [30.0, 30.0, 60.0, 60.0]]),
                "labels": torch.tensor([1, 2]),
            }
        ]
        # 2 masks: instance 0 = all ones, instance 1 = all twos (distinguishable)
        masks_aug = torch.zeros(1, 2, 8, 8, dtype=torch.float32)
        masks_aug[0, 0] = 1.0
        masks_aug[0, 1] = 1.0  # will be removed with box 1

        result = unpack_boxes(boxes_aug, valid, targets, 100, 100, masks_aug=masks_aug)

        assert "masks" in result[0], "masks key must be present in output target"
        assert result[0]["masks"].shape[0] == 1, f"Expected 1 surviving mask, got {result[0]['masks'].shape[0]}"

    def test_masks_converted_to_bool(self):
        """Float masks > 0.5 threshold converted to bool in output."""
        from rfdetr.datasets.kornia_transforms import unpack_boxes

        boxes_aug = torch.tensor([[[5.0, 5.0, 25.0, 25.0]]])
        valid = torch.tensor([[True]])
        targets = [{"boxes": torch.tensor([[5.0, 5.0, 25.0, 25.0]]), "labels": torch.tensor([1])}]
        masks_aug = torch.full((1, 1, 8, 8), 0.8, dtype=torch.float32)  # float, all 0.8

        result = unpack_boxes(boxes_aug, valid, targets, 100, 100, masks_aug=masks_aug)

        assert result[0]["masks"].dtype == torch.bool, f"masks must be bool, got {result[0]['masks'].dtype}"
        assert result[0]["masks"].all(), "All values > 0.5 should be True after thresholding"

    def test_no_masks_aug_leaves_masks_key_unchanged(self):
        """When masks_aug=None, existing masks key in target is preserved as-is."""
        from rfdetr.datasets.kornia_transforms import unpack_boxes

        boxes_aug = torch.tensor([[[5.0, 5.0, 25.0, 25.0]]])
        valid = torch.tensor([[True]])
        original_mask = torch.ones(1, 8, 8, dtype=torch.bool)
        targets = [
            {
                "boxes": torch.tensor([[5.0, 5.0, 25.0, 25.0]]),
                "labels": torch.tensor([1]),
                "masks": original_mask,
            }
        ]

        result = unpack_boxes(boxes_aug, valid, targets, 100, 100, masks_aug=None)

        assert "masks" in result[0], "masks key must still be present when masks_aug=None"
        assert result[0]["masks"] is original_mask, "Original masks object must be preserved unchanged"


class TestGaussNoiseStdRangeWarning:
    """_make_gauss_noise warns when the configured std range is non-degenerate (GPU uses a fixed upper-bound std)."""

    @pytest.fixture(autouse=True)
    def _require_kornia(self):
        pytest.importorskip("kornia")

    def test_warns_for_unequal_std_range(self):
        """A non-degenerate std_range emits a divergence warning at build time."""
        from unittest import mock

        from rfdetr.datasets import kornia_transforms

        with mock.patch.object(kornia_transforms.logger, "warning") as mock_warning:
            kornia_transforms._make_gauss_noise({"std_range": (0.01, 0.05), "p": 0.5})

        mock_warning.assert_called_once()

    def test_no_warning_for_degenerate_std_range(self):
        """An equal-bound std_range matches the CPU path exactly and stays silent."""
        from unittest import mock

        from rfdetr.datasets import kornia_transforms

        with mock.patch.object(kornia_transforms.logger, "warning") as mock_warning:
            kornia_transforms._make_gauss_noise({"std_range": (0.05, 0.05), "p": 0.5})

        mock_warning.assert_not_called()


class TestToGrayDroppedParamsWarning:
    """_make_to_gray warns when passed method/num_output_channels, which have no Kornia equivalent."""

    @pytest.fixture(autouse=True)
    def _require_kornia(self):
        pytest.importorskip("kornia")

    def test_warns_for_method(self):
        """A non-default method emits a dropped-param warning at build time."""
        from unittest import mock

        from rfdetr.datasets import kornia_transforms

        with mock.patch.object(kornia_transforms.logger, "warning") as mock_warning:
            kornia_transforms._make_to_gray({"method": "max", "p": 0.5})

        mock_warning.assert_called_once()

    def test_warns_for_num_output_channels(self):
        """A non-default num_output_channels emits a dropped-param warning at build time."""
        from unittest import mock

        from rfdetr.datasets import kornia_transforms

        with mock.patch.object(kornia_transforms.logger, "warning") as mock_warning:
            kornia_transforms._make_to_gray({"num_output_channels": 1, "p": 0.5})

        mock_warning.assert_called_once()

    def test_no_warning_for_p_only(self):
        """A config using only p matches the CPU path's default behavior and stays silent."""
        from unittest import mock

        from rfdetr.datasets import kornia_transforms

        with mock.patch.object(kornia_transforms.logger, "warning") as mock_warning:
            kornia_transforms._make_to_gray({"p": 0.5})

        mock_warning.assert_not_called()


# ---------------------------------------------------------------------------
# TestResolveAugmentationBackend — the single resolution seam that maps backend
# strings (incl. sentinels/legacy aliases) to concrete AugmentationBackend members.
# ---------------------------------------------------------------------------


class TestResolveAugmentationBackend:
    """resolve_augmentation_backend maps backend strings to concrete AugmentationBackend members."""

    @pytest.mark.parametrize(
        "value",
        [
            pytest.param("auto", id="auto"),
            pytest.param("cpu", id="cpu"),
        ],
    )
    def test_falls_back_to_tv_when_no_optional_packages(self, value: str) -> None:
        """'cpu'/'auto' resolve to torchvision when neither Albumentations nor Kornia is installed."""
        from unittest.mock import patch

        from rfdetr.config import AugmentationBackend
        from rfdetr.datasets.kornia_transforms import resolve_augmentation_backend

        with (
            patch.object(AugmentationBackend, "_is_albu_available", return_value=False),
            patch.object(AugmentationBackend, "_is_kornia_available", return_value=False),
        ):
            assert resolve_augmentation_backend(value, has_cuda=False) == AugmentationBackend.TV

    @pytest.mark.parametrize(
        ("value", "expected"),
        [
            pytest.param("kornia", "kornia", id="kornia_passthrough"),
            pytest.param("gpu", "kornia", id="gpu_alias_to_kornia"),
            pytest.param("torchvision", "torchvision", id="torchvision_passthrough"),
            pytest.param("tv", "torchvision", id="tv_alias_to_torchvision"),
        ],
    )
    def test_explicit_backend_passthrough_regardless_of_cuda(self, value: str, expected: str) -> None:
        """Explicit concrete backends and their legacy aliases ('gpu', 'tv') resolve regardless of CUDA."""
        from rfdetr.config import AugmentationBackend
        from rfdetr.datasets.kornia_transforms import resolve_augmentation_backend

        assert resolve_augmentation_backend(value, has_cuda=False) == AugmentationBackend(expected)

    def test_albu_passthrough_when_installed(self) -> None:
        """Explicit legacy 'albu' resolves to ALBU when Albumentations is installed."""
        from unittest.mock import patch

        from rfdetr.config import AugmentationBackend
        from rfdetr.datasets.kornia_transforms import resolve_augmentation_backend

        with patch.object(AugmentationBackend, "_is_albu_available", return_value=True):
            assert resolve_augmentation_backend("albu", has_cuda=False) == AugmentationBackend.ALBU

    def test_albu_missing_raises_import_error(self) -> None:
        """Explicit 'albu' fails fast with an install hint when Albumentations is not installed."""
        from unittest.mock import patch

        from rfdetr.config import AugmentationBackend
        from rfdetr.datasets.kornia_transforms import resolve_augmentation_backend

        with (
            patch.object(AugmentationBackend, "_is_albu_available", return_value=False),
            pytest.raises(ImportError, match=r"rfdetr\[augment\]"),
        ):
            resolve_augmentation_backend("albu", has_cuda=False)


class TestResolveBackendForBuild:
    """resolve_backend_for_build combines the GPU readiness fail-fast with backend resolution."""

    def test_resolves_concrete_backend_without_cuda(self) -> None:
        """A concrete non-GPU backend resolves without requiring a CUDA device."""
        from rfdetr.config import AugmentationBackend
        from rfdetr.datasets.kornia_transforms import resolve_backend_for_build

        assert resolve_backend_for_build("torchvision", has_cuda=False) == AugmentationBackend.TV

    def test_gpu_without_cuda_raises_runtime_error(self) -> None:
        """An explicit GPU request fails fast when no CUDA device is available."""
        from rfdetr.datasets.kornia_transforms import resolve_backend_for_build

        with pytest.raises(RuntimeError, match="CUDA"):
            resolve_backend_for_build("gpu", has_cuda=False)

    def test_gpu_without_kornia_raises_import_error(self) -> None:
        """An explicit GPU request with CUDA but no Kornia fails fast with an install hint."""
        from unittest.mock import patch

        from rfdetr.config import AugmentationBackend
        from rfdetr.datasets.kornia_transforms import resolve_backend_for_build

        with (
            patch.object(AugmentationBackend, "_is_kornia_available", return_value=False),
            pytest.raises(ImportError, match=r"rfdetr\[augment\]"),
        ):
            resolve_backend_for_build("gpu", has_cuda=True)


class TestPerspectiveFactory:
    """`Perspective` on the Kornia backend (issue #1252).

    Perspective preserves output resolution, and the DataModule carries the batch padding mask through the same Kornia
    sequence. Transforms that resize output remain unsupported because this path only transports fixed-size batches.
    """

    @pytest.mark.parametrize("scale", [(0.05, 0.2), 0.2], ids=["range", "scalar"])
    def test_distribution_divergence_is_always_reported(self, scale) -> None:
        """The GPU path draws uniformly where the CPU path draws a half-normal, so every config diverges.

        This holds for a scalar too: albumentations reads ``0.2`` as ``sigma`` in ``(0, 0.2)`` and samples
        ``abs(N(0, sigma))``, so even an "exact" request is not the same distribution Kornia produces.
        """
        from unittest import mock

        from rfdetr.datasets import kornia_transforms

        with mock.patch.object(kornia_transforms.logger, "warning") as warn:
            kornia_transforms.build_kornia_pipeline({"Perspective": {"scale": scale}}, 560)

        messages = [call[0][0] for call in warn.call_args_list]
        assert any("Perspective" in m and "abs(N(0, sigma))" in m for m in messages), messages

    @pytest.mark.parametrize(
        "key,value",
        [
            ("fit_output", True),
            ("interpolation", 1),
            ("mask_interpolation", 0),
            ("border_mode", 0),
            ("fill", 0),
            ("fill_mask", 0),
        ],
    )
    def test_unmappable_options_are_reported_not_silently_dropped(self, key, value) -> None:
        """Kornia's RandomPerspective exposes only distortion_scale and p; the rest must not vanish quietly."""
        from unittest import mock

        from rfdetr.datasets import kornia_transforms

        with mock.patch.object(kornia_transforms.logger, "warning") as warn:
            kornia_transforms.build_kornia_pipeline({"Perspective": {key: value}}, 560)

        messages = [call[0][0] % call[0][1:] if len(call[0]) > 1 else call[0][0] for call in warn.call_args_list]
        assert any("ignores" in m and key in m for m in messages), messages

    def test_keep_size_false_is_refused_not_ignored(self) -> None:
        """keep_size=False changes the output resolution, which this pipeline cannot express."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        with pytest.raises(ValueError, match="keep_size=False"):
            build_kornia_pipeline({"Perspective": {"keep_size": False}}, 560)

    def test_output_keeps_the_input_resolution(self) -> None:
        """The property the whole mapping rests on: image height and width survive the transform."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        img = torch.rand(2, 3, 64, 64)
        boxes = torch.tensor([[[8.0, 8.0, 40.0, 40.0]], [[4.0, 4.0, 20.0, 20.0]]])

        pipeline = build_kornia_pipeline({"Perspective": {"scale": 0.3, "p": 1.0}}, 560)
        img_out, _ = pipeline(img, boxes)

        assert img_out.shape[-2:] == img.shape[-2:]

    def test_boxes_follow_the_warp(self) -> None:
        """A geometric transform that left the boxes where they were would silently mislabel every image."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        img = torch.rand(1, 3, 64, 64)
        boxes = torch.tensor([[[8.0, 8.0, 40.0, 40.0]]])

        pipeline = build_kornia_pipeline({"Perspective": {"scale": 0.4, "p": 1.0}}, 560)
        _, boxes_out = pipeline(img, boxes)

        assert not torch.allclose(boxes_out, boxes), "boxes must be warped with the image"

    def test_padding_mask_follows_the_perspective_warp(self) -> None:
        """Batch padding stays an auxiliary mask under the same Perspective parameters as image and boxes."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        torch.manual_seed(7)
        image = torch.zeros(1, 3, 64, 64)
        image[:, :, :48, :48] = 1.0
        boxes = torch.tensor([[[4.0, 4.0, 44.0, 44.0]]])
        instance_mask = torch.zeros(1, 1, 64, 64)
        instance_mask[:, :, 12:36, 12:36] = 1.0
        padding_mask = torch.ones(1, 1, 64, 64)
        padding_mask[:, :, :48, :48] = 0.0
        auxiliary_masks = torch.cat((instance_mask, padding_mask), dim=1)

        pipeline = build_kornia_pipeline({"Perspective": {"scale": 0.4, "p": 1.0}}, 560, with_masks=True)
        image_aug, boxes_aug, auxiliary_masks_aug = pipeline(image, boxes, auxiliary_masks)
        instance_mask_aug = auxiliary_masks_aug[:, :1]
        padding_mask_aug = auxiliary_masks_aug[:, 1:]

        assert image_aug.shape == image.shape
        assert boxes_aug.shape == boxes.shape
        assert padding_mask_aug.shape == padding_mask.shape
        assert instance_mask_aug.shape == instance_mask.shape
        assert not torch.equal(padding_mask_aug, padding_mask)
        assert not torch.equal(instance_mask_aug, instance_mask)
        bright_pixels = image_aug[:, 0] > 0.99
        assert not padding_mask_aug[:, 0].to(torch.bool)[bright_pixels].any()

    @pytest.mark.parametrize("name", ["RandomCrop", "CenterCrop", "RandomResizedCrop"])
    def test_crop_names_from_1252_remain_unsupported(self, name) -> None:
        """Guard for the reason Perspective ships alone: the crops resize, so they are still rejected."""
        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline

        with pytest.raises(ValueError, match="Unknown augmentation key"):
            build_kornia_pipeline({name: {"height": 32, "width": 32}}, 560)