Сегментация

Сегментация — это задача компьютерного зрения, направленная на разбиение изображения на семантически значимые области. В отличие от классификации, где модель присваивает метку всему изображению, сегментация работает на уровне пикселей, что позволяет определять объекты и их границы с высокой точностью. В TensorFlow.js сегментация реализуется с использованием моделей, обученных на архитектурах типа U-Net, DeepLab или Mask R-CNN.

Подготовка данных для сегментации

Для успешного обучения или инференса модели сегментации необходимо корректное представление данных. Основные моменты:

  • Изображения должны быть нормализованы. Обычно пиксели масштабируются в диапазон [0, 1] или [-1, 1].
  • Маски сегментации представляют собой изображения той же размерности, что и исходные, где каждому пикселю соответствует индекс класса. Формат маски обычно Uint8Array или Tensor с целочисленными значениями.
  • Формат тензора: для подачи на вход модели изображения конвертируются в tf.tensor4d([batch, height, width, channels]). Маски сегментации аналогично подаются в виде тензора с shape [batch, height, width, 1].

Пример преобразования изображения и маски в тензоры:

const imageTensor = tf.browser.fromPixels(imageElement).toFloat().div(tf.scalar(255));
const maskTensor = tf.browser.fromPixels(maskElement, 1).toInt();
const input = imageTensor.expandDims(0); // [1, height, width, 3]
const target = maskTensor.expandDims(0);  // [1, height, width, 1]

Архитектуры моделей для сегментации

U-Net

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

DeepLab

DeepLab использует атросивные сверточные слои (dilated convolutions) и pyramid pooling для захвата контекста на разных масштабах. В TensorFlow.js модели DeepLab чаще всего предоставляются в предобученном виде через tfjs-models:

import * as deeplab from '@tensorflow-models/deeplab';

const model = await deeplab.load({base: 'pascal', quantizationBytes: 2});
const segmentation = await model.segment(imageElement);

segmentation возвращает объект с картой классов (segmentation.data) и информацией о цветовой маске.

Mask R-CNN

Mask R-CNN позволяет выполнять объектно-ориентированную сегментацию, где каждая отдельная сущность получает маску и bounding box. В отличие от U-Net, здесь сегментация делается для каждого объекта отдельно.

Подходы к инференсу

Существует два основных режима работы с сегментационными моделями:

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

Постобработка результатов

После получения маски сегментации необходимо часто применять постобработку для повышения качества:

  • Сглаживание границ с помощью морфологических операций (dilate, erode) или Gaussian blur.
  • Приведение к исходному размеру изображения через tf.image.resizeBilinear.
  • Преобразование индексов классов в цветовую маску для визуализации:
const colorMap = {0: [0,0,0], 1: [255,0,0], 2: [0,255,0]};
const rgbMask = tf.tidy(() => {
  return segmentationTensor.squeeze().cast('int32').map(idx => tf.tensor(colorMap[idx]));
});

Тренировка моделей сегментации

Процесс тренировки включает:

  • Функцию потерь: чаще всего используется sparseCategoricalCrossentropy для многоклассовой сегментации.
  • Метрики: точность на уровне пикселя (pixel accuracy) и Intersection over Union (IoU).
  • Аугментации: горизонтальные/вертикальные отражения, повороты, изменение яркости/контрастности, чтобы увеличить разнообразие тренировочных данных.

Пример компиляции модели:

model.compile({
  optimizer: tf.train.adam(0.001),
  loss: 'sparseCategoricalCrossentropy',
  metrics: ['accuracy']
});

Оптимизация и ускорение

TensorFlow.js позволяет ускорять сегментацию через:

  • WebGL backend для GPU-вычислений.
  • Квантование модели (quantizationBytes: 1 или 2) для уменьшения размера и ускорения инференса.
  • Батчинг изображений при инференсе на статических наборах.

Визуализация сегментации

В браузере визуализация осуществляется через Canvas API. Маска сегментации комбинируется с исходным изображением:

const canvas = document.getElementById('canvas');
const ctx = canvas.getContext('2d');

ctx.drawImage(imageElement, 0, 0);
ctx.globalAlpha = 0.5;
ctx.putImageData(segmentationImageData, 0, 0);

Прозрачность позволяет видеть исходное изображение под наложенной цветовой маской.

Оглавление