[WASM] Support integer inputs for bilinear resizing and crop and resizing.

tfjs-backend-wasm/src/kernels/Cast.ts

@@ -28,8 +28,9 @@ interface CastAttrs extends NamedAttrMap {
   dtype: DataType;
 }
 
-function cast(
-    args: {inputs: CastInputs, attrs: CastAttrs, backend: BackendWasm}) {
+export function cast(
+    args: {inputs: CastInputs, attrs: CastAttrs, backend: BackendWasm}):
+    TensorInfo {
   const {inputs: {x}, attrs: {dtype}, backend} = args;
   const out = backend.makeOutput(x.shape, dtype);
   const inVals = backend.typedArrayFromHeap(x);

tfjs-backend-wasm/src/kernels/CropAndResize.ts

@@ -18,6 +18,7 @@
 import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
+import {cast} from './Cast';
 
 interface CropAndResizeInputs extends NamedTensorInfoMap {
   images: TensorInfo;
@@ -71,11 +72,19 @@ function cropAndResize(args: {
   const [cropHeight, cropWidth] = cropSize as [number, number];
   const outShape = [numBoxes, cropHeight, cropWidth, images.shape[3]];
 
-  const imagesId = backend.dataIdMap.get(images.dataId).id;
+  let imagesData = backend.dataIdMap.get(images.dataId);
+  let castedData;
+  if (images.dtype !== 'float32') {
+    castedData =
+        cast({backend, inputs: {x: images}, attrs: {dtype: 'float32'}});
+    imagesData = backend.dataIdMap.get(castedData.dataId);
+  }
+
+  const imagesId = imagesData.id;
   const boxesId = backend.dataIdMap.get(boxes.dataId).id;
   const boxIndId = backend.dataIdMap.get(boxInd.dataId).id;
 
-  const out = backend.makeOutput(outShape, images.dtype);
+  const out = backend.makeOutput(outShape, 'float32');
   const outId = backend.dataIdMap.get(out.dataId).id;
 
   const imagesShapeBytes = new Uint8Array(new Int32Array(images.shape).buffer);
@@ -85,6 +94,11 @@ function cropAndResize(args: {
       cropWidth,
       InterpolationMethod[method as {} as keyof typeof InterpolationMethod],
       extrapolationValue, outId);
+
+  if (castedData != null) {
+    backend.disposeData(castedData.dataId);
+  }
+
   return out;
 }
 

tfjs-backend-wasm/src/kernels/ResizeBilinear.ts

@@ -19,6 +19,8 @@ import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo, util} from
 
 import {BackendWasm} from '../backend_wasm';
 
+import {cast} from './Cast';
+
 interface ResizeBilinearInputs extends NamedTensorInfoMap {
   x: TensorInfo;
 }
@@ -48,7 +50,7 @@ function setup(backend: BackendWasm): void {
   ]);
 }
 
-function cropAndResize(args: {
+function resizeBilinear(args: {
   backend: BackendWasm,
   inputs: ResizeBilinearInputs,
   attrs: ResizeBilinearAttrs
@@ -60,9 +62,15 @@ function cropAndResize(args: {
   const [batch, oldHeight, oldWidth, numChannels] = x.shape;
   const outShape = [batch, newHeight, newWidth, numChannels];
 
-  const xId = backend.dataIdMap.get(x.dataId).id;
+  let xData = backend.dataIdMap.get(x.dataId);
+  let castedData;
+  if (xData.dtype !== 'float32') {
+    castedData = cast({backend, inputs: {x}, attrs: {dtype: 'float32'}});
+    xData = backend.dataIdMap.get(castedData.dataId);
+  }
+  const xId = xData.id;
 
-  const out = backend.makeOutput(outShape, x.dtype);
+  const out = backend.makeOutput(outShape, 'float32');
   if (util.sizeFromShape(x.shape) === 0) {
     return out;
   }
@@ -71,12 +79,17 @@ function cropAndResize(args: {
   wasmResizeBilinear(
       xId, batch, oldHeight, oldWidth, numChannels, newHeight, newWidth,
       alignCorners ? 1 : 0, outId);
+
+  if (castedData != null) {
+    backend.disposeData(castedData.dataId);
+  }
+
   return out;
 }
 
 registerKernel({
   kernelName: 'ResizeBilinear',
   backendName: 'wasm',
   setupFunc: setup,
-  kernelFunc: cropAndResize
+  kernelFunc: resizeBilinear
 });

tfjs-core/src/backends/cpu/backend_cpu.ts

@@ -3561,8 +3561,8 @@ export class MathBackendCPU extends KernelBackend {
     const numBoxes = boxes.shape[0];
 
     const [cropHeight, cropWidth] = cropSize;
-    const output = ops.buffer(
-        [numBoxes, cropHeight, cropWidth, numChannels], images.dtype);
+    const output =
+        ops.buffer([numBoxes, cropHeight, cropWidth, numChannels], 'float32');
 
     const boxVals = this.readSync(boxes.dataId) as TypedArray;
     const boxIndVals = this.readSync(boxIndex.dataId) as TypedArray;

tfjs-core/src/backends/webgl/backend_webgl.ts

@@ -2195,7 +2195,7 @@ export class MathBackendWebGL extends KernelBackend {
         new ResizeBilinearPackedProgram(
             x.shape, newHeight, newWidth, alignCorners) :
         new ResizeBilinearProgram(x.shape, newHeight, newWidth, alignCorners);
-    return this.compileAndRun(program, [x]);
+    return this.compileAndRun(program, [x], 'float32');
   }
 
   resizeBilinearBackprop(dy: Tensor4D, x: Tensor4D, alignCorners: boolean):
@@ -2260,7 +2260,7 @@ export class MathBackendWebGL extends KernelBackend {
       extrapolationValue: number): Tensor4D {
     const program = new CropAndResizeProgram(
         image.shape, boxes.shape, cropSize, method, extrapolationValue);
-    return this.compileAndRun(program, [image, boxes, boxIndex]);
+    return this.compileAndRun(program, [image, boxes, boxIndex], 'float32');
   }
 
   depthToSpace(x: Tensor4D, blockSize: number, dataFormat: 'NHWC'|'NCHW'):

tfjs-core/src/ops/image_ops.ts

@@ -271,7 +271,7 @@ function cropAndResize_(
     method?: 'bilinear'|'nearest',
     extrapolationValue?: number,
     ): Tensor4D {
-  const $image = convertToTensor(image, 'image', 'cropAndResize', 'float32');
+  const $image = convertToTensor(image, 'image', 'cropAndResize');
   const $boxes = convertToTensor(boxes, 'boxes', 'cropAndResize', 'float32');
   const $boxInd = convertToTensor(boxInd, 'boxInd', 'cropAndResize', 'int32');
   method = method || 'bilinear';