onnxruntime/js/common/lib/tensor-conversion-impl.ts

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

import {TensorToDataUrlOptions, TensorToImageDataOptions} from './tensor-conversion.js';
import {Tensor} from './tensor.js';

/**
 * implementation of Tensor.toDataURL()
 */
export const tensorToDataURL = (tensor: Tensor, options?: TensorToDataUrlOptions): string => {
  const canvas = typeof document !== 'undefined' ? document.createElement('canvas') : (new OffscreenCanvas(1, 1));
  canvas.width = tensor.dims[3];
  canvas.height = tensor.dims[2];
  const pixels2DContext =
      canvas.getContext('2d') as (CanvasRenderingContext2D | OffscreenCanvasRenderingContext2D | null);

  if (pixels2DContext != null) {
    // Default values for height and width & format
    let width: number;
    let height: number;
    if (options?.tensorLayout !== undefined && options.tensorLayout === 'NHWC') {
      width = tensor.dims[2];
      height = tensor.dims[3];
    } else {  // Default layout is NCWH
      width = tensor.dims[3];
      height = tensor.dims[2];
    }

    const inputformat = options?.format !== undefined ? options.format : 'RGB';

    const norm = options?.norm;
    let normMean: [number, number, number, number];
    let normBias: [number, number, number, number];
    if (norm === undefined || norm.mean === undefined) {
      normMean = [255, 255, 255, 255];
    } else {
      if (typeof (norm.mean) === 'number') {
        normMean = [norm.mean, norm.mean, norm.mean, norm.mean];
      } else {
        normMean = [norm.mean[0], norm.mean[1], norm.mean[2], 0];
        if (norm.mean[3] !== undefined) {
          normMean[3] = norm.mean[3];
        }
      }
    }
    if (norm === undefined || norm.bias === undefined) {
      normBias = [0, 0, 0, 0];
    } else {
      if (typeof (norm.bias) === 'number') {
        normBias = [norm.bias, norm.bias, norm.bias, norm.bias];
      } else {
        normBias = [norm.bias[0], norm.bias[1], norm.bias[2], 0];
        if (norm.bias[3] !== undefined) {
          normBias[3] = norm.bias[3];
        }
      }
    }

    const stride = height * width;
    // Default pointer assignments
    let rTensorPointer = 0, gTensorPointer = stride, bTensorPointer = stride * 2, aTensorPointer = -1;

    // Updating the pointer assignments based on the input image format
    if (inputformat === 'RGBA') {
      rTensorPointer = 0;
      gTensorPointer = stride;
      bTensorPointer = stride * 2;
      aTensorPointer = stride * 3;
    } else if (inputformat === 'RGB') {
      rTensorPointer = 0;
      gTensorPointer = stride;
      bTensorPointer = stride * 2;
    } else if (inputformat === 'RBG') {
      rTensorPointer = 0;
      bTensorPointer = stride;
      gTensorPointer = stride * 2;
    }

    for (let i = 0; i < height; i++) {
      for (let j = 0; j < width; j++) {
        const R = ((tensor.data[rTensorPointer++] as number) - normBias[0]) * normMean[0];  // R value
        const G = ((tensor.data[gTensorPointer++] as number) - normBias[1]) * normMean[1];  // G value
        const B = ((tensor.data[bTensorPointer++] as number) - normBias[2]) * normMean[2];  // B value
        const A = aTensorPointer === -1 ?
            255 :
            ((tensor.data[aTensorPointer++] as number) - normBias[3]) * normMean[3];  // A value
        // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
        pixels2DContext.fillStyle = 'rgba(' + R + ',' + G + ',' + B + ',' + A + ')';
        pixels2DContext.fillRect(j, i, 1, 1);
      }
    }
    if ('toDataURL' in canvas) {
      return canvas.toDataURL();
    } else {
      throw new Error('toDataURL is not supported');
    }
  } else {
    throw new Error('Can not access image data');
  }
};

/**
 * implementation of Tensor.toImageData()
 */
export const tensorToImageData = (tensor: Tensor, options?: TensorToImageDataOptions): ImageData => {
  const pixels2DContext = typeof document !== 'undefined' ?
      document.createElement('canvas').getContext('2d') :
      new OffscreenCanvas(1, 1).getContext('2d') as OffscreenCanvasRenderingContext2D;
  let image: ImageData;
  if (pixels2DContext != null) {
    // Default values for height and width & format
    let width: number;
    let height: number;
    let channels: number;
    if (options?.tensorLayout !== undefined && options.tensorLayout === 'NHWC') {
      width = tensor.dims[2];
      height = tensor.dims[1];
      channels = tensor.dims[3];
    } else {  // Default layout is NCWH
      width = tensor.dims[3];
      height = tensor.dims[2];
      channels = tensor.dims[1];
    }
    const inputformat = options !== undefined ? (options.format !== undefined ? options.format : 'RGB') : 'RGB';

    const norm = options?.norm;
    let normMean: [number, number, number, number];
    let normBias: [number, number, number, number];
    if (norm === undefined || norm.mean === undefined) {
      normMean = [255, 255, 255, 255];
    } else {
      if (typeof (norm.mean) === 'number') {
        normMean = [norm.mean, norm.mean, norm.mean, norm.mean];
      } else {
        normMean = [norm.mean[0], norm.mean[1], norm.mean[2], 255];
        if (norm.mean[3] !== undefined) {
          normMean[3] = norm.mean[3];
        }
      }
    }
    if (norm === undefined || norm.bias === undefined) {
      normBias = [0, 0, 0, 0];
    } else {
      if (typeof (norm.bias) === 'number') {
        normBias = [norm.bias, norm.bias, norm.bias, norm.bias];
      } else {
        normBias = [norm.bias[0], norm.bias[1], norm.bias[2], 0];
        if (norm.bias[3] !== undefined) {
          normBias[3] = norm.bias[3];
        }
      }
    }

    const stride = height * width;
    if (options !== undefined) {
      if (options.format !== undefined && (channels === 4 && options.format !== 'RGBA') ||
          (channels === 3 && (options.format !== 'RGB' && options.format !== 'BGR'))) {
        throw new Error('Tensor format doesn\'t match input tensor dims');
      }
    }

    // Default pointer assignments
    const step = 4;
    let rImagePointer = 0, gImagePointer = 1, bImagePointer = 2, aImagePointer = 3;
    let rTensorPointer = 0, gTensorPointer = stride, bTensorPointer = stride * 2, aTensorPointer = -1;

    // Updating the pointer assignments based on the input image format
    if (inputformat === 'RGBA') {
      rTensorPointer = 0;
      gTensorPointer = stride;
      bTensorPointer = stride * 2;
      aTensorPointer = stride * 3;
    } else if (inputformat === 'RGB') {
      rTensorPointer = 0;
      gTensorPointer = stride;
      bTensorPointer = stride * 2;
    } else if (inputformat === 'RBG') {
      rTensorPointer = 0;
      bTensorPointer = stride;
      gTensorPointer = stride * 2;
    }

    image = pixels2DContext.createImageData(width, height);

    for (let i = 0; i < height * width;
         rImagePointer += step, gImagePointer += step, bImagePointer += step, aImagePointer += step, i++) {
      image.data[rImagePointer] = ((tensor.data[rTensorPointer++] as number) - normBias[0]) * normMean[0];  // R value
      image.data[gImagePointer] = ((tensor.data[gTensorPointer++] as number) - normBias[1]) * normMean[1];  // G value
      image.data[bImagePointer] = ((tensor.data[bTensorPointer++] as number) - normBias[2]) * normMean[2];  // B value
      image.data[aImagePointer] = aTensorPointer === -1 ?
          255 :
          ((tensor.data[aTensorPointer++] as number) - normBias[3]) * normMean[3];  // A value
    }

  } else {
    throw new Error('Can not access image data');
  }
  return image;
};