init

.gitignore

@@ -0,0 +1,12 @@
+.venv/
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.coverage
+dist/
+build/
+runs/
+mlruns/
+mlflow.db
+mlflow.db-shm
+mlflow.db-wal

README.md

@@ -0,0 +1,60 @@
+# YOLO Train TUI
+
+Терминальный интерфейс для обучения моделей Ultralytics YOLO с автоматической
+регистрацией параметров, метрик и артефактов в MLflow.
+
+## Возможности
+
+- задачи `detect`, `segment`, `classify`, `pose` и `obb`;
+- локальные пути, YAML-конфигурации и официальные имена моделей/датасетов;
+- настройка эпох, размера изображения, batch, устройства, workers и patience;
+- настройка цветовых и геометрических аугментаций, flip, Mosaic, MixUp,
+  CutMix, copy-paste, erasing и AutoAugment;
+- обучение в фоновом потоке, прогресс по эпохам, журнал и мягкая остановка;
+- встроенная интеграция Ultralytics ↔ MLflow;
+- локальное MLflow-хранилище по умолчанию или внешний tracking server.
+
+## Установка и запуск
+
+```bash
+uv sync
+uv run yolo-train-tui
+```
+
+Также приложение можно запустить как модуль:
+
+```bash
+uv run -m yolo_tui
+```
+
+При первом использовании официального имени модели (например, `yolo11n.pt`)
+Ultralytics автоматически скачает веса. Для полностью локальной работы укажите
+путь к уже загруженному `.pt` или `.yaml` файлу.
+
+## Датасеты
+
+Для `detect`, `segment`, `pose` и `obb` укажите путь к YAML-файлу датасета.
+Для `classify` укажите каталог с подкаталогами `train`, `test`/`val`, внутри
+которых изображения разложены по классам.
+
+## MLflow
+
+По умолчанию метаданные записываются в локальную SQLite-базу `./mlflow.db`,
+сервер для обучения не требуется. Артефакты сохраняются локально средствами MLflow.
+Открыть интерфейс просмотра:
+
+```bash
+uv run mlflow ui --backend-store-uri sqlite:///mlflow.db
+```
+
+Затем откройте `http://127.0.0.1:5000`. Для удаленного MLflow-сервера включите
+MLflow в TUI и замените Tracking URI на адрес вида `http://mlflow.example:5000`.
+
+## Проверка
+
+```bash
+uv run pytest
+```
+
+Ultralytics распространяется по лицензии AGPL-3.0; для закрытых коммерческих
+продуктов проверьте условия Enterprise-лицензии Ultralytics.

pyproject.toml

@@ -0,0 +1,31 @@
+[project]
+name = "yolo-train-tui"
+version = "0.1.0"
+description = "Terminal UI for training Ultralytics YOLO models with MLflow tracking"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "mlflow>=3.0",
+    "textual>=1.0",
+    "ultralytics>=8.3",
+]
+
+[project.scripts]
+yolo-train-tui = "yolo_tui.app:main"
+
+[dependency-groups]
+dev = [
+    "pytest>=8.3",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/yolo_tui"]
+
+[tool.pytest.ini_options]
+addopts = "-q"
+testpaths = ["tests"]
+

src/yolo_tui/__init__.py

@@ -0,0 +1,8 @@
+"""YOLO Train TUI package."""
+
+from .config import MlflowConfig, TrainingConfig
+
+__all__ = ["MlflowConfig", "TrainingConfig"]
+
+__version__ = "0.1.0"
+

src/yolo_tui/__main__.py

@@ -0,0 +1,6 @@
+from .app import main
+
+
+if __name__ == "__main__":
+    main()
+

src/yolo_tui/app.py

@@ -0,0 +1,518 @@
+from __future__ import annotations
+
+from typing import Any
+
+from textual import on, work
+from textual.app import App, ComposeResult
+from textual.containers import Container, Horizontal, Vertical, VerticalScroll
+from textual.widgets import (
+    Button,
+    Footer,
+    Header,
+    Input,
+    Label,
+    ProgressBar,
+    RichLog,
+    Select,
+    Static,
+    Switch,
+)
+
+from .config import (
+    AugmentationConfig,
+    MlflowConfig,
+    SUPPORTED_AUTO_AUGMENT_POLICIES,
+    SUPPORTED_COPY_PASTE_MODES,
+    SUPPORTED_TASKS,
+    TrainingConfig,
+)
+from .trainer import TrainingEvent, TrainingRunner
+
+
+class Field(Vertical):
+    def __init__(self, label: str, control: Any, *, classes: str = "") -> None:
+        super().__init__(classes=f"field {classes}".strip())
+        self.label_text = label
+        self.control = control
+
+    def compose(self) -> ComposeResult:
+        yield Label(self.label_text)
+        yield self.control
+
+
+class YoloTrainApp(App[None]):
+    TITLE = "YOLO Train Studio"
+    SUB_TITLE = "Ultralytics + MLflow"
+
+    CSS = """
+    Screen {
+        background: #0b1020;
+        color: #dbe7ff;
+    }
+
+    Header {
+        background: #111a33;
+        color: #f5f8ff;
+    }
+
+    #workspace {
+        height: 1fr;
+        layout: horizontal;
+        padding: 1 2;
+    }
+
+    #config-pane {
+        width: 46%;
+        min-width: 48;
+        height: 100%;
+        margin-right: 2;
+        padding: 0 1 2 1;
+        border: round #314268;
+        background: #0e162b;
+    }
+
+    #run-pane {
+        width: 1fr;
+        height: 100%;
+        padding: 1 2;
+        border: round #314268;
+        background: #0e162b;
+    }
+
+    .section-title {
+        height: 2;
+        margin-top: 1;
+        color: #78a9ff;
+        text-style: bold;
+    }
+
+    .field {
+        height: auto;
+        margin-bottom: 1;
+    }
+
+    .field Label {
+        height: 1;
+        margin-left: 1;
+        color: #9fb1d1;
+    }
+
+    .field Input, .field Select {
+        width: 100%;
+    }
+
+    .row {
+        height: auto;
+    }
+
+    .row .field {
+        width: 1fr;
+        margin-right: 1;
+    }
+
+    .row .field:last-child {
+        margin-right: 0;
+    }
+
+    .toggle-row {
+        height: 3;
+        align-vertical: middle;
+    }
+
+    .toggle-row Label {
+        width: 1fr;
+        color: #dbe7ff;
+    }
+
+    .toggle-row Switch {
+        width: auto;
+    }
+
+    #status-card {
+        height: auto;
+        min-height: 5;
+        padding: 1 2;
+        margin-bottom: 1;
+        border-left: thick #5b8def;
+        background: #131f3b;
+    }
+
+    #status-title {
+        color: #78a9ff;
+        text-style: bold;
+    }
+
+    #progress {
+        margin: 1 0;
+    }
+
+    #actions {
+        height: 3;
+        margin-bottom: 1;
+    }
+
+    #actions Button {
+        width: 1fr;
+        margin-right: 1;
+    }
+
+    #actions Button:last-child {
+        margin-right: 0;
+    }
+
+    #log-title {
+        height: 2;
+        margin-top: 1;
+        color: #9fb1d1;
+        text-style: bold;
+    }
+
+    #log {
+        height: 1fr;
+        border: round #253455;
+        background: #090f1e;
+        padding: 0 1;
+    }
+
+    #hint {
+        height: auto;
+        margin-top: 1;
+        color: #7384a3;
+    }
+
+    Footer {
+        background: #111a33;
+    }
+    """
+
+    BINDINGS = [
+        ("ctrl+s", "start_training", "Запустить"),
+        ("ctrl+x", "stop_training", "Остановить"),
+        ("q", "quit", "Выход"),
+    ]
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.runner = TrainingRunner()
+        self._running = False
+
+    def compose(self) -> ComposeResult:
+        yield Header(show_clock=True)
+        with Container(id="workspace"):
+            with VerticalScroll(id="config-pane"):
+                yield Static("Модель и данные", classes="section-title")
+                yield Field(
+                    "Тип задачи",
+                    Select(
+                        [(task.capitalize(), task) for task in SUPPORTED_TASKS],
+                        value="detect",
+                        id="task",
+                        allow_blank=False,
+                    ),
+                )
+                yield Field(
+                    "Модель — путь, .pt/.yaml или официальное имя",
+                    Input(value="yolo11n.pt", placeholder="/models/best.pt", id="model"),
+                )
+                yield Field(
+                    "Датасет — путь к YAML/каталогу или имя",
+                    Input(value="coco8.yaml", placeholder="/data/dataset.yaml", id="dataset"),
+                )
+
+                yield Static("Параметры обучения", classes="section-title")
+                with Horizontal(classes="row"):
+                    yield Field("Эпохи", Input(value="100", type="integer", id="epochs"))
+                    yield Field("Размер", Input(value="640", type="integer", id="image-size"))
+                    yield Field("Batch", Input(value="16", type="integer", id="batch-size"))
+                with Horizontal(classes="row"):
+                    yield Field("Device", Input(placeholder="cpu, 0, 0,1", id="device"))
+                    yield Field("Workers", Input(value="8", type="integer", id="workers"))
+                    yield Field("Patience", Input(value="100", type="integer", id="patience"))
+                with Horizontal(classes="row"):
+                    yield Field("Каталог результатов", Input(value="runs/train", id="project"))
+                    yield Field("Имя запуска", Input(placeholder="experiment-01", id="run-name"))
+
+                yield Static("Аугментация", classes="section-title")
+                with Horizontal(classes="toggle-row"):
+                    yield Label("Передавать свои параметры аугментации в Ultralytics")
+                    yield Switch(value=True, id="augmentation-enabled")
+                with Horizontal(classes="row augmentation-field"):
+                    yield Field("HSV hue · 0…1", Input(value="0.015", id="hsv-h"))
+                    yield Field("HSV saturation · 0…1", Input(value="0.7", id="hsv-s"))
+                    yield Field("HSV brightness · 0…1", Input(value="0.4", id="hsv-v"))
+                with Horizontal(classes="row augmentation-field"):
+                    yield Field("Поворот · градусы", Input(value="0.0", id="degrees"))
+                    yield Field("Смещение · 0…1", Input(value="0.1", id="translate"))
+                    yield Field("Масштаб · 0…1", Input(value="0.5", id="scale"))
+                with Horizontal(classes="row augmentation-field"):
+                    yield Field("Сдвиг · градусы", Input(value="0.0", id="shear"))
+                    yield Field("Перспектива · 0…1", Input(value="0.0", id="perspective"))
+                    yield Field("Закрыть mosaic · эпох", Input(value="10", type="integer", id="close-mosaic"))
+                with Horizontal(classes="row augmentation-field"):
+                    yield Field("Flip вверх/вниз · 0…1", Input(value="0.0", id="flipud"))
+                    yield Field("Flip влево/вправо · 0…1", Input(value="0.5", id="fliplr"))
+                    yield Field("RGB ↔ BGR · 0…1", Input(value="0.0", id="bgr"))
+                with Horizontal(classes="row augmentation-field"):
+                    yield Field("Mosaic · 0…1", Input(value="1.0", id="mosaic"))
+                    yield Field("MixUp · 0…1", Input(value="0.0", id="mixup"))
+                    yield Field("CutMix · 0…1", Input(value="0.0", id="cutmix"))
+                with Horizontal(classes="row augmentation-field"):
+                    yield Field("Copy-paste · 0…1", Input(value="0.0", id="copy-paste"))
+                    yield Field("Erasing · 0…1", Input(value="0.4", id="erasing"))
+                    yield Field(
+                        "Режим copy-paste · segment",
+                        Select(
+                            [(mode.capitalize(), mode) for mode in SUPPORTED_COPY_PASTE_MODES],
+                            value="flip",
+                            id="copy-paste-mode",
+                            allow_blank=False,
+                        ),
+                    )
+                yield Field(
+                    "AutoAugment · classify",
+                    Select(
+                        [(policy.capitalize(), policy) for policy in SUPPORTED_AUTO_AUGMENT_POLICIES],
+                        value="randaugment",
+                        id="auto-augment",
+                        allow_blank=False,
+                    ),
+                    classes="augmentation-field",
+                )
+
+                yield Static("MLflow", classes="section-title")
+                with Horizontal(classes="toggle-row"):
+                    yield Label("Записывать метрики, параметры и артефакты")
+                    yield Switch(value=True, id="mlflow-enabled")
+                yield Field(
+                    "Tracking URI",
+                    Input(value="sqlite:///mlflow.db", placeholder="http://127.0.0.1:5000", id="tracking-uri"),
+                    classes="mlflow-field",
+                )
+                with Horizontal(classes="row mlflow-field"):
+                    yield Field("Эксперимент", Input(value="yolo-tui", id="experiment-name"))
+                    yield Field("MLflow run", Input(placeholder="необязательно", id="mlflow-run-name"))
+
+            with Vertical(id="run-pane"):
+                with Vertical(id="status-card"):
+                    yield Static("ГОТОВО К ЗАПУСКУ", id="status-title")
+                    yield Static("Проверьте параметры и начните обучение.", id="status-text")
+                yield ProgressBar(total=100, show_eta=True, id="progress")
+                with Horizontal(id="actions"):
+                    yield Button("▶  Начать обучение", variant="primary", id="start-button")
+                    yield Button("■  Остановить", variant="error", id="stop-button", disabled=True)
+                yield Static("Журнал", id="log-title")
+                yield RichLog(id="log", markup=True, wrap=True, highlight=False)
+                yield Static(
+                    "MLflow работает локально без сервера. Просмотр: uv run mlflow ui --backend-store-uri sqlite:///mlflow.db",
+                    id="hint",
+                )
+        yield Footer()
+
+    def on_mount(self) -> None:
+        self.query_one("#progress", ProgressBar).update(progress=0)
+        self._write_log("[dim]Интерфейс готов. Обучение еще не запускалось.[/dim]")
+
+    @on(Switch.Changed, "#mlflow-enabled")
+    def toggle_mlflow(self, event: Switch.Changed) -> None:
+        for widget_id in ("tracking-uri", "experiment-name", "mlflow-run-name"):
+            self.query_one(f"#{widget_id}", Input).disabled = not event.value
+
+    @on(Switch.Changed, "#augmentation-enabled")
+    def toggle_augmentation(self, event: Switch.Changed) -> None:
+        for control in self.query(".augmentation-field Input"):
+            control.disabled = not event.value
+        for control in self.query(".augmentation-field Select"):
+            control.disabled = not event.value
+
+    @on(Button.Pressed, "#start-button")
+    def start_pressed(self) -> None:
+        self.action_start_training()
+
+    @on(Button.Pressed, "#stop-button")
+    def stop_pressed(self) -> None:
+        self.action_stop_training()
+
+    def action_start_training(self) -> None:
+        if self._running:
+            self.notify("Обучение уже выполняется.", severity="warning")
+            return
+        try:
+            config = self._read_config()
+            config.validate()
+        except ValueError as exc:
+            self.notify(str(exc), title="Проверьте параметры", severity="error")
+            return
+
+        self._set_running(True)
+        progress = self.query_one("#progress", ProgressBar)
+        progress.update(total=config.epochs, progress=0)
+        self.query_one("#status-title", Static).update("ПОДГОТОВКА")
+        self.query_one("#status-text", Static).update("Загружаю модель и датасет…")
+        self._write_log(
+            f"[bold #78a9ff]Запуск:[/] задача={config.task}, модель={config.model}, датасет={config.dataset}"
+        )
+        if config.mlflow.enabled:
+            self._write_log(
+                f"[dim]MLflow: {config.mlflow.tracking_uri} · эксперимент {config.mlflow.experiment_name}[/dim]"
+            )
+        if config.augmentation.enabled:
+            self._write_log(
+                f"[dim]Аугментация: mosaic={config.augmentation.mosaic}, "
+                f"mixup={config.augmentation.mixup}, fliplr={config.augmentation.fliplr}[/dim]"
+            )
+        self._train_in_background(config)
+
+    def action_stop_training(self) -> None:
+        if not self._running:
+            return
+        self.runner.request_stop()
+        self.query_one("#status-title", Static).update("ОСТАНОВКА")
+        self.query_one("#status-text", Static).update("Завершаю текущую эпоху…")
+        self.query_one("#stop-button", Button).disabled = True
+        self._write_log("[yellow]Запрошена остановка обучения.[/yellow]")
+
+    @work(thread=True, exclusive=True, group="yolo-training")
+    def _train_in_background(self, config: TrainingConfig) -> None:
+        try:
+            output_dir = self.runner.train(
+                config,
+                lambda event: self.app.call_from_thread(self._handle_training_event, event),
+            )
+        except Exception as exc:  # errors must be surfaced in the TUI, not hidden in a worker
+            self.app.call_from_thread(self._training_failed, exc)
+        else:
+            self.app.call_from_thread(self._training_finished, output_dir)
+
+    def _handle_training_event(self, event: TrainingEvent) -> None:
+        styles = {
+            "info": "#9fb1d1",
+            "started": "#78a9ff",
+            "epoch": "#b7c9e8",
+            "warning": "yellow",
+            "success": "green",
+        }
+        self._write_log(f"[{styles.get(event.kind, 'white')}]{event.message}[/]")
+        if event.kind == "started":
+            self.query_one("#status-title", Static).update("ОБУЧЕНИЕ")
+            self.query_one("#status-text", Static).update(f"Выполняется 0 из {event.total_epochs} эпох.")
+        elif event.kind == "epoch":
+            self.query_one("#progress", ProgressBar).update(
+                total=event.total_epochs or None,
+                progress=event.epoch,
+            )
+            self.query_one("#status-text", Static).update(event.message)
+
+    def _training_finished(self, output_dir: Any | None) -> None:
+        stopped = self.runner.stop_requested
+        self._set_running(False)
+        if stopped:
+            title = "ОСТАНОВЛЕНО"
+            message = "Обучение остановлено. Уже сохраненные checkpoints не удалены."
+            style = "yellow"
+        else:
+            title = "ГОТОВО"
+            message = "Обучение успешно завершено."
+            style = "green"
+            progress = self.query_one("#progress", ProgressBar)
+            progress.update(progress=progress.total)
+        self.query_one("#status-title", Static).update(title)
+        self.query_one("#status-text", Static).update(message)
+        if output_dir:
+            self._write_log(f"[{style}]Результаты: {output_dir}[/]")
+        self.notify(message, severity="warning" if stopped else "information")
+
+    def _training_failed(self, error: Exception) -> None:
+        self._set_running(False)
+        self.query_one("#status-title", Static).update("ОШИБКА")
+        self.query_one("#status-text", Static).update(str(error))
+        self._write_log(f"[bold red]Ошибка: {error}[/bold red]")
+        self.notify(str(error), title="Обучение не запущено", severity="error", timeout=10)
+
+    def _set_running(self, running: bool) -> None:
+        self._running = running
+        self.query_one("#start-button", Button).disabled = running
+        self.query_one("#stop-button", Button).disabled = not running
+
+    def _read_config(self) -> TrainingConfig:
+        task = self.query_one("#task", Select).value
+        if task not in SUPPORTED_TASKS:
+            raise ValueError("Выберите тип задачи YOLO.")
+        return TrainingConfig(
+            dataset=self._input("dataset"),
+            model=self._input("model"),
+            task=task,
+            epochs=self._integer("epochs", "Эпохи"),
+            image_size=self._integer("image-size", "Размер изображения"),
+            batch_size=self._integer("batch-size", "Batch"),
+            device=self._input("device"),
+            workers=self._integer("workers", "Workers"),
+            patience=self._integer("patience", "Patience"),
+            project=self._input("project"),
+            run_name=self._input("run-name"),
+            augmentation=AugmentationConfig(
+                enabled=self.query_one("#augmentation-enabled", Switch).value,
+                hsv_h=self._float("hsv-h", "HSV hue"),
+                hsv_s=self._float("hsv-s", "HSV saturation"),
+                hsv_v=self._float("hsv-v", "HSV brightness"),
+                degrees=self._float("degrees", "Поворот"),
+                translate=self._float("translate", "Смещение"),
+                scale=self._float("scale", "Масштаб"),
+                shear=self._float("shear", "Сдвиг"),
+                perspective=self._float("perspective", "Перспектива"),
+                flipud=self._float("flipud", "Flip вверх/вниз"),
+                fliplr=self._float("fliplr", "Flip влево/вправо"),
+                bgr=self._float("bgr", "RGB ↔ BGR"),
+                mosaic=self._float("mosaic", "Mosaic"),
+                mixup=self._float("mixup", "MixUp"),
+                cutmix=self._float("cutmix", "CutMix"),
+                copy_paste=self._float("copy-paste", "Copy-paste"),
+                copy_paste_mode=self._select(
+                    "copy-paste-mode",
+                    "режим copy-paste",
+                    SUPPORTED_COPY_PASTE_MODES,
+                ),
+                auto_augment=self._select(
+                    "auto-augment",
+                    "политику AutoAugment",
+                    SUPPORTED_AUTO_AUGMENT_POLICIES,
+                ),
+                erasing=self._float("erasing", "Erasing"),
+                close_mosaic=self._integer("close-mosaic", "Close mosaic"),
+            ),
+            mlflow=MlflowConfig(
+                enabled=self.query_one("#mlflow-enabled", Switch).value,
+                tracking_uri=self._input("tracking-uri"),
+                experiment_name=self._input("experiment-name"),
+                run_name=self._input("mlflow-run-name"),
+            ),
+        )
+
+    def _input(self, widget_id: str) -> str:
+        return self.query_one(f"#{widget_id}", Input).value.strip()
+
+    def _integer(self, widget_id: str, label: str) -> int:
+        value = self._input(widget_id)
+        try:
+            return int(value)
+        except ValueError as exc:
+            raise ValueError(f"Поле «{label}» должно быть целым числом.") from exc
+
+    def _float(self, widget_id: str, label: str) -> float:
+        value = self._input(widget_id).replace(",", ".")
+        try:
+            return float(value)
+        except ValueError as exc:
+            raise ValueError(f"Поле «{label}» должно быть числом.") from exc
+
+    def _select(self, widget_id: str, label: str, choices: tuple[str, ...]) -> Any:
+        value = self.query_one(f"#{widget_id}", Select).value
+        if value not in choices:
+            raise ValueError(f"Выберите {label}.")
+        return value
+
+    def _write_log(self, message: str) -> None:
+        self.query_one("#log", RichLog).write(message)
+
+
+def main() -> None:
+    YoloTrainApp().run()

src/yolo_tui/config.py

@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Literal
+
+
+YoloTask = Literal["detect", "segment", "classify", "pose", "obb"]
+AutoAugmentPolicy = Literal["randaugment", "autoaugment", "augmix"]
+CopyPasteMode = Literal["flip", "mixup"]
+SUPPORTED_TASKS: tuple[YoloTask, ...] = (
+    "detect",
+    "segment",
+    "classify",
+    "pose",
+    "obb",
+)
+SUPPORTED_AUTO_AUGMENT_POLICIES: tuple[AutoAugmentPolicy, ...] = (
+    "randaugment",
+    "autoaugment",
+    "augmix",
+)
+SUPPORTED_COPY_PASTE_MODES: tuple[CopyPasteMode, ...] = ("flip", "mixup")
+
+
+@dataclass(frozen=True, slots=True)
+class AugmentationConfig:
+    """Training augmentation overrides accepted by Ultralytics."""
+
+    enabled: bool = True
+    hsv_h: float = 0.015
+    hsv_s: float = 0.7
+    hsv_v: float = 0.4
+    degrees: float = 0.0
+    translate: float = 0.1
+    scale: float = 0.5
+    shear: float = 0.0
+    perspective: float = 0.0
+    flipud: float = 0.0
+    fliplr: float = 0.5
+    bgr: float = 0.0
+    mosaic: float = 1.0
+    mixup: float = 0.0
+    cutmix: float = 0.0
+    copy_paste: float = 0.0
+    copy_paste_mode: CopyPasteMode = "flip"
+    auto_augment: AutoAugmentPolicy = "randaugment"
+    erasing: float = 0.4
+    close_mosaic: int = 10
+
+    def validate(self) -> None:
+        if not self.enabled:
+            return
+
+        fractions = {
+            "HSV hue": self.hsv_h,
+            "HSV saturation": self.hsv_s,
+            "HSV brightness": self.hsv_v,
+            "Translate": self.translate,
+            "Scale": self.scale,
+            "Perspective": self.perspective,
+            "Flip up/down": self.flipud,
+            "Flip left/right": self.fliplr,
+            "BGR": self.bgr,
+            "Mosaic": self.mosaic,
+            "MixUp": self.mixup,
+            "CutMix": self.cutmix,
+            "Copy-paste": self.copy_paste,
+            "Erasing": self.erasing,
+        }
+        for label, value in fractions.items():
+            if not 0.0 <= value <= 1.0:
+                raise ValueError(f"Параметр «{label}» должен быть от 0 до 1.")
+        if self.degrees < 0:
+            raise ValueError("Угол поворота не может быть отрицательным.")
+        if self.shear < 0:
+            raise ValueError("Угол сдвига не может быть отрицательным.")
+        if self.close_mosaic < 0:
+            raise ValueError("Close mosaic не может быть отрицательным.")
+        if self.copy_paste_mode not in SUPPORTED_COPY_PASTE_MODES:
+            raise ValueError(f"Неизвестный режим copy-paste: {self.copy_paste_mode}.")
+        if self.auto_augment not in SUPPORTED_AUTO_AUGMENT_POLICIES:
+            raise ValueError(f"Неизвестная политика AutoAugment: {self.auto_augment}.")
+
+    def train_kwargs(self) -> dict[str, str | int | float]:
+        if not self.enabled:
+            return {}
+        return {
+            "hsv_h": self.hsv_h,
+            "hsv_s": self.hsv_s,
+            "hsv_v": self.hsv_v,
+            "degrees": self.degrees,
+            "translate": self.translate,
+            "scale": self.scale,
+            "shear": self.shear,
+            "perspective": self.perspective,
+            "flipud": self.flipud,
+            "fliplr": self.fliplr,
+            "bgr": self.bgr,
+            "mosaic": self.mosaic,
+            "mixup": self.mixup,
+            "cutmix": self.cutmix,
+            "copy_paste": self.copy_paste,
+            "copy_paste_mode": self.copy_paste_mode,
+            "auto_augment": self.auto_augment,
+            "erasing": self.erasing,
+            "close_mosaic": self.close_mosaic,
+        }
+
+
+@dataclass(frozen=True, slots=True)
+class MlflowConfig:
+    enabled: bool = True
+    tracking_uri: str = "sqlite:///mlflow.db"
+    experiment_name: str = "yolo-tui"
+    run_name: str = ""
+
+    def validate(self) -> None:
+        if self.enabled and not self.tracking_uri.strip():
+            raise ValueError("Укажите URI хранилища MLflow.")
+        if self.enabled and not self.experiment_name.strip():
+            raise ValueError("Укажите название эксперимента MLflow.")
+
+
+@dataclass(frozen=True, slots=True)
+class TrainingConfig:
+    dataset: str
+    model: str
+    task: YoloTask = "detect"
+    epochs: int = 100
+    image_size: int = 640
+    batch_size: int = 16
+    device: str = ""
+    workers: int = 8
+    patience: int = 100
+    project: str = "runs/train"
+    run_name: str = ""
+    augmentation: AugmentationConfig = field(default_factory=AugmentationConfig)
+    mlflow: MlflowConfig = field(default_factory=MlflowConfig)
+
+    def validate(self) -> None:
+        if not self.dataset.strip():
+            raise ValueError("Укажите путь или имя датасета.")
+        if not self.model.strip():
+            raise ValueError("Укажите путь или имя модели.")
+        if self.task not in SUPPORTED_TASKS:
+            raise ValueError(f"Неизвестный тип задачи: {self.task}.")
+        if self.epochs < 1:
+            raise ValueError("Количество эпох должно быть не меньше 1.")
+        if self.image_size < 32:
+            raise ValueError("Размер изображения должен быть не меньше 32.")
+        if self.batch_size == 0 or self.batch_size < -1:
+            raise ValueError("Batch должен быть положительным числом или -1 для автоподбора.")
+        if self.workers < 0:
+            raise ValueError("Количество workers не может быть отрицательным.")
+        if self.patience < 0:
+            raise ValueError("Patience не может быть отрицательным.")
+        self.augmentation.validate()
+        self.mlflow.validate()
+
+    def train_kwargs(self) -> dict[str, str | int | float | bool]:
+        """Convert the form values to arguments accepted by YOLO.train()."""
+        values: dict[str, str | int | float | bool] = {
+            "data": self.dataset.strip(),
+            "epochs": self.epochs,
+            "imgsz": self.image_size,
+            "batch": self.batch_size,
+            "workers": self.workers,
+            "patience": self.patience,
+            "project": self.project.strip() or "runs/train",
+            # Ultralytics' tqdm output would otherwise paint over Textual's screen.
+            # Epoch metrics are sent to the in-app log by callbacks instead.
+            "verbose": False,
+        }
+        if self.device.strip():
+            values["device"] = self.device.strip()
+        if self.run_name.strip():
+            values["name"] = self.run_name.strip()
+        values.update(self.augmentation.train_kwargs())
+        return values

src/yolo_tui/trainer.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+import os
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from threading import Event, Lock
+from typing import Any
+
+from .config import MlflowConfig, TrainingConfig
+
+
+@dataclass(frozen=True, slots=True)
+class TrainingEvent:
+    kind: str
+    message: str
+    epoch: int = 0
+    total_epochs: int = 0
+
+
+EventHandler = Callable[[TrainingEvent], None]
+
+
+@contextmanager
+def mlflow_environment(config: MlflowConfig) -> Iterator[None]:
+    """Temporarily expose the settings expected by Ultralytics' MLflow callback."""
+    keys = {
+        "MLFLOW_TRACKING_URI": config.tracking_uri.strip(),
+        "MLFLOW_EXPERIMENT_NAME": config.experiment_name.strip(),
+        "MLFLOW_RUN": config.run_name.strip(),
+        "MLFLOW_KEEP_RUN_ACTIVE": "False",
+    }
+    previous = {key: os.environ.get(key) for key in keys}
+    try:
+        if config.enabled:
+            for key, value in keys.items():
+                if value:
+                    os.environ[key] = value
+                else:
+                    os.environ.pop(key, None)
+        yield
+    finally:
+        for key, value in previous.items():
+            if value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = value
+
+
+class TrainingRunner:
+    """Owns a single YOLO training run and exposes cooperative cancellation."""
+
+    def __init__(self) -> None:
+        self._model: Any | None = None
+        self._state_lock = Lock()
+        self._stop_requested = Event()
+
+    def request_stop(self) -> None:
+        self._stop_requested.set()
+        with self._state_lock:
+            trainer = getattr(self._model, "trainer", None)
+            if trainer is not None:
+                trainer.stop = True
+
+    @property
+    def stop_requested(self) -> bool:
+        return self._stop_requested.is_set()
+
+    def train(self, config: TrainingConfig, on_event: EventHandler) -> Path | None:
+        config.validate()
+        self._stop_requested.clear()
+
+        on_event(TrainingEvent("info", "Загрузка Ultralytics и подготовка модели…"))
+        from ultralytics import YOLO, settings
+
+        settings.update({"mlflow": config.mlflow.enabled})
+
+        with mlflow_environment(config.mlflow):
+            model = YOLO(config.model.strip(), task=config.task)
+            with self._state_lock:
+                self._model = model
+
+            model.add_callback("on_train_start", self._on_train_start(on_event))
+            model.add_callback("on_train_epoch_end", self._on_epoch_end(on_event))
+            model.add_callback("on_train_end", self._on_train_end(on_event))
+
+            try:
+                model.train(**config.train_kwargs())
+                trainer = getattr(model, "trainer", None)
+                save_dir = getattr(trainer, "save_dir", None)
+                return Path(save_dir) if save_dir else None
+            finally:
+                with self._state_lock:
+                    self._model = None
+
+    def _on_train_start(self, on_event: EventHandler) -> Callable[[Any], None]:
+        def callback(trainer: Any) -> None:
+            total = int(getattr(getattr(trainer, "args", None), "epochs", 0))
+            on_event(TrainingEvent("started", "Обучение началось.", 0, total))
+            if self._stop_requested.is_set():
+                trainer.stop = True
+
+        return callback
+
+    def _on_epoch_end(self, on_event: EventHandler) -> Callable[[Any], None]:
+        def callback(trainer: Any) -> None:
+            epoch = int(getattr(trainer, "epoch", 0)) + 1
+            total = int(getattr(getattr(trainer, "args", None), "epochs", 0))
+            metrics = getattr(trainer, "metrics", {}) or {}
+            summary = self._metrics_summary(metrics)
+            message = f"Эпоха {epoch}/{total} завершена"
+            if summary:
+                message += f" · {summary}"
+            on_event(TrainingEvent("epoch", message, epoch, total))
+            if self._stop_requested.is_set():
+                trainer.stop = True
+
+        return callback
+
+    def _on_train_end(self, on_event: EventHandler) -> Callable[[Any], None]:
+        def callback(trainer: Any) -> None:
+            if self._stop_requested.is_set():
+                on_event(TrainingEvent("warning", "Обучение остановлено пользователем."))
+            else:
+                on_event(TrainingEvent("success", "Ultralytics завершил обучение."))
+
+        return callback
+
+    @staticmethod
+    def _metrics_summary(metrics: dict[str, Any]) -> str:
+        result: list[str] = []
+        for key, value in list(metrics.items())[:3]:
+            try:
+                result.append(f"{key.split('/')[-1]}={float(value):.4g}")
+            except (TypeError, ValueError):
+                continue
+        return " · ".join(result)

tests/test_app.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import asyncio
+
+from textual.widgets import Button, Input, Select, Switch
+
+from yolo_tui.app import YoloTrainApp
+
+
+def test_app_mounts_with_expected_defaults() -> None:
+    async def exercise() -> None:
+        app = YoloTrainApp()
+        async with app.run_test(size=(140, 45)):
+            assert app.query_one("#task", Select).value == "detect"
+            assert app.query_one("#model", Input).value == "yolo11n.pt"
+            assert app.query_one("#dataset", Input).value == "coco8.yaml"
+            assert app.query_one("#augmentation-enabled", Switch).value is True
+            assert app.query_one("#mosaic", Input).value == "1.0"
+            assert app.query_one("#auto-augment", Select).value == "randaugment"
+            assert app.query_one("#mlflow-enabled", Switch).value is True
+            assert app.query_one("#start-button", Button).disabled is False
+            assert app.query_one("#stop-button", Button).disabled is True
+
+    asyncio.run(exercise())
+
+
+def test_augmentation_fields_follow_switch() -> None:
+    async def exercise() -> None:
+        app = YoloTrainApp()
+        async with app.run_test(size=(140, 45)) as pilot:
+            switch = app.query_one("#augmentation-enabled", Switch)
+            switch.value = False
+            await pilot.pause()
+
+            assert app.query_one("#mosaic", Input).disabled is True
+            assert app.query_one("#auto-augment", Select).disabled is True
+
+    asyncio.run(exercise())
+
+
+def test_mlflow_fields_follow_switch() -> None:
+    async def exercise() -> None:
+        app = YoloTrainApp()
+        async with app.run_test(size=(140, 45)) as pilot:
+            switch = app.query_one("#mlflow-enabled", Switch)
+            switch.value = False
+            await pilot.pause()
+
+            assert app.query_one("#tracking-uri", Input).disabled is True
+            assert app.query_one("#experiment-name", Input).disabled is True
+
+    asyncio.run(exercise())

tests/test_config.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import os
+
+import pytest
+
+from yolo_tui.config import AugmentationConfig, MlflowConfig, TrainingConfig
+from yolo_tui.trainer import TrainingRunner, mlflow_environment
+
+
+def test_train_kwargs_omit_optional_empty_values() -> None:
+    config = TrainingConfig(
+        dataset=" dataset.yaml ",
+        model=" model.pt ",
+        augmentation=AugmentationConfig(enabled=False),
+    )
+
+    assert config.train_kwargs() == {
+        "data": "dataset.yaml",
+        "epochs": 100,
+        "imgsz": 640,
+        "batch": 16,
+        "workers": 8,
+        "patience": 100,
+        "project": "runs/train",
+        "verbose": False,
+    }
+
+
+def test_augmentation_kwargs_are_passed_to_ultralytics() -> None:
+    config = TrainingConfig(
+        dataset="dataset.yaml",
+        model="model.pt",
+        augmentation=AugmentationConfig(mosaic=0.5, mixup=0.2, close_mosaic=5),
+    )
+
+    kwargs = config.train_kwargs()
+
+    assert kwargs["mosaic"] == 0.5
+    assert kwargs["mixup"] == 0.2
+    assert kwargs["close_mosaic"] == 5
+    assert kwargs["auto_augment"] == "randaugment"
+
+
+@pytest.mark.parametrize("field", ["mosaic", "fliplr", "erasing", "perspective"])
+def test_augmentation_probabilities_are_validated(field: str) -> None:
+    augmentation = AugmentationConfig(**{field: 1.1})
+
+    with pytest.raises(ValueError, match="от 0 до 1"):
+        augmentation.validate()
+
+
+@pytest.mark.parametrize("batch", [0, -2])
+def test_invalid_batch_is_rejected(batch: int) -> None:
+    config = TrainingConfig(dataset="dataset.yaml", model="model.pt", batch_size=batch)
+
+    with pytest.raises(ValueError, match="Batch"):
+        config.validate()
+
+
+def test_mlflow_environment_is_restored(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("MLFLOW_TRACKING_URI", "previous")
+    config = MlflowConfig(
+        enabled=True,
+        tracking_uri="http://127.0.0.1:5000",
+        experiment_name="test-experiment",
+        run_name="test-run",
+    )
+
+    with mlflow_environment(config):
+        assert os.environ["MLFLOW_TRACKING_URI"] == config.tracking_uri
+        assert os.environ["MLFLOW_EXPERIMENT_NAME"] == config.experiment_name
+        assert os.environ["MLFLOW_RUN"] == config.run_name
+
+    assert os.environ["MLFLOW_TRACKING_URI"] == "previous"
+    assert "MLFLOW_EXPERIMENT_NAME" not in os.environ
+    assert "MLFLOW_RUN" not in os.environ
+
+
+def test_metrics_summary_skips_non_numeric_values() -> None:
+    summary = TrainingRunner._metrics_summary(
+        {"metrics/mAP50": 0.81234, "label": "invalid", "val/loss": 0.12345}
+    )
+
+    assert summary == "mAP50=0.8123 · loss=0.1235"

tests/test_trainer.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from types import ModuleType, SimpleNamespace
+from typing import Any
+
+from yolo_tui.config import MlflowConfig, TrainingConfig
+from yolo_tui.trainer import TrainingEvent, TrainingRunner
+
+
+def test_runner_wires_yolo_callbacks_and_returns_output(
+    monkeypatch: Any, tmp_path: Path
+) -> None:
+    settings_updates: list[dict[str, bool]] = []
+    constructed: list[tuple[str, str]] = []
+    train_arguments: list[dict[str, Any]] = []
+
+    class FakeSettings:
+        def update(self, values: dict[str, bool]) -> None:
+            settings_updates.append(values)
+
+    class FakeYOLO:
+        def __init__(self, model: str, task: str) -> None:
+            constructed.append((model, task))
+            self.callbacks: dict[str, Any] = {}
+            self.trainer = SimpleNamespace(
+                args=SimpleNamespace(epochs=2),
+                epoch=0,
+                metrics={},
+                stop=False,
+                save_dir=tmp_path / "run",
+            )
+
+        def add_callback(self, name: str, callback: Any) -> None:
+            self.callbacks[name] = callback
+
+        def train(self, **kwargs: Any) -> None:
+            train_arguments.append(kwargs)
+            self.callbacks["on_train_start"](self.trainer)
+            for epoch in range(2):
+                self.trainer.epoch = epoch
+                self.trainer.metrics = {"metrics/mAP50": 0.5 + epoch / 10}
+                self.callbacks["on_train_epoch_end"](self.trainer)
+            self.callbacks["on_train_end"](self.trainer)
+
+    fake_ultralytics = ModuleType("ultralytics")
+    fake_ultralytics.YOLO = FakeYOLO  # type: ignore[attr-defined]
+    fake_ultralytics.settings = FakeSettings()  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, "ultralytics", fake_ultralytics)
+
+    config = TrainingConfig(
+        dataset="dataset.yaml",
+        model="model.pt",
+        task="pose",
+        epochs=2,
+        mlflow=MlflowConfig(enabled=False),
+    )
+    events: list[TrainingEvent] = []
+
+    output = TrainingRunner().train(config, events.append)
+
+    assert output == tmp_path / "run"
+    assert constructed == [("model.pt", "pose")]
+    assert settings_updates == [{"mlflow": False}]
+    assert train_arguments[0]["data"] == "dataset.yaml"
+    assert train_arguments[0]["verbose"] is False
+    assert [event.kind for event in events] == ["info", "started", "epoch", "epoch", "success"]