Add Draw API (#7628)

FEATURE

* add draw

* add tests

* add comments

* lint

* upd

* typo

* cmt

* rename contextOption to canvasOption

* Update tfjs-core/src/ops/browser.ts

Co-authored-by: Matthew Soulanille <matthew@soulanille.net>

* Update tfjs-core/src/ops/draw_test.ts

Co-authored-by: Matthew Soulanille <matthew@soulanille.net>

* warn once

* lint

* unify options

* typo

* tune

* tune

* tune

* typo

---------

Co-authored-by: Matthew Soulanille <matthew@soulanille.net>

Author: Linchenn

tfjs-backend-cpu/BUILD.bazel

@@ -41,7 +41,7 @@ tfjs_web_test(
     ],
     headless = False,
     presubmit_browsers = [
-        "bs_safari_mac",
+        "bs_chrome_mac",
     ],
 )
 

tfjs-backend-cpu/src/kernels/Draw.ts

@@ -0,0 +1,96 @@
+/**
+ * @license
+ * Copyright 2023 Google LLC.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {Draw, DrawAttrs, DrawInputs, KernelConfig, KernelFunc, TypedArray} from '@tensorflow/tfjs-core';
+import {TensorInfo} from '@tensorflow/tfjs-core';
+
+import {MathBackendCPU} from '../backend_cpu';
+
+export function draw(
+    args: {inputs: DrawInputs, backend: MathBackendCPU, attrs: DrawAttrs}):
+    TensorInfo {
+  const {inputs, backend, attrs} = args;
+  const {image} = inputs;
+  const {canvas, options} = attrs;
+  const {contextOptions, imageOptions} = options || {};
+  const alpha = imageOptions ?.alpha || 1;
+
+  const contextType = contextOptions ?.contextType || '2d';
+  if (contextType !== '2d') {
+    throw new Error(`Context type ${
+        contextOptions.contextType} is not supported by the CPU backend.`);
+  }
+  const ctx = canvas.getContext(contextType,
+    contextOptions?.contextAttributes || {}) as CanvasRenderingContext2D ;
+  if (ctx == null) {
+    throw new Error(`Could not get the context with ${contextType} type.`);
+  }
+
+  const [height, width] = image.shape.slice(0, 2);
+  const depth = image.shape.length === 2 ? 1 : image.shape[2];
+  const data = backend.data.get(image.dataId).values as TypedArray;
+  const multiplier = image.dtype === 'float32' ? 255 : 1;
+  const bytes = new Uint8ClampedArray(width * height * 4);
+
+  for (let i = 0; i < height * width; ++i) {
+    const rgba = [0, 0, 0, 255 * alpha];
+
+    for (let d = 0; d < depth; d++) {
+      const value = data[i * depth + d];
+
+      if (image.dtype === 'float32') {
+        if (value < 0 || value > 1) {
+          throw new Error(
+              `Tensor values for a float32 Tensor must be in the ` +
+              `range [0 - 1] but encountered ${value}.`);
+        }
+      } else if (image.dtype === 'int32') {
+        if (value < 0 || value > 255) {
+          throw new Error(
+              `Tensor values for a int32 Tensor must be in the ` +
+              `range [0 - 255] but encountered ${value}.`);
+        }
+      }
+
+      if (depth === 1) {
+        rgba[0] = value * multiplier;
+        rgba[1] = value * multiplier;
+        rgba[2] = value * multiplier;
+      } else {
+        rgba[d] = value * multiplier;
+      }
+    }
+
+    const j = i * 4;
+    bytes[j + 0] = Math.round(rgba[0]);
+    bytes[j + 1] = Math.round(rgba[1]);
+    bytes[j + 2] = Math.round(rgba[2]);
+    bytes[j + 3] = Math.round(rgba[3]);
+  }
+
+  canvas.width = width;
+  canvas.height = height;
+  const imageData = new ImageData(bytes, width, height);
+  ctx.putImageData(imageData, 0, 0);
+  return image;
+}
+
+export const drawConfig: KernelConfig = {
+  kernelName: Draw,
+  backendName: 'cpu',
+  kernelFunc: draw as unknown as KernelFunc
+};

tfjs-backend-cpu/src/register_all_kernels.ts

@@ -70,6 +70,7 @@ import {diagConfig} from './kernels/Diag';
 import {dilation2DConfig} from './kernels/Dilation2D';
 import {dilation2DBackpropFilterConfig} from './kernels/Dilation2DBackpropFilter';
 import {dilation2DBackpropInputConfig} from './kernels/Dilation2DBackpropInput';
+import {drawConfig} from './kernels/Draw';
 import {einsumConfig} from './kernels/Einsum';
 import {eluConfig} from './kernels/Elu';
 import {eluGradConfig} from './kernels/EluGrad';
@@ -244,6 +245,7 @@ const kernelConfigs: KernelConfig[] = [
   dilation2DConfig,
   dilation2DBackpropFilterConfig,
   dilation2DBackpropInputConfig,
+  drawConfig,
   einsumConfig,
   eluConfig,
   eluGradConfig,

tfjs-backend-webgl/src/setup_test.ts

@@ -36,7 +36,7 @@ const customInclude = (testName: string) => {
     'isBrowser: false', 'dilation gradient',
     'throws when index is out of bound',
     // otsu tests for threshold op is failing on windows
-    'method otsu'
+    'method otsu', 'Draw on 2d context'
   ];
   for (const subStr of toExclude) {
     if (testName.includes(subStr)) {

tfjs-backend-webgpu/src/setup_test.ts

@@ -124,6 +124,12 @@ const TEST_FILTERS: TestFilter[] = [
       'canvas and image match',  // Failing on Linux
     ],
   },
+  {
+    startsWith: 'Draw',
+    excludes: [
+      'on 2d context',
+    ]
+  },
   {
     startsWith: 'sign ',
     excludes: [

tfjs-core/src/kernel_names.ts

@@ -22,7 +22,7 @@ import {NamedTensorInfoMap} from './kernel_registry';
 import {ExplicitPadding} from './ops/conv_util';
 import {Activation} from './ops/fused_types';
 import {TensorInfo} from './tensor_info';
-import {DataType, PixelData} from './types';
+import {DataType, DrawOptions, PixelData} from './types';
 
 export const Abs = 'Abs';
 export type AbsInputs = UnaryInputs;
@@ -335,6 +335,13 @@ export const Dilation2DBackpropFilter = 'Dilation2DBackpropFilter';
 export type Dilation2DBackpropFilterInputs =
     Pick<NamedTensorInfoMap, 'x'|'filter'|'dy'>;
 
+export const Draw = 'Draw';
+export type DrawInputs = Pick<NamedTensorInfoMap, 'image'>;
+export interface DrawAttrs {
+  canvas: HTMLCanvasElement;
+  options?: DrawOptions;
+}
+
 export const RealDiv = 'RealDiv';
 export type RealDivInputs = BinaryInputs;
 

tfjs-core/src/ops/browser.ts

@@ -17,18 +17,19 @@
 
 import {ENGINE} from '../engine';
 import {env} from '../environment';
-import {FromPixels, FromPixelsAttrs, FromPixelsInputs} from '../kernel_names';
+import {Draw, DrawAttrs, DrawInputs, FromPixels, FromPixelsAttrs, FromPixelsInputs} from '../kernel_names';
 import {getKernel, NamedAttrMap} from '../kernel_registry';
 import {Tensor, Tensor2D, Tensor3D} from '../tensor';
 import {NamedTensorMap} from '../tensor_types';
 import {convertToTensor} from '../tensor_util_env';
-import {PixelData, TensorLike} from '../types';
+import {DrawOptions, ImageOptions, PixelData, TensorLike} from '../types';
 
 import {cast} from './cast';
 import {op} from './operation';
 import {tensor3d} from './tensor3d';
 
 let fromPixels2DContext: CanvasRenderingContext2D;
+let hasToPixelsWarned = false;
 
 /**
  * Creates a `tf.Tensor` from an image.
@@ -145,9 +146,8 @@ function fromPixels_(
               'Reason: OffscreenCanvas Context2D rendering is not supported.');
         }
       } else {
-        fromPixels2DContext =
-            document.createElement('canvas').getContext(
-                '2d', {willReadFrequently: true});
+        fromPixels2DContext = document.createElement('canvas').getContext(
+            '2d', {willReadFrequently: true});
       }
     }
     fromPixels2DContext.canvas.width = width;
@@ -269,6 +269,33 @@ export async function fromPixelsAsync(
   return fromPixels_(inputs, numChannels);
 }
 
+function validateImgTensor(img: Tensor2D|Tensor3D) {
+  if (img.rank !== 2 && img.rank !== 3) {
+    throw new Error(
+        `toPixels only supports rank 2 or 3 tensors, got rank ${img.rank}.`);
+  }
+  const depth = img.rank === 2 ? 1 : img.shape[2];
+
+  if (depth > 4 || depth === 2) {
+    throw new Error(
+        `toPixels only supports depth of size ` +
+        `1, 3 or 4 but got ${depth}`);
+  }
+
+  if (img.dtype !== 'float32' && img.dtype !== 'int32') {
+    throw new Error(
+        `Unsupported type for toPixels: ${img.dtype}.` +
+        ` Please use float32 or int32 tensors.`);
+  }
+}
+
+function validateImageOptions(imageOptions: ImageOptions) {
+  const alpha = imageOptions ?.alpha || 1;
+  if (alpha > 1 || alpha < 0) {
+    throw new Error(`Alpha value ${alpha} is suppoed to be in range [0 - 1].`);
+  }
+}
+
 /**
  * Draws a `tf.Tensor` of pixel values to a byte array or optionally a
  * canvas.
@@ -299,25 +326,10 @@ export async function toPixels(
     $img = cast(originalImgTensor, 'int32');
     originalImgTensor.dispose();
   }
-  if ($img.rank !== 2 && $img.rank !== 3) {
-    throw new Error(
-        `toPixels only supports rank 2 or 3 tensors, got rank ${$img.rank}.`);
-  }
+  validateImgTensor($img);
+
   const [height, width] = $img.shape.slice(0, 2);
   const depth = $img.rank === 2 ? 1 : $img.shape[2];
-
-  if (depth > 4 || depth === 2) {
-    throw new Error(
-        `toPixels only supports depth of size ` +
-        `1, 3 or 4 but got ${depth}`);
-  }
-
-  if ($img.dtype !== 'float32' && $img.dtype !== 'int32') {
-    throw new Error(
-        `Unsupported type for toPixels: ${$img.dtype}.` +
-        ` Please use float32 or int32 tensors.`);
-  }
-
   const data = await $img.data();
   const multiplier = $img.dtype === 'float32' ? 255 : 1;
   const bytes = new Uint8ClampedArray(width * height * 4);
@@ -359,6 +371,13 @@ export async function toPixels(
   }
 
   if (canvas != null) {
+    if (!hasToPixelsWarned) {
+      console.warn(
+          'tf.browser.toPixels is not efficient to draw tensor on canvas. ' +
+          'Please try tf.browser.draw instead.');
+      hasToPixelsWarned = true;
+    }
+
     canvas.width = width;
     canvas.height = height;
     const ctx = canvas.getContext('2d');
@@ -371,4 +390,44 @@ export async function toPixels(
   return bytes;
 }
 
+/**
+ * Draws a `tf.Tensor` to a canvas.
+ *
+ * When the dtype of the input is 'float32', we assume values in the range
+ * [0-1]. Otherwise, when input is 'int32', we assume values in the range
+ * [0-255].
+ *
+ * @param image The tensor to draw on the canvas. Must match one of
+ * these shapes:
+ *   - Rank-2 with shape `[height, width`]: Drawn as grayscale.
+ *   - Rank-3 with shape `[height, width, 1]`: Drawn as grayscale.
+ *   - Rank-3 with shape `[height, width, 3]`: Drawn as RGB with alpha set in
+ *     `imageOptions` (defaults to 1, which is opaque).
+ *   - Rank-3 with shape `[height, width, 4]`: Drawn as RGBA.
+ * @param canvas The canvas to draw to.
+ * @param options The configuration arguments for image to be drawn and the
+ *     canvas to draw to.
+ *
+ * @doc {heading: 'Browser', namespace: 'browser'}
+ */
+export function draw(
+    image: Tensor2D|Tensor3D|TensorLike, canvas: HTMLCanvasElement,
+    options?: DrawOptions): void {
+  let $img = convertToTensor(image, 'img', 'draw');
+  if (!(image instanceof Tensor)) {
+    // Assume int32 if user passed a native array.
+    const originalImgTensor = $img;
+    $img = cast(originalImgTensor, 'int32');
+    originalImgTensor.dispose();
+  }
+  validateImgTensor($img);
+  validateImageOptions(options?.imageOptions);
+
+  const inputs: DrawInputs = {image: $img};
+  const attrs: DrawAttrs = {canvas, options};
+  ENGINE.runKernel(
+      Draw, inputs as unknown as NamedTensorMap,
+      attrs as unknown as NamedAttrMap);
+}
+
 export const fromPixels = /* @__PURE__ */ op({fromPixels_});