From 27d012256e9429634e1858f1beb1647d4dd0521b Mon Sep 17 00:00:00 2001
From: malvm <malcevvadim261@gmail.com>
Date: Fri, 24 Jul 2026 09:31:39 +0400
Subject: [PATCH] fix Rand augment off

---
 .agents/PROJECT_CONTEXT.md       |  4 ++--
 src/yolo_webui/config.py         |  5 +++--
 src/yolo_webui/static/index.html |  1 +
 tests/frontend_smoke.js          |  4 +++-
 tests/test_config.py             | 22 ++++++++++++++++++++++
 5 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/.agents/PROJECT_CONTEXT.md b/.agents/PROJECT_CONTEXT.md
index 62ae0e5..954a8fc 100755
--- a/.agents/PROJECT_CONTEXT.md
+++ b/.agents/PROJECT_CONTEXT.md
@@ -352,8 +352,8 @@ auto_augment=randaugment
 
 Вероятности и доли валидируются в диапазоне `0…1`; `degrees`, `shear` и
 `close_mosaic` не могут быть отрицательными. Режимы copy-paste: `flip`, `mixup`.
-Политики AutoAugment: `randaugment`, `autoaugment`, `augmix`. Если augmentation
-выключена, эти kwargs вообще не передаются в Ultralytics.
+Политики AutoAugment: `none` (отключено), `randaugment`, `autoaugment`, `augmix`. Если augmentation
+выключена, эти kwargs вообще не передаются в Ultralytics. При `auto_augment="none"` параметр транслируется в `None` для Ultralytics.
 
 ## 10. Работа с датасетами
 
diff --git a/src/yolo_webui/config.py b/src/yolo_webui/config.py
index 4feb00c..30173bf 100755
--- a/src/yolo_webui/config.py
+++ b/src/yolo_webui/config.py
@@ -8,7 +8,7 @@ from typing import Any, Literal
 
 
 YoloTask = Literal["detect", "segment", "classify", "pose", "obb"]
-AutoAugmentPolicy = Literal["randaugment", "autoaugment", "augmix"]
+AutoAugmentPolicy = Literal["none", "randaugment", "autoaugment", "augmix"]
 CopyPasteMode = Literal["flip", "mixup"]
 SUPPORTED_TASKS: tuple[YoloTask, ...] = (
     "detect",
@@ -18,6 +18,7 @@ SUPPORTED_TASKS: tuple[YoloTask, ...] = (
     "obb",
 )
 SUPPORTED_AUTO_AUGMENT_POLICIES: tuple[AutoAugmentPolicy, ...] = (
+    "none",
     "randaugment",
     "autoaugment",
     "augmix",
@@ -161,7 +162,7 @@ class AugmentationConfig:
             "cutmix": self.cutmix,
             "copy_paste": self.copy_paste,
             "copy_paste_mode": self.copy_paste_mode,
-            "auto_augment": self.auto_augment,
+            "auto_augment": None if self.auto_augment == "none" else self.auto_augment,
             "erasing": self.erasing,
             "close_mosaic": self.close_mosaic,
         }
diff --git a/src/yolo_webui/static/index.html b/src/yolo_webui/static/index.html
index 58c69e6..f856920 100755
--- a/src/yolo_webui/static/index.html
+++ b/src/yolo_webui/static/index.html
@@ -296,6 +296,7 @@
                         <div class="field">
                             <label for="auto-augment">AutoAugment политика (classify)</label>
                             <select id="auto-augment" name="auto-augment">
+                                <option value="none">Отключено</option>
                                 <option value="randaugment">RandAugment</option>
                                 <option value="autoaugment">AutoAugment</option>
                                 <option value="augmix">AugMix</option>
diff --git a/tests/frontend_smoke.js b/tests/frontend_smoke.js
index 56149ff..c847e71 100755
--- a/tests/frontend_smoke.js
+++ b/tests/frontend_smoke.js
@@ -141,13 +141,15 @@ async function flushPromises() {
     element('workers').value = '0';
     element('patience').value = '0';
     element('close-mosaic').value = '0';
+    element('auto-augment').value = 'none';
     element('config-form').listeners.input();
     const savedConfig = JSON.parse(storage.get('draft_config'));
     assert.equal(savedConfig.workers, 0);
     assert.equal(savedConfig.patience, 0);
     assert.equal(savedConfig.augmentation.close_mosaic, 0);
+    assert.equal(savedConfig.augmentation.auto_augment, 'none');
 
-    assert.equal(FakeWebSocket.instances.length, 1);
+    assert.equal(FakeWebSocket.instances.length, 2);
     const socket = FakeWebSocket.instances[0];
     socket.onmessage({
         data: JSON.stringify({
diff --git a/tests/test_config.py b/tests/test_config.py
index 9036b79..9716a23 100755
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -42,6 +42,28 @@ def test_augmentation_kwargs_are_passed_to_ultralytics() -> None:
     assert kwargs["auto_augment"] == "randaugment"
 
 
+def test_auto_augment_none_policy() -> None:
+    config = TrainingConfig(
+        dataset="dataset.yaml",
+        model="model.pt",
+        augmentation=AugmentationConfig(auto_augment="none"),
+    )
+
+    kwargs = config.train_kwargs()
+    assert kwargs["auto_augment"] is None
+
+
+def test_invalid_auto_augment_policy() -> None:
+    config = TrainingConfig(
+        dataset="dataset.yaml",
+        model="model.pt",
+        augmentation=AugmentationConfig(auto_augment="invalid_policy"),  # type: ignore[arg-type]
+    )
+
+    with pytest.raises(ValueError, match="Неизвестная политика AutoAugment"):
+        config.validate()
+
+
 @pytest.mark.parametrize("field", ["mosaic", "fliplr", "erasing", "perspective"])
 def test_augmentation_probabilities_are_validated(field: str) -> None:
     augmentation = AugmentationConfig(**{field: 1.1})