Чистое изображение: удаление бликов и отражений на фото с помощью Machine Learning

Решение нетривиальных задач computer vision и сбор синтетического датасета – Артём Топильский из команды Поиска и рекламных технологий Яндекса рассказывает о том, чему можно научиться, пока пишешь диплом.

От школы до Яндекса

К окончанию школы я увлекался физикой, и думал поступать куда-то по этому профилю, поэтому среди вузов рассматривал в том числе и МФТИ. Уже в 11-ом классе к нам пришёл новый учитель по информатике, который занимался подготовкой олимпиадников. Он поставил перед собой сложную задачу: за год научить нас базовым вещам C++, причём в олимпиадном формате. Это у него получилось: на одной из олимпиад я победил и решил попробовать себя в IT.  В МФТИ я выбрал «Прикладную математику и информатику», а при распределении пошёл на «Дискретную математику». Во время обучения на ФИВТе у нас были очень крутые преподаватели — команда А. М. Райгородского. Там я закончил бакалавриат, а в магистратуру решил поступить на кафедру анализа данных Яндекса, где были дисциплины из ШАДа.

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

Диплом писал по computer vision — на тему удаления отражений с изображений при помощи нейросетей. Такие модели нужны для быстрого и удобного удаления дефектов с фотографий. Второе очевидное применение – улучшение качества детекции и распознавания. Например, ты видишь понравившуюся вещь в витрине, но у тебя нет времени заходить в магазин, либо он закрыт. Прямо на ходу фотографируешь её через стекло, а позже прогоняешь через приложение для удаления бликов. Качество фотографии улучшается, и теперь вещь можно найти с помощью умной камеры.

Работа над дипломом

Летом после первого курса магистратуры к нам пришёл Алексей Озёрин из Лаборатории машинного интеллекта Яндекса и сказал: «Я нашёл классные прорывные статьи, но во время чтения не понятны тонкие, но важные детали реализации идей, потому что нет приложенного кода. Хотелось бы попробовать воспроизвести описанные методы: честно сделать то, что предлагают авторы, и проверить, работает ли это на практике». Одна из этих статей была про удаление отражений с изображения. Она мне особенно понравилась: раньше я уже работал с компьютерным зрением во время стажировки в Samsung Research. Там я два месяца занимался сверхразрешением – это немного другая задача, но некоторые концепты в них достаточно близки.

Задача сверхразрешения состоит в следующем: допустим, у тебя есть изображение с разрешением 1 мегапиксель, а тебе нужно улучшить качество, превратив 1 в 4, чтобы картинка стала более чёткой и детальной. В решении таких задач в последнее время очень популярны приёмы машинного обучения, в том числе свёрточные сети. И я подумал: почему бы не попробовать в аналогичной задаче применить навыки, которые у меня уже есть, например, вспомнить Research ML — теоретические исследования в ML, чтение и разбор статей и научных исследований. Он противопоставляется Production ML, для которого более важна явная польза бизнесу.

Сначала я должен был проверить, действительно ли метод из статьи можно воспроизвести. Параллельно с этим возникла идея: описанные свёрточные сети различной глубины и ширины можно изменять по размеру и по времени работы. Мы с Лёшей решили, что будет полезно исследовать задачу удаления отражений на предмет масштабирования и построить типичный граф размена качества на уменьшение и ускорение модели.

К сожалению, аналогичных источников по масштабированию нейросетей в задаче удаления отражений мы не нашли. С одной стороны, стандартные задачи классификации уже давно определены, конкретизированы и формализованы. За последние 10 лет было придумано много всего, нарощены огромные сети из сотен, даже тысяч слоёв,  и дальнейшее увеличение их размера не приводит к значительному росту качества — оно выходит на «плато». 

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

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

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

Представим себе модель, у которой есть несколько слоёв. Количество последовательных слоёв обычно называют высотой или глубиной, а количество параметров на каждом слое — шириной. Предположение было такое: чтобы оптимально разменивать качество на скорость и размер модели, надо построить trade-off качества и скорости или размера модели, так как изменение ширины или глубины должно происходить равномерно. 

Есть такая закономерность: если слишком увеличить глубину, то качество быстро увеличивается, а потом выходит на плато и больше не растёт. Или модель вообще перестаёт обучаться, если сеть очень глубокая. И аналогично: если наращивать только ширину, тоже очень быстро выходишь на плато. А если их наращивать равномерно, то качество получается лучше при сравнимом размере и скорости обработки данных.

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

Советы для тех, кто решил погрузиться в машинное обучение

Сейчас мои рабочие задачи с темой диплома напрямую не связаны — я решил, что нужно испробовать много областей, пока есть такая возможность. Но как и любой специалист, работающий с машинным обучением, я читаю статьи и слушаю выступления других исследователей. Ещё я подписан в Telegram на канал Kantor AI,  автор которого рассказывает про работу команд машинного обучения и образование в области анализа данных. Мне интересно читать про то, как изменился ML с тех пор, как я начал им заниматься, про особенности работы с большими данными.  

ТОП-выводов, которые я сделал во время работы над дипломом и которые до сих пор помогают мне в работе:

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

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

Важно уметь подбирать для модели оптимальную скорость обучения (learning rate) – то, насколько она по датасету исправляет себя с каждым шагом. Этот параметр очень важен, потому что если выставить скорость обучения слишком высокой, то модель начинает переобучаться на тренировочный датасет, и на реальных данных покажет плохое качество. Или вообще не обучится, потому что веса этой модели скачут из-за высокой скорости изменений. Очень низкую скорость обучения тоже ставить нельзя — тогда модель просто не обучится, будет видеть пример, но исправлять веса будет совсем незначительно, например, на 0,001%. Целой жизни не хватит дождаться, пока при такой скорости модель обучится. В своей статье «SIRR, не соизволите ли удалить отражение?» я, в том числе, рассказываю о том, как процесс нахождения оптимальной скорости обучения можно автоматизировать.

Говорить и писать о своей работе нужно понятно. Понятно для тех, кто тебя будет слушать или читать. Например, для написания статьи на Хабр я использовал много технических деталей и кода — «мяса», которое потом какой-нибудь студент сможет разобрать. А вот если бы пожилая соседка спросила меня, чем я занимаюсь, я бы показал ей пару фотографий и объяснил, что научил искусственный интеллект убирать блики с таких изображений. Учиться этому важно ещё по одной причине: когда рассказываешь о своей работе разным людям, осознаёшь, какое значение она имеет для других.

поделиться
статьёй
поделиться
статьёй