YOLO нейросеть: как обучить модель обнаружения объектов с нуля

Что такое YOLO и почему она так популярна

YOLO (You Only Look Once) — это семейство нейронных сетей для обнаружения объектов на изображениях и видео в реальном времени. Главная особенность YOLO в том, что она анализирует весь кадр за один проход, а не сканирует его по частям. Это обеспечивает высокую скорость работы, позволяя обрабатывать десятки кадров в секунду даже на относительно скромном оборудовании.

YOLO решает несколько задач компьютерного зрения: детекцию (поиск объектов и их границ), сегментацию (разделение изображения на области, принадлежащие разным объектам), классификацию (определение того, что изображено на картинке), поиск ключевых точек тела (оценка позы человека) и трекинг объектов (отслеживание перемещения объектов в видеопотоке).

Популярность YOLO объясняется сочетанием скорости и приемлемой точности. Модель подходит для беспилотных автомобилей, систем видеонаблюдения, робототехники, дополненной реальности и многих других приложений, где важна оперативность.

Как работает YOLO: принцип «одного взгляда»

В отличие от более старых подходов, которые сначала выделяли потенциальные области интереса, а затем классифицировали каждую из них, YOLO обрабатывает изображение целиком. Нейросеть делит картинку на сетку, например 7×7 или 13×13. Каждая ячейка сетки отвечает за обнаружение объектов, центр которых попадает в эту ячейку.

Для каждой ячейки YOLO предсказывает несколько ограничивающих рамок (bounding boxes) разного размера и формы, а также вероятность наличия объекта внутри каждой рамки и вероятности принадлежности к каждому из известных классов. После этого применяется фильтр Non-Maximum Suppression (NMS), который удаляет дублирующиеся рамки, оставляя только наиболее уверенные предсказания.

Такой подход позволяет YOLO работать очень быстро — модель не тратит время на перебор множества кандидатов. Однако за скорость приходится платить некоторой потерей точности при обнаружении мелких объектов или объектов, расположенных близко друг к другу.

Установка YOLO и подготовка окружения

Для работы с YOLO (на примере версии YOLOv8) необходимо установить библиотеку Ultralytics. Это можно сделать одной командой:

pip install ultralytics

После установки все необходимые модули будут доступны. Для работы с изображениями и видео также потребуется библиотека OpenCV, которая обычно устанавливается вместе с Ultralytics. Если её нет, установите отдельно:

pip install opencv-python

YOLO поддерживает работу на CPU, но для обучения и быстрой обработки рекомендуется использовать GPU (NVIDIA CUDA). Убедитесь, что у вас установлены драйверы и CUDA Toolkit, если планируете задействовать видеокарту.

Для начала можно загрузить предобученную модель, например yolov8n.pt — самую лёгкую и быструю версию. Она будет скачана автоматически при первом запуске скрипта.

Структура датасета для обучения YOLO

Чтобы обучить YOLO распознавать новые объекты, необходимо подготовить собственный датасет. Данные должны быть размечены в формате YOLO — это текстовые файлы, где для каждого изображения указаны класс и координаты ограничивающей рамки в нормализованном виде.

Формат одной строки в файле разметки:

Где:

• class_id — целочисленный идентификатор класса (начиная с 0).
• x_center, y_center — координаты центра рамки, нормализованные относительно ширины и высоты изображения (значения от 0 до 1).
• width, height — ширина и высота рамки, также нормализованные.

Структура папок датасета должна быть следующей:

dataset/
  images/
    train/
    val/
  labels/
    train/
    val/

В папках images/train и images/val хранятся изображения, а в labels/train и labels/val — соответствующие текстовые файлы разметки с теми же именами, но с расширением .txt.

Также необходимо создать файл конфигурации датасета (например, dataset.yaml) со следующим содержанием:

train: dataset/images/train
val: dataset/images/val
nc: 3  # количество классов
names: ['cat', 'dog', 'bird']  # названия классов

Процесс обучения модели YOLO

Обучение YOLO выполняется с помощью модуля Ultralytics. Пример команды для запуска обучения:

from ultralytics import YOLO

# Загрузка предобученной модели (можно взять yolov8n.pt, yolov8s.pt и т.д.)
model = YOLO('yolov8n.pt')

# Запуск обучения
results = model.train(data='dataset.yaml', epochs=100, imgsz=640, batch=16)

Параметры обучения:

• data — путь к файлу конфигурации датасета.
• epochs — количество эпох (полных проходов по датасету).
• imgsz — размер входного изображения (обычно 640).
• batch — размер батча (зависит от объёма видеопамяти).

Во время обучения модель будет выводить метрики: потери (loss), точность (precision), полноту (recall) и mAP (mean Average Precision). После завершения обучения лучшие веса сохраняются в файл runs/detect/train/weights/best.pt.

Если вы хотите дообучить модель на своих данных, начиная с предобученных весов, это значительно ускорит процесс и повысит качество, особенно при небольшом датасете.

Метрики оценки качества: IoU, Precision, Recall, mAP

Для оценки того, насколько хорошо обучена модель, используются стандартные метрики компьютерного зрения.

IoU (Intersection over Union) — отношение площади пересечения предсказанной рамки с истинной к площади их объединения. Значение от 0 до 1, где 1 — идеальное совпадение. Обычно порогом успешного обнаружения считают IoU > 0.5.

Precision (точность) — доля правильных положительных предсказаний среди всех положительных предсказаний модели. Высокая точность означает, что модель редко ошибается, но может пропускать объекты.

Recall (полнота) — доля найденных истинных объектов среди всех объектов, которые есть на изображении. Высокая полнота означает, что модель находит почти всё, но может давать ложные срабатывания.

AP (Average Precision) — усреднённое значение точности при разных порогах полноты. Рассчитывается для каждого класса отдельно.

mAP (mean Average Precision) — среднее AP по всем классам. Это основная метрика, по которой сравнивают модели YOLO. Например, mAP@0.5 — это mAP при пороге IoU 0.5, а mAP@0.5:0.95 — среднее mAP при порогах от 0.5 до 0.95 с шагом 0.05.

Выбор версии YOLO: от nano до extra-large

Семейство YOLOv8 включает несколько предобученных моделей, различающихся размером, скоростью и точностью:

• YOLOv8n (nano) — самая лёгкая и быстрая модель. Подходит для мобильных устройств и встраиваемых систем. Точность ниже, но скорость максимальна.
• YOLOv8s (small) — небольшой баланс между скоростью и точностью.
• YOLOv8m (medium) — средняя модель, хороша для большинства задач.
• YOLOv8l (large) — более точная, но требует больше ресурсов.
• YOLOv8x (extra-large) — самая точная и самая медленная. Используется, когда качество важнее скорости.

Для сегментации доступны модели с суффиксом -seg, например yolov8n-seg.pt. Они работают медленнее, но позволяют получать маски объектов.

Выбор модели зависит от вашего оборудования и требований к скорости. Если вы работаете на слабом ПК или в реальном времени, выбирайте nano или small. Если точность критична и есть мощная видеокарта — используйте large или extra-large.

Практические примеры: детекция, сегментация, трекинг

После обучения модель можно использовать для решения различных задач.

Детекция объектов на изображении — самый простой случай. Загружаем модель, передаём изображение и получаем рамки с классами:

from ultralytics import YOLO
model = YOLO('best.pt')
results = model('test.jpg')
results[0].show()  # показать результат

Детекция в видеопотоке — обрабатываем каждый кадр видео. Для этого можно использовать OpenCV для чтения видео и записи результата. Важно учитывать, что скорость обработки должна быть не ниже частоты кадров исходного видео.

Сегментация — позволяет не только найти объект, но и выделить его контур. Для этого используется модель с суффиксом -seg. Пример удаления фона: находим маску человека, заменяем фон на однотонный или другое изображение.

Трекинг объектов — YOLO может отслеживать перемещение объектов между кадрами, присваивая каждому уникальный идентификатор. Это полезно для подсчёта людей, анализа трафика или наблюдения за животными. Трекинг реализован в Ultralytics и активируется параметром persist=True.

Типичные ошибки при обучении и способы их избежать

При обучении YOLO новички часто сталкиваются с рядом проблем.

Недостаточный размер датасета. Для качественного обучения нужно хотя бы несколько сотен размеченных изображений на каждый класс. Если данных мало, используйте аугментацию (повороты, масштабирование, изменение яркости) — Ultralytics поддерживает это автоматически.

Неправильная разметка. Ошибки в координатах рамок или неверные классы приводят к плохому обучению. Проверяйте разметку визуально перед запуском.

Дисбаланс классов. Если одних объектов много, а других — единицы, модель будет хуже распознавать редкие классы. Постарайтесь сбалансировать датасет или использовать взвешенные функции потерь.

Слишком маленькое изображение. Если объекты мелкие, увеличьте размер входного изображения (imgsz=1280) или используйте модель с более высоким разрешением.

Переобучение. Если модель отлично работает на тренировочных данных, но плохо на новых — увеличьте аугментацию, добавьте регуляризацию или уменьшите количество эпох.

Неправильный выбор модели. Для простых задач достаточно nano или small, для сложных — large или extra-large. Не используйте тяжёлую модель без необходимости.

Вопросы и ответы