Colorization Transformer — это нейросетевая модель, которая окрашивает черно-белые изображения и основывается на механизме внимания. Модель принимает на вход черно-белое изображение. Процесс окрашивания проходит в три этапа. Сначала авторегрессионная Transformer-модель производить окраску в низком разрешении. Архитектура модели использует условные Transformer слои для эффективной обработки изображений. Затем две параллельные нейросети улучшают разрешение и окраску изображения. По результатам экспериментов, модель обходит state-of-the-art архитектуры для раскраски изображений. В качестве метрики использовали FID и результаты опроса на платформе Mechanical Turk. В более чем 60% участники опроса предпочитали результат модели реальному изображению.
Зачем это нужно
Окраска изображения — это комплексная задача, которая требует от модели понимания семантики сцены и базовых знаний об окружающем мире.Одними из применений таких моделей являются цветовая коррекция существующих цветных изображений и восстановление цвета у черно-белых изображений и видеозаписей.
Детали архитектуры
Исследователи утверждают, что Colorization Transformer — это первый случай использования Transformer-модели для задачи окраски изображения в высоком разрешении (256 × 256). Они предлагают условные Transformer-слои для колоризации изображения в низком разрешении, а затем последовательное улучшение разрешения.
Подход состоит из трех отдельных моделей, каждая из которых обучается независимо от остальных:
- Авторегрессионный колоризатор: Axial Transformer с условными Transformer-слоями;
- Модель апсемплинга цвета
- Модель пространственного апсемплинга
Модели апсемплинга в ColTran являются слоями механизма внимания, которые параллельно апсемплят цвет и разрешение изображения. Каждый блок внимания состоит из следующих операций: нормализация слоев → механизм внимания → многослойный перцептрон.