Object Detection или обнаружение объектов - это одна из ключевых задач компьютерного зрения и обработки видео и изображений. Эта функция используется в самых разных приложениях и сферах, начиная с охранных систем и заканчивая управлением беспилотного транспорта. Поэтому важно, чтобы обнаружение происходило в реальном времени.
Одна из самых популярных архитектур для решения этой задачи - YOLO (You Only Look Once). Это алгоритм обнаружения объектов в реальном времени, который представляет собой одну глубокую сверточную нейронную сеть. Она разбивает входное изображение на набор ячеек сетки, поэтому, в отличие от классификации изображений или обнаружения лиц, каждая ячейка сетки в алгоритме YOLO будет иметь связанный вектор на выходе, который сообщает нам:
- Есть ли объект в конкретной ячейке сетки.
- Класс этого объекта.
- Предполагаемая ограничивающая рамка для этого объекта (местоположение).
Конечно, есть и другие архитектуры, например, Fast R-CNN или Faster R-CNN. Однако, они требуют тысячи прогнозов для одного изображения (в каждом окне). Из-за этого страдает скорость выполнения: YOLO даже третьей версии будет в 1000 раз быстрее, чем R-CNN и в 100 раз быстрее, чем Fast R-CNN.
Поэтому предлагаем рассмотреть порядок действий при работе именно с YOLOv3. А для этого нам понадобятся OpenCV и PyTorch. Погнали.
1. Установка и настройка
Прежде чем мы погрузимся в код, давайте установим необходимые библиотеки. Если вы хотите использовать код PyTorch, перейдите на эту страницу для установки:
pip3 install opencv- python numpy matplotlib
Импорт необходимых модулей:
Определяем некоторые переменные и параметры, которые нам понадобятся:
config_path и weights_path представляют конфигурацию модели и соответствующие предварительно обученные веса. Метки - это список всех классов для различных объектов обнаружения. Они будут выделяться разными цветами, поэтому мы генерируем случайные цвета.
Все необходимые файлы есть в этом репозитории. Файл с весами придется скачать отдельно, так как он весит 237 мб и не вмещается в github. Приведенный ниже код загружает модель:
2. Подготовка изображений
Загрузим примеры изображений из этого репозитория:
Затем нужно нормализовать, масштабировать и изменить изображение, чтобы оно подходило в качестве входных данных для нейронной сети:
Это нормализует значения пикселей в диапазоне от 0 до 1, изменит размер изображения до (416х416) и изменит его форму, давайте посмотрим:
Выход:
3. Получение прогнозов
Загружаем полученное изображение в нейронную сеть, чтобы получить прогноз на выходе:
Это позволит извлечь выходные данные нейронной сети и вывести общее время, затраченное на вывод:
Time took: 1. 54s
Вы можете возразить: 1,5 секунды - это очень медленно. Мы используем для получения вывода наш ЦП, поэтому позже задействуем PyTorch. Но даже в таком виде, результат намного больше, чем с тем же R-CNN.
Теперь нам нужно перебрать выходные данные нейронной сети и отбросить любой объект, уровень достоверности которого меньше, чем параметр CONFIDENCE. Ранее мы указали значение 0,5 или 50%.
Теперь наша нейронка будет циклически прогонять все предсказания и сохранять только те объекты, которые точнее всего соответствуют основным параметрам. Наш вектор обнаружения:
print(detection. shape)
Выход:
(85,)
Для каждого предсказания объекта есть вектор 8. Первые четыре значения представляют расположение объекта: координаты (x, y) для центральной точки, а также ширину и высоту ограничивающего прямоугольника. Остальные числа соответствуют меткам объекта, так как это набор данных COCO, он имеет 80 категорий объектов.
Например, если обнаруженный объект - человек, первое значение в векторе длины 80 должно быть 1. Все остальные значения должны быть 0. Второе число для велосипеда, Третье для автомобиля. И так вплоть до 80-го объекта.
Вот почему мы используем функцию np.argmax() для получения идентификатора класса, поскольку она возвращает индекс максимального значения из этого вектора длины.
4. Создание прямоугольников выделения
Осталось совсем немного. Сейчас мы зададим форму, цвет и метки для выделяемых объектов:
Добавим изображение и посмотрим, что выйдет:
cv2. imwrite( filename + "_yolo3." + ext, image)
Тут мы столкнулись с проблемой: двойные рамки выделения. Чтобы исправить это, нам придется прибегнуть к немаксимальному подавлению (Non-maximal Suppression).
5. Non-maximal Suppression
Немаксимальное подавление - это метод, который подавляет перекрывающиеся ограничивающие рамки с низкой вероятностью обнаружения объекта. В основном это достигается в два этапа:
- Он выбирает рамку, у которой более высокая оценка предсказания
- Затем он сопоставляет его всем остальным рамкам и удаляет те, которые имеют высокий IoU .
IoU (Intersection over Union) - это метод позволяет понять, насколько схожи две разные ограничивающие рамки.
Чем выше IoU, тем ближе ограничивающие рамки:
- если IoU = 1, значит, что две рамки идентичны
- если IoU = 0, значит, что они даже не пересекаются
Мы будем использовать пороговое значение IoU = 0,5. SCORE_THRESHOLD устранит прямоугольник с меньшей оценкой прогноза.
Теперь попробуем повторить наложение рамок:
Вы можете использовать cv2.imshow ("image", image), чтобы показать изображение, но мы просто сохраним его на диск:
cv2. imwrite( filename + "_yolo3." + ext, image)
Если у ваших изображений высокое разрешение, убедитесь, что вы увеличили параметр font_scale.
6. Настройка PyTorch
Если вы хотите задействовать графический процессор для вывода, то нужно использовать библиотеку PyTorch, поддерживающую вычисления CUDA. Вот нужные фрагменты кода: darknet.py и utils.py.
Примечание. Для приведенного выше кода требуются файлы darknet.py и utils.py в текущем каталоге. Также должен быть установлен PyTorch (рекомендуется ускорение на GPU).
Заключение
Теперь вы знаете, как создать модель обнаружения объектов YOLO с помощью OpenCV и PyTorch. Если вы хотите пойти дальше и добиваться определения в реальном времен, то ниже мы оставим фрагмент кода и пример того, как это работает.
Обратите внимание, что у YOLO есть некоторые проблемы с обнаружением небольших сгруппированных объектов. Чем ближе и мельче объекты, тем выше вероятность неправильного выделения. Это сравнительно небольшая ложка дегтя, ведь этот способ демонстрирует отличную скорость. Важно лишь верно подобрать сферу применения.
Другие наши статьи: