UNet-Segmentation

Передо мной стояла следующая задача: есть обучающая выборка, которая состоит из нескольких тысяч фотографий автомобилей. Необходимо обучить нейронную сеть отделять объект автомобиль от фона.

Задача сегментации

Сегментация – задача разбиения изображений на области, соответствующие различным объектам.

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

Структура сети UNet

За основу была взята архитектура UNet. Эта сверточная нейронная сеть, созданная в 2015 году для сегментации биомедицинских изображений в отделении Computer Science Фрайбургского университета. Архитектура сети представляет собой последовательность слоёв свёртка+пулинг, которые сначала уменьшают пространственное разрешение картинки, а потом увеличивают его, предварительно объединив с данными картинки и пропустив через другие слои свёртки. Сеть выполняет роль своеобразного фильтра.

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

Обучение

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

Граница разделения вычисляется с использованием морфологических операций. Затем вычисляется карта весовых коэффициентов:

где wc — карта весов для балансировки частот классов, d1 — расстояние до границы ближайшей ячейки, а d2 — расстояние до границы второй ближайшей ячейки.

Загрузка и подготовка данных

Для обучения модели был использован dataset с различными изображениями и файл с бинарными масками.

И маски для данных картинок

Код для импорта данных:

from google.colab import drive
drive.mount('/content/gdrive', force_remount=True)
!cp /content/gdrive/'My Drive'/data.zip .
!unzip data.zip

import numpy as np
import matplotlib.pyplot as plt
import cv2
import pandas as pd

Более подробно в Application 1 - Import data.

Создание архитектуры

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

inp = Input(shape=(256, 256, 3)) 

conv_1_1 = Conv2D(32, (3, 3), padding='same')(inp) # слой свертки
conv_1_1 = Activation('relu')(conv_1_1) # активация свертки

conv_1_2 = Conv2D(32, (3, 3), padding='same')(conv_1_1)
conv_1_2 = Activation('relu')(conv_1_2)

pool_1 = MaxPooling2D(2)(conv_1_2) # пулинг

Во второй части обучения делаем увеличение изображения и параллельную конкатенацию блоков одинаковой размерности. Здесь применим функции UpSampling2D с билинейной интерполяцией для увеличения изображения. Реализуем блок из 2х сверток с их последующей активацией.

up_1 = UpSampling2D(2, interpolation='bilinear')(pool_4) # апсамплинг 
conc_1 = Concatenate()([conv_4_2, up_1]) # конкатенация с тензором пулинга той же размерности

conv_up_1_1 = Conv2D(256, (3, 3), padding='same')(conc_1) # свертка уже сконкатенированного блока
conv_up_1_1 = Activation('relu')(conv_up_1_1)

conv_up_1_2 = Conv2D(256, (3, 3), padding='same')(conv_up_1_1)
conv_up_1_2 = Activation('relu')(conv_up_1_2)

На последнем этапе выполним свертку изображения с 1м каналом для определения конечной маски и активацию с параметром Sigmoid для получения вероятностей пикселей маски.

conv_up_4_2 = Conv2D(1, (3, 3), padding='same')(conv_up_4_1) # свертка с 1м каналом, для вывода маски изображения
result = Activation('sigmoid')(conv_up_4_2)

Более подробно в Application 2 - Create UNet.

Обучение модели

Создадим модель функцией Model.

model = Model(inputs=inp, outputs=result)

Проведем компиляцию модели, в качестве функции ошибки берем кросс-энтропию.

model.compile(adam, 'binary_crossentropy')

batch_size = 16
model.fit_generator(keras_generator(train_df, batch_size),
          steps_per_epoch=100, # кол-во батчей
          epochs=100, # кол-во эпох обучения 
          verbose=1, # вывод результата
          callbacks=callbacks, # наши листы весов
          validation_data=keras_generator(val_df, batch_size), # отдельный генератор для валидационной выборки
          validation_steps=50,
          class_weight=None,
          max_queue_size=10,
          workers=1,
          use_multiprocessing=False,
          shuffle=True,
          initial_epoch=0)
         
pred = model.predict(x)# функция прогнозирования маски

Процесс обучения модели:

Более подробно в Application 3 - Running Model Training.

Полученные результаты работы нейронной сети

Некоторые из полученных данных сегментации. Нашей задачей было выделить маску машины по изображению (т.е. отделить объект машина от фона). Как реализует это полученная нейронная сеть: