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

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

Сейчас мы попытаемся решить похожую задачу. В данном кейсе будут проанализированы спутниковые снимки на предмет определения на них географических объектов, таких как реки, поля, дома, дороги и леса. Для решения таких задач используется сверточная нейронная сеть. Одной из распространенных её архитектур является модель 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 часто используется для сегментации изображений и выдает хорошие результаты.