extractFaceTensors.js 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839
  1. import { __awaiter, __generator } from "tslib";
  2. import * as tf from '@tensorflow/tfjs-core';
  3. import { FaceDetection } from '../classes/FaceDetection';
  4. import { isTensor3D, isTensor4D } from '../utils';
  5. /**
  6. * Extracts the tensors of the image regions containing the detected faces.
  7. * Useful if you want to compute the face descriptors for the face images.
  8. * Using this method is faster then extracting a canvas for each face and
  9. * converting them to tensors individually.
  10. *
  11. * @param imageTensor The image tensor that face detection has been performed on.
  12. * @param detections The face detection results or face bounding boxes for that image.
  13. * @returns Tensors of the corresponding image region for each detected face.
  14. */
  15. export function extractFaceTensors(imageTensor, detections) {
  16. return __awaiter(this, void 0, void 0, function () {
  17. return __generator(this, function (_a) {
  18. if (!isTensor3D(imageTensor) && !isTensor4D(imageTensor)) {
  19. throw new Error('extractFaceTensors - expected image tensor to be 3D or 4D');
  20. }
  21. if (isTensor4D(imageTensor) && imageTensor.shape[0] > 1) {
  22. throw new Error('extractFaceTensors - batchSize > 1 not supported');
  23. }
  24. return [2 /*return*/, tf.tidy(function () {
  25. var _a = imageTensor.shape.slice(isTensor4D(imageTensor) ? 1 : 0), imgHeight = _a[0], imgWidth = _a[1], numChannels = _a[2];
  26. var boxes = detections.map(function (det) { return det instanceof FaceDetection
  27. ? det.forSize(imgWidth, imgHeight).box
  28. : det; })
  29. .map(function (box) { return box.clipAtImageBorders(imgWidth, imgHeight); });
  30. var faceTensors = boxes.map(function (_a) {
  31. var x = _a.x, y = _a.y, width = _a.width, height = _a.height;
  32. return tf.slice3d(imageTensor.as3D(imgHeight, imgWidth, numChannels), [y, x, 0], [height, width, numChannels]);
  33. });
  34. return faceTensors;
  35. })];
  36. });
  37. });
  38. }
  39. //# sourceMappingURL=extractFaceTensors.js.map