Сегментация изображения с помощью нейронных сетей на примере практической задачи



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

Ключевые слова: сегментация изображений, нейронные сети.

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

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

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

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

Для того, чтобы решить описанную задачу, необходимо:

1) Выполнить задачу сегментации, а именно: подобрать датасет с фотографиями комнат и соответствующими им изображениями-масками, на которых сегментированы стены и построить модель нейронной сети с архитектурой, соответствующей задаче, и обучить ее сегментации;

2) Разработать алгоритм, который выполнял бы перекраску изображения.

Для проведения исследования в данной работе был использован язык программирования Python 3.8, с использованием библиотек: Keras 2.4.3, Numpy 1.20.1, Matplotlib 3.3.4, Scikit-image 0.18.1.

Был найден датасет [1], в котором содержались 1800 фотографий с соответствующих им масками. Помимо стен, на масках изначально были размечены еще 22 различных класса объектов, которые не были нужны для выполнения нашей задачи, поэтому они были объединены в один.

Рис. 1. Пример фотографии и маски из датасета

Датасет был размечен недостаточно корректно, поэтому до построения модели были исключены 400 фотографий и масок к ним.

В качестве архитектуры была взята U-Net (ее преимущества были описаны в статье [2]). Сначала была использована классическая по количеству слоев U-Net, которая описана, например, в книге [3].

Далее, по очереди обучались модели архитектуры U-Net с общими параметрами:

Таблица 1

Общие параметры моделей

Функция активации свертки	Изменение разрешения при свертке	Функция активации на выходе	Оптимизатор	Функция потерь	Метрика
ReLU	Нет	Sigmoid	Adam	Binary crossentropy	Accuracy

В Таблице 2 представлены параметры обучения каждой модели, которые менялись в процессе исследования.

Таблица 2

Переменные параметры

№	Размер выборки	Размер батча	Шаг обучения	Число нейронов	Аугмен-тация	Разрешение	Кол-во эпох
1	1400	20	0.001	1 941105	Нет	128x128	50
2	1000	25	0.001	1941105	Нет	128x128	70
3	640	32	0.005	1941105	Нет	128x128	70
4	640	32	0.003	1941105	Да	128x128	70
5	640	32	0.003	1941105	Да	512x512	70
6	640	32	0.003	31110001	Да	512x512	70

При аугментации варьировались параметры: приближение, поворот, яркость и горизонтальный поворот (зеркальное отображение). Для каждой модели из датасета 90 % изображений и масок были включены в обучающий набор, а остальные 10 % — в валидационный. В итоге, 6-я модель имела на последней эпохе наивысшую точность: 86 %, а также визуально лучше всех производила сегментацию на тестовых фотографиях.

Рис. 2. Пример работы 6-й модели нейронной сети

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

(1)

где — текущий пиксел нового изображения, color — цвет, задаваемый пользователем, original — цвет текущего пиксела исходного изображения,

— цветовая компонента.

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

Литература:

1. RMRC Reconstruction Meets Recognition Challenge 2014. [Электронный ресурс]. URL: https://cs.nyu.edu/~silberman/rmrc2014/indoor.php (дата обращения: 11.06.2021).
2. U-Net: нейросеть для сегментации изображений. [Электронный ресурс]. URL: https://neurohive.io/ru/vidy-nejrosetej/u-net-image-segmentation/ (дата обращения: 11.06.2021).
3. S. Pattanayak. Pro Deep Learning with TensorFlow [Книга]. — Berkeley, CA: Apress, 2017. — 398 с.
4. Аугментация (augmentation, “раздутие”) данных для обучения нейронной сети на примере печатных символов. [Электронный ресурс]. URL: https://habr.com/ru/company/smartengines/blog/264677/ (дата обращения: 11.06.2021).