yolo-trainer

Репозиторий для обучения YOLO-моделей на пользовательских датасетах через ultralytics.

Требования

• Python 3.10+;
• CUDA (опционально; рекомендуется для torch из requirements.txt).

Установка

python -m venv .venv
.venv\Scripts\activate
pip install -r requirements.txt

Структура проекта

.
├── datasets/
│   └── <DatasetName>/
│       ├── data.yaml
│       ├── train/
│       │   ├── images/
│       │   └── labels/
│       ├── valid/
│       │   ├── images/
│       │   └── labels/
│       └── test/
│           ├── images/
│           └── labels/
├── runs/
├── requirements.txt
└── README.md

Формат датасета

Разметка должна быть в YOLO-формате:

<class_id> <x_center> <y_center> <width> <height>

Координаты нормализованы в диапазоне 0..1.

Пример data.yaml:

train: train/images
val: valid/images
test: test/images

nc: 2
names: ['class_0', 'class_1']

nc должен совпадать с количеством классов в names.

Обучение

Пример запуска:

yolo task=detect mode=train model=yolo26n.pt data=./datasets/DatasetName/data.yaml epochs=200 imgsz=320 batch=32 degrees=180 scale=0.5 mosaic=1.0 flipud=0.5 fliplr=0.5 hsv_h=0.015 hsv_s=0.7 hsv_v=0.4 name=DatasetName

Основные параметры:

• task=detect - задача детекции объектов;
• mode=train - режим обучения;
• model=yolo26n.pt - стартовые веса модели;
• data=.../data.yaml - конфигурация датасета;
• epochs=200 - количество эпох;
• imgsz=320 - размер входного изображения;
• batch=32 - размер batch;
• degrees=180 - случайный поворот изображения;
• scale=0.5 - случайное масштабирование;
• mosaic=1.0 - вероятность mosaic-аугментации;
• flipud=0.5 - вероятность вертикального отражения;
• fliplr=0.5 - вероятность горизонтального отражения;
• hsv_h=0.015 - изменение оттенка;
• hsv_s=0.7 - изменение насыщенности;
• hsv_v=0.4 - изменение яркости;
• name=DatasetName - имя запуска в runs/detect/.

Результаты

После обучения артефакты сохраняются в:

runs/detect/<name>/

Основные файлы:

• weights/best.pt - лучшие веса по метрике валидации.
• weights/last.pt - веса последней эпохи.
• results.csv - метрики по эпохам.
• results.png - графики обучения.
• confusion_matrix.png - матрица ошибок.

Валидация

yolo task=detect mode=val model=runs/detect/DatasetName/weights/best.pt data=./datasets/DatasetName/data.yaml imgsz=320

Предсказание

yolo task=detect mode=predict model=runs/detect/DatasetName/weights/best.pt source=./datasets/DatasetName/test/images imgsz=320

Экспорт модели

yolo export model=runs/detect/DatasetName/weights/best.pt format=onnx imgsz=320