yolo-helpers
Version:
Helper functions to use models converted from YOLO in browser and Node.js
283 lines (282 loc) • 10.5 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.decodeSegment = decodeSegment;
exports.decodeSegmentSync = decodeSegmentSync;
exports.combineMask = combineMask;
exports.hasOverlap = hasOverlap;
function getMaskShape(args) {
if ('mask_shape' in args)
return args.mask_shape;
if ('input_shape' in args) {
return {
width: args.input_shape.width / 4,
height: args.input_shape.height / 4,
};
}
throw new Error('missing mask_shape or input_shape');
}
/**
* tensorflow output: boxes [batch, features, channel] and masks [batch, height, width, channel]
*
* box features:
* - 4: x, y, width, height
* - num_classes: class confidence
* - 32: channel coefficients
*
* segmentation mask:
* - 0 for background, 1 for object
* - 32 channels, correspond to the 32 channel coefficients in the bounding box
*
* The x, y, width, height are in pixel unit, NOT normalized in the range of [0, 1].
* The the pixel units are scaled to the input_shape.
*
* The confidence are already normalized between 0 to 1.
*/
async function decodeSegment(args) {
let { tf, num_classes, maxOutputSize, iouThreshold, scoreThreshold } = args;
let num_channels = args.num_channels ?? 32;
let { width: mask_width, height: mask_height } = getMaskShape(args);
let boxes_length = 4 + num_classes + num_channels;
// e.g. 1x116x8400
let batches_boxes = args.output_boxes;
if (batches_boxes[0].length === 0) {
// no a single batch
return [];
}
if (batches_boxes[0].length !== boxes_length) {
throw new Error(`boxes_data[batch].length must be ${boxes_length}`);
}
let num_boxes = batches_boxes[0][0].length;
// e.g. 1x160x160x32
let batches_masks = args.output_masks;
if (batches_masks[0].length !== mask_height) {
throw new Error(`masks_data[batch].length must be ${mask_height}`);
}
if (batches_masks[0][0].length !== mask_width) {
throw new Error(`masks_data[batch][y].length must be ${mask_width}`);
}
if (batches_masks[0][0][0].length !== num_channels) {
throw new Error(`masks_data[batch][y][x].length must be ${num_channels}`);
}
if (batches_boxes.length !== batches_masks.length) {
throw new Error('boxes_data and masks_data must have the same length');
}
let result = [];
let batch_size = batches_boxes.length;
for (let batch = 0; batch < batch_size; batch++) {
// 116x8400
let batch_boxes = batches_boxes[batch];
// 160x160x32
let batch_masks = batches_masks[batch];
let boxes = [];
let cls_scores = [];
let cls_indices = [];
for (let box_index = 0; box_index < num_boxes; box_index++) {
let x = batch_boxes[0][box_index];
let y = batch_boxes[1][box_index];
let width = batch_boxes[2][box_index];
let height = batch_boxes[3][box_index];
let x1 = x - width / 2;
let y1 = y - height / 2;
let x2 = x + width / 2;
let y2 = y + height / 2;
let cls_score = batch_boxes[4][box_index];
let cls_index = 0;
for (let i = 1; i < num_classes; i++) {
let score = batch_boxes[4 + i][box_index];
if (score > cls_score) {
cls_score = score;
cls_index = i;
}
}
boxes.push([x1, y1, x2, y2]);
cls_scores.push(cls_score);
cls_indices.push(cls_index);
}
let box_indices;
if (maxOutputSize) {
let box_indices_tensor = await tf.image.nonMaxSuppressionAsync(boxes, cls_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_boxes[0][box_index];
let y = batch_boxes[1][box_index];
let width = batch_boxes[2][box_index];
let height = batch_boxes[3][box_index];
let class_index = cls_indices[box_index];
let confidence = batch_boxes[4 + class_index][box_index];
let all_confidences = new Array(num_classes);
for (let i = 0; i < num_classes; i++) {
all_confidences[i] = batch_boxes[4 + i][box_index];
}
let mask_coefficients = new Array(num_channels);
for (let i = 0; i < num_channels; i++) {
mask_coefficients[i] = batch_boxes[4 + num_classes + i][box_index];
}
bounding_boxes.push({
x,
y,
width,
height,
class_index,
confidence,
all_confidences,
mask_coefficients,
});
}
result.push({
bounding_boxes,
masks: batch_masks,
});
}
return result;
}
/**
* Sync version of `decodeSegment`.
*/
function decodeSegmentSync(args) {
let { tf, num_classes, maxOutputSize, iouThreshold, scoreThreshold } = args;
let num_channels = args.num_channels ?? 32;
let { width: mask_width, height: mask_height } = getMaskShape(args);
let boxes_length = 4 + num_classes + num_channels;
// e.g. 1x116x8400
let batches_boxes = args.output_boxes;
if (batches_boxes[0].length === 0) {
// no a single batch
return [];
}
if (batches_boxes[0].length !== boxes_length) {
throw new Error(`boxes_data[batch].length must be ${boxes_length}`);
}
let num_boxes = batches_boxes[0][0].length;
// e.g. 1x160x160x32
let batches_masks = args.output_masks;
if (batches_masks[0].length !== mask_height) {
throw new Error(`masks_data[batch].length must be ${mask_height}`);
}
if (batches_masks[0][0].length !== mask_width) {
throw new Error(`masks_data[batch][y].length must be ${mask_width}`);
}
if (batches_masks[0][0][0].length !== num_channels) {
throw new Error(`masks_data[batch][y][x].length must be ${num_channels}`);
}
if (batches_boxes.length !== batches_masks.length) {
throw new Error('boxes_data and masks_data must have the same length');
}
let result = [];
let batch_size = batches_boxes.length;
for (let batch = 0; batch < batch_size; batch++) {
// 116x8400
let batch_boxes = batches_boxes[batch];
// 160x160x32
let batch_masks = batches_masks[batch];
let boxes = [];
let cls_scores = [];
let cls_indices = [];
for (let box_index = 0; box_index < num_boxes; box_index++) {
let x = batch_boxes[0][box_index];
let y = batch_boxes[1][box_index];
let width = batch_boxes[2][box_index];
let height = batch_boxes[3][box_index];
let x1 = x - width / 2;
let y1 = y - height / 2;
let x2 = x + width / 2;
let y2 = y + height / 2;
let cls_score = batch_boxes[4][box_index];
let cls_index = 0;
for (let i = 1; i < num_classes; i++) {
let score = batch_boxes[4 + i][box_index];
if (score > cls_score) {
cls_score = score;
cls_index = i;
}
}
boxes.push([x1, y1, x2, y2]);
cls_scores.push(cls_score);
cls_indices.push(cls_index);
}
let box_indices;
if (maxOutputSize) {
box_indices = tf.tidy(() => {
let box_indices_tensor = tf.image.nonMaxSuppression(boxes, cls_scores, maxOutputSize, iouThreshold, scoreThreshold);
return box_indices_tensor.arraySync();
});
}
else {
box_indices = Array.from({ length: num_boxes }, (_, i) => i);
}
let bounding_boxes = [];
for (let box_index of box_indices) {
let x = batch_boxes[0][box_index];
let y = batch_boxes[1][box_index];
let width = batch_boxes[2][box_index];
let height = batch_boxes[3][box_index];
let class_index = cls_indices[box_index];
let confidence = batch_boxes[4 + class_index][box_index];
let all_confidences = new Array(num_classes);
for (let i = 0; i < num_classes; i++) {
all_confidences[i] = batch_boxes[4 + i][box_index];
}
let mask_coefficients = new Array(num_channels);
for (let i = 0; i < num_channels; i++) {
mask_coefficients[i] = batch_boxes[4 + num_classes + i][box_index];
}
bounding_boxes.push({
x,
y,
width,
height,
class_index,
confidence,
all_confidences,
mask_coefficients,
});
}
result.push({
bounding_boxes,
masks: batch_masks,
});
}
return result;
}
/**
* @description final mask = mask coefficients * mask channels
*/
function combineMask(bounding_box,
/** e.g. [mask_height, mask_width, 32] for 32 channels of masks */
masks) {
let mask_coefficients = bounding_box.mask_coefficients;
let mask_height = masks.length;
let mask_width = masks[0].length;
let num_channels = masks[0][0].length;
if (num_channels != mask_coefficients.length) {
throw new Error(`expect ${num_channels} mask coefficients, but got ${mask_coefficients.length}`);
}
let final_mask = new Array(mask_height);
for (let h = 0; h < mask_height; h++) {
final_mask[h] = new Array(mask_width);
for (let w = 0; w < mask_width; w++) {
let acc = 0;
for (let i = 0; i < num_channels; i++) {
acc += mask_coefficients[i] * masks[h][w][i];
}
acc = sigmoid(acc);
final_mask[h][w] = acc;
}
}
return final_mask;
}
function sigmoid(x) {
return 1 / (1 + Math.exp(-x));
}
function hasOverlap(a, b) {
return !(a.right < b.left ||
a.left > b.right ||
a.bottom < b.top ||
a.top > b.bottom);
}