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

Привет, Хабр! Меня зовут Денис Смирнов, я старший инженер по разработке ПО искусственного интеллекта KVADRA AI в YADRO. В этой статье расскажу, как мы искали нейросеть для автоматической цветокоррекции на планшете KVADRA_T, почему выбрали MSLTNet и как довели ее до работы на реальном устройстве. Покажу, с какими проблемами столкнулись при запуске на CPU, GPU и NPU, как боролись с артефактами и в итоге ускорили обработку фотографии 12 Мп до 0,45 секунды.

Что нужно от будущего решения

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

Основная задача цветокоррекции в нашем случае — сделать фотографии с планшета более яркими и выразительными. Исходная цветопередача была нормальной: небо оставалось синим, трава — зеленой, люди выглядели естественно. Но самим фотографиям не хватало насыщенности и контраста. Поэтому хотелось получить автоматический режим, который сам корректирует яркость, насыщенность и контраст, примерно как фильтр «Яркий» в мобильных галереях.

Главное техническое ограничение — скорость обработки, поэтому мы выбрали нейросетевой подход. Алгоритм должен работать непосредственно на планшете, а не на сервере. 

Целевым показателем выбрали примерно 1 секунду на один кадр. В экспериментах использовались снимки вплоть до 12 Мп, поэтому модель должна оставаться достаточно быстрой и на больших изображениях. Из-за этого важен и размер самой модели. Большие могли давать хороший результат, но работали слишком долго.

При этом модель должна не только усиливать цвета, но и делать это осторожно. Слишком сильная коррекция может привести к перенасыщенным цветам или неестественным оттенкам кожи. Еще одно важное требование — отсутствие артефактов: модель не должна дорисовывать новые детали или заметно менять содержимое кадра. Это одна из причин, почему при исследовании больше внимания уделяли небольшим моделям, работающим в основном с цветом и пикселями изображения.

Кроме того, модель нужно было запустить на аппаратной платформе планшета с NPU. Поэтому учитывалась возможность конвертации из PyTorch в более подходящие для устройства форматы и runtime — ONNX и TensorFlow Lite, а также поддержка отдельных операций сети на CPU, GPU и NPU.

Шесть моделей для цветокоррекции и один победитель

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

→ MIRNet оказалась самой медленной из всех рассмотренных моделей. При этом она больше подходит для фотографий, снятых при плохом освещении, поэтому для обычного улучшения снимков ее использовать было не очень удобно.

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

→ Pix2PixHD давала один из самых приятных результатов по цветам и яркости. Фотографии действительно выглядели заметно лучше, но сеть работала слишком медленно — 49.5443 секунды. Для планшета такая скорость уже была неприемлемой, особенно если учитывать фотографии большого разрешения.

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

→ IAT показала средний результат. Она была неплохой и по скорости, и по качеству, но ни в одном из этих параметров не превосходила остальные модели. Время работы — 9.0251 с.

Лучше всего себя показала MSLTNet. Это была самая легкая сеть среди рассмотренных — около 8 тысяч параметров.

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

Особенности модели MSLTNet

MSLTNet — очень легкая нейросеть, в ней всего около 8 тысяч параметров. При этом она рассчитана на обработку изображений высокого разрешения, вплоть до 4K, практически в реальном времени. Сеть умеет корректировать яркость и экспозицию, но при этом не требует тяжелых вычислений на полном разрешении изображения.

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

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

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

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

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

Первые эксперименты с запуском модели

После выбора MSLTNet начали проверять, как она работает уже на железе, близком по характеристикам к планшету. Для первых экспериментов использовали Mac, а затем одноплатный компьютер с похожим SoC, запуская модель через ONNX. Эта платформа использовалась скорее как модельная машина: по вычислительным характеристикам она была близка к целевому устройству, поэтому на ней было удобно проверять идеи до полноценного запуска на планшете. Там удалось довольно быстро получить время обработки меньше одной секунды на кадр. Стало понятно, что сама модель достаточно быстрая, а основная проблема дальше будет связана уже с переносом на планшет и выбором подходящего инференс-движка.

Еще одна проблема была связана с оператором GridSample, который не поддерживал инференс-движок для запуска модели на NPU или на TensorFlow Lite.

Из-за этого модель пришлось разделить: основная часть выполнялась через выбранный runtime, а GridSample обрабатывался отдельно. При использовании NPU для этого оператора применялась кастомная реализация на CPU, поэтому часть вычислений не могла выполняться на нейроускорителе. Кроме того, такой подход требовал дополнительных трансферов данных между памятью NPU и CPU, что также увеличивало общее время обработки.

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

3 Мп* — ускоренный режим: сеть запускается на изображении 3 Мп, но затем результат масштабируется до 12 Мп, после чего к нему подмешивается канал яркости исходного изображения на 12 Мп с весом около 30%. Это частично сохраняет детали исходного кадра при существенно меньших вычислительных затратах, чем при запуске сети непосредственно на 12 Мп.

Уже хорошо, но можно лучше 

После первых тестов стало понятно, что вариант с TensorFlow Lite (TFLite) на GPU уже дает хорошую скорость, около 0,8 секунды на кадр при 3 Мп, но качество все еще было не таким, как при обработке изображения в полном разрешении.

Чтобы сохранить исходное качество, MSLTNet запустили на полном кадре 12 Мп. Для этого часть слоев в начале и в конце сети переписали так, чтобы они выполнялись на CPU, а основная часть модели продолжала работать на GPU.

Такой вариант CPU + GPU оказался заметно лучше. Причина в том, что TFLite неэффективно работает со слоями, где используются большие тензоры размером около 12 Мп × 3 канала. Такие операции находились в начале и в конце сети и становились узким местом. 

Чтобы убрать это ограничение, первые и последние слои перенесли на CPU и оптимизировали их с помощью ARM NEON, а основную часть модели оставили на GPU. На полном разрешении 12 Мп время обработки сократилось с 2,3 до 1,3 секунды на кадр, при этом удалось сохранить качество изображения без уменьшения его до 3 Мп.

В итоге TFLite с совместной работой CPU и GPU стал лучшим вариантом на этом этапе: он давал почти нужную скорость и позволял обрабатывать фотографию в полном разрешении без потери деталей.

Понимаю, что можно уже запутаться в наших приключениях, ниже — общая таблица с техническими характеристиками.

Наслаждаемся первыми результатами

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

Природа

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

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

Телесные оттенки

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

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

Предметы

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

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

Полосы на насыщенных изображениях

Во время тестов обнаружилась еще одна проблема: если почти весь кадр был заполнен одним очень насыщенным цветом, после обработки на изображении могли появляться заметные горизонтальные полосы. Причина оказалась связана с размером 3D-объема, который используется в операции GridSample.

Изначально для GridSample использовался объем размером (6, 16, 16). После экспериментов его уменьшили до (6, 4, 4), и полосы практически исчезли. При этом на обычных фотографиях качество обработки почти не изменилось — это отдельно проверили на наборе снимков с KVADRA.

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

Все отлично, но надо переделать: собственная реализация модели

На последнем этапе MSLTNet полностью переписали под CPU, чтобы не зависеть от TFLite и GPU. Основные вычислительные ядра оптимизировали с помощью распараллеливания и NEON-инструкций.

Основная часть алгоритма теперь работает в FP16, а для x86 оставили поддержку FP32. Это дало заметный прирост скорости: раньше вариант TFLite CPU + GPU обрабатывал кадр 12 Мп примерно за 1,3 секунды.

Кастомная реализация на CPU + NEON сократила это время до 0,6 секунды на кадр.

Готовый алгоритм интегрировали в KvadraOS для приложений камеры и галереи. Для тестирования использовали демонстрационное приложение, в котором можно сравнивать оригинальное и обработанное изображения, перемещая границу между ними. Там же измерялось реальное время работы алгоритма непосредственно на планшете.

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

Параллельно продолжили ускорять сам алгоритм. Время обработки кадра 12 Мп удалось сократить с 0,6 до 0,45 секунды.

Основной прирост получили за счет оптимизации пред- и постобработки, уменьшения потребления памяти примерно на 150 МБ, а также объединения операций свертки и Resize, чтобы сократить число промежуточных вычислений и обращений к памяти.

А что дальше? 

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

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

И напоследок

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

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

Если хочется не только следить за развитием AI, но и участвовать в нем напрямую, сейчас есть несколько открытых позиций для специалистов с разным профилем — от разработки и архитектуры до качества и продуктового дизайна. В команде ищут:

• AI Tech Lead (LLM),

• Руководителя группы разработки (AI, LLM),

• Архитектора систем ИИ,

• Старшего продуктового дизайнера.