Время прочтения: 4 мин.

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