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

Для решения таких задач используется сверточная нейронная сеть. Одной из распространенных её архитектур является модель U-Net. На вход нейронной сети подается изображение, и далее создается маска, которая будет определять объекты из разных классов на изображении.

Исходный датасет состоял из снимков со спутника. Ниже можно увидеть пример одного из таких изображений:

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

Для начала необходимо выбрать область, на которой будет проходить обучение нейронной сети. То есть выбираем патч рандомным образом и формируем для него маску. На данном патче будет тренироваться модель. Функция на Python выглядит следующим образом:

def get_rand_patch(img, mask, s=160): assert len(img.shape) == 3 and img.shape[0] > sz and img.shape[1] > sz and img.shape[0:2] == mask.shape[0:2] x = random.randint(0, img.shape[0] - s) y = random.randint(0, img.shape[1] - s) patch_img = img[x:(x + s), y:(y + s)] patch_mask = mask[x:(x + s), y:(y + s)] return patch_img, patch_mask

Сверточная нейронная сеть состоит из четырех шагов: Convolution, Max Pooling, Flattening и Full Connection. Таким образом, ниже будет представлена функция на Python, в которой подробно описано построение модели U-Net для рассматриваемого кейса. На шаге Convolution было взято количество фильтров 32 размеров 160 на 160. Также был использован ReLu Layer, который избавил feature map от отрицательных значений и превратил их в нули. Таким образом, были получены новые Rectified feature maps. Параметр, отвечающий за padding в функции стоит ‘same’, что означает обрамление входного изображения нулями для контроля размера feature map.

def unet_model(classes=5, size=160, channels=8, filters=32, factor=2, conv=True, weights=[0.2, 0.3, 0.1, 0.1, 0.3]): number_of_filters = filters input_1 = Input((size, size, channels)) convolutional_1 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')(input_1) convolutional_1 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')( convolutional_1)

Далее следует Max Pooling, где указывается матрица размером 2 на 2. Среди данных значений будет выбираться максимальное число, чтобы уменьшить размерность Rectified feature map.

pooling_1 = MaxPooling2D(pool_size=(2, 2))(convolutional_1)

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

number_of_filters *= factor pooling_1 = BatchNormalization()(pooling_1)

Ниже представлена вся модель U-Net, строящаяся по аналогичному принципу.

convolutional_2 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')(pooling_1) convolutional_2 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')( convolutional_2) pooling_2 = MaxPooling2D(pool_size=(2, 2))( convolutional_2) number_of_filters *= factor pooling_2 = BatchNormalization()(pool2) convolutional_3 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')(pooling_2) convolutional_3 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')( convolutional_3) pooling_3 = MaxPooling2D(pool_size=(2, 2))(convolutional_3) number_of_filters *= factor pooling_3 = BatchNormalization()(pooling_3) convolutional_4 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')(pooling_3) convolutional_4 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')( convolutional_4) pooling_4 = MaxPooling2D(pool_size=(2, 2))( convolutional_4) number_of_filters *= factor convolutional_5 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')(pooling_4) convolutional_5 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')(convolutional_5) number_of_filters //= factor if conv: up_6 = concatenate([Conv2DTranspose(number_of_filters, (2, 2), strides=(2, 2), padding='same')(convolutional_5), convolutional_4]) else: up_6 = concatenate([UpSampling2D(size=(2, 2))(convolutional_5), convolutional_4]) up_6 = BatchNormalization()(up6) convolutional_6 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')(up_6) convolutional_6 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')(convolutional_6) number_of_filters //= factor if conv: up_7 = concatenate([Conv2DTranspose(number_of_filters, (2, 2), strides=(2, 2), padding='same')(convolutional_6), convolutional_3]) else: up_7 = concatenate([UpSampling2D(size=(2, 2))(convolutional_6), convolutional_3]) up_7 = BatchNormalization()(up_7) convolutional_7 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')(up_7) convolutional_7 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')(convolutional_7) number_of_filters //= factor if conv: up_8 = concatenate([Conv2DTranspose(number_of_filters, (2, 2), strides=(2, 2), padding='same')(convolutional_7), convolutional_2]) else: up_8 = concatenate([UpSampling2D(size=(2, 2))(convolutional_7), convolutional_2]) up_8 = BatchNormalization()(up_8) convolutional_8 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')(up_8) convolutional_8 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')(convolutional_8) number_of_filters //= factor if conv: up_9 = concatenate([Conv2DTranspose(number_of_filters, (2, 2), strides=(2, 2), padding='same')(convolutional_8), convolutional_1]) else: up_9 = concatenate([UpSampling2D(size=(2, 2))(convolutional_8), convolutional_1]) convolutional_9 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')(up_9) convolutional_9 = Conv2D(number_of_filters, (3, 3), activation='relu', padding='same')(convolutional_9) convolutional_10 = Conv2D(classes, (1, 1), activation='sigmoid')(convolutional_9) model = Model(inputs=input_1, outputs=convolutional_10)

Далее на этапе Full Connection выбираем функцию оптимизации Adam, которая будет минимизировать ошибку наших предсказаний.

model.compile(optimizer=Adam(), loss=K.sum(K.mean(K.binary_crossentropy(y_true, y_pred), axis=[0, 1, 2]) * K.constant(weights) ) return model

Метрикой качества данной модели была выбрана logloss, так как она часто используется в таких задачах. После обучения модели она достигла примерно 0,15.

Таким образом, в результате реализации данной модели удалось получить следующую картинку.

По результатам работы модели можно заметить, что лучше всего распознаются поля и реки. Для распознавания других объектов модели для обучения требуется больше времени, итераций, размеров тренировочной и валидационной выборки, для чего необходим графический процессор. Модель U-Net часто используется для сегментации изображений и выдает хорошие результаты.