@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
51 lines • 8.04 kB
JavaScript
/**
* @license
* Copyright 2020 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 { NonMaxSuppressionV4 } from '@tensorflow/tfjs-core';
import { parseResultStruct } from './NonMaxSuppression_util';
let wasmFunc;
function setup(backend) {
wasmFunc = backend.wasm.cwrap(NonMaxSuppressionV4, 'number', // Result*
[
'number',
'number',
'number',
'number',
'number',
'bool', // padToMaxOutputSize
]);
}
function nonMaxSuppressionV4(args) {
const { backend, inputs, attrs } = args;
const { iouThreshold, maxOutputSize, scoreThreshold, padToMaxOutputSize } = 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, padToMaxOutputSize);
const { pSelectedIndices, selectedSize, pSelectedScores, pValidOutputs } = parseResultStruct(backend, resOffset);
// Since we are not using scores for V4, we have to delete it from the heap.
backend.wasm._free(pSelectedScores);
const selectedIndicesTensor = backend.makeOutput([selectedSize], 'int32', pSelectedIndices);
const validOutputsTensor = backend.makeOutput([], 'int32', pValidOutputs);
return [selectedIndicesTensor, validOutputsTensor];
}
export const nonMaxSuppressionV4Config = {
kernelName: NonMaxSuppressionV4,
backendName: 'wasm',
setupFunc: setup,
kernelFunc: nonMaxSuppressionV4,
};
//# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoiTm9uTWF4U3VwcHJlc3Npb25WNC5qcyIsInNvdXJjZVJvb3QiOiIiLCJzb3VyY2VzIjpbIi4uLy4uLy4uLy4uLy4uLy4uL3RmanMtYmFja2VuZC13YXNtL3NyYy9rZXJuZWxzL05vbk1heFN1cHByZXNzaW9uVjQudHMiXSwibmFtZXMiOltdLCJtYXBwaW5ncyI6IkFBQUE7Ozs7Ozs7Ozs7Ozs7OztHQWVHO0FBRUgsT0FBTyxFQUEyQixtQkFBbUIsRUFBa0UsTUFBTSx1QkFBdUIsQ0FBQztBQUlySixPQUFPLEVBQUMsaUJBQWlCLEVBQUMsTUFBTSwwQkFBMEIsQ0FBQztBQUUzRCxJQUFJLFFBR3NDLENBQUM7QUFFM0MsU0FBUyxLQUFLLENBQUMsT0FBb0I7SUFDakMsUUFBUSxHQUFHLE9BQU8sQ0FBQyxJQUFJLENBQUMsS0FBSyxDQUN6QixtQkFBbUIsRUFDbkIsUUFBUSxFQUFHLFVBQVU7SUFDckI7UUFDRSxRQUFRO1FBQ1IsUUFBUTtRQUNSLFFBQVE7UUFDUixRQUFRO1FBQ1IsUUFBUTtRQUNSLE1BQU0sRUFBSyxxQkFBcUI7S0FDakMsQ0FBQyxDQUFDO0FBQ1QsQ0FBQztBQUVELFNBQVMsbUJBQW1CLENBQUMsSUFJNUI7SUFDQyxNQUFNLEVBQUMsT0FBTyxFQUFFLE1BQU0sRUFBRSxLQUFLLEVBQUMsR0FBRyxJQUFJLENBQUM7SUFDdEMsTUFBTSxFQUFDLFlBQVksRUFBRSxhQUFhLEVBQUUsY0FBYyxFQUFFLGtCQUFrQixFQUFDLEdBQ25FLEtBQUssQ0FBQztJQUNWLE1BQU0sRUFBQyxLQUFLLEVBQUUsTUFBTSxFQUFDLEdBQUcsTUFBTSxDQUFDO0lBRS9CLE1BQU0sT0FBTyxHQUFHLE9BQU8sQ0FBQyxTQUFTLENBQUMsR0FBRyxDQUFDLEtBQUssQ0FBQyxNQUFNLENBQUMsQ0FBQyxFQUFFLENBQUM7SUFDdkQsTUFBTSxRQUFRLEdBQUcsT0FBTyxDQUFDLFNBQVMsQ0FBQyxHQUFHLENBQUMsTUFBTSxDQUFDLE1BQU0sQ0FBQyxDQUFDLEVBQUUsQ0FBQztJQUV6RCxNQUFNLFNBQVMsR0FBRyxRQUFRLENBQ3RCLE9BQU8sRUFBRSxRQUFRLEVBQUUsYUFBYSxFQUFFLFlBQVksRUFBRSxjQUFjLEVBQzlELGtCQUFrQixDQUFDLENBQUM7SUFFeEIsTUFBTSxFQUFDLGdCQUFnQixFQUFFLFlBQVksRUFBRSxlQUFlLEVBQUUsYUFBYSxFQUFDLEdBQ2xFLGlCQUFpQixDQUFDLE9BQU8sRUFBRSxTQUFTLENBQUMsQ0FBQztJQUUxQyw0RUFBNEU7SUFDNUUsT0FBTyxDQUFDLElBQUksQ0FBQyxLQUFLLENBQUMsZUFBZSxDQUFDLENBQUM7SUFFcEMsTUFBTSxxQkFBcUIsR0FDdkIsT0FBTyxDQUFDLFVBQVUsQ0FBQyxDQUFDLFlBQVksQ0FBQyxFQUFFLE9BQU8sRUFBRSxnQkFBZ0IsQ0FBQyxDQUFDO0lBRWxFLE1BQU0sa0JBQWtCLEdBQUcsT0FBTyxDQUFDLFVBQVUsQ0FBQyxFQUFFLEVBQUUsT0FBTyxFQUFFLGFBQWEsQ0FBQyxDQUFDO0lBRTFFLE9BQU8sQ0FBQyxxQkFBcUIsRUFBRSxrQkFBa0IsQ0FBQyxDQUFDO0FBQ3JELENBQUM7QUFFRCxNQUFNLENBQUMsTUFBTSx5QkFBeUIsR0FBaUI7SUFDckQsVUFBVSxFQUFFLG1CQUFtQjtJQUMvQixXQUFXLEVBQUUsTUFBTTtJQUNuQixTQUFTLEVBQUUsS0FBSztJQUNoQixVQUFVLEVBQUUsbUJBQTRDO0NBQ3pELENBQUMiLCJzb3VyY2VzQ29udGVudCI6WyIvKipcbiAqIEBsaWNlbnNlXG4gKiBDb3B5cmlnaHQgMjAyMCBHb29nbGUgTExDLiBBbGwgUmlnaHRzIFJlc2VydmVkLlxuICogTGljZW5zZWQgdW5kZXIgdGhlIEFwYWNoZSBMaWNlbnNlLCBWZXJzaW9uIDIuMCAodGhlIFwiTGljZW5zZVwiKTtcbiAqIHlvdSBtYXkgbm90IHVzZSB0aGlzIGZpbGUgZXhjZXB0IGluIGNvbXBsaWFuY2Ugd2l0aCB0aGUgTGljZW5zZS5cbiAqIFlvdSBtYXkgb2J0YWluIGEgY29weSBvZiB0aGUgTGljZW5zZSBhdFxuICpcbiAqIGh0dHA6Ly93d3cuYXBhY2hlLm9yZy9saWNlbnNlcy9MSUNFTlNFLTIuMFxuICpcbiAqIFVubGVzcyByZXF1aXJlZCBieSBhcHBsaWNhYmxlIGxhdyBvciBhZ3JlZWQgdG8gaW4gd3JpdGluZywgc29mdHdhcmVcbiAqIGRpc3RyaWJ1dGVkIHVuZGVyIHRoZSBMaWNlbnNlIGlzIGRpc3RyaWJ1dGVkIG9uIGFuIFwiQVMgSVNcIiBCQVNJUyxcbiAqIFdJVEhPVVQgV0FSUkFOVElFUyBPUiBDT05ESVRJT05TIE9GIEFOWSBLSU5ELCBlaXRoZXIgZXhwcmVzcyBvciBpbXBsaWVkLlxuICogU2VlIHRoZSBMaWNlbnNlIGZvciB0aGUgc3BlY2lmaWMgbGFuZ3VhZ2UgZ292ZXJuaW5nIHBlcm1pc3Npb25zIGFuZFxuICogbGltaXRhdGlvbnMgdW5kZXIgdGhlIExpY2Vuc2UuXG4gKiA9PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PVxuICovXG5cbmltcG9ydCB7S2VybmVsQ29uZmlnLCBLZXJuZWxGdW5jLCBOb25NYXhTdXBwcmVzc2lvblY0LCBOb25NYXhTdXBwcmVzc2lvblY0QXR0cnMsIE5vbk1heFN1cHByZXNzaW9uVjRJbnB1dHMsIFRlbnNvckluZm99IGZyb20gJ0B0ZW5zb3JmbG93L3RmanMtY29yZSc7XG5cbmltcG9ydCB7QmFja2VuZFdhc219IGZyb20gJy4uL2JhY2tlbmRfd2FzbSc7XG5cbmltcG9ydCB7cGFyc2VSZXN1bHRTdHJ1Y3R9IGZyb20gJy4vTm9uTWF4U3VwcHJlc3Npb25fdXRpbCc7XG5cbmxldCB3YXNtRnVuYzogKFxuICAgIGJveGVzSWQ6IG51bWJlciwgc2NvcmVzSWQ6IG51bWJlciwgbWF4T3V0cHV0U2l6ZTogbnVtYmVyLFxuICAgIGlvdVRocmVzaG9sZDogbnVtYmVyLCBzY29yZVRocmVzaG9sZDogbnVtYmVyLFxuICAgIHBhZFRvTWF4T3V0cHV0U2l6ZTogYm9vbGVhbikgPT4gbnVtYmVyO1xuXG5mdW5jdGlvbiBzZXR1cChiYWNrZW5kOiBCYWNrZW5kV2FzbSk6IHZvaWQge1xuICB3YXNtRnVuYyA9IGJhY2tlbmQud2FzbS5jd3JhcChcbiAgICAgIE5vbk1heFN1cHByZXNzaW9uVjQsXG4gICAgICAnbnVtYmVyJywgIC8vIFJlc3VsdCpcbiAgICAgIFtcbiAgICAgICAgJ251bWJlcicsICAvLyBib3hlc0lkXG4gICAgICAgICdudW1iZXInLCAgLy8gc2NvcmVzSWRcbiAgICAgICAgJ251bWJlcicsICAvLyBtYXhPdXRwdXRTaXplXG4gICAgICAgICdudW1iZXInLCAgLy8gaW91VGhyZXNob2xkXG4gICAgICAgICdudW1iZXInLCAgLy8gc2NvcmVUaHJlc2hvbGRcbiAgICAgICAgJ2Jvb2wnLCAgICAvLyBwYWRUb01heE91dHB1dFNpemVcbiAgICAgIF0pO1xufVxuXG5mdW5jdGlvbiBub25NYXhTdXBwcmVzc2lvblY0KGFyZ3M6IHtcbiAgYmFja2VuZDogQmFja2VuZFdhc20sXG4gIGlucHV0czogTm9uTWF4U3VwcHJlc3Npb25WNElucHV0cyxcbiAgYXR0cnM6IE5vbk1heFN1cHByZXNzaW9uVjRBdHRyc1xufSk6IFRlbnNvckluZm9bXSB7XG4gIGNvbnN0IHtiYWNrZW5kLCBpbnB1dHMsIGF0dHJzfSA9IGFyZ3M7XG4gIGNvbnN0IHtpb3VUaHJlc2hvbGQsIG1heE91dHB1dFNpemUsIHNjb3JlVGhyZXNob2xkLCBwYWRUb01heE91dHB1dFNpemV9ID1cbiAgICAgIGF0dHJzO1xuICBjb25zdCB7Ym94ZXMsIHNjb3Jlc30gPSBpbnB1dHM7XG5cbiAgY29uc3QgYm94ZXNJZCA9IGJhY2tlbmQuZGF0YUlkTWFwLmdldChib3hlcy5kYXRhSWQpLmlkO1xuICBjb25zdCBzY29yZXNJZCA9IGJhY2tlbmQuZGF0YUlkTWFwLmdldChzY29yZXMuZGF0YUlkKS5pZDtcblxuICBjb25zdCByZXNPZmZzZXQgPSB3YXNtRnVuYyhcbiAgICAgIGJveGVzSWQsIHNjb3Jlc0lkLCBtYXhPdXRwdXRTaXplLCBpb3VUaHJlc2hvbGQsIHNjb3JlVGhyZXNob2xkLFxuICAgICAgcGFkVG9NYXhPdXRwdXRTaXplKTtcblxuICBjb25zdCB7cFNlbGVjdGVkSW5kaWNlcywgc2VsZWN0ZWRTaXplLCBwU2VsZWN0ZWRTY29yZXMsIHBWYWxpZE91dHB1dHN9ID1cbiAgICAgIHBhcnNlUmVzdWx0U3RydWN0KGJhY2tlbmQsIHJlc09mZnNldCk7XG5cbiAgLy8gU2luY2Ugd2UgYXJlIG5vdCB1c2luZyBzY29yZXMgZm9yIFY0LCB3ZSBoYXZlIHRvIGRlbGV0ZSBpdCBmcm9tIHRoZSBoZWFwLlxuICBiYWNrZW5kLndhc20uX2ZyZWUocFNlbGVjdGVkU2NvcmVzKTtcblxuICBjb25zdCBzZWxlY3RlZEluZGljZXNUZW5zb3IgPVxuICAgICAgYmFja2VuZC5tYWtlT3V0cHV0KFtzZWxlY3RlZFNpemVdLCAnaW50MzInLCBwU2VsZWN0ZWRJbmRpY2VzKTtcblxuICBjb25zdCB2YWxpZE91dHB1dHNUZW5zb3IgPSBiYWNrZW5kLm1ha2VPdXRwdXQoW10sICdpbnQzMicsIHBWYWxpZE91dHB1dHMpO1xuXG4gIHJldHVybiBbc2VsZWN0ZWRJbmRpY2VzVGVuc29yLCB2YWxpZE91dHB1dHNUZW5zb3JdO1xufVxuXG5leHBvcnQgY29uc3Qgbm9uTWF4U3VwcHJlc3Npb25WNENvbmZpZzogS2VybmVsQ29uZmlnID0ge1xuICBrZXJuZWxOYW1lOiBOb25NYXhTdXBwcmVzc2lvblY0LFxuICBiYWNrZW5kTmFtZTogJ3dhc20nLFxuICBzZXR1cEZ1bmM6IHNldHVwLFxuICBrZXJuZWxGdW5jOiBub25NYXhTdXBwcmVzc2lvblY0IGFzIHVua25vd24gYXMgS2VybmVsRnVuYyxcbn07XG4iXX0=