Machine Learning, Нейронные сети

Использование YOLOv5 для задачи детекции

Время прочтения: 7 мин.

Ликбез по нейронным сетям

Определение. Нейронная сеть (или Искусственная нейронная сеть) представляет собой математическую модель, а также её программное или аппаратное воплощение, построенную по принципу организации и функционирования биологических нейронных сетей — сетей нервных клеток живого организма

Помните, как в школе мы чертили график функции по точкам. Нам была известна сама функция, например парабола   и для того, чтобы построить график, мы вместо   подставляли какое-либо число и получали соответствующее значение  . Имея эти два значения, можно было ставить точку. Затем брали другое значение  , снова подставляли, получали новый   и ставили следующую точку. Так делали несколько раз и по точкам строили график. Можно сказать, что мы по функции получали точки (график)

Нейронные сети решают обратную задачу: по точкам находят функцию. Причем, можно найти (а точнее сказать приблизить) любую функцию с любой точностью, но это только теоретически 🙂

Процесс нахождения функции по точкам происходит путем минимизации ошибки, т.е. сводится к минимуму «расстояние» между значениями  , предсказываемых нейронной сетью и значениями  , которые мы наблюдаем.

YOLO

Под архитектурой нейронной сети понимается её устройство — последовательность нейронов и связей между ними.

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

Один из способов решения задачи детекции заключается в разбиении изображения на квадратные области, затем классификация этих областей на наличие объекта и классификация самого объекта. Таким образом, изображение просматривается дважды (один раз для определения областей, где есть объект, второй — для классификации этого объекта.) Этот способ работает долго и требует больших затрат вычислительных мощностей.

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

YOLOv5

YOLOv5 – усовершенствованная пятая версия YOLO, реализованная на фреймворке PyTorch. YOLOv5 встроена в одноименный модуль для Python3, который можно установить с pypi. Этот модуль предоставляет очень хорошую инфраструктуру как для обучения модели, так и для тестирования с построением графиков всех ключевых показателей.

Обучение YOLOv5

Для обучения yolov5 нам потребуется:

  1. сама библиотека yolov5 (есть на pypi и на github)
  2. Предобученные веса (можно найти на github)
  3. файл .yaml (в нем указываются пути до тренировочной и валидационной выборок, а также количество классов и их метки)
  4. собственно датасет
  5. GPU с драйвером CUDA, совместимым с версией pytorch (Совместимость версий pytorch и CUDA можно посмотреть на сайте )

Вот так выглядит файл .yaml:

Структура каталога с датасетом:

Датасет должен быть разделен на две папки: train (тренировочная выборка) и val (валидационная). В каждой папке лежат еще две папки: images (картинки) и labels – папка с текстовыми файлами, содержащими метки объектов на этих картинках в формате YOLO. Формат подразумевает, что каждая строка текстового файла представляется в виде:

n, x, y, w, h, где

n — номер класса объекта

— относительная координата bounding box’а объекта по оси Ox

y — относительная координата bounding box’а объекта по оси Oy

— относительная ширина bounding box’а объекта

— относительная высота bounding box’а объекта

Пример: 0 0.397396 0.467014 0.0635417 0.100694

Перед запуском обучения поясню параметры, с которыми запускаемся:

  • —img — размер изображения, подаваемого на вход yolo. Это значит, что размер исходного изображения преобразуется к виду n x n, где n — число, введенное пользователем после ключа —img. ВАЖНО: n должно быть кратно 32 (это связано с архитектурой yolo)
  • —batch — размер батча, то есть количество картинок, одновременно подаваемых на вход yolo
  • —epoch — количество эпох для обучения
  • —data — путь до .yaml файлика
  • —weights — путь до файла весов (нужны предобученные веса, чтобы не обучаться с нуля)
  • —name — куда сохраняем результаты обучения
  • —workers — количество потоков, на которых запускаем обучение (нужно для увеличения скорости обучения).

Теперь можно запускаться 🙂

!python yolov5/train.py --img 896 --batch 8 --epoch 300 --data ./dataset.yaml --weights ./yolov5x.pt --name result --workers 2

Теперь у нас есть обученная нейронная сеть, результаты обучения которой сохраняются в папку runs/<то что ввели в параметре —name><номер запуска>. В этой папке очень много разной полезной информации об обучении. Советую посмотреть.

Детекция

Для запуска детекции объектов обученной нейронной сетью нужно запустить файл detect.py. Но перед этим несколько слов о параметрах

Параметры —img, —weights, —name вам уже известны. Скажу только про —weights. Здесь указывается путь до весов, которые сохранились после обучения.

Вообще в папке ./runs/result/weights хранятся 2 файлика весов:

  1. best.pt — лучший вариант модели за все время обучения
  2. last.pt — модель на последней эпохе

Замечание
Лучшая модель не всегда (почти никогда) последняя!

Параметр —save-text обеспечивает нам сохранение меток детектированных объектов в формате YOLO (что это за формат — объяснялось выше)

Определение
Confidence – число от 0 до 1, характеризующее «уверенность» модели в том, что детектирован объект или детектирован объект определенного класса

Параметр —conf-thres позволяет установить пороговое значение для confidence модели. Все объекты, confidence которых ниже этого значения не считаются объектами.

Параметр —source позволяет указать путь до картинок, на которых будет осуществляться детекция. Также можно указать путь до видео или целой папки с видео. ЗАПУСК

!python yolov5/detect.py --img 896 --weights ./runs/result/weights/best.pt --name result_detect --save-txt --conf-thres 0.15 --source ./dataset/images/val

В результате получаются такие изображения с обозначением детектированных объектов и значениями confidence.

Пример обучения

В качестве реального примера приведем процесс обучения YOLO на датасете COCO128. Этот датасет содержит 128 тренировочных и 128 валидационных изображений с объектами 80 различных классов.

Вот так выглядит вывод во время обучения.

Здесь представлена информация об оптимизаторе, тренировочном и валидационном датасетах, а также параметры обучения – такие как размер входного изображения, количество потоков, директория с результатами обучения.

Далее следует информация о процессе самого обучения. Здесь следует обратить внимание на столбцы P, R, mAP@.5 и mAP@.5:.95 – в них записаны значения метрик precision, recall и двух вариантов метрики mAP соответственно. Подробнее про эти метрики можно почитать тут.

Также по окончании обучения приводятся значения этих метрик по каждому классу для модели с наилучшими весами.

А в директории runs/train/exp/ можно увидеть графики precision, recall, f1, а также ещё много другой полезной информации

Более подробно этот пример разобран здесь.

Выводы

В заключение хочу сказать, что модуль yolov5 – это очень удобный инструмент, позволяющий обучать и тестировать нейронные сети архитектуры YOLOv5. Благодаря богатому функционалу модуль позволяет достаточно быстро решить задачу детекции, а также визуализировать процесс обучения.

Советуем почитать