Семантическая и инстанс-сегментация

ONNX Runtime Web (ORT Web) представляет собой высокопроизводительную JavaScript-библиотеку для запуска моделей в формате ONNX прямо в браузере или на Node.js. Она обеспечивает переносимость моделей, высокую скорость выполнения и совместимость с различными аппаратными ускорителями, такими как WebAssembly (WASM) и WebGPU. Основной объект для работы — Ort.InferenceSession, который управляет загрузкой модели, подготовкой входных данных и выполнением инференса.

import * as ort from 'onnxruntime-web';

const session = await ort.InferenceSession.create('model.onnx', {
    executionProviders: ['wasm', 'webgl']
});

Ключевые моменты:

  • executionProviders определяет, на каком движке будет выполняться модель: wasm для универсальной поддержки или webgl/webgpu для ускорения на GPU.
  • Загрузка модели происходит асинхронно, что важно учитывать при организации потока обработки данных.

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

Модели семантической и инстанс-сегментации требуют специфической предобработки. Основные шаги включают:

  1. Масштабирование и обрезка изображения: входное изображение приводится к размеру, ожидаемому моделью, например 512x512 или 640x640 пикселей.
  2. Нормализация: значения пикселей преобразуются в диапазон [0,1] или нормализуются по среднему и стандартному отклонению.
  3. Форматирование: изображение преобразуется в тензор типа Float32Array с размерностью [1, C, H, W], где C — число каналов, H и W — высота и ширина.

Пример подготовки входного тензора:

function preprocessImage(imageData, width, height) {
    const canvas = document.createElement('canvas');
    canvas.width = width;
    canvas.height = height;
    const ctx = canvas.getContext('2d');
    ctx.drawImage(imageData, 0, 0, width, height);
    const imageArray = ctx.getImageData(0, 0, width, height).data;

    const floatArray = new Float32Array(width * height * 3);
    for (let i = 0; i < width * height; i++) {
        floatArray[i] = imageArray[i * 4] / 255;       // R
        floatArray[i + width * height] = imageArray[i * 4 + 1] / 255; // G
        floatArray[i + 2 * width * height] = imageArray[i * 4 + 2] / 255; // B
    }
    return new ort.Tensor('float32', floatArray, [1, 3, height, width]);
}

Запуск инференса

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

const inputTensor = preprocessImage(imgElement, 512, 512);

const feeds = { input: inputTensor };
const results = await session.run(feeds);

const outputTensor = results['output'];

Особенности:

  • Семантическая сегментация обычно возвращает тензор [1, num_classes, H, W], где каждый канал соответствует вероятностям для класса.
  • Инстанс-сегментация может возвращать несколько выходных тензоров: boxes, masks, labels и scores.

Обработка результатов семантической сегментации

Для визуализации семантической карты необходимо преобразовать вероятности в метки классов:

function getSegmentationMap(outputTensor) {
    const [batch, numClasses, height, width] = outputTensor.dims;
    const data = outputTensor.data;
    const segmentationMap = new Uint8ClampedArray(height * width);

    for (let h = 0; h < height; h++) {
        for (let w = 0; w < width; w++) {
            let maxProb = -Infinity;
            let classIdx = -1;
            for (let c = 0; c < numClasses; c++) {
                const idx = c * height * width + h * width + w;
                if (data[idx] > maxProb) {
                    maxProb = data[idx];
                    classIdx = c;
                }
            }
            segmentationMap[h * width + w] = classIdx;
        }
    }
    return segmentationMap;
}

Для визуализации можно сопоставить каждому классу цвет и отрисовать на canvas.

Инстанс-сегментация: работа с масками и объектами

Инстанс-сегментация требует дополнительной постобработки:

  1. Выбор объектов по порогу вероятности (score_threshold).
  2. Применение масок к исходному изображению с учетом координат ограничивающих рамок.
  3. Слияние нескольких масок с прозрачностью для корректной визуализации перекрывающихся объектов.

Пример обработки одного объекта:

function applyMask(ctx, maskTensor, box, color) {
    const [height, width] = maskTensor.dims;
    const data = maskTensor.data;

    ctx.fillStyle = color;
    ctx.globalAlpha = 0.5;

    for (let y = 0; y < height; y++) {
        for (let x = 0; x < width; x++) {
            if (data[y * width + x] > 0.5) {
                ctx.fillRect(box[0] + x, box[1] + y, 1, 1);
            }
        }
    }

    ctx.globalAlpha = 1.0;
}

Особенности работы с инстанс-сегментацией:

  • Маски часто имеют меньший размер, чем исходное изображение, и требуют масштабирования до координат box.
  • Многоканальные выходные тензоры позволяют эффективно извлекать объекты и их принадлежность к классам.

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

  • Выбор execution provider: webgpu обеспечивает максимальную скорость на поддерживаемых устройствах. wasm универсален, но медленнее.
  • Минимизация пересылки данных: использовать ImageData напрямую, избегать лишнего копирования в массивы.
  • Батчинг изображений: если модель поддерживает пакетный инференс, несколько изображений можно обрабатывать одновременно для ускорения.

Совместимость с браузером и Node.js

ORT Web поддерживает работу:

  • В браузере: через script тег или ES-модули, поддерживает WebAssembly и WebGPU.
  • В Node.js: через onnxruntime-web или onnxruntime-node с идентичным API, но с использованием нативного ускорения CPU/GPU.

Особое внимание требуется при работе с WebGPU в браузере: проверяется доступность через navigator.gpu и конфигурируются соответствующие execution providers.

Оглавление