Use augmentation disable intent for crop branch

Co-authored-by: Codex <codex@openai.com>

Author: Borda

src/rfdetr/datasets/coco.py

@@ -294,7 +294,7 @@ def _build_train_resize_config(
     *,
     square: bool,
     max_size: Optional[int] = None,
-    include_crop_branch: bool = True,
+    disable_augmentations: bool = False,
 ) -> List[Dict[str, Any]]:
     """Build the training resize pipeline as an Albumentations config list.
 
@@ -317,15 +317,14 @@ def _build_train_resize_config(
             optional long-side cap.
         max_size: Maximum long-side size for non-square resizes.  Defaults to
             ``1333`` when *square* is ``False``.
-        include_crop_branch: If ``False``, omit the resize-and-crop branch so the
-            pipeline always uses Option A (direct resize). Useful when objects
-            of interest can be partially or fully cropped out of frame.
-            Defaults to ``True`` (preserves the original two-branch behavior).
+        disable_augmentations: If ``True``, omit the resize-and-crop branch so
+            explicitly disabled augmentations do not still randomly crop images.
+            Defaults to ``False`` (preserves the original two-branch behavior).
 
     Returns:
-        A single-element list. When ``include_crop_branch`` is ``True`` (default)
-        the entry wraps a ``OneOf`` over both branches; when ``False`` the
-        entry is Option A directly.
+        A single-element list. By default the entry wraps a ``OneOf`` over both
+        branches; when augmentations are disabled, the entry is Option A
+        directly.
     """
     if square:
         option_a: Dict[str, Any] = {
@@ -371,14 +370,14 @@ def _build_train_resize_config(
             }
         }
 
-    if not include_crop_branch:
+    if disable_augmentations:
         return [option_a]
 
     return [{"OneOf": {"transforms": [option_a, option_b]}}]
 
 
-def _crop_branch_enabled(aug_config: Optional[Dict[str, Any]]) -> bool:
-    """Decide whether the training resize pipeline keeps its resize-and-crop branch.
+def _augmentations_disabled(aug_config: Optional[Dict[str, Any]]) -> bool:
+    """Decide whether the user explicitly disabled augmentations.
 
     ``aug_config={}`` is an explicit request to disable augmentations; it also
     drops the resize-and-crop branch. ``aug_config=None`` (the default) and any
@@ -390,8 +389,8 @@ def _crop_branch_enabled(aug_config: Optional[Dict[str, Any]]) -> bool:
             "augmentations; images will not be randomly cropped. Pass aug_config=None to keep "
             "the default resize pipeline."
         )
-        return False
-    return True
+        return True
+    return False
 
 
 def make_coco_transforms(
@@ -468,9 +467,13 @@ def make_coco_transforms(
 
     if image_set == "train":
         resolved_aug_config = aug_config if aug_config is not None else AUG_CONFIG
-        include_crop_branch = _crop_branch_enabled(aug_config)
         resize_wrappers = AlbumentationsWrapper.from_config(
-            _build_train_resize_config(scales, square=False, max_size=1333, include_crop_branch=include_crop_branch)
+            _build_train_resize_config(
+                scales,
+                square=False,
+                max_size=1333,
+                disable_augmentations=_augmentations_disabled(aug_config),
+            )
         )
         pipeline = [*resize_wrappers]
         if not gpu_postprocess:
@@ -556,9 +559,12 @@ def make_coco_transforms_square_div_64(
 
     if image_set == "train":
         resolved_aug_config = aug_config if aug_config is not None else AUG_CONFIG
-        include_crop_branch = _crop_branch_enabled(aug_config)
         resize_wrappers = AlbumentationsWrapper.from_config(
-            _build_train_resize_config(scales, square=True, include_crop_branch=include_crop_branch)
+            _build_train_resize_config(
+                scales,
+                square=True,
+                disable_augmentations=_augmentations_disabled(aug_config),
+            )
         )
         pipeline = [*resize_wrappers]
         if not gpu_postprocess:

tests/datasets/test_coco.py

@@ -454,36 +454,36 @@ class TestAugConfigDisablesCrop:
     @pytest.mark.parametrize(
         "aug_config,expected",
         [
-            pytest.param(None, True, id="none_keeps_crop"),
-            pytest.param({}, False, id="empty_disables_crop"),
-            pytest.param({"HorizontalFlip": {"p": 0.5}}, True, id="nonempty_keeps_crop"),
+            pytest.param(None, False, id="none_keeps_augmentations"),
+            pytest.param({}, True, id="empty_disables_augmentations"),
+            pytest.param({"HorizontalFlip": {"p": 0.5}}, False, id="nonempty_keeps_augmentations"),
         ],
     )
-    def test_crop_branch_enabled(self, aug_config, expected):
-        """_crop_branch_enabled maps aug_config to the crop-branch decision."""
-        from rfdetr.datasets.coco import _crop_branch_enabled
+    def test_augmentations_disabled(self, aug_config, expected):
+        """_augmentations_disabled maps aug_config to the explicit disable decision."""
+        from rfdetr.datasets.coco import _augmentations_disabled
 
-        assert _crop_branch_enabled(aug_config) is expected
+        assert _augmentations_disabled(aug_config) is expected
 
     def test_empty_aug_config_warns(self):
         """Passing aug_config={} logs a warning about the dropped resize-and-crop branch."""
         from unittest.mock import patch
 
-        from rfdetr.datasets.coco import _crop_branch_enabled
+        from rfdetr.datasets.coco import _augmentations_disabled
 
         with patch("rfdetr.datasets.coco.logger") as mock_logger:
-            _crop_branch_enabled({})
+            _augmentations_disabled({})
         mock_logger.warning.assert_called_once()
 
     @pytest.mark.parametrize(
         "aug_config,expected",
         [
-            pytest.param(None, True, id="none_keeps_crop"),
-            pytest.param({}, False, id="empty_disables_crop"),
+            pytest.param(None, False, id="none_keeps_augmentations"),
+            pytest.param({}, True, id="empty_disables_augmentations"),
         ],
     )
-    def test_make_coco_transforms_forwards_crop_decision(self, aug_config, expected):
-        """make_coco_transforms passes the aug_config-derived value to _build_train_resize_config."""
+    def test_make_coco_transforms_forwards_disable_augmentations(self, aug_config, expected):
+        """make_coco_transforms passes the aug_config-derived disable decision to resize config."""
         from unittest.mock import patch
 
         from rfdetr.datasets.coco import make_coco_transforms
@@ -497,4 +497,4 @@ def test_make_coco_transforms_forwards_crop_decision(self, aug_config, expected)
         ):
             make_coco_transforms("train", 640, aug_config=aug_config)
 
-        assert mock_build.call_args.kwargs["include_crop_branch"] is expected
+        assert mock_build.call_args.kwargs["disable_augmentations"] is expected

tests/datasets/test_coco_resize_config.py

@@ -254,15 +254,15 @@ def test_square_option_b_unchanged(self, scales, square):
 
 
 class TestBuildTrainResizeConfigCropBranch:
-    """include_crop_branch=False drops the resize-and-crop branch so only Option A runs."""
+    """disable_augmentations=True drops the resize-and-crop branch so only Option A runs."""
 
     @pytest.mark.parametrize(
         "square",
         [pytest.param(True, id="square"), pytest.param(False, id="nonsquare")],
     )
-    def test_include_crop_branch_false_drops_crop_branch(self, square):
+    def test_disable_augmentations_drops_crop_branch(self, square):
         """No RandomCrop/RandomSizedCrop appears in either square or non-square pipelines."""
-        result = _build_train_resize_config([480, 640], square=square, include_crop_branch=False)
+        result = _build_train_resize_config([480, 640], square=square, disable_augmentations=True)
         flat = str(result)
         assert "RandomSizedCrop" not in flat
         assert "RandomCrop" not in flat