Детекция объектов: YOLOv8, SSD, DETR

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

Ключевые возможности:

  • Кроссплатформенность: работа как в браузере, так и в Node.js.
  • Поддержка WebAssembly (WASM) для широкого охвата браузеров.
  • Использование WebGPU для ускоренного инференса на GPU при наличии поддержки.
  • Полная поддержка моделей ONNX, включая сложные архитектуры детекции объектов.

Инициализация среды и загрузка модели

Для начала работы необходимо подключить библиотеку и создать сессию выполнения. В браузере это обычно делается через импорт скрипта или npm-пакет:

import * as ort from 'onnxruntime-web';

const session = await ort.InferenceSession.create('model.onnx', {
  executionProviders: ['wasm'], // или 'webgl', 'webgpu' при необходимости
});

Параметры сессии:

  • executionProviders — массив движков, которые можно использовать, например: 'wasm', 'webgl', 'webgpu'.
  • graphOptimizationLevel — уровень оптимизации графа модели ('disabled', 'basic', 'extended', 'all').

Использование оптимизации графа позволяет ускорить инференс, особенно для больших моделей детекции.

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

Модели детекции объектов ожидают тензоры фиксированного размера, чаще всего с каналами RGB и нормализованными значениями. Для YOLOv8 и SSD вход обычно выглядит как [1, 3, H, W].

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

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

  for (let i = 0; i < width * height; i++) {
    data[i * 3 + 0] = imageData.data[i * 4 + 0] / 255.0; // R
    data[i * 3 + 1] = imageData.data[i * 4 + 1] / 255.0; // G
    data[i * 3 + 2] = imageData.data[i * 4 + 2] / 255.0; // B
  }

  return new ort.Tensor('float32', data, [1, 3, height, width]);
}

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

Инференс модели

После подготовки тензора можно выполнить инференс. В ONNX Runtime Web это делается с помощью метода run, передавая объект с именами входов модели.

const inputTensor = preprocessImage(imgElement, 640, 640);
const feeds = { input: inputTensor };

const results = await session.run(feeds);

Результат модели — объект, где ключи соответствуют выходам сети. Для YOLOv8 и SSD это часто координаты ограничивающих рамок (boxes), вероятности (scores) и классы (labels). DETR возвращает результаты в виде прямых предсказаний с использованием механизма трансформера, включая позиции и идентификаторы объектов.

Постобработка детекции

После получения сырых выходов требуется интерпретировать их в человекочитаемом формате. Основные этапы:

  1. Фильтрация по порогу уверенности: удаление объектов с низкой вероятностью.
  2. Нормализация координат: преобразование из относительных значений [0,1] в пиксели изображения.
  3. Non-Maximum Suppression (NMS): устранение перекрывающихся рамок с меньшей вероятностью.

Пример NMS для YOLO/SSD:

function nonMaxSuppression(boxes, scores, iouThreshold) {
  const indices = [];
  const sorted = scores
    .map((score, idx) => ({ score, idx }))
    .sort((a, b) => b.score - a.score);

  while (sorted.length > 0) {
    const { idx } = sorted.shift();
    indices.push(idx);

    sorted = sorted.filter(({ idx: i }) => {
      const iou = computeIoU(boxes[idx], boxes[i]);
      return iou < iouThreshold;
    });
  }

  return indices;
}

function computeIoU(boxA, boxB) {
  const x1 = Math.max(boxA[0], boxB[0]);
  const y1 = Math.max(boxA[1], boxB[1]);
  const x2 = Math.min(boxA[2], boxB[2]);
  const y2 = Math.min(boxA[3], boxB[3]);

  const intersection = Math.max(0, x2 - x1) * Math.max(0, y2 - y1);
  const areaA = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1]);
  const areaB = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1]);

  return intersection / (areaA + areaB - intersection);
}

Оптимизация инференса

  • Выбор движка: WebGPU обеспечивает значительное ускорение на современных устройствах. WebAssembly подходит для совместимости с любыми браузерами.
  • Предобработка и батчинг: выполнение предобработки через Web Workers снижает блокировку основного потока.
  • Кэширование сессии: повторное использование объекта InferenceSession для нескольких изображений исключает повторную инициализацию модели.
  • Сжатие модели: использование ONNX с оптимизированными весами и квантованием уменьшает объем памяти и ускоряет загрузку.

Особенности моделей

YOLOv8:

  • Высокая скорость и точность на небольших изображениях.
  • Требует специфической постобработки — NMS и порог уверенности.

SSD:

  • Баланс между скоростью и точностью.
  • Предсказывает ограничивающие рамки через якоря (anchors), которые требуют декодирования.

DETR:

  • Использует архитектуру трансформера, не нуждается в NMS, так как сеть сама генерирует уникальные предсказания.
  • Постобработка сводится к фильтрации по вероятности и масштабированию координат.

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

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

function drawDetections(canvas, detections) {
  const ctx = canvas.getContext('2d');
  ctx.lineWidth = 2;
  ctx.font = '18px Arial';
  ctx.strokeStyle = 'red';
  ctx.fillStyle = 'red';

  detections.forEach(det => {
    const [x1, y1, x2, y2] = det.box;
    ctx.strokeRect(x1, y1, x2 - x1, y2 - y1);
    ctx.fillText(`${det.label} ${Math.round(det.score * 100)}%`, x1, y1 - 5);
  });
}

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