YOLO-World: распознавание произвольного числа объектов с высокой точностью и скоростью

Всем привет! Буквально несколько дней назад была представлена новая модель семейства Yolo. Ее основная фишка заключается в том, что в отличие от своих старших братьев, она способна распознавать на изображении фактически любые объекты (которые интересуют человека) без предварительного обучения и делает она это в real-time режиме! Звучит неплохо, не так ли?

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

Типы задач детекции

Думаю, многие из читателей, которые уже немного знакомы со сферой Computer Vision, представляя задачу детекции, подразумевают традиционный object detection. При таком подходе модель способна детектировать на изображениях строго заданный заранее список объектов, на котором она обучалась:

6dfb4a569293c7a9e736563e2960a9bc.png

Другим типом, который в последние годы набирает популярность в виду большей гибкости, а также зачастую отличной работы без предварительного дообучения, является, так называемый open vocabulary object detection. Основная идея детекторов, решающих задачу распознавания с «открытым словарем», является использование не самих целочисленных меток классов, а эмбеддингов (векторных представлений) названий этих классов. Благодаря этому они могут находить даже те классы, которые не были явно заданы заранее, а также хорошо работать с фразами. Например, мы можем заставить модель искать не просто котов, а конкретную породу, даже если ее не было в обучении. Такие детекторы способны находить практически неограниченное число классов объектов на изображении:

4395f1244669c38e521567c015512e2a.png

Здесь стоит отметить, что, говоря неограниченное, мы подразумеваем, что человек в любом случае должен определить список объектов, которые он хочет распознавать, но этот список может быть сколь угодно большим (фактически конечным).  К примеру, таким детекторам можно передать 21 тысячу наименований классов из полного датасета ImageNet и они действительно будут пытаться распознать каждый из них!

Как уже было отмечено в начале, герой данной статьи Yolo-world как раз-таки относится к детекторам, способным решать более интересную вторую задачу. Но авторы статьи пошли дальше и предложили использовать так называемый prompt-then-detect подход. Это означает, что если ранее мы один раз создавали большой список интересующих нас слов (online vocabulary, векторы в котором статичны), то теперь мы формируем так называемые промты, которые кодируются в offline vocabulary и уже эти эмбеддинги идут далее по пайплайну. Такой подход позволяет уменьшить количество вычислений для каждого ввода, а также обеспечивает гибкость, корректируя словарь по мере необходимости.

3ce1ad0bbf03849fdf171fdb3844cce7.png

Итак, осознав потенциал новой модели, думаю многим не терпится перейти к разбору самой архитектуры и понять, в чем же ее основные особенности. Скорее приступим!

Архитектура

Как и многие современные нейросетевые архитектуры, YOLO-World можно разбить на множество отдельных блоков. Давайте подробнее остановимся на некоторых из них.

3642203d63c0a3907bbfe9c36e19729a.png

YOLO Detector

Здесь все просто: для извлечения фичей изображения используется относительно новая YOLOv8. Она же в свою очередь содержит backbone Darknet в качестве кодировщика и сеть PAN для формирования multi-scale фичей.

 Text Encoder

Для получения эмбеддингов текста используется уже хорошо зарекомендовавший себя CLIP, а именно его трансформер для кодировки текста. Кодировщик CLIP формирует векторы текстов так, чтобы их можно было хорошо сопоставить с соответствующими им векторными представлениями изображений (высокое косинусное сходство).

Re-parameterizable Vision-Language PAN

Пожалуй, основной и главный строительный блок всей архитектуры. Состоит он из top-down и bottom-up частей, внутри которых производится сопоставление ранее извлеченных текстовых эмбеддингов и multi-scale фичей изображения.

Блок внутри себя включает 2 основных компонента: CSPLayer и Image-Pooling Attention. Говоря обывательским языком, первый пытается добавить языковую информацию в элементы изображения, а второй наоборот, заложить информацию с изображения в текстовые эмбеддинги:

663185d1699ead60d29fee72a905dc51.png

После RepVL-PAN следуют блоки Box Head и Text Contrastive Head. Если первый, очевидно, предсказывает ограничительные рамки объектов, то второй их эмбеддинги (на основании близости объект-текст).

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

В первом приближении это все про саму модель! Для того, чтобы не усложнять изложения обзора, я не стал приводить формул и строгих математических выкладок. Желающих разобраться во всех деталях и тонкостях я отправляю к прочтению оригинальной статьи. Там же можно ознакомиться с тем, как все это обучается с помощью Region-Text Contrastive Loss, а также найти описание множества экспериментов по дообучению под конкретные задачи на разных датасетах для сравнения с предыдущими решениями.

Точность и скорость работы

Конечно, архитектура — это хорошо, но какие же метрики показывает новая SOTA в распознавании объектов? Здесь авторы не оставили нас без Speed-Accuracy графика, который без лишней скромности заявляет о повышении производительности в 20 раз, при как минимум сопоставимой точности (mAP), измеренной на датасете LVIS:

10247bcc54932688b707bf46ad1cae39.png

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

Качество распознавания вполне объяснимо главным блоком (RepVL-PAN), который использует многоуровневое кросс-модальное объединение фичей (тексты и картинки).

Заключение

В конце хочется отдельно обозначить особенности нового YOLO-World детектора:

  • Способен распознавать неограниченное число объектов (в том числе фразы) из коробки

  • Самая большая large модель показывает real-time скорость на инференсе (первая сеть с таким показателем для задачи OVD)

  • Использует под капотом как устоявшиеся архитектуры (Yolov8, CLIP), так и потенциально многообещающие новые (RepVL-PAN).

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

Ну и конечно же, у тех, кто дочитал до конца, мог остаться вопрос:, а есть ли открытый код? К счастью, ответ положительный! Правда на момент написания статьи код с инференсом еще не доступен, да и тот, что есть может быть еще сыроватым. Но веса моделей уже выложены в открытый доступ, так что особо любопытные могут смело экспериментировать)

Полезные ссылки:

© Habrahabr.ru