Постобработка сегментации: разбор масок и полигонов

ONNX Runtime Web (ORT Web) предоставляет возможность запускать модели машинного обучения прямо в браузере с использованием JavaScript. Основная цель — обеспечить выполнение предобученных моделей ONNX с высокой производительностью на стороне клиента. Для работы с сегментацией это особенно важно, поскольку результаты модели обычно представляют собой многоканальные тензоры, которые необходимо правильно интерпретировать и визуализировать.

Подключение библиотеки осуществляется через пакет onnxruntime-web, который можно установить через npm:

import * as ort from 'onnxruntime-web';

После этого создаётся сессия для конкретной модели ONNX:

const session = await ort.InferenceSession.create('model.onnx');

Формат выходных данных сегментационной модели

Сегментационные модели в ONNX обычно возвращают тензоры с размерностью [1, C, H, W], где:

  • 1 — размер батча;
  • C — количество классов или каналов;
  • H и W — высота и ширина изображения.

Каждый канал содержит вероятности принадлежности пикселя к соответствующему классу. В некоторых моделях выход может быть [1, H, W], где значения пикселей уже содержат индексы классов.

Преобразование тензора в карту масок

Для дальнейшей работы необходимо преобразовать многоканальный тензор в карту масок. Это делается с помощью операции argmax по канальному измерению:

function getSegmentationMask(outputTensor) {
    const [batch, channels, height, width] = outputTensor.dims;
    const data = outputTensor.data;
    const mask = new Uint8Array(height * width);

    for (let h = 0; h < height; h++) {
        for (let w = 0; w < width; w++) {
            let maxVal = -Infinity;
            let classIdx = 0;
            for (let c = 0; c < channels; c++) {
                const value = data[c * height * width + h * width + w];
                if (value > maxVal) {
                    maxVal = value;
                    classIdx = c;
                }
            }
            mask[h * width + w] = classIdx;
        }
    }
    return { mask, width, height };
}

Генерация полигонов из маски

Для визуализации объектов часто требуется преобразовать растровую маску в полигоны. Наиболее распространённый метод — алгоритм Marching Squares, который позволяет выделить контуры классов.

Пример генерации контуров:

import { findContours } from 'marchingsquares';

function maskToPolygons(maskData) {
    const { mask, width, height } = maskData;
    const polygonsByClass = {};

    // Для каждого класса
    const classes = [...new Set(mask)];
    classes.forEach(cls => {
        const binaryMask = [];
        for (let i = 0; i < mask.length; i++) {
            binaryMask[i] = mask[i] === cls ? 1 : 0;
        }
        const contours = findContours(binaryMask, width);
        polygonsByClass[cls] = contours.map(contour => contour.map(([x, y]) => ({ x, y })));
    });

    return polygonsByClass;
}

Особенности обработки

  1. Шумовые пиксели: иногда одиночные пиксели могут образовывать маленькие контуры. Рекомендуется фильтровать полигоны по площади.
  2. Масштабирование координат: если маска была получена на уменьшенной версии изображения, координаты полигонов необходимо масштабировать до исходного размера.
  3. Аппроксимация контуров: для упрощения геометрии полигонов можно использовать алгоритмы типа Ramer–Douglas–Peucker.

Визуализация результатов

После генерации полигонов их можно отобразить на HTML Canvas:

function drawPolygons(ctx, polygonsByClass, colors) {
    Object.entries(polygonsByClass).forEach(([cls, polygons]) => {
        ctx.fillStyle = colors[cls] || 'rgba(255,0,0,0.5)';
        polygons.forEach(polygon => {
            ctx.beginPath();
            polygon.forEach((point, index) => {
                if (index === 0) ctx.moveTo(point.x, point.y);
                else ctx.lineTo(point.x, point.y);
            });
            ctx.closePath();
            ctx.fill();
        });
    });
}

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

  • WebGL-поддержка: ORT Web может использовать WebGL для ускорения вычислений на GPU браузера.
  • TypedArrays: обработка масок через Uint8Array или Float32Array минимизирует накладные расходы.
  • Деление на блоки: при работе с очень большими изображениями целесообразно делить тензор на блоки и обрабатывать их параллельно.

Дополнительные аспекты

  • Поддержка нескольких моделей: ORT Web позволяет одновременно загружать несколько моделей, что полезно при комплексной сегментации с разными уровнями детализации.
  • Асинхронные вычисления: методы run и create возвращают промисы, что обеспечивает неблокирующее выполнение в интерфейсе пользователя.
  • Кэширование моделей: можно хранить сессии в памяти для многократного использования, избегая повторной загрузки больших файлов ONNX.

Эта последовательность операций — от получения тензора до визуализации полигонов — обеспечивает полный цикл постобработки сегментационных моделей в браузере с использованием ONNX Runtime Web. Она сочетает точность распознавания с гибкостью в представлении и оптимизацией вычислений.

Оглавление