Восстановление цвета на изображении — это одна из открытых задач в компьютерном зрении. Исследователи из National Tsing Hua University и Virginia Tech обучили модель, которая учитывает распознанные instance-маски объектов при восстановлении цвета. Предложенная архитектура обходит существующие подходы.
Проблема существующих подходов
Предыдущие подходы для восстановления цветов изображения предсказывают цвета изображения напрямую для всего изображения сразу. Такие модели хуже справляются с входными изображениями, которые содержат множество объектов. Основной причиной является тот факт, что модели обучаются восстанавливать цвета сразу на всем изображении. Это приводит к тому, что нейросеть не может эффективно локализовать объекты и учитывать семантику на уровне объектов.
Архитектура подхода
Модель принимает на вход черно-белое изображение. На первом этапе изображение проходит через нейросеть для распознавания объектов. Затем вырезаются все распознанные объекты. Каждый из распознанных объектов поступает в модель, которая восстанавливает цвета на изображении. Чтобы объединить восстановленные части изображения в одно, используется fusion модуль. Fusion модуль обеспечивает консистентность цветов на всем изображении.
Оценка работы модели
Исследователи оценили обучили модель на COCO-Stuff и сравнили с текущими state-of-the-art подходами на трех датасетах. Ниже видно, что по количественным метрикам предложенная нейросеть работает лучше, чем state-of-the-art.
