yolo-helpers/dist/yolo-box/common.js

Helper functions to use models converted from YOLO in browser and Node.js

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.decodeBox = decodeBox;
exports.decodeBoxSync = decodeBoxSync;
/**
 * tensorflow output: [batch, features, instances]
 * features:
 * - 4: x, y, width, height
 * - num_classes: class confidence
 *
 * e.g. 1x84x8400 for 1 batch of 8400 instances with 80 classes
 *
 * The confidence are already normalized between 0 to 1.
 */
async function decodeBox(args) {
    let { tf, num_classes, maxOutputSize, iouThreshold, scoreThreshold } = args;
    let length = 4 + num_classes;
    // e.g. 1x84x8400
    let batches = args.output;
    if (batches[0].length === 0) {
        // no a single batch
        return [];
    }
    if (batches[0].length !== length) {
        throw new Error(`data[batch].length must be ${length}`);
    }
    let num_boxes = batches[0][0].length;
    let result = [];
    for (let batch of batches) {
        // e.g. 84x8400
        let boxes = [];
        let scores = [];
        let cls_indices = [];
        for (let box_index = 0; box_index < num_boxes; box_index++) {
            let x = batch[0][box_index];
            let y = batch[1][box_index];
            let width = batch[2][box_index];
            let height = batch[3][box_index];
            let x1 = x - width / 2;
            let y1 = y - height / 2;
            let x2 = x + width / 2;
            let y2 = y + height / 2;
            let box_score = batch[4][box_index];
            let cls_index = 0;
            for (let i = 1; i < num_classes; i++) {
                let cls_score = batch[4 + i][box_index];
                if (cls_score > box_score) {
                    box_score = cls_score;
                    cls_index = i;
                }
            }
            boxes.push([x1, y1, x2, y2]);
            scores.push(box_score);
            cls_indices.push(cls_index);
        }
        let box_indices;
        if (maxOutputSize) {
            let box_indices_tensor = await tf.image.nonMaxSuppressionAsync(boxes, scores, maxOutputSize, iouThreshold, scoreThreshold);
            box_indices = await box_indices_tensor.array();
            box_indices_tensor.dispose();
        }
        else {
            box_indices = Array.from({ length: num_boxes }, (_, i) => i);
        }
        let bounding_boxes = [];
        for (let box_index of box_indices) {
            let x = batch[0][box_index];
            let y = batch[1][box_index];
            let width = batch[2][box_index];
            let height = batch[3][box_index];
            let class_index = cls_indices[box_index];
            let confidence = batch[4 + class_index][box_index];
            let all_confidences = new Array(num_classes);
            for (let i = 0; i < num_classes; i++) {
                all_confidences[i] = batch[4 + i][box_index];
            }
            bounding_boxes.push({
                x,
                y,
                width,
                height,
                class_index,
                confidence,
                all_confidences,
            });
        }
        result.push(bounding_boxes);
    }
    return result;
}
/**
 * Sync version of `decodeBox`.
 */
function decodeBoxSync(args) {
    let { tf, num_classes, maxOutputSize, iouThreshold, scoreThreshold } = args;
    let length = 4 + num_classes;
    // e.g. 1x84x8400
    let batches = args.output;
    if (batches[0].length === 0) {
        // no a single batch
        return [];
    }
    if (batches[0].length !== length) {
        throw new Error(`data[batch].length must be ${length}`);
    }
    let num_boxes = batches[0][0].length;
    let result = [];
    for (let batch of batches) {
        // e.g. 84x8400
        let boxes = [];
        let scores = [];
        let cls_indices = [];
        for (let box_index = 0; box_index < num_boxes; box_index++) {
            let x = batch[0][box_index];
            let y = batch[1][box_index];
            let width = batch[2][box_index];
            let height = batch[3][box_index];
            let x1 = x - width / 2;
            let y1 = y - height / 2;
            let x2 = x + width / 2;
            let y2 = y + height / 2;
            let box_score = batch[4][box_index];
            let cls_index = 0;
            for (let i = 1; i < num_classes; i++) {
                let cls_score = batch[4 + i][box_index];
                if (cls_score > box_score) {
                    box_score = cls_score;
                    cls_index = i;
                }
            }
            boxes.push([x1, y1, x2, y2]);
            scores.push(box_score);
            cls_indices.push(cls_index);
        }
        let box_indices;
        if (maxOutputSize) {
            box_indices = tf.tidy(() => tf.image
                .nonMaxSuppression(boxes, scores, maxOutputSize, iouThreshold, scoreThreshold)
                .arraySync());
        }
        else {
            box_indices = Array.from({ length: num_boxes }, (_, i) => i);
        }
        let bounding_boxes = [];
        for (let box_index of box_indices) {
            let x = batch[0][box_index];
            let y = batch[1][box_index];
            let width = batch[2][box_index];
            let height = batch[3][box_index];
            let class_index = cls_indices[box_index];
            let confidence = batch[4 + class_index][box_index];
            let all_confidences = new Array(num_classes);
            for (let i = 0; i < num_classes; i++) {
                all_confidences[i] = batch[4 + i][box_index];
            }
            bounding_boxes.push({
                x,
                y,
                width,
                height,
                class_index,
                confidence,
                all_confidences,
            });
        }
        result.push(bounding_boxes);
    }
    return result;
}