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

Некоторые устройства могут использовать стереокамеры, способные извлекать информацию о глубине для сегментации изображения. Однако подход, описанный в этой статье, направлен на создание системы сегментации, которая сможет вывести желаемую информацию из одного изображения RGB. Это позволит нам применять портретные эффекты на фронтальной камере для более широкого набора устройств.

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

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

Наконец, последняя проблема - это обучающий набор данных. Чтобы добиться максимальной производительности, мы собрали набор данных для этой задачи и создали специальный конвейер дополнений.

Набор данных

Нейронные сети нуждаются в обучении, поэтому большая часть построения нейронных сетей этого типа требует обучения, поэтому основная часть построения алгоритмов этого типа - это набор данных. Мы хотели, чтобы карты сегментации были максимально точными, поэтому все изображения в нашем наборе данных были вручную помечены с помощью профессиональной программы редактирования. Кроме того, мы хотели разнообразить набор данных. Отсутствие разнообразия может привести к структурным ошибкам алгоритма. Например, если модель не может справиться с определенной позой, возможно, набор данных недостаточно репрезентативен. Мы использовали два способа улучшить качество нашего набора данных: увеличение и сбор дополнительных изображений. Количество изображений в нашем наборе данных постепенно увеличивается с течением времени. На данный момент у нас есть около 13 тысяч изображений в тренировочном наборе и еще тысяча в тестовом наборе с людьми в различных позах, включая фотографии в полный рост, портретные и селфи-фотографии.

Архитектура

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

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

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

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

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

Следующая ключевая особенность - разделимая свертка по глубине. По сравнению с обычными свертками он позволяет достичь того же качества с меньшими затратами на вычисления и меньшими требованиями к памяти. Мы добавляем остаточный блок с разделяемыми по глубине свертками в качестве базового блока в нашей архитектуре вместо исходных блоков Convolution + ReLU. Еще одна особенность нашей архитектуры - асимметрия. Декодер содержит больше блоков, чем кодировщик. Это позволяет нам достичь большей точности при тех же вычислительных затратах. Наконец, чтобы добиться вывода в реальном времени, мы постарались максимально сократить количество слоев и карт объектов.

Окончательная архитектура изображена на рисунке ниже.

Учебные приемы

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

Еще один момент, о котором мы хотим упомянуть, - это артефакты вывода. Нейронная сеть имеет тенденцию генерировать шахматные доски - сетку, которую можно увидеть после применения вычитания фона. Этот тип артефакта не влияет на показатель mIoU, но создает плохое визуальное восприятие. Чтобы решить эту проблему, мы добавляем дополнительные остаточные блоки к окончательной форме и даем дополнительное время на обучение.

Полученные результаты

Наконец, мы получаем модель для портретной сегментации, которая хорошо сбалансирована по качеству и скорости. Занимает 3,7 мб в формате fp32 onnx. Наш алгоритм можно перенести на большинство фреймворков, включая CoreML, MetalPerformanceShaders, SNPE, Huawei Kirin, OpenVINO. В частности, на iPhone 7 и Mi Mix 2S эта модель работает быстрее, чем пороговое значение 30 кадров в секунду. Конвейер включает в себя получение пиксельного буфера с камеры, предварительную обработку, логический вывод, постобработку и рендеринг конечного изображения с удаленным фоном на экране.

Пост скриптум

Представленная система сегментации портретов создана в сотрудничестве с моими замечательными коллегами. К сожалению, на Medium нет официального варианта совместного редактирования, но они упомянуты ниже:

Лебедев Антон, Константин Семьянов, Артур Чахвадзе - обучающий конвейер и сетевая архитектура
Роман Куцев, Павел Воропаев - управление процессом сбора и аннотирования данных
Максим Скороходов, Вячеслав Тарасов - интеграция с Android
Олег Пояганов, Андрей Володин - интеграция с iOS
MR - обзор этой статьи