20 lines
598 B
TypeScript
20 lines
598 B
TypeScript
import * as tf from '@tensorflow/tfjs-core';
|
|
|
|
import { env } from '../env';
|
|
import { isTensor4D } from '../utils';
|
|
|
|
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
|
|
} |