face-api/src/dom/imageTensorToCanvas.ts

20 lines
615 B
TypeScript

import * as tf from '../../dist/tfjs.esm';
import { env } from '../env/index';
import { isTensor4D } from '../utils/index';
export async function imageTensorToCanvas(
imgTensor: tf.Tensor,
canvas?: HTMLCanvasElement,
): Promise<HTMLCanvasElement> {
const targetCanvas = canvas || env.getEnv().createCanvasElement();
const [height, width, numChannels] = imgTensor.shape.slice(isTensor4D(imgTensor) ? 1 : 0);
const imgTensor3D = tf.tidy(() => imgTensor.as3D(height, width, numChannels).toInt());
await tf.browser.toPixels(imgTensor3D, targetCanvas);
imgTensor3D.dispose();
return targetCanvas;
}