Семантическая сегментация: архитектура U-Net

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

Архитектура U-Net

U-Net является одной из наиболее популярных архитектур для сегментации. Она обладает характерной U-образной формой и состоит из двух основных частей: сжимающего (encoder) и расширяющего (decoder) блоков.

Сжимающая часть (Encoder)

Сжимающая часть отвечает за извлечение признаков и уменьшение пространственных размеров изображения. Она построена на последовательности:

  1. Сверточные слои (Conv2D) Каждый блок содержит два сверточных слоя с функцией активации ReLU. Свертки позволяют выявлять локальные признаки изображения: края, текстуры, формы объектов.

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

Пример блока в Keras:

const convBlock = (input, filters) => {
    let x = keras.layers.conv2d({filters: filters, kernel_size: 3, activation: 'relu', padding: 'same'}).apply(input);
    x = keras.layers.conv2d({filters: filters, kernel_size: 3, activation: 'relu', padding: 'same'}).apply(x);
    return x;
};

const pool = keras.layers.maxPooling2d({pool_size: [2, 2]}).apply(convBlock);

Расширяющая часть (Decoder)

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

  1. Обратная свертка (Conv2DTranspose) Этот слой увеличивает размерность карты признаков, выполняя операцию обратной подвыборки.

  2. Конкатенация (Concatenate) Для сохранения детальной информации на каждом уровне используются skip connections: карты признаков из encoder соединяются с соответствующими уровнями decoder.

  3. Сверточные слои После объединения выполняются сверточные операции для уточнения признаков и формирования окончательной карты сегментации.

Пример блока в Keras.js:

const upConvBlock = (input, skipConnection, filters) => {
    let x = keras.layers.conv2dTranspose({filters: filters, kernel_size: 2, strides: 2, padding: 'same'}).apply(input);
    x = keras.layers.concatenate().apply([x, skipConnection]);
    x = keras.layers.conv2d({filters: filters, kernel_size: 3, activation: 'relu', padding: 'same'}).apply(x);
    x = keras.layers.conv2d({filters: filters, kernel_size: 3, activation: 'relu', padding: 'same'}).apply(x);
    return x;
};

Финальный слой

Для многоклассовой сегментации используется Conv2D с количеством фильтров, равным числу классов, и активацией softmax. Для бинарной сегментации применяется один фильтр с активацией sigmoid.

const output = keras.layers.conv2d({filters: numClasses, kernel_size: 1, activation: 'softmax'}).apply(lastDecoderLayer);

Интеграция Keras.js в браузере

Модели, обученные в Python Keras, экспортируются в формат .json с бинарными весами. В Keras.js они загружаются и выполняются на стороне клиента:

const model = new KerasJS.Model({
    filepath: 'model.json',
    gpu: true
});

await model.ready();
const outputData = await model.predict({input_1: inputTensor});

Формат входных данных

Модель ожидает тензор размерности [1, height, width, channels]. Для изображений RGB это [1, 128, 128, 3]. Перед подачей на модель изображение нормализуется в диапазон [0,1].

const normalizeImage = img => {
    const floatData = new Float32Array(img.length);
    for (let i = 0; i < img.length; i++) {
        floatData[i] = img[i] / 255.0;
    }
    return floatData;
};

Оптимизация производительности

  • Использование WebGL через параметр gpu: true ускоряет вычисления.
  • Сжатие модели и уменьшение глубины U-Net снижает задержку загрузки и скорость предсказания.
  • Мини-батчи и предварительная обработка изображений позволяют плавно интегрировать сегментацию в интерактивные веб-приложения.

Особенности применения

U-Net особенно эффективен для задач с ограниченным количеством обучающих данных благодаря skip connections, сохраняющим пространственные признаки. В Keras.js эта архитектура обеспечивает точную сегментацию прямо в браузере без необходимости серверной обработки, что критично для интерактивных визуализаторов и веб-приложений с ограниченной пропускной способностью сети.