@tensorflow/tfjs-backend-wasm
Version:
This package adds a WebAssembly backend to TensorFlow.js. It currently supports the following models from our [models](https://github.com/tensorflow/tfjs-models) repo: - BlazeFace - BodyPix - CocoSSD - Face landmarks detection - HandPose - KNN classifier
52 lines • 8.04 kB
JavaScript
/**
* @license
* Copyright 2019 Google LLC. All Rights Reserved.
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
* =============================================================================
*/
import { NonMaxSuppressionV5 } from '@tensorflow/tfjs-core';
import { parseResultStruct } from './NonMaxSuppression_util';
let wasmFunc;
function setup(backend) {
wasmFunc = backend.wasm.cwrap(NonMaxSuppressionV5, 'number', // Result*
[
'number',
'number',
'number',
'number',
'number',
'number', // softNmsSigma
]);
}
function kernelFunc(args) {
const { backend, inputs, attrs } = args;
const { iouThreshold, maxOutputSize, scoreThreshold, softNmsSigma } = attrs;
const { boxes, scores } = inputs;
const boxesId = backend.dataIdMap.get(boxes.dataId).id;
const scoresId = backend.dataIdMap.get(scores.dataId).id;
const resOffset = wasmFunc(boxesId, scoresId, maxOutputSize, iouThreshold, scoreThreshold, softNmsSigma);
const { pSelectedIndices, selectedSize, pSelectedScores, pValidOutputs } = parseResultStruct(backend, resOffset);
// Since we are not using validOutputs for V5, we have to delete it from the
// heap.
backend.wasm._free(pValidOutputs);
const selectedIndicesTensor = backend.makeOutput([selectedSize], 'int32', pSelectedIndices);
const selectedScoresTensor = backend.makeOutput([selectedSize], 'float32', pSelectedScores);
return [selectedIndicesTensor, selectedScoresTensor];
}
export const nonMaxSuppressionV5Config = {
kernelName: NonMaxSuppressionV5,
backendName: 'wasm',
setupFunc: setup,
kernelFunc: kernelFunc,
};
//# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoiTm9uTWF4U3VwcHJlc3Npb25WNS5qcyIsInNvdXJjZVJvb3QiOiIiLCJzb3VyY2VzIjpbIi4uLy4uLy4uLy4uLy4uLy4uL3RmanMtYmFja2VuZC13YXNtL3NyYy9rZXJuZWxzL05vbk1heFN1cHByZXNzaW9uVjUudHMiXSwibmFtZXMiOltdLCJtYXBwaW5ncyI6IkFBQUE7Ozs7Ozs7Ozs7Ozs7OztHQWVHO0FBRUgsT0FBTyxFQUEyQixtQkFBbUIsRUFBa0UsTUFBTSx1QkFBdUIsQ0FBQztBQUlySixPQUFPLEVBQUMsaUJBQWlCLEVBQUMsTUFBTSwwQkFBMEIsQ0FBQztBQUUzRCxJQUFJLFFBR1UsQ0FBQztBQUVmLFNBQVMsS0FBSyxDQUFDLE9BQW9CO0lBQ2pDLFFBQVEsR0FBRyxPQUFPLENBQUMsSUFBSSxDQUFDLEtBQUssQ0FDekIsbUJBQW1CLEVBQ25CLFFBQVEsRUFBRyxVQUFVO0lBQ3JCO1FBQ0UsUUFBUTtRQUNSLFFBQVE7UUFDUixRQUFRO1FBQ1IsUUFBUTtRQUNSLFFBQVE7UUFDUixRQUFRLEVBQUcsZUFBZTtLQUMzQixDQUFDLENBQUM7QUFDVCxDQUFDO0FBRUQsU0FBUyxVQUFVLENBQUMsSUFJbkI7SUFDQyxNQUFNLEVBQUMsT0FBTyxFQUFFLE1BQU0sRUFBRSxLQUFLLEVBQUMsR0FBRyxJQUFJLENBQUM7SUFDdEMsTUFBTSxFQUFDLFlBQVksRUFBRSxhQUFhLEVBQUUsY0FBYyxFQUFFLFlBQVksRUFBQyxHQUFHLEtBQUssQ0FBQztJQUMxRSxNQUFNLEVBQUMsS0FBSyxFQUFFLE1BQU0sRUFBQyxHQUFHLE1BQU0sQ0FBQztJQUUvQixNQUFNLE9BQU8sR0FBRyxPQUFPLENBQUMsU0FBUyxDQUFDLEdBQUcsQ0FBQyxLQUFLLENBQUMsTUFBTSxDQUFDLENBQUMsRUFBRSxDQUFDO0lBQ3ZELE1BQU0sUUFBUSxHQUFHLE9BQU8sQ0FBQyxTQUFTLENBQUMsR0FBRyxDQUFDLE1BQU0sQ0FBQyxNQUFNLENBQUMsQ0FBQyxFQUFFLENBQUM7SUFFekQsTUFBTSxTQUFTLEdBQUcsUUFBUSxDQUN0QixPQUFPLEVBQUUsUUFBUSxFQUFFLGFBQWEsRUFBRSxZQUFZLEVBQUUsY0FBYyxFQUM5RCxZQUFZLENBQUMsQ0FBQztJQUVsQixNQUFNLEVBQUMsZ0JBQWdCLEVBQUUsWUFBWSxFQUFFLGVBQWUsRUFBRSxhQUFhLEVBQUMsR0FDbEUsaUJBQWlCLENBQUMsT0FBTyxFQUFFLFNBQVMsQ0FBQyxDQUFDO0lBRTFDLDRFQUE0RTtJQUM1RSxRQUFRO0lBQ1IsT0FBTyxDQUFDLElBQUksQ0FBQyxLQUFLLENBQUMsYUFBYSxDQUFDLENBQUM7SUFFbEMsTUFBTSxxQkFBcUIsR0FDdkIsT0FBTyxDQUFDLFVBQVUsQ0FBQyxDQUFDLFlBQVksQ0FBQyxFQUFFLE9BQU8sRUFBRSxnQkFBZ0IsQ0FBQyxDQUFDO0lBQ2xFLE1BQU0sb0JBQW9CLEdBQ3RCLE9BQU8sQ0FBQyxVQUFVLENBQUMsQ0FBQyxZQUFZLENBQUMsRUFBRSxTQUFTLEVBQUUsZUFBZSxDQUFDLENBQUM7SUFFbkUsT0FBTyxDQUFDLHFCQUFxQixFQUFFLG9CQUFvQixDQUFDLENBQUM7QUFDdkQsQ0FBQztBQUVELE1BQU0sQ0FBQyxNQUFNLHlCQUF5QixHQUFpQjtJQUNyRCxVQUFVLEVBQUUsbUJBQW1CO0lBQy9CLFdBQVcsRUFBRSxNQUFNO0lBQ25CLFNBQVMsRUFBRSxLQUFLO0lBQ2hCLFVBQVUsRUFBRSxVQUFtQztDQUNoRCxDQUFDIiwic291cmNlc0NvbnRlbnQiOlsiLyoqXG4gKiBAbGljZW5zZVxuICogQ29weXJpZ2h0IDIwMTkgR29vZ2xlIExMQy4gQWxsIFJpZ2h0cyBSZXNlcnZlZC5cbiAqIExpY2Vuc2VkIHVuZGVyIHRoZSBBcGFjaGUgTGljZW5zZSwgVmVyc2lvbiAyLjAgKHRoZSBcIkxpY2Vuc2VcIik7XG4gKiB5b3UgbWF5IG5vdCB1c2UgdGhpcyBmaWxlIGV4Y2VwdCBpbiBjb21wbGlhbmNlIHdpdGggdGhlIExpY2Vuc2UuXG4gKiBZb3UgbWF5IG9idGFpbiBhIGNvcHkgb2YgdGhlIExpY2Vuc2UgYXRcbiAqXG4gKiBodHRwOi8vd3d3LmFwYWNoZS5vcmcvbGljZW5zZXMvTElDRU5TRS0yLjBcbiAqXG4gKiBVbmxlc3MgcmVxdWlyZWQgYnkgYXBwbGljYWJsZSBsYXcgb3IgYWdyZWVkIHRvIGluIHdyaXRpbmcsIHNvZnR3YXJlXG4gKiBkaXN0cmlidXRlZCB1bmRlciB0aGUgTGljZW5zZSBpcyBkaXN0cmlidXRlZCBvbiBhbiBcIkFTIElTXCIgQkFTSVMsXG4gKiBXSVRIT1VUIFdBUlJBTlRJRVMgT1IgQ09ORElUSU9OUyBPRiBBTlkgS0lORCwgZWl0aGVyIGV4cHJlc3Mgb3IgaW1wbGllZC5cbiAqIFNlZSB0aGUgTGljZW5zZSBmb3IgdGhlIHNwZWNpZmljIGxhbmd1YWdlIGdvdmVybmluZyBwZXJtaXNzaW9ucyBhbmRcbiAqIGxpbWl0YXRpb25zIHVuZGVyIHRoZSBMaWNlbnNlLlxuICogPT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT1cbiAqL1xuXG5pbXBvcnQge0tlcm5lbENvbmZpZywgS2VybmVsRnVuYywgTm9uTWF4U3VwcHJlc3Npb25WNSwgTm9uTWF4U3VwcHJlc3Npb25WNUF0dHJzLCBOb25NYXhTdXBwcmVzc2lvblY1SW5wdXRzLCBUZW5zb3JJbmZvfSBmcm9tICdAdGVuc29yZmxvdy90ZmpzLWNvcmUnO1xuXG5pbXBvcnQge0JhY2tlbmRXYXNtfSBmcm9tICcuLi9iYWNrZW5kX3dhc20nO1xuXG5pbXBvcnQge3BhcnNlUmVzdWx0U3RydWN0fSBmcm9tICcuL05vbk1heFN1cHByZXNzaW9uX3V0aWwnO1xuXG5sZXQgd2FzbUZ1bmM6XG4gICAgKGJveGVzSWQ6IG51bWJlciwgc2NvcmVzSWQ6IG51bWJlciwgbWF4T3V0cHV0U2l6ZTogbnVtYmVyLFxuICAgICBpb3VUaHJlc2hvbGQ6IG51bWJlciwgc2NvcmVUaHJlc2hvbGQ6IG51bWJlciwgc29mdE5tc1NpZ21hOiBudW1iZXIpID0+XG4gICAgICAgIG51bWJlcjtcblxuZnVuY3Rpb24gc2V0dXAoYmFja2VuZDogQmFja2VuZFdhc20pOiB2b2lkIHtcbiAgd2FzbUZ1bmMgPSBiYWNrZW5kLndhc20uY3dyYXAoXG4gICAgICBOb25NYXhTdXBwcmVzc2lvblY1LFxuICAgICAgJ251bWJlcicsICAvLyBSZXN1bHQqXG4gICAgICBbXG4gICAgICAgICdudW1iZXInLCAgLy8gYm94ZXNJZFxuICAgICAgICAnbnVtYmVyJywgIC8vIHNjb3Jlc0lkXG4gICAgICAgICdudW1iZXInLCAgLy8gbWF4T3V0cHV0U2l6ZVxuICAgICAgICAnbnVtYmVyJywgIC8vIGlvdVRocmVzaG9sZFxuICAgICAgICAnbnVtYmVyJywgIC8vIHNjb3JlVGhyZXNob2xkXG4gICAgICAgICdudW1iZXInLCAgLy8gc29mdE5tc1NpZ21hXG4gICAgICBdKTtcbn1cblxuZnVuY3Rpb24ga2VybmVsRnVuYyhhcmdzOiB7XG4gIGJhY2tlbmQ6IEJhY2tlbmRXYXNtLFxuICBpbnB1dHM6IE5vbk1heFN1cHByZXNzaW9uVjVJbnB1dHMsXG4gIGF0dHJzOiBOb25NYXhTdXBwcmVzc2lvblY1QXR0cnNcbn0pOiBUZW5zb3JJbmZvW10ge1xuICBjb25zdCB7YmFja2VuZCwgaW5wdXRzLCBhdHRyc30gPSBhcmdzO1xuICBjb25zdCB7aW91VGhyZXNob2xkLCBtYXhPdXRwdXRTaXplLCBzY29yZVRocmVzaG9sZCwgc29mdE5tc1NpZ21hfSA9IGF0dHJzO1xuICBjb25zdCB7Ym94ZXMsIHNjb3Jlc30gPSBpbnB1dHM7XG5cbiAgY29uc3QgYm94ZXNJZCA9IGJhY2tlbmQuZGF0YUlkTWFwLmdldChib3hlcy5kYXRhSWQpLmlkO1xuICBjb25zdCBzY29yZXNJZCA9IGJhY2tlbmQuZGF0YUlkTWFwLmdldChzY29yZXMuZGF0YUlkKS5pZDtcblxuICBjb25zdCByZXNPZmZzZXQgPSB3YXNtRnVuYyhcbiAgICAgIGJveGVzSWQsIHNjb3Jlc0lkLCBtYXhPdXRwdXRTaXplLCBpb3VUaHJlc2hvbGQsIHNjb3JlVGhyZXNob2xkLFxuICAgICAgc29mdE5tc1NpZ21hKTtcblxuICBjb25zdCB7cFNlbGVjdGVkSW5kaWNlcywgc2VsZWN0ZWRTaXplLCBwU2VsZWN0ZWRTY29yZXMsIHBWYWxpZE91dHB1dHN9ID1cbiAgICAgIHBhcnNlUmVzdWx0U3RydWN0KGJhY2tlbmQsIHJlc09mZnNldCk7XG5cbiAgLy8gU2luY2Ugd2UgYXJlIG5vdCB1c2luZyB2YWxpZE91dHB1dHMgZm9yIFY1LCB3ZSBoYXZlIHRvIGRlbGV0ZSBpdCBmcm9tIHRoZVxuICAvLyBoZWFwLlxuICBiYWNrZW5kLndhc20uX2ZyZWUocFZhbGlkT3V0cHV0cyk7XG5cbiAgY29uc3Qgc2VsZWN0ZWRJbmRpY2VzVGVuc29yID1cbiAgICAgIGJhY2tlbmQubWFrZU91dHB1dChbc2VsZWN0ZWRTaXplXSwgJ2ludDMyJywgcFNlbGVjdGVkSW5kaWNlcyk7XG4gIGNvbnN0IHNlbGVjdGVkU2NvcmVzVGVuc29yID1cbiAgICAgIGJhY2tlbmQubWFrZU91dHB1dChbc2VsZWN0ZWRTaXplXSwgJ2Zsb2F0MzInLCBwU2VsZWN0ZWRTY29yZXMpO1xuXG4gIHJldHVybiBbc2VsZWN0ZWRJbmRpY2VzVGVuc29yLCBzZWxlY3RlZFNjb3Jlc1RlbnNvcl07XG59XG5cbmV4cG9ydCBjb25zdCBub25NYXhTdXBwcmVzc2lvblY1Q29uZmlnOiBLZXJuZWxDb25maWcgPSB7XG4gIGtlcm5lbE5hbWU6IE5vbk1heFN1cHByZXNzaW9uVjUsXG4gIGJhY2tlbmROYW1lOiAnd2FzbScsXG4gIHNldHVwRnVuYzogc2V0dXAsXG4gIGtlcm5lbEZ1bmM6IGtlcm5lbEZ1bmMgYXMgdW5rbm93biBhcyBLZXJuZWxGdW5jLFxufTtcbiJdfQ==