Sapiens: фундаментальная CV-модель для задач с людьми

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

Итак, у нас есть семейство моделей, которое из коробки поможет решить «четыре фундаментальные задачи человека в CV» (цитата авторов) и давайте посмотрим, что же там есть есть и как это работает.

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

Демонстрация возможностей модели, картиночку подбирал сам, надеюсь вам понравится

Демонстрация возможностей модели, картиночку подбирал сам, надеюсь вам понравится

И так, Sapiens.

Авторы нацелились дать инструменты для следующих задач: определение позы, сегментация частей тела, оценка глубины и прогнозирование нормали поверхности. Кажется, будто бы Meta сознательно нацелилась удешевить все, что касается виртуальной реальности.

Вместо слов — видео с демонстрацией от авторов.

Датасет и претрейн

Датасет закрытый и откуда у меты такие объемы изображений с людьми можно только догадываться シ. Первичная версия состояла из 1B изображений, которую подвергли двойной фильтрации: сначала отфильтровали по наличию неподходящих элементов (текст, водяные знаки, подвергнутые сильной обработке сами фото и так далее), затем прошлись bbox-детектором на людей и убрали те фото, где людей не было или они занимали небольшую часть фото. Оставили только те, где сам бокс больше 300 px и уверенность в нем 0.9. Ну, а затем свели все к квадратной размерности.

В итоге получился датасет из 300M фотографий, на котором обучали модель: маскированный автоэнкодер (masked autoencoders / MAE), когда модель учится восстанавливать исходное изображение по частично видимым сегментам.

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

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

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

Модель обучалась на качестве 1k, что позволило сделать скачок по качеству и найти много новых деталей. Если данные поступают ниже качеством, то модель ожидаемо деградирует, но со слов авторов — они целились в определенный уровень качества и задачи поддерживать качество хуже — не было.

Подстройка под конкретные задачи

Определение позы

В таких задачах используются ключевые точки (keypoints, «skeleton») для выделения частей тела, в результате можно понять в какой позе он находится, и как следующий шаг — понять что именно он делает. Такая штука используется, например, в спорте, виртуальной реальности и в анализе видеопотока на определенные действия (например, странное поведение или падение).

Примеры определения позы

Примеры определения позы

Здесь авторы завели аж 308 ключевых точек для определения позы. 243 пришлось на лицо (не просто глаза/нос/губы/язык и так далее, а все, что может влиять на «выражение лица») и 40 на кисти рук. И это прям очень неплохо.

Демонстрация ключевых точек - из публикации

Демонстрация ключевых точек — из публикации

Из интересного: модель расставляет точки даже для перекрытых объектов, «доугадывая» их расположение.

Определение перекрытых кистей рук

Определение перекрытых кистей рук

Сегментация частей телаhuman parsing»)

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

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

Human parsing - сегментация частей тела, не очень понятно к какой части тела относится ракетка, но это мелочи

Human parsing — сегментация частей тела, не очень понятно к какой части тела относится ракетка, но это мелочи

Оценка глубины

Задача оценки расстояния до объекта. Чем ярче объект, тем ближе он расположен к объективу и наоборот. Сферы применений: автономный транспорт, AR/VR, 3D-моделирование из 2D, медицина, жестовое управление устройствами и всякое-разное другое.

Оценка расстояния до объекта

Оценка расстояния до объекта

Оценка нормалей поверхности

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

Пример вывода модели

Пример вывода модели

Выводы и ссылки

Модель хороша! Авторы заявляют, что файнтюнинг под ваши дальнейшие хотелки делается легко, но на данный момент есть инструкция только для нормалей, но очень скоро обещают и все остальные. Но здесь стоит дополнительно добавить, что модель построена исключительно на людях и только на них.

Код и веса выложены, в том числе демо-скрипты чтобы поиграться самому на локалке (во избежание проблем разворачиваться лучше через conda). А вот датасет — нет и выложен не будет.

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

Из полезного:
Код и веса
Публикация (paper)
Демки hugginface
Готовый notebook для локалки

Мои другие статьи:

Разбор SAM2 через колено в голову или революция в разметке видео

Офис Apple в Москве: как я с нуля стал экспертом и попал на приватную вечеринку для разработчиков

Спасибо и всем хороших исследований!

© Habrahabr.ru