Function transform

tfjs-backend-webgpu/src/kernels/Transform.ts:23–65

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

import {TransformProgram} from '../transform_webgpu';

export function transform(args: {
  inputs: TransformInputs,
  backend: WebGPUBackend,
  attrs: TransformAttrs
}): TensorInfo {
  const {inputs, backend, attrs} = 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] as [number, number, number, number];

  const program = new TransformProgram(outShape);
  const interpolationModeId = interpolation === 'nearest' ? 1 : 2;
  let fillModeId: number;
  switch (fillMode) {
    case 'constant':
      fillModeId = 1;
      break;
    case 'reflect':
      fillModeId = 2;
      break;
    case 'wrap':
      fillModeId = 3;
      break;
    case 'nearest':
      fillModeId = 4;
      break;
    default:
      fillModeId = 1;
      break;
  }
  const uniformData = [
    {type: 'int32', data: [interpolationModeId]},
    {type: 'int32', data: [fillModeId]}, {type: 'float32', data: [fillValue]}
  ];
  return backend.runWebGPUProgram(
      program, [image, transforms], 'float32', uniformData);
}

export const transformConfig: KernelConfig = {
  kernelName: Transform,