Function transform

tfjs-backend-cpu/src/kernels/Transform.ts:22–116

(args: {
  inputs: TransformInputs,
  attrs: TransformAttrs,
  backend: MathBackendCPU
})

import {MathBackendCPU} from '../backend_cpu';

export function transform(args: {
  inputs: TransformInputs,
  attrs: TransformAttrs,
  backend: MathBackendCPU
}): TensorInfo {
  const {inputs, attrs, backend} = args;
  const {image, transforms} = inputs;
  const {interpolation, fillMode, fillValue, outputShape} = attrs;

  const [batch, imageHeight, imageWidth, numChannels] = image.shape;
  const [outHeight, outWidth] =
      outputShape != null ? outputShape : [imageHeight, imageWidth];
  const outShape = [batch, outHeight, outWidth, numChannels];

  const inStrides = util.computeStrides(image.shape);
  const batchInStride = inStrides[0];
  const rowInStride = inStrides[1];
  const colInStride = inStrides[2];

  const outStrides = util.computeStrides(outShape);
  const batchOutStride = outStrides[0];
  const rowOutStride = outStrides[1];
  const colOutStride = outStrides[2];

  const outVals = util.getTypedArrayFromDType(
      image.dtype as NumericDataType, util.sizeFromShape(outShape));

  outVals.fill(fillValue);

  const imageVals = backend.data.get(image.dataId).values as TypedArray;
  const transformVals =
      backend.data.get(transforms.dataId).values as TypedArray;

  // Ref TF implementation:
  // https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/image/image_ops.h
  for (let b = 0; b < batch; ++b) {
    const transform = transforms.shape[0] === 1 ?
        transformVals :
        transformVals.subarray(b * 8, b * 8 + 8);

    for (let outY = 0; outY < outHeight; ++outY) {
      for (let outX = 0; outX < outWidth; ++outX) {
        for (let channel = 0; channel < numChannels; ++channel) {
          let val;

          const projection = transform[6] * outX + transform[7] * outY + 1;

          if (projection === 0) {
            // Return the fill value for infinite coordinates,
            // which are outside the input image
            continue;
          }

          const inX =
              (transform[0] * outX + transform[1] * outY + transform[2]) /
              projection;
          const inY =
              (transform[3] * outX + transform[4] * outY + transform[5]) /