Fix training lifecycle and dataset splitting bugs
This commit is contained in:
parent
06deeb25d1
commit
3a5f502735
11 changed files with 661 additions and 132 deletions
73
.agents/PROJECT_ISSUES.md
Normal file
73
.agents/PROJECT_ISSUES.md
Normal file
|
|
@ -0,0 +1,73 @@
|
||||||
|
# Исправленные проблемы проекта YOLO Train TUI
|
||||||
|
|
||||||
|
Дата исправления и повторной проверки: 2026-07-17
|
||||||
|
|
||||||
|
## Итог
|
||||||
|
|
||||||
|
Все 11 ранее зафиксированных дефектов исправлены и покрыты регрессионными
|
||||||
|
проверками.
|
||||||
|
|
||||||
|
| ID | Приоритет | Статус | Исправление |
|
||||||
|
|---|---|---|---|
|
||||||
|
| BUG-001 | Критический | Исправлено | `subprocess_runner.main()` возвращает код, а `SystemExit` создаётся только снаружи обрабатывающего блока |
|
||||||
|
| BUG-002 | Высокий | Исправлено | Перед каждым запуском `prepare_run()` сбрасывает состояние остановки |
|
||||||
|
| BUG-003 | Высокий | Исправлено | Родитель отправляет кооперативный сигнал; принудительный `kill()` используется только после таймаута |
|
||||||
|
| BUG-004 | Высокий | Исправлено | Запрос, сделанный до готовности subprocess, сохраняется и доставляется после маркера `READY` |
|
||||||
|
| BUG-005 | Высокий | Исправлено | Detection-style авторазбиение запрещено для `classify` в UI и конфигурации |
|
||||||
|
| BUG-006 | Средний | Исправлено | `.yaml`/`.yml` разбираются через `yaml.safe_load()`, поле `names` валидируется |
|
||||||
|
| BUG-007 | Средний | Исправлено | Датасет с одним изображением отклоняется с понятной ошибкой |
|
||||||
|
| BUG-008 | Средний | Исправлено | Изображения и метки ищутся рекурсивно с сохранением вложенных путей |
|
||||||
|
| BUG-009 | Средний | Исправлено | Каждый результат создаётся в уникальном `.yolo-tui/splits/<id>` без перезаписи пользовательского `split/` |
|
||||||
|
| BUG-010 | Низкий | Исправлено | Traceback выводится в журнал TUI; абсолютный путь другого пользователя удалён |
|
||||||
|
| BUG-011 | Низкий | Исправлено | Явно указанный отсутствующий или некорректный файл классов вызывает точную ошибку без fallback |
|
||||||
|
|
||||||
|
## Жизненный цикл обучения
|
||||||
|
|
||||||
|
- Дочерний процесс устанавливает обработчики остановки и только затем печатает
|
||||||
|
`__YOLO_TUI_READY__`.
|
||||||
|
- Если пользователь нажал «Остановить» раньше, родитель запоминает запрос и
|
||||||
|
отправляет его после получения маркера готовности.
|
||||||
|
- Дочерний `TrainingRunner` устанавливает `trainer.stop = True`; Ultralytics
|
||||||
|
останавливается между пакетами данных, затем выполняет штатную финализацию и
|
||||||
|
завершающие callbacks.
|
||||||
|
- Если процесс не завершился за 30 секунд, используется принудительный fallback.
|
||||||
|
- После завершения ссылка на subprocess и таймер очищаются; перед следующим
|
||||||
|
запуском флаг остановки сбрасывается.
|
||||||
|
|
||||||
|
## Работа с датасетами
|
||||||
|
|
||||||
|
- Текущий splitter предназначен для задач `detect`, `segment`, `pose` и `obb`
|
||||||
|
со структурой `images/` + `labels/`.
|
||||||
|
- Для `classify` требуется готовый каталог с `train`/`val` и подкаталогами
|
||||||
|
классов; несовместимый переключатель в UI отключён.
|
||||||
|
- Поддерживаются `classes.txt`, `.yaml` и `.yml`; YAML может хранить `names` как
|
||||||
|
список или словарь с последовательными ID от 0.
|
||||||
|
- Явный путь к классам считается обязательным и не заменяется автопоиском при
|
||||||
|
опечатке или ошибке формата.
|
||||||
|
- Split требует минимум два изображения и рекурсивно обрабатывает вложенные
|
||||||
|
каталоги.
|
||||||
|
- Файлы каждого запуска создаются эксклюзивно в отдельном управляемом каталоге.
|
||||||
|
|
||||||
|
## Проверка
|
||||||
|
|
||||||
|
Выполнены команды:
|
||||||
|
|
||||||
|
```text
|
||||||
|
uv run pytest -q
|
||||||
|
uv run python -m compileall -q src tests
|
||||||
|
git diff --check
|
||||||
|
```
|
||||||
|
|
||||||
|
Результат: `38 passed`; ошибок компиляции и форматирования diff нет.
|
||||||
|
|
||||||
|
Регрессионные тесты проверяют:
|
||||||
|
|
||||||
|
1. успешный и ошибочный коды `subprocess_runner.main()`;
|
||||||
|
2. сброс остановки между запусками;
|
||||||
|
3. кооперативный сигнал вместо немедленного `terminate()`;
|
||||||
|
4. доставку раннего запроса после готовности subprocess;
|
||||||
|
5. запрет авторазбиения для `classify`;
|
||||||
|
6. пользовательские YAML-файлы классов и ошибочный явный путь;
|
||||||
|
7. датасеты из одного и двух изображений;
|
||||||
|
8. вложенные изображения и метки;
|
||||||
|
9. сохранность пользовательского каталога `split/` и уникальность результатов.
|
||||||
|
|
@ -329,9 +329,26 @@ class YoloTrainApp(App[None]):
|
||||||
|
|
||||||
@on(Switch.Changed, "#split-enabled")
|
@on(Switch.Changed, "#split-enabled")
|
||||||
def toggle_split(self, event: Switch.Changed) -> None:
|
def toggle_split(self, event: Switch.Changed) -> None:
|
||||||
|
if self.query_one("#task", Select).value == "classify" and event.value:
|
||||||
|
self.query_one("#split-enabled", Switch).value = False
|
||||||
|
self.notify(
|
||||||
|
"Для classify укажите готовый каталог с train/val по классам.",
|
||||||
|
severity="warning",
|
||||||
|
)
|
||||||
|
return
|
||||||
for control in self.query(".split-field Input"):
|
for control in self.query(".split-field Input"):
|
||||||
control.disabled = not event.value
|
control.disabled = not event.value
|
||||||
|
|
||||||
|
@on(Select.Changed, "#task")
|
||||||
|
def task_changed(self, event: Select.Changed) -> None:
|
||||||
|
split_switch = self.query_one("#split-enabled", Switch)
|
||||||
|
is_classify = event.value == "classify"
|
||||||
|
if is_classify and split_switch.value:
|
||||||
|
split_switch.value = False
|
||||||
|
split_switch.disabled = is_classify
|
||||||
|
for control in self.query(".split-field Input"):
|
||||||
|
control.disabled = is_classify or not split_switch.value
|
||||||
|
|
||||||
@on(Switch.Changed, "#mlflow-enabled")
|
@on(Switch.Changed, "#mlflow-enabled")
|
||||||
def toggle_mlflow(self, event: Switch.Changed) -> None:
|
def toggle_mlflow(self, event: Switch.Changed) -> None:
|
||||||
for widget_id in ("tracking-uri", "experiment-name", "mlflow-run-name"):
|
for widget_id in ("tracking-uri", "experiment-name", "mlflow-run-name"):
|
||||||
|
|
@ -363,6 +380,7 @@ class YoloTrainApp(App[None]):
|
||||||
self.notify(str(exc), title="Проверьте параметры", severity="error")
|
self.notify(str(exc), title="Проверьте параметры", severity="error")
|
||||||
return
|
return
|
||||||
|
|
||||||
|
self.runner.prepare_run()
|
||||||
self._set_running(True)
|
self._set_running(True)
|
||||||
progress = self.query_one("#progress", ProgressBar)
|
progress = self.query_one("#progress", ProgressBar)
|
||||||
progress.update(total=config.epochs, progress=0)
|
progress.update(total=config.epochs, progress=0)
|
||||||
|
|
@ -387,7 +405,7 @@ class YoloTrainApp(App[None]):
|
||||||
return
|
return
|
||||||
self.runner.request_stop()
|
self.runner.request_stop()
|
||||||
self.query_one("#status-title", Static).update("ОСТАНОВКА")
|
self.query_one("#status-title", Static).update("ОСТАНОВКА")
|
||||||
self.query_one("#status-text", Static).update("Завершаю текущую эпоху…")
|
self.query_one("#status-text", Static).update("Корректно останавливаю обучение…")
|
||||||
self.query_one("#stop-button", Button).disabled = True
|
self.query_one("#stop-button", Button).disabled = True
|
||||||
self._write_log("[yellow]Запрошена остановка обучения.[/yellow]")
|
self._write_log("[yellow]Запрошена остановка обучения.[/yellow]")
|
||||||
|
|
||||||
|
|
@ -415,7 +433,7 @@ class YoloTrainApp(App[None]):
|
||||||
text=True,
|
text=True,
|
||||||
bufsize=1,
|
bufsize=1,
|
||||||
)
|
)
|
||||||
self.runner.set_subprocess(process)
|
self.runner.set_subprocess(process, ready=False)
|
||||||
|
|
||||||
output_dir = None
|
output_dir = None
|
||||||
|
|
||||||
|
|
@ -427,7 +445,9 @@ class YoloTrainApp(App[None]):
|
||||||
if not line_str:
|
if not line_str:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
if line_str.startswith("__YOLO_TUI_EVENT__:"):
|
if line_str == "__YOLO_TUI_READY__":
|
||||||
|
self.runner.mark_subprocess_ready()
|
||||||
|
elif line_str.startswith("__YOLO_TUI_EVENT__:"):
|
||||||
try:
|
try:
|
||||||
event_data = json.loads(line_str[len("__YOLO_TUI_EVENT__:"):])
|
event_data = json.loads(line_str[len("__YOLO_TUI_EVENT__:"):])
|
||||||
event = TrainingEvent(
|
event = TrainingEvent(
|
||||||
|
|
@ -446,6 +466,7 @@ class YoloTrainApp(App[None]):
|
||||||
|
|
||||||
process.wait()
|
process.wait()
|
||||||
rc = process.returncode
|
rc = process.returncode
|
||||||
|
self.runner.clear_subprocess()
|
||||||
|
|
||||||
if rc == 0:
|
if rc == 0:
|
||||||
self.app.call_from_thread(self._training_finished, output_dir)
|
self.app.call_from_thread(self._training_finished, output_dir)
|
||||||
|
|
@ -458,13 +479,21 @@ class YoloTrainApp(App[None]):
|
||||||
Exception("Процесс обучения завершился с ошибкой. Проверьте логи выше."),
|
Exception("Процесс обучения завершился с ошибкой. Проверьте логи выше."),
|
||||||
)
|
)
|
||||||
|
|
||||||
except BaseException as exc:
|
except Exception as exc:
|
||||||
import traceback
|
import traceback
|
||||||
|
if process is not None:
|
||||||
try:
|
try:
|
||||||
with open("/Users/vadim/.gemini/antigravity/brain/ab216120-71db-4a55-a2bf-2eff6a9caf26/error.log", "w", encoding="utf-8") as f:
|
if process.poll() is None:
|
||||||
traceback.print_exc(file=f)
|
process.kill()
|
||||||
|
process.wait(timeout=5)
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
self.runner.clear_subprocess()
|
||||||
|
details = escape(traceback.format_exc())
|
||||||
|
self.app.call_from_thread(
|
||||||
|
self._write_log,
|
||||||
|
f"[red]{details}[/red]",
|
||||||
|
)
|
||||||
self.app.call_from_thread(self._training_failed, exc)
|
self.app.call_from_thread(self._training_failed, exc)
|
||||||
finally:
|
finally:
|
||||||
if temp_config_path and os.path.exists(temp_config_path):
|
if temp_config_path and os.path.exists(temp_config_path):
|
||||||
|
|
|
||||||
|
|
@ -157,6 +157,11 @@ class TrainingConfig:
|
||||||
raise ValueError("Укажите путь или имя модели.")
|
raise ValueError("Укажите путь или имя модели.")
|
||||||
if self.task not in SUPPORTED_TASKS:
|
if self.task not in SUPPORTED_TASKS:
|
||||||
raise ValueError(f"Неизвестный тип задачи: {self.task}.")
|
raise ValueError(f"Неизвестный тип задачи: {self.task}.")
|
||||||
|
if self.task == "classify" and self.split.enabled:
|
||||||
|
raise ValueError(
|
||||||
|
"Автоматическое разделение доступно только для YOLO-датасетов "
|
||||||
|
"с папками images/labels и не поддерживает задачу classify."
|
||||||
|
)
|
||||||
if self.epochs < 1:
|
if self.epochs < 1:
|
||||||
raise ValueError("Количество эпох должно быть не меньше 1.")
|
raise ValueError("Количество эпох должно быть не меньше 1.")
|
||||||
if self.image_size < 32:
|
if self.image_size < 32:
|
||||||
|
|
|
||||||
|
|
@ -2,138 +2,217 @@ from __future__ import annotations
|
||||||
|
|
||||||
import random
|
import random
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
from typing import Any
|
||||||
|
from uuid import uuid4
|
||||||
|
|
||||||
import yaml
|
import yaml
|
||||||
|
|
||||||
|
|
||||||
|
def _normalize_names(names: Any, source: Path) -> dict[int, str]:
|
||||||
|
if isinstance(names, list):
|
||||||
|
items = enumerate(names)
|
||||||
|
elif isinstance(names, dict):
|
||||||
|
normalized_items: list[tuple[int, Any]] = []
|
||||||
|
seen: set[int] = set()
|
||||||
|
for raw_key, value in names.items():
|
||||||
|
if isinstance(raw_key, bool):
|
||||||
|
raise ValueError(f"Некорректный ID класса в '{source}': {raw_key!r}.")
|
||||||
|
try:
|
||||||
|
class_id = int(raw_key)
|
||||||
|
except (TypeError, ValueError) as exc:
|
||||||
|
raise ValueError(
|
||||||
|
f"Некорректный ID класса в '{source}': {raw_key!r}."
|
||||||
|
) from exc
|
||||||
|
if class_id in seen:
|
||||||
|
raise ValueError(f"Повторяющийся ID класса {class_id} в '{source}'.")
|
||||||
|
seen.add(class_id)
|
||||||
|
normalized_items.append((class_id, value))
|
||||||
|
items = normalized_items
|
||||||
|
else:
|
||||||
|
raise ValueError(
|
||||||
|
f"Поле names в '{source}' должно быть списком или словарём."
|
||||||
|
)
|
||||||
|
|
||||||
|
result: dict[int, str] = {}
|
||||||
|
for class_id, raw_name in items:
|
||||||
|
if isinstance(raw_name, (dict, list)):
|
||||||
|
name = ""
|
||||||
|
else:
|
||||||
|
name = str(raw_name).strip()
|
||||||
|
if not name:
|
||||||
|
raise ValueError(f"Пустое имя класса {class_id} в '{source}'.")
|
||||||
|
result[class_id] = name
|
||||||
|
|
||||||
|
if not result:
|
||||||
|
raise ValueError(f"Список классов в '{source}' пуст.")
|
||||||
|
expected_ids = list(range(len(result)))
|
||||||
|
if sorted(result) != expected_ids:
|
||||||
|
raise ValueError(
|
||||||
|
f"ID классов в '{source}' должны идти подряд, начиная с 0."
|
||||||
|
)
|
||||||
|
return dict(sorted(result.items()))
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_text_classes(path: Path) -> dict[int, str]:
|
||||||
|
with path.open("r", encoding="utf-8") as classes_file:
|
||||||
|
lines = [line.strip() for line in classes_file if line.strip()]
|
||||||
|
return _normalize_names(lines, path)
|
||||||
|
|
||||||
|
|
||||||
|
def _load_yaml(path: Path) -> Any:
|
||||||
|
try:
|
||||||
|
with path.open("r", encoding="utf-8") as yaml_file:
|
||||||
|
return yaml.safe_load(yaml_file)
|
||||||
|
except yaml.YAMLError as exc:
|
||||||
|
raise ValueError(f"Не удалось разобрать YAML классов '{path}': {exc}.") from exc
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_yaml_classes(path: Path) -> dict[int, str]:
|
||||||
|
data = _load_yaml(path)
|
||||||
|
if not isinstance(data, dict) or "names" not in data:
|
||||||
|
raise ValueError(f"В YAML-файле '{path}' отсутствует поле names.")
|
||||||
|
return _normalize_names(data["names"], path)
|
||||||
|
|
||||||
|
|
||||||
def _parse_classes_file(path: Path) -> dict[int, str]:
|
def _parse_classes_file(path: Path) -> dict[int, str]:
|
||||||
with open(path, "r", encoding="utf-8") as f:
|
suffix = path.suffix.lower()
|
||||||
lines = [line.strip() for line in f if line.strip()]
|
if suffix == ".txt":
|
||||||
return {i: name for i, name in enumerate(lines)}
|
return _parse_text_classes(path)
|
||||||
|
if suffix in {".yaml", ".yml"}:
|
||||||
|
return _parse_yaml_classes(path)
|
||||||
|
raise ValueError(
|
||||||
|
f"Неподдерживаемый формат файла классов '{path}'. "
|
||||||
|
"Используйте .txt, .yaml или .yml."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def read_classes(dataset_dir: Path, custom_classes_path: str) -> dict[int, str]:
|
def read_classes(dataset_dir: Path, custom_classes_path: str) -> dict[int, str]:
|
||||||
# 1. Custom path specified by user
|
# An explicit path is authoritative: typos and malformed files must not fall back.
|
||||||
if custom_classes_path.strip():
|
if custom_classes_path.strip():
|
||||||
path = Path(custom_classes_path.strip())
|
path = Path(custom_classes_path.strip()).expanduser()
|
||||||
if path.exists():
|
if not path.exists():
|
||||||
return _parse_classes_file(path)
|
raise ValueError(f"Указанный файл классов '{path}' не существует.")
|
||||||
|
if not path.is_file():
|
||||||
# 2. classes.txt in dataset_dir root
|
raise ValueError(f"Путь к классам '{path}' не является файлом.")
|
||||||
path = dataset_dir / "classes.txt"
|
|
||||||
if path.exists():
|
|
||||||
return _parse_classes_file(path)
|
|
||||||
|
|
||||||
# 3. classes.txt inside dataset_dir/labels
|
|
||||||
path = dataset_dir / "labels" / "classes.txt"
|
|
||||||
if path.exists():
|
|
||||||
return _parse_classes_file(path)
|
|
||||||
|
|
||||||
# 4. Any .yaml file in dataset_dir (excluding split/dataset.yaml)
|
|
||||||
yaml_files = list(dataset_dir.glob("*.yaml"))
|
|
||||||
yaml_files = [f for f in yaml_files if f.name != "dataset.yaml"]
|
|
||||||
if yaml_files:
|
|
||||||
try:
|
try:
|
||||||
with open(yaml_files[0], "r", encoding="utf-8") as f:
|
return _parse_classes_file(path)
|
||||||
data = yaml.safe_load(f)
|
except OSError as exc:
|
||||||
if isinstance(data, dict) and "names" in data:
|
raise ValueError(f"Не удалось прочитать файл классов '{path}': {exc}.") from exc
|
||||||
names = data["names"]
|
|
||||||
if isinstance(names, dict):
|
|
||||||
return {int(k): str(v) for k, v in names.items()}
|
|
||||||
elif isinstance(names, list):
|
|
||||||
return {i: str(v) for i, v in enumerate(names)}
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
|
|
||||||
# 5. Fallback: Scan label files to determine number of classes and use class_i names
|
for path in (dataset_dir / "classes.txt", dataset_dir / "labels" / "classes.txt"):
|
||||||
class_ids = set()
|
if path.is_file():
|
||||||
|
return _parse_text_classes(path)
|
||||||
|
|
||||||
|
yaml_files = sorted(
|
||||||
|
(*dataset_dir.glob("*.yaml"), *dataset_dir.glob("*.yml")),
|
||||||
|
key=lambda item: item.name,
|
||||||
|
)
|
||||||
|
for path in yaml_files:
|
||||||
|
try:
|
||||||
|
data = _load_yaml(path)
|
||||||
|
except (OSError, ValueError):
|
||||||
|
continue
|
||||||
|
if isinstance(data, dict) and "names" in data:
|
||||||
|
return _normalize_names(data["names"], path)
|
||||||
|
|
||||||
|
# Infer the range from every nested label file when no class list is available.
|
||||||
|
class_ids: set[int] = set()
|
||||||
labels_dir = dataset_dir / "labels"
|
labels_dir = dataset_dir / "labels"
|
||||||
if labels_dir.exists():
|
if labels_dir.exists():
|
||||||
for txt_file in labels_dir.glob("*.txt"):
|
for txt_file in labels_dir.rglob("*.txt"):
|
||||||
if txt_file.name == "classes.txt":
|
if txt_file.name == "classes.txt":
|
||||||
continue
|
continue
|
||||||
try:
|
try:
|
||||||
with open(txt_file, "r", encoding="utf-8") as f:
|
with txt_file.open("r", encoding="utf-8") as label_file:
|
||||||
for line in f:
|
for line in label_file:
|
||||||
parts = line.strip().split()
|
parts = line.strip().split()
|
||||||
if parts:
|
if parts:
|
||||||
class_ids.add(int(parts[0]))
|
class_ids.add(int(parts[0]))
|
||||||
except Exception:
|
except (OSError, ValueError):
|
||||||
pass
|
continue
|
||||||
|
|
||||||
if class_ids:
|
if class_ids:
|
||||||
max_id = max(class_ids)
|
max_id = max(class_ids)
|
||||||
return {i: f"class_{i}" for i in range(max_id + 1)}
|
return {class_id: f"class_{class_id}" for class_id in range(max_id + 1)}
|
||||||
|
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
"Не удалось найти список классов. Пожалуйста, создайте файл classes.txt "
|
"Не удалось найти список классов. Создайте classes.txt или YAML с полем "
|
||||||
"в корневой папке датасета или укажите путь к нему."
|
"names в корневой папке датасета либо укажите путь к нему."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _write_new(path: Path, content: str) -> None:
|
||||||
|
with path.open("x", encoding="utf-8") as output_file:
|
||||||
|
output_file.write(content)
|
||||||
|
|
||||||
|
|
||||||
def split_dataset(
|
def split_dataset(
|
||||||
dataset_dir: str, train_ratio: float, classes_path: str
|
dataset_dir: str, train_ratio: float, classes_path: str
|
||||||
) -> tuple[int, int, str]:
|
) -> tuple[int, int, str]:
|
||||||
base_dir = Path(dataset_dir.strip()).absolute()
|
if not dataset_dir.strip():
|
||||||
|
raise ValueError("Укажите каталог датасета.")
|
||||||
|
if not 0.1 <= train_ratio <= 0.95:
|
||||||
|
raise ValueError("Доля обучающей выборки должна быть от 0.1 до 0.95.")
|
||||||
|
|
||||||
|
base_dir = Path(dataset_dir.strip()).expanduser().absolute()
|
||||||
images_dir = base_dir / "images"
|
images_dir = base_dir / "images"
|
||||||
labels_dir = base_dir / "labels"
|
labels_dir = base_dir / "labels"
|
||||||
|
|
||||||
if not base_dir.exists():
|
if not base_dir.is_dir():
|
||||||
raise ValueError(f"Каталог датасета '{base_dir}' не существует.")
|
raise ValueError(f"Каталог датасета '{base_dir}' не существует.")
|
||||||
if not images_dir.exists():
|
if not images_dir.is_dir():
|
||||||
raise ValueError(f"Папка с изображениями '{images_dir}' не найдена.")
|
raise ValueError(f"Папка с изображениями '{images_dir}' не найдена.")
|
||||||
if not labels_dir.exists():
|
if not labels_dir.is_dir():
|
||||||
raise ValueError(f"Папка с разметкой '{labels_dir}' не найдена.")
|
raise ValueError(f"Папка с разметкой '{labels_dir}' не найдена.")
|
||||||
|
|
||||||
# Find images
|
|
||||||
valid_extensions = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}
|
valid_extensions = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}
|
||||||
image_files = [
|
image_files = sorted(
|
||||||
f for f in images_dir.iterdir()
|
(
|
||||||
if f.is_file() and f.suffix.lower() in valid_extensions
|
path
|
||||||
]
|
for path in images_dir.rglob("*")
|
||||||
|
if path.is_file() and path.suffix.lower() in valid_extensions
|
||||||
|
),
|
||||||
|
key=lambda item: item.as_posix(),
|
||||||
|
)
|
||||||
|
|
||||||
if not image_files:
|
if not image_files:
|
||||||
raise ValueError(f"В папке '{images_dir}' не найдено изображений.")
|
raise ValueError(f"В папке '{images_dir}' не найдено изображений.")
|
||||||
|
if len(image_files) == 1:
|
||||||
|
raise ValueError(
|
||||||
|
"Для разделения нужно минимум 2 изображения; найдено: 1."
|
||||||
|
)
|
||||||
|
|
||||||
# Shuffle deterministically using a fixed seed
|
|
||||||
rng = random.Random(42)
|
rng = random.Random(42)
|
||||||
rng.shuffle(image_files)
|
rng.shuffle(image_files)
|
||||||
|
|
||||||
split_idx = int(len(image_files) * train_ratio)
|
split_idx = int(len(image_files) * train_ratio)
|
||||||
if split_idx == 0:
|
split_idx = max(1, min(split_idx, len(image_files) - 1))
|
||||||
split_idx = 1
|
|
||||||
if split_idx >= len(image_files):
|
|
||||||
split_idx = len(image_files) - 1
|
|
||||||
|
|
||||||
train_images = image_files[:split_idx]
|
train_images = image_files[:split_idx]
|
||||||
val_images = image_files[split_idx:]
|
val_images = image_files[split_idx:]
|
||||||
|
|
||||||
split_dir = base_dir / "split"
|
# Resolve classes before creating output so invalid input leaves no partial split.
|
||||||
split_dir.mkdir(parents=True, exist_ok=True)
|
classes = read_classes(base_dir, classes_path)
|
||||||
|
|
||||||
|
relative_split_dir = Path(".yolo-tui") / "splits" / uuid4().hex
|
||||||
|
split_dir = base_dir / relative_split_dir
|
||||||
|
split_dir.mkdir(parents=True, exist_ok=False)
|
||||||
|
|
||||||
train_txt_path = split_dir / "train.txt"
|
train_txt_path = split_dir / "train.txt"
|
||||||
val_txt_path = split_dir / "val.txt"
|
val_txt_path = split_dir / "val.txt"
|
||||||
|
|
||||||
with open(train_txt_path, "w", encoding="utf-8") as f:
|
|
||||||
for img in train_images:
|
|
||||||
f.write(f"{img}\n")
|
|
||||||
|
|
||||||
with open(val_txt_path, "w", encoding="utf-8") as f:
|
|
||||||
for img in val_images:
|
|
||||||
f.write(f"{img}\n")
|
|
||||||
|
|
||||||
# Load class mapping
|
|
||||||
classes = read_classes(base_dir, classes_path)
|
|
||||||
|
|
||||||
# Write data.yaml equivalent
|
|
||||||
dataset_yaml_path = split_dir / "dataset.yaml"
|
dataset_yaml_path = split_dir / "dataset.yaml"
|
||||||
|
|
||||||
|
_write_new(train_txt_path, "".join(f"{image}\n" for image in train_images))
|
||||||
|
_write_new(val_txt_path, "".join(f"{image}\n" for image in val_images))
|
||||||
|
|
||||||
dataset_data = {
|
dataset_data = {
|
||||||
"path": str(base_dir),
|
"path": str(base_dir),
|
||||||
"train": f"split/{train_txt_path.name}",
|
"train": (relative_split_dir / train_txt_path.name).as_posix(),
|
||||||
"val": f"split/{val_txt_path.name}",
|
"val": (relative_split_dir / val_txt_path.name).as_posix(),
|
||||||
"names": classes,
|
"names": classes,
|
||||||
}
|
}
|
||||||
|
_write_new(
|
||||||
with open(dataset_yaml_path, "w", encoding="utf-8") as f:
|
dataset_yaml_path,
|
||||||
yaml.safe_dump(dataset_data, f, allow_unicode=True, sort_keys=False)
|
yaml.safe_dump(dataset_data, allow_unicode=True, sort_keys=False),
|
||||||
|
)
|
||||||
|
|
||||||
return len(train_images), len(val_images), str(dataset_yaml_path)
|
return len(train_images), len(val_images), str(dataset_yaml_path)
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,10 @@
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
|
import signal
|
||||||
import sys
|
import sys
|
||||||
|
import traceback
|
||||||
|
from collections.abc import Iterator, Sequence
|
||||||
|
from contextlib import contextmanager
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
# Force headless Matplotlib to avoid any thread/process GUI issues
|
# Force headless Matplotlib to avoid any thread/process GUI issues
|
||||||
|
|
@ -9,21 +13,47 @@ os.environ["MPLBACKEND"] = "Agg"
|
||||||
from yolo_tui.config import TrainingConfig
|
from yolo_tui.config import TrainingConfig
|
||||||
from yolo_tui.trainer import TrainingEvent, TrainingRunner
|
from yolo_tui.trainer import TrainingEvent, TrainingRunner
|
||||||
|
|
||||||
def main():
|
|
||||||
if len(sys.argv) < 2:
|
|
||||||
print("Usage: python -m yolo_tui.subprocess_runner <config_json_path>", file=sys.stderr)
|
|
||||||
sys.exit(1)
|
|
||||||
|
|
||||||
config_path = sys.argv[1]
|
@contextmanager
|
||||||
|
def _stop_signal_handlers(runner: TrainingRunner) -> Iterator[None]:
|
||||||
|
previous: dict[signal.Signals, signal.Handlers] = {}
|
||||||
|
|
||||||
|
def request_stop(_signum: int, _frame: object) -> None:
|
||||||
|
runner.request_stop()
|
||||||
|
|
||||||
|
for signum in (signal.SIGTERM, signal.SIGINT):
|
||||||
|
previous[signum] = signal.getsignal(signum)
|
||||||
|
signal.signal(signum, request_stop)
|
||||||
try:
|
try:
|
||||||
with open(config_path, "r", encoding="utf-8") as f:
|
yield
|
||||||
config_dict = json.load(f)
|
finally:
|
||||||
config = TrainingConfig.from_dict(config_dict)
|
for signum, handler in previous.items():
|
||||||
except Exception as e:
|
signal.signal(signum, handler)
|
||||||
print(f"Error loading config: {e}", file=sys.stderr)
|
|
||||||
sys.exit(1)
|
|
||||||
|
|
||||||
|
|
||||||
|
def main(argv: Sequence[str] | None = None) -> int:
|
||||||
|
args = list(sys.argv[1:] if argv is None else argv)
|
||||||
|
if not args:
|
||||||
|
print(
|
||||||
|
"Usage: python -m yolo_tui.subprocess_runner <config_json_path>",
|
||||||
|
file=sys.stderr,
|
||||||
|
)
|
||||||
|
return 1
|
||||||
|
|
||||||
|
config_path = Path(args[0])
|
||||||
runner = TrainingRunner()
|
runner = TrainingRunner()
|
||||||
|
runner.prepare_run()
|
||||||
|
|
||||||
|
with _stop_signal_handlers(runner):
|
||||||
|
# The parent waits for this marker before sending a cooperative signal.
|
||||||
|
print("__YOLO_TUI_READY__", flush=True)
|
||||||
|
try:
|
||||||
|
with config_path.open("r", encoding="utf-8") as config_file:
|
||||||
|
config_dict = json.load(config_file)
|
||||||
|
config = TrainingConfig.from_dict(config_dict)
|
||||||
|
except Exception as exc:
|
||||||
|
print(f"Error loading config: {exc}", file=sys.stderr)
|
||||||
|
return 1
|
||||||
|
|
||||||
def handle_event(event: TrainingEvent) -> None:
|
def handle_event(event: TrainingEvent) -> None:
|
||||||
event_dict = {
|
event_dict = {
|
||||||
|
|
@ -39,11 +69,10 @@ def main():
|
||||||
output_dir = runner.train(config, handle_event)
|
output_dir = runner.train(config, handle_event)
|
||||||
if output_dir:
|
if output_dir:
|
||||||
print(f"__YOLO_TUI_RESULT__:{output_dir}", flush=True)
|
print(f"__YOLO_TUI_RESULT__:{output_dir}", flush=True)
|
||||||
sys.exit(0)
|
return 0
|
||||||
except BaseException as e:
|
except Exception:
|
||||||
import traceback
|
|
||||||
traceback.print_exc()
|
traceback.print_exc()
|
||||||
sys.exit(1)
|
return 1
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
main()
|
raise SystemExit(main())
|
||||||
|
|
|
||||||
|
|
@ -1,11 +1,12 @@
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import os
|
import os
|
||||||
|
import signal
|
||||||
from collections.abc import Callable, Iterator
|
from collections.abc import Callable, Iterator
|
||||||
from contextlib import contextmanager
|
from contextlib import contextmanager
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from threading import Event, Lock
|
from threading import Event, RLock, Timer
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from .config import MlflowConfig, TrainingConfig
|
from .config import MlflowConfig, TrainingConfig
|
||||||
|
|
@ -51,29 +52,66 @@ def mlflow_environment(config: MlflowConfig) -> Iterator[None]:
|
||||||
class TrainingRunner:
|
class TrainingRunner:
|
||||||
"""Owns a single YOLO training run and exposes cooperative cancellation."""
|
"""Owns a single YOLO training run and exposes cooperative cancellation."""
|
||||||
|
|
||||||
|
FORCE_STOP_TIMEOUT_SECONDS = 30.0
|
||||||
|
|
||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
self._model: Any | None = None
|
self._model: Any | None = None
|
||||||
self._state_lock = Lock()
|
self._state_lock = RLock()
|
||||||
self._stop_requested = Event()
|
self._stop_requested = Event()
|
||||||
self._subprocess: Any | None = None
|
self._subprocess: Any | None = None
|
||||||
|
self._subprocess_ready = False
|
||||||
|
self._force_stop_timer: Timer | None = None
|
||||||
|
|
||||||
def set_subprocess(self, process: Any) -> None:
|
def prepare_run(self) -> None:
|
||||||
|
"""Reset cancellation state before starting a new training run."""
|
||||||
|
with self._state_lock:
|
||||||
|
timer = self._force_stop_timer
|
||||||
|
self._force_stop_timer = None
|
||||||
|
self._subprocess_ready = False
|
||||||
|
self._stop_requested.clear()
|
||||||
|
if timer is not None:
|
||||||
|
timer.cancel()
|
||||||
|
|
||||||
|
def set_subprocess(self, process: Any, *, ready: bool = False) -> None:
|
||||||
|
"""Register the child process without losing an earlier stop request."""
|
||||||
with self._state_lock:
|
with self._state_lock:
|
||||||
self._subprocess = process
|
self._subprocess = process
|
||||||
|
self._subprocess_ready = ready
|
||||||
|
stop_requested = self._stop_requested.is_set()
|
||||||
|
if stop_requested:
|
||||||
|
if ready:
|
||||||
|
self._send_cooperative_stop(process)
|
||||||
|
self._schedule_force_stop(process)
|
||||||
|
|
||||||
|
def mark_subprocess_ready(self) -> None:
|
||||||
|
"""Mark the child signal handler as ready and deliver any pending stop."""
|
||||||
|
with self._state_lock:
|
||||||
|
process = self._subprocess
|
||||||
|
self._subprocess_ready = process is not None
|
||||||
|
stop_requested = self._stop_requested.is_set()
|
||||||
|
if process is not None and stop_requested:
|
||||||
|
self._send_cooperative_stop(process)
|
||||||
|
self._schedule_force_stop(process)
|
||||||
|
|
||||||
def clear_subprocess(self) -> None:
|
def clear_subprocess(self) -> None:
|
||||||
with self._state_lock:
|
with self._state_lock:
|
||||||
self._subprocess = None
|
self._subprocess = None
|
||||||
|
self._subprocess_ready = False
|
||||||
|
timer = self._force_stop_timer
|
||||||
|
self._force_stop_timer = None
|
||||||
|
if timer is not None:
|
||||||
|
timer.cancel()
|
||||||
|
|
||||||
def request_stop(self) -> None:
|
def request_stop(self) -> None:
|
||||||
self._stop_requested.set()
|
self._stop_requested.set()
|
||||||
with self._state_lock:
|
with self._state_lock:
|
||||||
if self._subprocess is not None:
|
process = self._subprocess
|
||||||
try:
|
process_ready = self._subprocess_ready
|
||||||
self._subprocess.terminate()
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
trainer = getattr(self._model, "trainer", None)
|
trainer = getattr(self._model, "trainer", None)
|
||||||
|
if process is not None:
|
||||||
|
if process_ready:
|
||||||
|
self._send_cooperative_stop(process)
|
||||||
|
self._schedule_force_stop(process)
|
||||||
if trainer is not None:
|
if trainer is not None:
|
||||||
trainer.stop = True
|
trainer.stop = True
|
||||||
|
|
||||||
|
|
@ -81,9 +119,41 @@ class TrainingRunner:
|
||||||
def stop_requested(self) -> bool:
|
def stop_requested(self) -> bool:
|
||||||
return self._stop_requested.is_set()
|
return self._stop_requested.is_set()
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _send_cooperative_stop(process: Any) -> None:
|
||||||
|
try:
|
||||||
|
process.send_signal(signal.SIGTERM)
|
||||||
|
except (AttributeError, OSError, ProcessLookupError):
|
||||||
|
try:
|
||||||
|
process.terminate()
|
||||||
|
except (AttributeError, OSError, ProcessLookupError):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def _schedule_force_stop(self, process: Any) -> None:
|
||||||
|
with self._state_lock:
|
||||||
|
if self._subprocess is not process or self._force_stop_timer is not None:
|
||||||
|
return
|
||||||
|
timer = Timer(
|
||||||
|
self.FORCE_STOP_TIMEOUT_SECONDS,
|
||||||
|
self._force_stop,
|
||||||
|
args=(process,),
|
||||||
|
)
|
||||||
|
timer.daemon = True
|
||||||
|
self._force_stop_timer = timer
|
||||||
|
timer.start()
|
||||||
|
|
||||||
|
def _force_stop(self, process: Any) -> None:
|
||||||
|
with self._state_lock:
|
||||||
|
if self._subprocess is not process:
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
if process.poll() is None:
|
||||||
|
process.kill()
|
||||||
|
except (AttributeError, OSError, ProcessLookupError):
|
||||||
|
pass
|
||||||
|
|
||||||
def train(self, config: TrainingConfig, on_event: EventHandler) -> Path | None:
|
def train(self, config: TrainingConfig, on_event: EventHandler) -> Path | None:
|
||||||
config.validate()
|
config.validate()
|
||||||
self._stop_requested.clear()
|
|
||||||
|
|
||||||
train_args = config.train_kwargs()
|
train_args = config.train_kwargs()
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -98,3 +98,21 @@ def test_read_config_ignores_disabled_split() -> None:
|
||||||
assert config.split.train_ratio == DatasetSplitConfig(enabled=False).train_ratio
|
assert config.split.train_ratio == DatasetSplitConfig(enabled=False).train_ratio
|
||||||
|
|
||||||
asyncio.run(exercise())
|
asyncio.run(exercise())
|
||||||
|
|
||||||
|
|
||||||
|
def test_classify_disables_detection_style_split() -> None:
|
||||||
|
async def exercise() -> None:
|
||||||
|
app = YoloTrainApp()
|
||||||
|
async with app.run_test(size=(140, 45)) as pilot:
|
||||||
|
app.query_one("#split-enabled", Switch).value = True
|
||||||
|
await pilot.pause()
|
||||||
|
|
||||||
|
app.query_one("#task", Select).value = "classify"
|
||||||
|
await pilot.pause()
|
||||||
|
|
||||||
|
assert app.query_one("#split-enabled", Switch).value is False
|
||||||
|
assert app.query_one("#split-enabled", Switch).disabled is True
|
||||||
|
assert app.query_one("#split-ratio", Input).disabled is True
|
||||||
|
assert app.query_one("#split-classes", Input).disabled is True
|
||||||
|
|
||||||
|
asyncio.run(exercise())
|
||||||
|
|
|
||||||
|
|
@ -96,3 +96,15 @@ def test_dataset_split_ratio_is_validated(ratio: float) -> None:
|
||||||
config = DatasetSplitConfig(enabled=True, train_ratio=ratio)
|
config = DatasetSplitConfig(enabled=True, train_ratio=ratio)
|
||||||
with pytest.raises(ValueError, match="Доля обучающей выборки"):
|
with pytest.raises(ValueError, match="Доля обучающей выборки"):
|
||||||
config.validate()
|
config.validate()
|
||||||
|
|
||||||
|
|
||||||
|
def test_classification_rejects_detection_style_auto_split() -> None:
|
||||||
|
config = TrainingConfig(
|
||||||
|
dataset="classification-dataset",
|
||||||
|
model="model.pt",
|
||||||
|
task="classify",
|
||||||
|
split=DatasetSplitConfig(enabled=True),
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="classify"):
|
||||||
|
config.validate()
|
||||||
|
|
|
||||||
|
|
@ -1,7 +1,7 @@
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import os
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
import yaml
|
import yaml
|
||||||
|
|
||||||
|
|
@ -16,6 +16,32 @@ def test_read_classes_custom_path(tmp_path: Path) -> None:
|
||||||
assert classes == {0: "classA", 1: "classB"}
|
assert classes == {0: "classA", 1: "classB"}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("names", "expected"),
|
||||||
|
[
|
||||||
|
(["cat", "dog"], {0: "cat", 1: "dog"}),
|
||||||
|
({0: "cat", 1: "dog"}, {0: "cat", 1: "dog"}),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_read_classes_custom_yaml(
|
||||||
|
tmp_path: Path, names: object, expected: dict[int, str]
|
||||||
|
) -> None:
|
||||||
|
custom_file = tmp_path / "custom_classes.yaml"
|
||||||
|
custom_file.write_text(
|
||||||
|
yaml.safe_dump({"names": names}, allow_unicode=True),
|
||||||
|
encoding="utf-8",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert read_classes(tmp_path, str(custom_file)) == expected
|
||||||
|
|
||||||
|
|
||||||
|
def test_missing_custom_classes_path_does_not_fall_back(tmp_path: Path) -> None:
|
||||||
|
(tmp_path / "classes.txt").write_text("fallback\n", encoding="utf-8")
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="не существует"):
|
||||||
|
read_classes(tmp_path, str(tmp_path / "typo.yaml"))
|
||||||
|
|
||||||
|
|
||||||
def test_read_classes_root_classes_txt(tmp_path: Path) -> None:
|
def test_read_classes_root_classes_txt(tmp_path: Path) -> None:
|
||||||
classes_file = tmp_path / "classes.txt"
|
classes_file = tmp_path / "classes.txt"
|
||||||
classes_file.write_text("class0\nclass1\nclass2\n", encoding="utf-8")
|
classes_file.write_text("class0\nclass1\nclass2\n", encoding="utf-8")
|
||||||
|
|
@ -84,13 +110,15 @@ def test_split_dataset_flow(tmp_path: Path) -> None:
|
||||||
with open(yaml_path, "r", encoding="utf-8") as f:
|
with open(yaml_path, "r", encoding="utf-8") as f:
|
||||||
data = yaml.safe_load(f)
|
data = yaml.safe_load(f)
|
||||||
assert data["path"] == str(tmp_path)
|
assert data["path"] == str(tmp_path)
|
||||||
assert data["train"] == "split/train.txt"
|
relative_output_dir = Path(yaml_path).parent.relative_to(tmp_path)
|
||||||
assert data["val"] == "split/val.txt"
|
assert data["train"] == (relative_output_dir / "train.txt").as_posix()
|
||||||
|
assert data["val"] == (relative_output_dir / "val.txt").as_posix()
|
||||||
assert data["names"] == {0: "dummy_class"}
|
assert data["names"] == {0: "dummy_class"}
|
||||||
|
|
||||||
# Verify lists content
|
# Verify lists content
|
||||||
train_list = (tmp_path / "split" / "train.txt").read_text(encoding="utf-8").strip().split("\n")
|
output_dir = Path(yaml_path).parent
|
||||||
val_list = (tmp_path / "split" / "val.txt").read_text(encoding="utf-8").strip().split("\n")
|
train_list = (output_dir / "train.txt").read_text(encoding="utf-8").strip().split("\n")
|
||||||
|
val_list = (output_dir / "val.txt").read_text(encoding="utf-8").strip().split("\n")
|
||||||
|
|
||||||
assert len(train_list) == 3
|
assert len(train_list) == 3
|
||||||
assert len(val_list) == 1
|
assert len(val_list) == 1
|
||||||
|
|
@ -98,3 +126,69 @@ def test_split_dataset_flow(tmp_path: Path) -> None:
|
||||||
# Verify paths are absolute
|
# Verify paths are absolute
|
||||||
assert Path(train_list[0]).is_absolute()
|
assert Path(train_list[0]).is_absolute()
|
||||||
assert Path(val_list[0]).is_absolute()
|
assert Path(val_list[0]).is_absolute()
|
||||||
|
|
||||||
|
|
||||||
|
def test_split_dataset_rejects_single_image_without_writing_output(
|
||||||
|
tmp_path: Path,
|
||||||
|
) -> None:
|
||||||
|
images_dir = tmp_path / "images"
|
||||||
|
labels_dir = tmp_path / "labels"
|
||||||
|
images_dir.mkdir()
|
||||||
|
labels_dir.mkdir()
|
||||||
|
(images_dir / "only.jpg").write_bytes(b"")
|
||||||
|
(labels_dir / "only.txt").write_text("0 0.5 0.5 1 1\n", encoding="utf-8")
|
||||||
|
(tmp_path / "classes.txt").write_text("item\n", encoding="utf-8")
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="минимум 2"):
|
||||||
|
split_dataset(str(tmp_path), 0.8, "")
|
||||||
|
|
||||||
|
assert not (tmp_path / ".yolo-tui").exists()
|
||||||
|
|
||||||
|
|
||||||
|
def test_split_dataset_finds_nested_images_and_labels(tmp_path: Path) -> None:
|
||||||
|
images_dir = tmp_path / "images" / "day"
|
||||||
|
labels_dir = tmp_path / "labels" / "day"
|
||||||
|
images_dir.mkdir(parents=True)
|
||||||
|
labels_dir.mkdir(parents=True)
|
||||||
|
for index in range(2):
|
||||||
|
(images_dir / f"nested-{index}.jpg").write_bytes(b"")
|
||||||
|
(labels_dir / f"nested-{index}.txt").write_text(
|
||||||
|
"2 0.5 0.5 0.2 0.2\n",
|
||||||
|
encoding="utf-8",
|
||||||
|
)
|
||||||
|
|
||||||
|
train_count, val_count, yaml_path = split_dataset(str(tmp_path), 0.8, "")
|
||||||
|
|
||||||
|
assert (train_count, val_count) == (1, 1)
|
||||||
|
data = yaml.safe_load(Path(yaml_path).read_text(encoding="utf-8"))
|
||||||
|
assert data["names"] == {0: "class_0", 1: "class_1", 2: "class_2"}
|
||||||
|
listed_images = "".join(
|
||||||
|
(Path(yaml_path).parent / filename).read_text(encoding="utf-8")
|
||||||
|
for filename in ("train.txt", "val.txt")
|
||||||
|
)
|
||||||
|
assert "images/day/nested-0.jpg" in listed_images
|
||||||
|
assert "images/day/nested-1.jpg" in listed_images
|
||||||
|
|
||||||
|
|
||||||
|
def test_split_dataset_preserves_existing_split_files(tmp_path: Path) -> None:
|
||||||
|
images_dir = tmp_path / "images"
|
||||||
|
labels_dir = tmp_path / "labels"
|
||||||
|
images_dir.mkdir()
|
||||||
|
labels_dir.mkdir()
|
||||||
|
for index in range(2):
|
||||||
|
(images_dir / f"image-{index}.jpg").write_bytes(b"")
|
||||||
|
(labels_dir / f"image-{index}.txt").write_text(
|
||||||
|
"0 0.5 0.5 0.2 0.2\n",
|
||||||
|
encoding="utf-8",
|
||||||
|
)
|
||||||
|
(tmp_path / "classes.txt").write_text("item\n", encoding="utf-8")
|
||||||
|
user_split = tmp_path / "split"
|
||||||
|
user_split.mkdir()
|
||||||
|
(user_split / "train.txt").write_text("user data\n", encoding="utf-8")
|
||||||
|
|
||||||
|
first_yaml = Path(split_dataset(str(tmp_path), 0.5, "")[2])
|
||||||
|
second_yaml = Path(split_dataset(str(tmp_path), 0.5, "")[2])
|
||||||
|
|
||||||
|
assert (user_split / "train.txt").read_text(encoding="utf-8") == "user data\n"
|
||||||
|
assert first_yaml.parent != second_yaml.parent
|
||||||
|
assert user_split not in first_yaml.parents
|
||||||
|
|
|
||||||
63
tests/test_subprocess_runner.py
Normal file
63
tests/test_subprocess_runner.py
Normal file
|
|
@ -0,0 +1,63 @@
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from yolo_tui import subprocess_runner
|
||||||
|
|
||||||
|
|
||||||
|
class FakeRunner:
|
||||||
|
def __init__(self, *, error: Exception | None = None) -> None:
|
||||||
|
self.error = error
|
||||||
|
self.prepared = False
|
||||||
|
self.stop_requested = False
|
||||||
|
|
||||||
|
def prepare_run(self) -> None:
|
||||||
|
self.prepared = True
|
||||||
|
|
||||||
|
def request_stop(self) -> None:
|
||||||
|
self.stop_requested = True
|
||||||
|
|
||||||
|
def train(self, config: Any, on_event: Any) -> Path:
|
||||||
|
if self.error is not None:
|
||||||
|
raise self.error
|
||||||
|
return Path("/tmp/successful-run")
|
||||||
|
|
||||||
|
|
||||||
|
def _write_config(tmp_path: Path) -> Path:
|
||||||
|
config_path = tmp_path / "config.json"
|
||||||
|
config_path.write_text(
|
||||||
|
json.dumps({"dataset": "dataset.yaml", "model": "model.pt"}),
|
||||||
|
encoding="utf-8",
|
||||||
|
)
|
||||||
|
return config_path
|
||||||
|
|
||||||
|
|
||||||
|
def test_main_returns_zero_after_successful_training(
|
||||||
|
monkeypatch: Any, tmp_path: Path, capsys: Any
|
||||||
|
) -> None:
|
||||||
|
runner = FakeRunner()
|
||||||
|
monkeypatch.setattr(subprocess_runner, "TrainingRunner", lambda: runner)
|
||||||
|
|
||||||
|
return_code = subprocess_runner.main([str(_write_config(tmp_path))])
|
||||||
|
|
||||||
|
output = capsys.readouterr()
|
||||||
|
assert return_code == 0
|
||||||
|
assert runner.prepared is True
|
||||||
|
assert "__YOLO_TUI_READY__" in output.out
|
||||||
|
assert "__YOLO_TUI_RESULT__:/tmp/successful-run" in output.out
|
||||||
|
assert "Traceback" not in output.err
|
||||||
|
|
||||||
|
|
||||||
|
def test_main_returns_one_when_training_raises(
|
||||||
|
monkeypatch: Any, tmp_path: Path, capsys: Any
|
||||||
|
) -> None:
|
||||||
|
runner = FakeRunner(error=RuntimeError("training failed"))
|
||||||
|
monkeypatch.setattr(subprocess_runner, "TrainingRunner", lambda: runner)
|
||||||
|
|
||||||
|
return_code = subprocess_runner.main([str(_write_config(tmp_path))])
|
||||||
|
|
||||||
|
output = capsys.readouterr()
|
||||||
|
assert return_code == 1
|
||||||
|
assert "RuntimeError: training failed" in output.err
|
||||||
|
|
@ -1,6 +1,7 @@
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import sys
|
import sys
|
||||||
|
import signal
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from types import ModuleType, SimpleNamespace
|
from types import ModuleType, SimpleNamespace
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
@ -9,6 +10,25 @@ from yolo_tui.config import MlflowConfig, TrainingConfig
|
||||||
from yolo_tui.trainer import TrainingEvent, TrainingRunner
|
from yolo_tui.trainer import TrainingEvent, TrainingRunner
|
||||||
|
|
||||||
|
|
||||||
|
class FakeProcess:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.signals: list[int] = []
|
||||||
|
self.terminate_calls = 0
|
||||||
|
self.kill_calls = 0
|
||||||
|
|
||||||
|
def send_signal(self, signum: int) -> None:
|
||||||
|
self.signals.append(signum)
|
||||||
|
|
||||||
|
def terminate(self) -> None:
|
||||||
|
self.terminate_calls += 1
|
||||||
|
|
||||||
|
def kill(self) -> None:
|
||||||
|
self.kill_calls += 1
|
||||||
|
|
||||||
|
def poll(self) -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
def test_runner_wires_yolo_callbacks_and_returns_output(
|
def test_runner_wires_yolo_callbacks_and_returns_output(
|
||||||
monkeypatch: Any, tmp_path: Path
|
monkeypatch: Any, tmp_path: Path
|
||||||
) -> None:
|
) -> None:
|
||||||
|
|
@ -66,3 +86,40 @@ def test_runner_wires_yolo_callbacks_and_returns_output(
|
||||||
assert train_arguments[0]["data"] == "dataset.yaml"
|
assert train_arguments[0]["data"] == "dataset.yaml"
|
||||||
assert train_arguments[0]["verbose"] is False
|
assert train_arguments[0]["verbose"] is False
|
||||||
assert [event.kind for event in events] == ["info", "started", "epoch", "epoch", "success"]
|
assert [event.kind for event in events] == ["info", "started", "epoch", "epoch", "success"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_prepare_run_clears_previous_stop_request() -> None:
|
||||||
|
runner = TrainingRunner()
|
||||||
|
runner.request_stop()
|
||||||
|
assert runner.stop_requested is True
|
||||||
|
|
||||||
|
runner.prepare_run()
|
||||||
|
|
||||||
|
assert runner.stop_requested is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_early_stop_is_delivered_after_subprocess_ready() -> None:
|
||||||
|
runner = TrainingRunner()
|
||||||
|
process = FakeProcess()
|
||||||
|
runner.request_stop()
|
||||||
|
|
||||||
|
runner.set_subprocess(process, ready=False)
|
||||||
|
assert process.signals == []
|
||||||
|
|
||||||
|
runner.mark_subprocess_ready()
|
||||||
|
|
||||||
|
assert process.signals == [signal.SIGTERM]
|
||||||
|
assert process.terminate_calls == 0
|
||||||
|
runner.clear_subprocess()
|
||||||
|
|
||||||
|
|
||||||
|
def test_ready_subprocess_receives_cooperative_signal_not_terminate() -> None:
|
||||||
|
runner = TrainingRunner()
|
||||||
|
process = FakeProcess()
|
||||||
|
runner.set_subprocess(process, ready=True)
|
||||||
|
|
||||||
|
runner.request_stop()
|
||||||
|
|
||||||
|
assert process.signals == [signal.SIGTERM]
|
||||||
|
assert process.terminate_calls == 0
|
||||||
|
runner.clear_subprocess()
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue