Началось все с того, что осенью прошлого года я общался со своим другом-коллегой на кухне. Ему хотелось ускорить виртуальные фоны, а у меня было немного времени.
Виртуальный фон – это важная фича Толка, которая позволяет обрезать все сзади тебя и вклеить одну фотографию. Например, вот фотография той самой кухни, где все и обсуждалось:
А результат наложения обычно такой:
Ускорение дело благородное, потому что у среднего пользователя компьютер совсем «картофельный» в плане производительности. Но даже если нет, то нам всё равно хочется, чтобы Толк не нагревал без дела коленки и садил батарейку ноутбука быстрее, чем нужно.
Как устроен виртуальный фон
На текущий момент в Толке зашита нейросеть, умеющая разделять каждый пиксель того, что видит камера, на две категории: фон и «человек в кадре». Нейросеть должна работать на любом ПК, выдавая приличное число кадров в секунду. Чтобы попасть в требования реального времени, нейросеть не может быть очень большой и сложной. Поэтому она работает на уменьшенном кадре. На выходе у нее тоже не супер много пикселей.
Вот так выглядит выход нейросети на самом деле:
После этого данные с нейросети попадают на волшебный Joint Bilaterial Filter, про который будет рассказано дальше. Он делает немного магии и получает вот такую маску:
Если потом смешать по этой маске картинки, то получится нужный результат:
Общая схема вот такая:
Joint Bilaterial Filter
Как видно из предыдущих картинок, именно этот алгоритм позволяет сглаживать грубую маску. Что, в свою очередь, нужно для облегчения нейросети. Но сам алгоритм билатеральной фильтрации не самый быстрый вычислительно, особенно если у вас 4К камера.
Попробуем разобрат��, как он устроен. Для этого немного откатимся к более простым фильтрам. Начнем с блура.
Gaussian Blur
Фильтр Gaussian Blur тоже есть в Толке. Он включается, когда нужно размыть задний фон. Например, вот тут:
Он устроен следующим образом. Нам нужно для каждой точки картинки взять окрестные точки, помноженные на определенные «веса» влияния. Веса в таком случае называют kernel или «ядром», а саму операцию называют «свёрткой с ядром» или 2d convolution.
В гаусовом блуре в качестве весов принято брать дискретизацию гауссовой кривой, прямо как тут:
На самом деле, если вместо гауссианы взять пирамиду, работать тоже будет. Но откуда тут гауссиана? По больше части так принято, исходя из её свойств.
Мы помним, что гауссиана имеет параметр – сигма. Тут он задает степень размытия, то есть по сути, насколько дальние точки мы будем учитывать.
Гауссов фильтр можно математически описать для каждой точки ещё как «сумму по окрестным точкам с учетом весов», где вес – это расстояние между координатами точки и координатами окрестной точки, подставленное в гауссиану (G):
Разница под двойными палками – это и есть вариант записи расстояния. В данном случае его можно заменить на корень из разницы квадратов по координатам.
Проблема с этим фильтром заключается в том, что он не сохраняет чётких границ. Другими словами, он размазывает всё. Поэтому для пролечивания этого на сцену выходит другой игрок.
Billaterial filter
Идея его в том, чтобы учитывать еще и границы, суммируя пиксели на основании того, насколько они похожи. А записывается Billaterial filter следующим образом:
Что поменялось? Добавился второй вес, который учитывает кроме обычного расстояния еще и так называемое «цветовое расстояние». То есть сейчас мы стараемся учитывать только похожие по цвету пиксели.
Ещё появилась явная нормировка, так как теперь не учитывается часть пикселей и теряется в яркости картинки. А чтобы вернуть уровень яркости, нужно посчитать, сколько мы потеряли в ней, и умножить на обратное.
Как это выглядит в деле:
При этом у нас есть теперь два параметра сигмы: сигма по расстоянию и сигма по цвету.
Joint Billaterial filter
Итак, в маске, которую нам отдала нейросеть и которую хотим наложить на видео, отсутствует информация о цвете, поэтому за «цветовым» расстоянием мы идем в исходный кадр изображения. Такой фильтр называется уже Joint, ведь мы используем два входа. Joint Billaterial filter находит примене��ие ещё в HDR фотографии, когда делается два снимка – со вспышкой и без – и на базе двух фото получается что-то более качественное.
А в нашей задаче мы используем информацию о цвете из исходного кадра, чтобы восстановить гладкую границу маски.
А где тут тормоза?
Не только нейросеть тормозная, но Bilaterial Filter сам по себе тоже тормозной. Для каждого кадра, для каждого пикселя нужно посчитать сумму по соседним пикселям. Вычисления у нас происходят на GPU и самая дорогая операция – вычитка цвета исходной картинки по координатам. Таких операции N^2 на каждый пиксель картинки.
Почему так? Потому что кадр запускается весь в обсчёт на GPU параллельно. Расчет каждого пикселя проходит независимо, но как только тебе потребовался цвет того или иного пикселя, ты должен его считать. А память находится далеко, и она общая – находится уже не в GPU, а в микросхемах рядом. И тут если не повезло попасть в кеш, придется ждать свою очередь.
Тот же Gaussian Filter сильно быстрее, потому что его можно наложить в два этапа: по горизонтали и по вертикали. Другими словами, 2D фильтр гаусса эквивалентен двум проходам 1D фильтра. Это уже N.
Но у билатерального фильтра такого свойства нет. И это его основная проблема.
Как мы их побороли?
Логически! Дело в том, что на самом деле необязательно фильтровать всю картинку. Ведь каждый пиксель на видео можно записать в 3 группы:
Точно фон
Точно человек
Может это человек, а может и фон
На основании этой идеи в коде возник своебразный хак. Так, мы делаем из маски картинку еще более низкого разрешения, исходя из правил:
Если все пиксели вокруг пикселя (включая его) – человек, то это точно человек
Если все пиксели рядом и сам пиксель – фон, то это точно фон
В коде расчета фильтра мы больше не считаем билатеральный фильтр, а принимаем решение заочно, если у нас точно человек (или точно фон). То есть мы не суммируем и не считаем, если нет явной границы.
Вот как работает разбивка на группы:
Результаты
Стало точно быстрее. Теперь вместо N^2 или N получилось что-то посередине – всё зависит от длины границы на видео.
Грубая оценка до и после по нагрузке на gpu на 720p
На 1080p профита еще больше.
Фича уже давно на проде!
Выводы
Оптимизация специфична по применению для целей Толка. В общем случае билатеральный фильтр нужно считать по всему изображению, но тут фильтруется именно маска, а маска нам интересна только на границе.
Если взглянуть внимательно на общий алгоритм и учесть в нем, что мы решаем уже не совсем общую задачу, а частную, то можно неплохо так выиграть на скорости работы.
