padToSquare.js 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243
  1. "use strict";
  2. Object.defineProperty(exports, "__esModule", { value: true });
  3. var tf = require("@tensorflow/tfjs-core");
  4. /**
  5. * Pads the smaller dimension of an image tensor with zeros, such that width === height.
  6. *
  7. * @param imgTensor The image tensor.
  8. * @param isCenterImage (optional, default: false) If true, add an equal amount of padding on
  9. * both sides of the minor dimension oof the image.
  10. * @returns The padded tensor with width === height.
  11. */
  12. function padToSquare(imgTensor, isCenterImage) {
  13. if (isCenterImage === void 0) { isCenterImage = false; }
  14. return tf.tidy(function () {
  15. var _a = imgTensor.shape.slice(1), height = _a[0], width = _a[1];
  16. if (height === width) {
  17. return imgTensor;
  18. }
  19. var dimDiff = Math.abs(height - width);
  20. var paddingAmount = Math.round(dimDiff * (isCenterImage ? 0.5 : 1));
  21. var paddingAxis = height > width ? 2 : 1;
  22. var createPaddingTensor = function (paddingAmount) {
  23. var paddingTensorShape = imgTensor.shape.slice();
  24. paddingTensorShape[paddingAxis] = paddingAmount;
  25. return tf.fill(paddingTensorShape, 0);
  26. };
  27. var paddingTensorAppend = createPaddingTensor(paddingAmount);
  28. var remainingPaddingAmount = dimDiff - paddingTensorAppend.shape[paddingAxis];
  29. var paddingTensorPrepend = isCenterImage && remainingPaddingAmount
  30. ? createPaddingTensor(remainingPaddingAmount)
  31. : null;
  32. var tensorsToStack = [
  33. paddingTensorPrepend,
  34. imgTensor,
  35. paddingTensorAppend
  36. ]
  37. .filter(function (t) { return !!t; })
  38. .map(function (t) { return t.toFloat(); });
  39. return tf.concat(tensorsToStack, paddingAxis);
  40. });
  41. }
  42. exports.padToSquare = padToSquare;
  43. //# sourceMappingURL=padToSquare.js.map