yolo-trainer

A project template for preparing datasets and training YOLO models with ultralytics.

⚙️ Requirements

• Python 3.10+;
• NVIDIA GPU and a driver with CUDA support.

requirements.txt uses PyTorch wheels from the https://download.pytorch.org/whl/cu130 index. If you need a different CUDA version, replace the index with a suitable one from PyTorch.

🛠️ Installation

1. Clone the repository:

   git clone https://github.com/fibit/yolo-trainer
   cd yolo-trainer

2. Create a virtual environment:

   python -m venv .venv

3. Activate the virtual environment:

   .venv\Scripts\activate    # Windows
   source .venv/bin/activate # Linux

4. Install the dependencies:

   pip install -r requirements.txt

5. Configure ultralytics for this project (see the 🧩 Configuring Ultralytics section).

🧩 Configuring Ultralytics

On first launch, ultralytics creates a settings.json file in the user directory (~/.config/Ultralytics/ on Linux, ~/Library/Application Support/Ultralytics/ on macOS, %APPDATA%\Ultralytics\ on Windows) and records the absolute paths datasets_dir, weights_dir, and runs_dir of the project in which the first launch occurred. If there are several projects, model weights and training results will end up in the folders of the project that configured ultralytics first.

To isolate the settings inside the project, run the following once from the project root after cloning.

Linux/macOS:

mkdir -p .config
export YOLO_CONFIG_DIR=.config
yolo settings datasets_dir=datasets weights_dir=weights runs_dir=runs

Windows (PowerShell):

mkdir .config -Force
$env:YOLO_CONFIG_DIR = ".config"
yolo settings datasets_dir=datasets weights_dir=weights runs_dir=runs

This gives you the following:

• settings.json is stored in .config/Ultralytics/ inside the project, and the global settings of the user directory are not used;
• datasets_dir, weights_dir, and runs_dir are set to relative paths, so when commands are run from the project root, pretrained weights are downloaded to weights/, and the results of training, validation, and prediction are saved to this project's runs/.

The YOLO_CONFIG_DIR variable is active within the current terminal session. In a new terminal, set it again:

• Linux/macOS: export YOLO_CONFIG_DIR=.config
• Windows (PowerShell): $env:YOLO_CONFIG_DIR = ".config"

All commands below are run from the project root with YOLO_CONFIG_DIR set.

📁 Structure

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

🧾 Dataset format

The annotations must be in YOLO format:

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

The coordinates are normalized to the 0..1 range.

Example data.yaml:

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

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

nc must match the number of classes in names.

🧱 Creating a dataset

You can create an empty dataset structure with the command:

python dataset.py DatasetName

The script will create the train/, valid/, and test/ folders with images/ and labels/ subfolders, as well as a data.yaml file. Before training, fill in the nc and names values in data.yaml.

🚀 Training

Start training from the project root after activating the virtual environment and configuring Ultralytics:

yolo task=detect mode=train model=weights/yolo26n.pt data=./datasets/DatasetName/data.yaml epochs=100 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

Replace DatasetName with the name of the dataset folder in datasets/, and model=weights/yolo26n.pt with the desired model or path to the weights. If the specified weights file does not exist, it will be downloaded automatically to the project's weights/.

Main parameters:

• task=detect - object detection task;
• mode=train - training mode;
• model=weights/yolo26n.pt - initial model weights (downloaded to the project's weights/);
• data=.../data.yaml - dataset configuration;
• epochs=100 - number of epochs;
• imgsz=320 - input image size;
• batch=32 - batch size;
• degrees=180 - random image rotation;
• scale=0.5 - random scaling;
• mosaic=1.0 - mosaic augmentation probability;
• flipud=0.5 - vertical flip probability;
• fliplr=0.5 - horizontal flip probability;
• hsv_h=0.015 - hue change;
• hsv_s=0.7 - saturation change;
• hsv_v=0.4 - brightness change;
• name=DatasetName - run name in runs/detect/.

📊 Results

After training, the artifacts are saved to:

runs/detect/<name>/

Main files:

• weights/best.pt - best weights by validation metric;
• weights/last.pt - weights of the last epoch;
• results.csv - metrics per epoch;
• results.png - training plots;
• confusion_matrix.png - confusion matrix.

🧪 Validation

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

🎯 Prediction

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