license import { env , KernelConfig , KernelFunc , TensorInfo } from '@tensorflow/tfjs-core' ; import { FromPixels , FromPixelsAttrs , FromPixelsInputs } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { TextureUsage } from '../tex_util' ; import { FromPixelsProgram } from './FromPixels_utils/from_pixels_gpu' ; import { FromPixelsPackedProgram } from './FromPixels_utils/from_pixels_packed_gpu' ; export const fromPixelsConfig = { kernelName : FromPixels , backendName : 'webgl' , kernelFunc : fromPixels as { } as KernelFunc , } ; let fromPixels2DContext ; function fromPixels ( args ) { const { inputs , backend , attrs } = args ; let { pixels } = inputs ; const { numChannels } = attrs ; const isVideo = typeof ( HTMLVideoElement ) !== 'undefined' && pixels instanceof HTMLVideoElement ; const isImage = typeof ( HTMLImageElement ) !== 'undefined' && pixels instanceof HTMLImageElement ; const [ width , height ] = isVideo ? [ ( pixels as HTMLVideoElement ) . videoWidth , ( pixels as HTMLVideoElement ) . videoHeight ] : [ pixels . width , pixels . height ] ; const texShape = [ height , width ] ; const outShape = [ height , width , numChannels ] ; if ( isImage || isVideo ) { if ( fromPixels2DContext == null ) { fromPixels2DContext = document . createElement ( 'canvas' ) . getContext ( '2d' ) ; } fromPixels2DContext . canvas . width = width ; fromPixels2DContext . canvas . height = height ; fromPixels2DContext . drawImage ( pixels as HTMLVideoElement | HTMLImageElement | ImageBitmap , 0 , 0 , width , height ) ; pixels = fromPixels2DContext . canvas ; } const tempPixelHandle = backend . makeTensorInfo ( texShape , 'int32' ) ; backend . texData . get ( tempPixelHandle . dataId ) . usage = TextureUsage . PIXELS ; backend . gpgpu . uploadPixelDataToTexture ( backend . getTexture ( tempPixelHandle . dataId ) , pixels as ImageData ) ; const program = env ( ) . getBool ( 'WEBGL_PACK' ) ? new FromPixelsPackedProgram ( outShape ) : new FromPixelsProgram ( outShape ) ; const res = backend . runWebGLProgram ( program , [ tempPixelHandle ] , 'int32' ) ; backend . disposeData ( tempPixelHandle . dataId ) ; return res ; }