Оценка ключевых точек и pose estimation

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

Инициализация среды ONNX Runtime Web

Для работы с ORT Web необходимо создать сессию модели. Основные шаги включают загрузку модели в формате ONNX, настройку сессии и подготовку входных данных.

import * as ort from 'onnxruntime-web';

// Создание сессии
const session = await ort.InferenceSession.create('pose_model.onnx', {
  executionProviders: ['wasm'] // Использование WebAssembly для кросс-платформенной совместимости
});

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

  • executionProviders: определяет, как будет выполняться модель (WebAssembly, WebGL). WebGL обеспечивает аппаратное ускорение через GPU, что критично для реального времени.
  • InferenceSession.create: асинхронный метод, возвращающий объект сессии, который используется для дальнейшего инференса.

Подготовка входных данных

Модели оценки поз обычно принимают изображения в формате тензора. Необходимые шаги включают:

  1. Изменение размера изображения под требования модели.
  2. Преобразование данных в формат Float32Array.
  3. Нормализация значений (например, деление на 255 и вычитание среднего значения).
import { createCanvas, loadImage } from 'canvas';

const image = await loadImage('person.jpg');
const canvas = createCanvas(256, 256); // Подгон под модель
const ctx = canvas.getContext('2d');
ctx.drawImage(image, 0, 0, 256, 256);

const imageData = ctx.getImageData(0, 0, 256, 256);
const data = Float32Array.from(imageData.data).map((v, i) => i % 4 !== 3 ? v / 255.0 : 0);

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

  • Игнорирование альфа-канала (i % 4 !== 3) важно для моделей, обученных на RGB.
  • Порядок каналов может зависеть от модели (RGB или BGR), необходимо проверить спецификацию модели.

Инференс и получение ключевых точек

После подготовки входных данных необходимо выполнить инференс и интерпретировать результаты:

const inputTensor = new ort.Tensor('float32', data, [1, 3, 256, 256]);
const feeds = { input: inputTensor };

const results = await session.run(feeds);
const keypoints = results['output'].data; // 'output' зависит от конкретной модели

Замечания по результатам:

  • Выходной тензор обычно имеет форму [1, num_keypoints, 3], где третья координата — confidence score.
  • Координаты могут быть нормализованы (0–1), требуется масштабирование обратно к исходному размеру изображения.

Визуализация ключевых точек

Для проверки и отладки удобно визуализировать ключевые точки поверх изображения.

keypoints.forEach(([x, y, confidence]) => {
  if (confidence > 0.5) {
    ctx.beginPath();
    ctx.arc(x * 256, y * 256, 3, 0, 2 * Math.PI);
    ctx.fillStyle = 'red';
    ctx.fill();
  }
});

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

  • Порог confidence позволяет фильтровать малонадежные предсказания.
  • Масштабирование координат необходимо, если модель выдает нормализованные значения.

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

  • WebGL Execution Provider: ускоряет инференс на GPU.
  • Размер батча: использование батчей для последовательных кадров снижает накладные расходы на подготовку тензоров.
  • Разделение вычислений и визуализации: выполнение инференса в requestAnimationFrame для плавного отображения видео.
async function processFrame(videoFrame) {
  const inputTensor = preprocess(videoFrame);
  const results = await session.run({ input: inputTensor });
  renderKeypoints(results['output'].data);
  requestAnimationFrame(() => processFrame(nextVideoFrame));
}

Поддержка различных моделей

ORT Web совместим с моделями, экспортированными в ONNX из PyTorch или TensorFlow. Для pose estimation часто используют:

  • OpenPose: высокая точность, но большие размеры модели.
  • MoveNet: легковесная модель, оптимизированная для браузеров.
  • MediaPipe Pose (ONNX): хорошая компромиссная точность и производительность.

Каждая модель имеет свои требования к форме входного тензора, нормализации и формату выходных ключевых точек.

Примеры использования на видео

Для работы с видеопотоком браузера:

const video = document.getElementById('video');
navigator.mediaDevices.getUserMedia({ video: true }).then(stream => {
  video.srcObject = stream;
  video.play();
});

video.addEventListener('loadeddata', () => {
  function analyzeFrame() {
    processFrame(video);
    requestAnimationFrame(analyzeFrame);
  }
  analyzeFrame();
});

Важные моменты:

  • Предварительная подготовка кадра через canvas.
  • Асинхронный инференс не блокирует основной поток UI.
  • Для плавного отображения рекомендуется использовать requestAnimationFrame.

Особенности масштабирования и post-processing

  • Преобразование координат обратно к исходному размеру изображения для точной наложенной визуализации.
  • Использование сглаживания координат по времени для уменьшения дрожания при видео.
  • Сохранение confidence scores для принятия решений о действиях на основе детектированных поз.

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

Оглавление