Обнаружение объектов является достаточно популярной задачей компьютерного зрения, которая включает в себя идентификацию и обнаружение объектов на изображениях или видео. Данная задача является частью многих приложений, например, таких как беспилотные автомобили, робототехника, видеонаблюдение и т. д. За прошедшие годы разработано множество алгоритмов и методов для поиска объектов на изображениях и их положениях. Наилучшее качество выполнения таких задач достигается при использовании сверточных нейронных сетей.

Одной из самых популярных архитектур нейронных сетей для таких задач, является YOLO (you only look once), созданная в 2015 году. С тех пор появилось довольно много версий данных алгоритмов. Последние выпуски сети предназначены для таких задач как распознавание, обнаружение и сегментация изображений. 

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

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

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

Ниже мы разберем примеры обнаружения объектов на моделях YOLOv5 и YOLOv8 в двух вариантах.

1. Предположим, что мы имеем одно или несколько изображений и наша задача состоит в обнаружении на нем максимального количества объектов, которое позволит модель.

2. Предположим, что у нас есть набор данных и мы хотим обучить модель YOLO именно на этом наборе.

Обе версии YOLO бесплатные от одного разработчика Ultralytics, хорошо показывают себя  сравнительных тестах и реализованы на современной библиотеке глубокого обучения PyTorch. Модель YOLOv8 позиционируется как более новая и современная по сравнению с 5й версией, хотя обе версии активно развиваются. 

Начало работы с YOLO и подготовка датасета

Модели YOLO поставляются уже предварительно обученными на наборе данных COCO и экспортированы в файлы с расширением .pt. Существует три типа моделей (классификация, обнаружение, сегментация) и по 5 моделей разного размера для каждого типа:

Набор данных COCO представляет собой огромную коллекцию изображений 80 классов. Поэтому, если нет особых потребностей, то можно просто запустить уже готовую модель без предварительного обучения и оценивать результат. Посмотрим на названия классов набора COCO.

names: { 0: person,  1: bicycle, 2: car, 3: motorcycle, 4: airplane, 5: bus, 6: train, 7: truck, 8: boat,  9: traffic light, 10: fire hydrant, 11: stop sign, 12: parking meter, 13: bench, 14: bird,  15: cat, 16: dog, 17: horse, 18: sheep, 19: cow, 20: elephant, 21: bear, 22: zebra, 23: giraffe,  24: backpack, 25: umbrella, 26: handbag, 27: tie, 28: suitcase, 29: frisbee, 30: skis, 31: snowboard, 32: sports ball, 33: kite, 34: baseball bat, 35: baseball glove,36: skateboard, 37: surfboard, 38: tennis racket, 39: bottle, 40: wine glass, 41: cup, 42: fork, 43: knife, 44: spoon,  45: bowl, 46: banana, 47: apple, 48: sandwich, 49: orange, 50: broccoli, 51: carrot, 52: hot dog, 53: pizza, 54: donut, 55: cake, 56: chair, 57: couch, 58: potted plant, 59: bed, 60: dining table, 61: toilet, 62: tv, 63: laptop, 64: mouse, 65: remote, 66: keyboard, 67: cell phone, 68: microwave, 69: oven, 70: toaster, 71: sink, 72: refrigerator, 73: book, 74: clock, 75: vase, 76: scissors, 77: teddy bear, 78: hair drier, 79: toothbrush}

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

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

Структура набора данных для обучения моделей YOLO должна иметь такой формат

├── dataset
│   ├── test
│   │   ├── images
│   │   └── labels
│   ├── train
│   │   ├── images
│   │   └── labels
│   └── val
│       ├── images
│       └── labels
├── dataset.yaml
└── yolov5x.pt

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

{object_class_id} {x_center} {y_center} {width} {height}

Все координаты должны быть нормализованы так, чтобы они помещались в диапазон от 0 до 1. Для их расчета используются формулы

x_center = (box_x_left+box_x_width/2)/image_width
y_center = (box_y_top+box_y_height/2)/image_height
width = box_width/image_width
height = box_height/image_height

где box - ограничивающая рамка для объекта того класса, для которого мы производим вычисления координат.

Если на изображении несколько объектов, тогда текстовый файл аннотаций может выглядеть так

1 0.589869281 0.490361446 0.326797386 0.527710843
0 0.323529412 0.585542169 0.189542484 0.351807229

Первая строка содержит ограничивающую рамку для одного объекта, вторая строка содержит ограничивающую рамку для второго объекта. Конечно, файл может содержать и большее количество строк, в зависимости от количества объектов на изображении.

Текстовые файлы аннотаций должны иметь те же имена, что и файлы изображений, и иметь расширение .txt.

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

train: ../train/images
val: ../val/images
test: ../test/images
nc: 2
names: ['cat','dog']

В первых двух строчках указываются пути к изображениям и проверочного набора данных. Затем строка nc указывает количество классов, существующих в этих наборах данных, а имена представляют собой имена классов в правильном порядке. Индексы этих элементов - это числа, которые использовались при аннотировании изображений, и эти индексы будут возвращены моделью при обнаружении объектов с использованием метода прогнозирования. Файл .yaml будет передаваться при обучении модели.

И вот он главный вопрос, как получить такой набор данных?

• Взять готовый датасет в свободном доступе. Можно например, воспользоваться сайтом kaggle.com. Но для того, чтобы найти набор в правильном формате, придется потратить время. 

• Также большой выбор бесплатных наборов для компьютерного зрения можно найти здесь universe.roboflow.com. Чем удобен данный сервис, что датасет можно загружать в разных форматах.

Недостаток — много повторов и некачественной разметки.

• Собрать набор данных самостоятельно, который потом, к сожалению, придется размечать вручную. Или даже если есть готовый набор, но не аннотированный, разметить его придется. Хорошая новость, что для визуального аннотирования изображений для задач машинного обучения., есть множество различных и удобных сервисов. Например, можно просто ввести в поисковую систему что‑то вроде «программное обеспечение для аннотирования изображений для машинного обучения», чтобы получить их список. Есть также множество онлайн‑инструментов, которые могут выполнить всю эту работу. Разметку, конечно, необходимо проводить в зависимости от модели сети, какую мы будем использовать. Так как мы можем разметить рамками, точками, областями сегментации и так далее. Одним из замечательных онлайн‑инструментов для этого является Roboflow Annotate, подробнее можно посмотреть здесь Используя этот сервис, нужно просто загрузить изображения, нарисовать на них ограничивающие рамки и установить класс для каждой ограничивающей рамки. Затем инструмент автоматически создаст файлы аннотаций, разделит наши данные на наборы данных для обучения и проверки, создаст файл дескриптора YAML, а затем можно экспортировать и загрузить аннотированные данные в виде ZIP‑файла.

Обнаружение объектов

1. YOLOv8

Сначала установим пакет YOLOv8, для этого в Jupyter Notebook необходимо набрать команду

%pip install ultralytics

В пакете Ultralytics есть класс YOLO, который используется для создания нейросетевых моделей. Далее импортируем этот модуль

from ultralytics import YOLO

Теперь все готово для создания модели нейронной сети

model = YOLO('yolov8m.pt')

При первом запуске такого кода в текущую папку будет загружен файл yolov8m.pt с сервера Ultralytics, который содержит модель и предобученные веса и далее будет создан объект модели. Теперь можно обучать эту модель, обнаруживать объекты и экспортировать ее для использования в производстве. 

Для всех этих задач есть удобные методы:

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

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

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

Теперь мы можем взять любое изображение и обнаружить там объекты, которые были заложены в модель YOLO на этапе обучения.

results = model.predict('boy_dog.jpg')

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

result = results[0]

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

result.boxes

Результат

• cls — номер класса (один из 80 типов объектов с идентификатором от 0 до 79. Классы объектов COCO были определены выше).

• conf — уровень доверия модели для этого класса.

• xywh — координаты прямоугольника в формате xywh.

• xywhn — координаты прямоугольника в формате xywh, нормализованные по исходному размеру.

• xyxy — координаты прямоугольника в виде массива [x1,y1,x2,y2]

• xyxyn — координаты прямоугольника, нормализованные по исходному размеру изображения

То есть на нашем изображении найдено 2 класса с номерами 0 и 16, что соответствует классу человек и собака.

Можем переписать код и получить информацию в более удобном виде

def print_box(box):
    class_id, cords, conf = box
    print("Object type:", class_id)
    print("Coordinates:", cords)
    print("Probability:", conf)
    print("---")

[
    print_box([
        result.names[box.cls[0].item()],
        [round(x) for x in box.xyxy[0].tolist()],
        round(box.conf[0].item(), 2)
    ]) for box in result.boxes
]

Результат:

Если мы хотим получить именно изображения с уже ограничивающими рамками, это можно сделать следующим образом

from PIL import Image
for i, r in enumerate(results):
	img_bgr = r.plot()
	im_rgb = Image.fromarray(img_bgr[...,::-1])
	r.show()
	r.save(filename = f"results{i}.jpg")

2. YOLOv5

Если мы работаем с моделью YOLOv5, тогда необходимо просто склонировать репозиторий YOLOV с github и также работать с сетью. В этом случае последовательность действий такая. 

1. Устанавливаем каталог, куда будем клонировать репозиторий

%cd D:/yolov5

2. Клонируем репозиторий с github

!git clone https://github.com/ultralytics/yolov5.git

3. Импортируем библиотеку pytorch

import torch

4. Переходим в нашу папку и устанавливаем зависимости из requirements.txt

%cd D:/yolov5/yolov5
!pip install -r requirements.txt

Для задач обнаружения, обучения и экспорта модели используются следующие команды:

• Для обнаружения объектов набираем команду 

!python detect.py --weights yolov5x.pt --source img.jpg

Входное изображение может быть просто файлом, скриншотом, папкой, видео и т. д. Веса (weights) указываются в зависимости от модели с которой мы хотим работать (nano, small, medium, large, xlarge).

• Для обучения модели набираем команду

!python train.py --img 640 --epochs 10 --data coco128.yaml --weights yolov5s.pt

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

• Для экспорта модели набираем команду

!python export.py --weights yolov5s.pt --include torchscript 

Необходимо указать модель, которую будем экспортировать и формат. Данная команда экспортирует модель в формат torchscript.Также еще используются форматы: pytorch, onnx, openvino, engine, coreml и другие. 

Теперь перейдем к загрузке изображения и получению результатов. Используем команду detect.py, предполагая, что у нас уже установлена директория проекта.

!python detect.py --weights yolov5x.pt --source D:/yolov_test1.jpg

Аргументы командной строки следующие

--source: путь к входному файлу 

--weights: путь к файлу с весами yolov5, который хотим использовать для обнаружения, в данном случае yolov5x.pt

--project: необязательный аргумент, с помощью него указывается папка для сохранения результатов, которая в свою очередь будет содержать файл exp.

Получаем следующий результат

Здесь указана информация о том,какие параметры сети использовались (веса, набор данных на каком обучалась сеть, количество слоев в сети и т. д.) и какие результаты обнаружения получены (сколько и каких объектов обнаружено). Все результаты сохраняются в папку ...runs\detect\exp.  

Здесь даже не нужен никакой дополнительный код для отрисовки ограничивающих рамок или определения имени класса, за нас уже все сделано. Мы получили готовое изображение с обнаруженными классами. Если на вход подавать папку  с изображениями, в результате также в директории runs\detect\exp получим результаты обнаружения. 

Обучение

1. YOLOv8

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

Я взяла набор данных с сайта universe.roboflow по обнаружению трех цветов воздушных шаров. Набор удобен тем, что его структура каталогов организована как раз для обучения моделей YOLO. В наборе имеются папки train, val и test, а в каждой из этих папок есть папки labels и images, также имеется файл data.yaml.

В моем случае файл .yaml выглядит следующим образом

train: ../train/images
val: ../val/images
test: ../test/images

nc: 3
names: ['blue','green','red']

Всего имеется 744 изображения для обучения и 71 для валидации. Набор данных есть, теперь необходимо его загрузить в модель. Для обучения будем использовать метод train:

model.train(data = 'data.yaml', epochs = 20)

Предполагаем, что модель сети уже загружена с помощью команды

model = YOLO('yolov8m.pt')

Файл .yaml является единственным обязательным параметром, но также стоит задать количество эпох для обучения (по умолчанию равно 100). Конечно, настройки обучения включают в себя и другие гиперпараметры и конфигурации, но на них мы не будем останавливаться.. Эти настройки влияют на производительность, скорость и точность модели. Ключевые параметры тренировки включают batch size, learning rate (скорость обучения), weight_decay (параметр регуляризации L2). Также, значительно повлиять на процесс обучения может выбор оптимизатора, функции потерь и состава набора обучающих данных может. Каждый цикл обучения состоит из этапа обучения и этапа поверки. 

На этапе обучения метод train выполняет следующие действия. 

• Извлекает случайную партию изображений из набора обучающих данных (количество изображений указывается с помощью параметра batch size). 

• Пропускает эти изображения через модель и получает результирующие ограничивающие рамки всех обнаруженных объектов и их классов. 

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

• Результат функции потерь передается оптимизатору для корректировки весов модели на основе величины ошибки в правильном направлении, чтобы уменьшить ошибку в следующем цикле. По умолчанию используется оптимизатор SGD (Stochastic Gradient Descent), но можно попробовать другие, например Adam, чтобы увидеть разницу. 

На этапе проверки метод train выполняет следующее:

• Извлекает изображения из набора проверочных данных. 

• Пропускает их через модель и получает обнаруженные ограничивающие рамки для этих изображений. 

• Сравнивает полученный результат с истинными значениями для этих изображений из текстовых файлов аннотаций. 

• Вычисляет точность модели на основе разницы между фактическими и ожидаемыми результатами.

Ход и результат для каждой эпохи отображаются на экране.

….

Фаза обучения включает в себя расчет величины ошибки в функции потерь, поэтому наиболее ценными метриками здесь являются box_loss и cls_loss. 

• box_loss показывает количество ошибок в обнаруженных ограничивающих рамках.

• cls_loss показывает количество ошибок в обнаруженных классах объектов. 

Почему потери разбиваются на несколько показателей? Потому что модель могла правильно определить ограничивающую рамку вокруг объекта, но неправильно определить класс объекта в этой рамке. Наиболее ценным показателем качества является mAP50-95, который представляет собой метрику Mean Average Precision (mAP). Данная метрика показывает на сколько безошибочно нейронная сеть смогла обнаружить объекты на изображениях из валидационной выборки. Если модель обучается и совершенствуется, точность должна расти от эпохи к эпохе.

Если после последней эпохи мы не получили приемлемой точности, тогда можно увеличить количество эпох. Также можно попробовать настроить другие параметры сеть. Для каждого случая необходимо подбирать свое решение.

Помимо этих метрик, во время обучения сети на диск записывается много статистики. Когда начинается обучение, создается  подпапка /runs/detect/train в текущей папке и после каждой эпохи в нее записываются разные файлы журналов. Стоит заглянуть в эту папку, в ней много интересной информации. Данная директория сохраняется в директории проекта, если дополнительно не была установлена какая-либо другая директория. Также обученная модель поле каждой эпохи записывается в файл /runs/detect/train/weights/last.pt, а модель с наивысшей точностью — в файл /runs/detect/train/weights/best.pt. Итак, после завершения обучения мы должны получить файл best.pt для его использования. 

Содержимое подпапки /runs/detect/train.

Содержимое файла results.png

Теперь посмотрим на наши результаты после обучения YOLOv8. Не забываем сначала загрузить новую модель.

Этот результат получен после 20 эпох обучения. Обучение заняло около 6 часов на компьютере с процессором AMD Ryzen 7 3700X 8-Core Processor 3.59 GHz, 32 ГБ оперативной памяти. На первом изображении пока пропущен 1 голубой шар и в 1 красном шаре модель немного сомневается, так как не очень высокая вероятность доверия, всего лишь 58%. Можно продолжить обучение или попробовать изменить параметры модели.

2. YOLOv5

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

--data: путь к датасету для обучения

--img: размер изображения, подаваемого на вход yolo

--epochs: количество эпох

--weights: стартовые веса для обучения

--batch: размер батча, то есть количество картинок, одновременно подаваемых на вход yolo

Существуют  другие настройки про которые можно почитать на странице проекта.

В данном случае я взяла набор данных с сайта universe.roboflow.com по доступности городской среды на улицах. Имеется 5 классов, на которые можно разделить препятствия, встречающиеся на улицах: затрудняет проезд (dificulta el paso), недоступно из-за строительства (inaccesible por obras), нужен пандус (necesita rampa), недоступный пандус (rampa inaccesible), неустойчивая площадка (suelo inestable).

Примеры изображений

Файл .yaml выглядит следующим образом

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

nc: 5
names: ['dificulta el paso', 'inaccesible por obras', 'necesita rampa', 'rampa inaccesible', 'suelo inestable']

Всего имеется 7126 изображений для обучения и 918 для валидации. Попробуем запустить обучение на 10 эпох просто для примера. Конечно, для обучения модели на таком наборе данных, необходимо большее количество эпох.

!python train.py --img 416 --batch 16 --epochs 10 --data C:/yolov5/data1/data.yaml --weights C:yolov5/data1/yolov5x.pt

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

Обучение проходит таким же образом, как и было описано выше при работе с моделью YOLOv8. Все результаты обучения также сохраняются в директорию runs/train/exp. При этом для каждого нового обучения создаются папки …./exp1, …/exp2 и т.д.

Содержимое подпапки runs/detect/train

Файл results.png выглядит следующим образом

А теперь посмотрим на результат обнаружения, только теперь с параметром –weights мы загружаем нашу модель best.py, которая была получена после обучения. Должна быть установлена директория откуда запускать скрипт detect.py.

!python detect.py --weights runs/train/exp8/weights/best.pt --source C:/streets1.jpg

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

Заключение

Целью статьи было познакомить читателей, которые только начинают работу в области компьютерного зрения и детекции объектов, с работой современных моделей нейронных сетей YOLOv5 и YOLOv8. То есть я не ставила цель решить какую-то сложную задачу или получить новые результаты, сравнить модели и так далее. Таким образом, мы рассмотрели задачу обнаружения объектов на предварительно обученных моделях YOLOv5 и YOLOv8, а также попробовали обучить модели для обнаружения объектов на своем датасете, обсудили вопрос создания своего датасета. При обучении таких моделей главное - это терпение, так как процесс обучения не быстрый, если вы не являетесь обладателем графического процессора и не арендуете облачный сервис с GPU).

---

В завершение приглашаем всех желающих на открытые уроки по компьютерному зрению в Otus:

• 10 июля: Семейство детекторов YOLO: от v1 до v10. Рассмотрим детекцию объектов, проследим эволюцию детекторов семейства YOLO от самой первой до самой актуальной версии (v10). Запись по ссылке.

• 23 июля: Human pose estimation — обзор подходов, моделей и решений. Вы узнаете: постановку задачи оценки позы человека, обзор основных направлений к решению этой задачи и какие есть популярные фреймворки. Запись по ссылке.