39 lines
1.6 KiB
JavaScript
39 lines
1.6 KiB
JavaScript
import * as tf from '@tensorflow/tfjs-core';
|
|
/**
|
|
* Pads the smaller dimension of an image tensor with zeros, such that width === height.
|
|
*
|
|
* @param imgTensor The image tensor.
|
|
* @param isCenterImage (optional, default: false) If true, add an equal amount of padding on
|
|
* both sides of the minor dimension oof the image.
|
|
* @returns The padded tensor with width === height.
|
|
*/
|
|
export function padToSquare(imgTensor, isCenterImage = false) {
|
|
return tf.tidy(() => {
|
|
const [height, width] = imgTensor.shape.slice(1);
|
|
if (height === width) {
|
|
return imgTensor;
|
|
}
|
|
const dimDiff = Math.abs(height - width);
|
|
const paddingAmount = Math.round(dimDiff * (isCenterImage ? 0.5 : 1));
|
|
const paddingAxis = height > width ? 2 : 1;
|
|
const createPaddingTensor = (paddingAmount) => {
|
|
const paddingTensorShape = imgTensor.shape.slice();
|
|
paddingTensorShape[paddingAxis] = paddingAmount;
|
|
return tf.fill(paddingTensorShape, 0);
|
|
};
|
|
const paddingTensorAppend = createPaddingTensor(paddingAmount);
|
|
const remainingPaddingAmount = dimDiff - paddingTensorAppend.shape[paddingAxis];
|
|
const paddingTensorPrepend = isCenterImage && remainingPaddingAmount
|
|
? createPaddingTensor(remainingPaddingAmount)
|
|
: null;
|
|
const tensorsToStack = [
|
|
paddingTensorPrepend,
|
|
imgTensor,
|
|
paddingTensorAppend
|
|
]
|
|
.filter(t => !!t)
|
|
.map((t) => t.toFloat());
|
|
return tf.concat(tensorsToStack, paddingAxis);
|
|
});
|
|
}
|
|
//# sourceMappingURL=padToSquare.js.map
|