Ускорение цветокоррекции на планшете до 0,45 секунды: опыт оптимизации MSLTNet

Автоматическая цветокоррекция на мобильном устройстве: задачи и вызовы

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

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

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

Исходная цветопередача планшета была нормальной: небо оставалось естественным синим, трава — зеленой, а лица людей сохраняли стандартный вид. Однако фотографиям не хватало сочности, выразительности и объёмного контраста. Требовалось создать автоматический режим улучшайзинга, который сам подтягивает яркость, контраст и насыщенность — по аналогии с профильными пресетами в iOS или Android.

Главное техническое ограничение заключалось в том, что алгоритм обязан работать локально на самом планшете, а не задействовать удаленные серверы. В качестве ориентира установили планку инференса около 1 секунды на кадр с разрешением вплоть до 12 мегапикселей. Это предъявляло жесткие рамки к размеру нейросети: тяжелые модели отсеивались из-за задержек, а слишком простые алгоритмы могли перенасыщать снимки или генерировать неприятные визуальные артефакты. Кроме того, модель требовалось безболезненно конвертировать из PyTorch в форматы ONNX или TensorFlow Lite для запуска на CPU, GPU и NPU мобильной платформы.

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

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

• MIRNet: показала худшие результаты по скорости. Кроме того, она заточена под съемку в условиях экстремально низкого освещения, поэтому для базового улучшения обычных дневных фото оказалась малопригодной.
• ZeroDCE: работала быстрее, но страдала той же проблемой — она ориентирована преимущественно на высветление темных кадров.
• Pix2PixHD: сформировала наилучшее и наиболее приятное визуальное качество с глубокими цветами. Однако время обработки одного снимка составило катастрофические 49,5 секунд, что совершенно недопустимо для мобильного гаджета.
• MSpecNet: выдавала приемлемую скорость, но на готовых фотографиях постоянно возникали заметные артефакты и странные искажения деталей.
• IAT: продемонстрировала средние результаты по всем фронтам — около 9 секунд на кадр, что не вписывалось в жесткие рамки целевого лимита.
• MSLTNet: стала безусловным победителем. Это суперлегкая нейросеть с объемом всего порядка 8 тысяч параметров. Она превзошла ближайших конкурентов по скорости на 1–2 порядка, показав прекрасный уровень корректировки яркости и цвета.

Особенности архитектуры модели MSLTNet

Секрет высокой эффективности MSLTNet кроется в её оригинальной структуре. Нейросеть не пытается перемалывать тяжелое изображение высокого разрешения (вплоть до 4K) единым сплошным блоком. В место этого кадр раскладывается с помощью пирамиды Лапласа на несколько частотных уровней:

• Низкочастотный слой: содержит информацию о глобальной яркости, контрасте и базовых цветах. Этот слой имеет низкое разрешение и обрабатывается через сетку двусторонней фильтрации (Bilateral Grid). Вычисления здесь требуют минимума ресурсов.
• Высокочастотные слои: отвечают за мелкие грани, текстуры и детали. К ним применяются точечные свертки 1×1, которые меняют значения цветовых каналов без сложных пространственных операций.

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

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

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

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

Для сохранения скорости на первых порах ввели промежуточный режим «3 Мп*». В нем сеть обрабатывала копию снимка на 3 Мп, после чего результат масштабировался до 12 Мп и смешивался с оригинальным каналом яркости с весом 30%. Это позволило уложиться в 0,8 секунды на кадр, но приводило к некоторой потере мелкого локального контраста.

Оптимизация инференса и совместная работа CPU и GPU

Полноценная обработка кадра 12 Мп через TensorFlow Lite на GPU изначально занимала 2,3 секунды. Анализ показал, что TFLite неэффективно справляется со слоями в начале и конце сети, где фигурируют гигантские тензоры (12 Мп × 3 канала).

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

Тестирование результатов на реальных снимках

После отладки гибридного режима алгоритм протестировали на разных категориях фотографий:

Обработка природных пейзажей

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

Коррекция телесных оттенков и портретов

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

Обработка предметов и студийных кадров

При съёмке отдельных предметов результат обработки иногда казался чересчур агрессивным. Решением со стороны пользовательского интерфейса стало добавление интерактивного слайдера регулировки степени воздействия (по аналогии со стандартными инструментами iOS). В будущем сеть дообучат автоматическому регулированию силы обработки под разные сцены.

Устранение артефактов и горизонтальных полос

На кадрах, заполненных единым насыщенным цветом, после обработки возникали видимые горизонтальные полосы. Причиной оказалась размерность 3D-тензора во входе GridSample (6, 16, 16). Экспериментально её уменьшили до (6, 4, 4), что полностью убрало артефакты полошения без потери качества на обычных фото.

Собственная кастомная реализация нейросети на CPU

Чтобы полностью отказаться от зависимостей и ограничений TFLite, нейросеть MSLTNet полностью переписали на C++ с ручной оптимизацией под процессоры ARM. Вычислительные ядра параллелились и оптимизировались инструкциями NEON, а формат данных перевели в FP16.

Это дало внушительный колоссальный задел по производительности:

• Время обработки снимка 12 Мп сократилось с 1,3 секунды до 0,6 секунды.
• Последующая оптимизация пред- и постобработки, слияние сверток с операциями Resize и оптимизация ОЗУ (минус 150 МБ) сократили время до рекордных 0,45 секунды!

Готовый алгоритм был успешно интегрирован в KvadraOS для встроенного приложения камеры и галереи.

Перспективы и обработка видеопотока

Следующая задача — перенос оставшихся вычислений на GPU и адаптация MSLTNet для видео. При независимой обработке каждого кадра возникает неприятное мерцание (flicker), так как коэффициенты Bilateral Grid слегка колебаются от кадра к кадру. Для решения проблемы разрабатывается система временного сглаживания параметров с учетом предыдущих кадров.

Итоги и выводы

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

Источник: habr.com