В рамках данной статьи мы обсудим новое zero-shot решение (то есть способное справляться с задачей без дополнительного обучения на данных из конкретного домена) задачи Visual Object Tracking под названием SAMURAI(SAM-based Unified and Robust zero-shot visual tracker with motion-Aware Instance-level memory). Эта модель продемонстрировала хороший перфоманс в задаче визуального трекинга, обойдя на нескольких бенчмарках своего прямого предка - SAM 2, а также многие supervised-решения(требующие дообучения под конкретный домен и задачу).

Прежде чем перейти непосредственно к SAMURAI, я нахожу полезным кратко ввести читателя в курс дела в области основных вех развития SAM-based моделей. Начнём с родоначальника - самого SAM.

SAM

SAM задумывалась как фундаментальная модель для сегментации изображений, которая выполняет сегментацию на основе промпта, которым может выступать текст, точка (указывает на объект для сегментации), bounding box или маска и без дообучения адаптируется под различные домены и задачи. SAM обучался на огромном датасете SA-1B, содержащем более миллиарда масок и видел большое количество разнообразных данных, что помогает ей эффективно работать с новыми данными. Это делает её универсальным инструментом, который можно использовать для получения новых данных.

Давайте теперь разберём из чего состоит SAM и как он работает

Конструкция SAM выглядит следующим образом:

После того как изображение подаётся в SAM, оно попадает в тяжёлый энкодер изображений. В качестве энкодера выступает предобученный Vision Transformer(ViT). Энкодер без особой потери в масштабируемости можно сделать потяжелее, поскольку эмбеддинги изображения достаточно получить лишь единожды при прогоне изображения через модель и их же можно несколько раз переиспользовать для последующих промптов.
Для получения эмбеддингов введённого пользователем промпта в архитектуре SAM имеется prompt encoder. В нём промпт в виде bounding box и точек кодируется с помощью позиционного энкодинга, который суммируется с обученными эмбеддингами для каждого типа промпта. Текстовый промпт кодируется с помощью текстового энкодера из модели CLIP, так как этот энкодер предназначен для перевода текста в векторное представление, которое может быть использовано совместно с визуальными данными. Промпт в виде маски с помощью свёрток приводится к размерности кодированного изображения и поэлементно складывается с ним.
Роль совмещения эмбеддингов изображения и промпта, а после получения итоговых предсказаний играет декодер. Кодированное изображение и промпт пропускается через два трансформерных блока, которые представляют собой слой self-attention в промпте и cross-attention между промптом и изображением в двух направлениях: от изображения к промпту и наоборот. После прохождения через трансформерные блоки происходит upsampling изображения до исходных размеров. Далее классификационная голова в виде линейного перцептрона(MLP) делает окончательное предсказание маски.

Однако, стоит учесть что во многих случаях промпт может быть довольно многозначным, как показано на примере ниже. Как бороться с подобными кейсами? Чтобы разрешать такие неопределённости, SAM генерирует не одну маску, а три. Это позволяет модели рассмотреть все возможные случаи, так как маски зачастую имеют не более трёх уровней глубины: целый объект, его часть и подчасть. Но как из трёх масок выбрать наиболее подходящую? Для этого у SAM есть дополнительная голова, которая отвечает как раз за это. Она рассчитывает score каждой маски - $s_{mask}$ . В качестве этого скора она учится предсказывать Intersection over Union(IoU) предсказанной и правдивой маски. В качестве результата модель выбирает маску с наибольшим скором.

SAM-2

SAM проявил себя отлично в задаче сегментации изображений, продемонстрировав высокую способность к обобщению. Однако, у него есть довольно существенное ограничение - SAM умеет работать только с изображениями. Разработчики SAM-2 пошли дальше, адаптировав SAM для работы с видео.

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

Итак, давайте взглянем на внутренности SAM-2

Основное нововведение SAM-2 - это банк памяти(memory_bank), в котором хранится последние N кадров из видео. Давайте рассмотрим подробнее

Энкодер изображений - тут идейно без изменений относительно SAM. Единственное, SAM-2 использует иерархический энкодер изображений. То есть он выдаёт на выходе признаки разных уровней, как высокоуровневые(ближе к входу модели), так и более низкоуровневые(ближе к выходу модели). О том, зачем это нужно, будет чуть ниже.
После энкодера, изображение попадает в memory attention. Memory attention представляет собой несколько блоков self-attention между элементами памяти, cross-attention между изображением и памятью и MLP. Именно на этом этапе изображение приобретает контекстную информацию из прошлых кадров и предсказаний.
Prompt encoder - без изменений относительно SAM
Mask decoder - по большей части аналогично без изменений относительно SAM, за тем исключением, что в видео целевой объект может скрываться из виду, поэтому необходимо добавить дополнительную голову, которая будет предсказывать скор, отвечающий за наличие объекта в кадре - $s_{obj}$ . Ещё одним нововведением являются skip-connections между разными уровнями иерархического энкодера и декодера(в обход memory attention) по аналогии с тем как это делается в архитектуре Unet(приложена ниже). С помощью этого информация о высоких пространственных разрешениях передаётся в декодер напрямую. Это помогает сохранить детали высокого разрешения, которые могут быть потеряны при проходе через энкодер.

Memory encoder - на этом этапе полученная на выходе из декодера маска прогоняется через свёрточные слои, затем суммируется поэлементно с закодированным изображением и, наконец, всё это "полируется" свёртками, после чего отправляется в memory bank.
Memory bank - банк памяти хранит последние N кадров, которые затем используются в memory attention для учёта контекста видео.

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

SAMURAI

Несмотря на то что SAM-2 показал хорошие результаты в задаче Visual Object Tracking, он также столкнулся с рядом трудностей. Во-первых модель до сих пор не очень хорошо чувствует себя в переполненных сценах с быстро движущимися, пропадающими из вида объектами, а также в сценах с похожими объектами(например, стая животных, где каждый экземпляр похож на остальных), так как SAM-2 часто отдаёт предпочтение визуальной схожести объектов, а не пространственной и временной согласованности. У модели отсутствует явная информация о движении, поэтому ей трудно различать похожие объекты. Ещё один источник ошибок - неэффективная конструкция памяти, в которой хранится только последние N кадров, вне зависимости от того, какого качества эти кадры. Такое нерациональное строение памяти может привести к распространению ошибок в видео.

Например:

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

На втором кейсе произошло распространение ошибки из-за того, что на нескольких кадрах объект было плохо видно на изображении => детекции получились с низкой уверенностью и неточные, а на последующих кадрах модель доставала из памяти эти кадры с неточной детекцией и на основе них делала следующее предсказание. Таким образом произошло распространение ошибки сквозь кадры и поведение модели приобрело высокую степень непредсказуемости. Настолько высокую, что вместо велосипеда она начала трекать автомобиль. Модель памяти, в которой мы будем не беспорядочно сохранять кадры в памяти и использовать их, а выбирать из них только качественные, которые не приведут к подобному результату, поможет решить подобные проблемы.

Рассмотрим архитектуру SAMURAI и разберём, какие нововведения помогли ей превзойти своего предшественника SAM-2.

Одно из важных нововведений SAMURAI - использование модели движения для задачи трекинга, основанной на фильтре Калмана.

На простом примере рассмотрим как работает фильтр Калмана:

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

В нашем случае "прибор", или "глаз" который следит за неким объектом - это сама модель детекции. А модель движения задействуется SAMURAI на этапе, когда изображение прошло через декодер и было предсказано M масок - она помогает выбрать наиболее корректную маску на основе информации о движении объекта.

Разберём теперь подробнее фильтр Калмана и то, как он помогает делать предсказания:

У нас есть вектор-состояние маски, который определяется следующим образом:

$x = [x, y, h, w, \dot{x}, \dot{y}, \dot{h}, \dot{w}]$

Это координаты и скорости bounding box'а маски, который определяется с помощью вычисления её максимальных и минимальных координат x и y.

Также вводится ковариационная матрица состояния $P_{t}$ , которая задаёт неопределённость состояния и обычно является диагональной, при условии что все компоненты нескоррелированы.

Также определяется используемая модель движения. Она задаётся оператором F преобразования из состояния $x_{t|t}$ в состояние $x_{t+1|t}$ . С помощью неё мы предсказываем следующее состояние нашего объекта:

$x_{t+1|t} = Fx_{t|t}$

В нашем случае это модель равномерного прямолинейного движения(это не значит что объект на протяжении всего видео движется равномерно прямолинейно. Это значит что он движется только на промежутке между кадрами, то есть измерениями).

Также предсказываем с помощью модели движения матрицу ковариации в следующий момент времени:

$P_{t+1|t}=FP_{t|t}F^{-1} + Q$ , где Q - матрица шума модели перехода, которая описывает неопределённость модели и обусловлена её неточностью. Чем меньше наша модель подходит для описания движения, тем этот шум больше. Например, если движение объекта не равномерное, а равноускоренное, шум будет больше. Значение этой матрицы зависит от задачи и подбирается эмпирически.

После получения предсказаний модели для каждой из M масок считаем следующий коэффициент:

$s_{kf} = IoU(x_{t+1|t}, M)$

Он представляет собой метрику IoU bounding box'a данной маски и предсказанного положения объекта. Это и есть коэффициент, который показывает качество предсказанной маски с точки зрения пространственной и временной согласованности. Окончательную маску выбираем следующим образом:

$M^{*} = \operatorname*{argmax}_{M}(\alpha_{kf} * s_{kf}(M_{i})+(1 - \alpha_{kf})*s_{mask}(M_{i}))$

Для каждой маски считаем взвешенную сумму коэффициентов $s_{kf}$ и $s_{mask}$ с коэффициентом $\alpha_{kf}$ , где $s_{mask}$ - оценка IoU с реальной маской, которая предсказывается самой моделью(как в SAM). В оригинальной статье $\alpha_{kf} = 0.2$ . Таким образом, модель при выборе итоговой маски из предсказанных, использует непосредственно информацию о движении объекта, что довольно сильно увеличивает качество предсказаний, особенно в сценах с множеством похожих объектов, ведь теперь у модели есть ещё один явный способ отличить объекты друг от друга помимо визуального сходства.

Полученную маску используем как новое измерение $z_{t+1}$ нашего "прибора" для уточнения состояния.

Считаем инновацию(невязку) между измерением и предсказанием

$y_{t+1}=z_{t+1}-Hx_{t+1|t}$

Где H - матрица - которая извлекает параметры из полного состояния $[x, y, h, w, \dot{x}, \dot{y}, \dot{h}, \dot{w}]$

Рассчитываем ковариацию невязок:

$S_{t+1}=HP_{t+1|t}H^{T}+R$ , где матрица R описывает шум измерений. Тоже оценивается эмпирически.

Вычисляем коэффициент Калмана:

$K_{t} = P_{t+1|t}H^{T}S_{t+1}^{-1}$

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

Теперь остаётся скорректировать состояние x и матрицу ковариации P:

$x_{t+1|t+1}=x_{t+1|t} + K_{t}(z_{t}-Hx_{t+1|t})$

$P_{t+1|t+1}=(I+K_{t}H)P_{t+1|t}$

И далее по-новой.

Память

Memory Encoder остался неизменным относительно SAM-2, но вот структура памяти претерпела изменения: ранее там просто хранились N последних кадров, теперь хранение происходит более умным образом, учитывая информацию о движении, наличии объекта в кадре и уверенности модели в предсказании. Мы позволяем модели выбрать, какие кадры хранить в памяти на основе 3-х оценок, о которых шла речь выше: $s_{kf}$ - коэффициент, показывающий, насколько данная маска соответствует модели движения, $s_{obj}$ - коэффициент, отвечающий за оценку наличия объекта в кадре и $s_{mask}$ - коэффициент, показывающий уверенность модели в предсказанной маске(предсказание IoU с правильной маской). Кадр считается хорошим тогда и только тогда, когда все три оценки превышают соответствующие пороги(эти пороги - гиперпараметры модели). В итоге в памяти хранится $N_{max}$ хороших кадров, которые выбраны моделью, это решает проблему распространения ошибки из-за хранения в памяти некачественных кадров.

Это все нововведения модели SAMURAI. Остальное осталось аналогично SAM-2.

Заключение

В данной статье мы рассмотрели эволюцию моделей SAM, начиная с оригинального SAM, разработанного для сегментации изображений, затем SAM-2, который адаптировали для работы с видео, и, наконец, SAMURAI — актуальное SOTA zero-shot решение для Visual Object Tracking.

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

Таким образом, SAMURAI демонстрирует значительный шаг вперёд в области zero-shot трекинга, показывая конкурентные результаты не только по сравнению с SAM-2, но и с многими supervised-методами.

У SAMURAI есть цель — zero-shot решение задачи Visual Object Tracking(VOT)