Постобработка детекции объектов: NMS, декодирование якорей

ONNX Runtime Web предоставляет инструменты для выполнения моделей глубокого обучения прямо в браузере на JavaScript. После получения сырых предсказаний от модели детекции объектов необходимо выполнить постобработку для преобразования этих данных в удобный для использования формат: координаты прямоугольников объектов, классы и вероятности. Ключевыми этапами постобработки являются декодирование якорей и Non-Maximum Suppression (NMS).

Декодирование якорей (Anchor Decoding)

Модели детекции, такие как SSD или YOLO с якорями, возвращают координаты объектов относительно заранее заданных якорей (anchor boxes). Эти предсказания обычно представляют собой смещения и масштабные коэффициенты относительно якорей, а не абсолютные координаты в изображении.

Формула декодирования:

[ x_{center} = x_{anchor} + dx w_{anchor}
y_{center} = y_{anchor} + dy h_{anchor}
w = w_{anchor} (dw)
h = h_{anchor} (dh)]

где:

  • (dx, dy, dw, dh) — выходы модели для конкретного якоря
  • (x_{anchor}, y_{anchor}, w_{anchor}, h_{anchor}) — параметры якоря
  • (x_{center}, y_{center}, w, h) — декодированные координаты прямоугольника объекта

После вычисления центра и размеров, прямоугольник можно преобразовать в привычный формат ([x_{min}, y_{min}, x_{max}, y_{max}]):

[ x_{min} = x_{center} - w/2, y_{min} = y_{center} - h/2
x_{max} = x_{center} + w/2, y_{max} = y_{center} + h/2]

На практике это часто реализуется в JavaScript с использованием массивов или TypedArray для повышения производительности:

function decodeBoxes(predictions, anchors) {
    const boxes = [];
    for (let i = 0; i < anchors.length; i++) {
        const [dx, dy, dw, dh] = predictions[i];
        const [ax, ay, aw, ah] = anchors[i];
        const cx = ax + dx * aw;
        const cy = ay + dy * ah;
        const w = aw * Math.exp(dw);
        const h = ah * Math.exp(dh);
        boxes.push([cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2]);
    }
    return boxes;
}

Non-Maximum Suppression (NMS)

После декодирования якорей модель может выдавать несколько перекрывающихся прямоугольников для одного объекта. Для удаления избыточных предсказаний применяется Non-Maximum Suppression (NMS).

Основная логика NMS:

  1. Сортировка всех предсказанных прямоугольников по вероятности объекта (score) по убыванию.
  2. Выбор прямоугольника с максимальной вероятностью как текущего.
  3. Вычисление IoU (Intersection over Union) между текущим прямоугольником и всеми оставшимися.
  4. Удаление всех прямоугольников, у которых IoU выше заданного порога (обычно 0.5).
  5. Повторение шагов 2–4 до тех пор, пока не останутся прямоугольники.

Формула IoU:

[ IoU = ]

Пример реализации на Jav * aScript:

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 interArea = Math.max(0, x2 - x1) * Math.max(0, y2 - y1);
    const boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1]);
    const boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1]);

    return interArea / (boxAArea + boxBArea - interArea);
}

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

    const picked = [];
    while (indices.length) {
        const current = indices.shift();
        picked.push(current);
        indices = indices.filter(idx => computeIoU(boxes[current], boxes[idx]) < iouThreshold);
    }
    return picked;
}

Оптимизация для ONNX Runtime Web

  • Использование Float32Array и векторных операций позволяет ускорить обработку больших массивов предсказаний.
  • Для моделей с большим числом якорей имеет смысл предварительно фильтровать предсказания по минимальному порогу вероятности до выполнения NMS, чтобы сократить количество вычислений.
  • Поддержка WebAssembly в ONNX Runtime Web обеспечивает значительное ускорение по сравнению с чистым JavaScript.
  • В современных браузерах можно использовать Web Workers для выполнения декодирования и NMS в отдельном потоке, предотвращая блокировку UI.

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

  1. Получение предсказаний от модели через ONNX Runtime Web: session.run({input}).
  2. Декодирование якорей с использованием функции decodeBoxes.
  3. Применение порогового фильтра по вероятности.
  4. Применение NMS для удаления перекрывающихся боксов.
  5. Формирование финального массива объектов с координатами и метками классов для визуализации или дальнейшей логики.

Этот подход обеспечивает точную и эффективную постобработку детекции объектов в браузере без необходимости передачи данных на сервер.

Оглавление