license import { browser , ENV , image , tensor1d , Tensor1D , tensor2d , Tensor2D , Tensor3D , Tensor4D , tidy , util } from '@tensorflow/tfjs-core' ; import { WebcamConfig } from '../types' ; import { LazyIterator } from './lazy_iterator' ; /**
 * Provide a stream of image tensors from webcam video stream. Only works in
 * browser environment.
 */ export class WebcamIterator extends LazyIterator < Tensor3D > { private isClosed = true ; private stream : MediaStream ; private resize = false ; private cropSize : [ number , number ] ; private cropBox : Tensor2D ; private cropBoxInd : Tensor1D ; private constructor ( protected readonly webcamVideoElement , protected readonly webcamConfig ) { super ( ) ; if ( this . needToResize ( ) ) { this . resize = true ; this . cropSize = [ this . webcamConfig . resizeHeight , this . webcamConfig . resizeWidth ] ; this . cropBoxInd = tensor1d ( [ 0 ] , 'int32' ) ; if ( this . webcamConfig . centerCrop ) { const widthCroppingRatio = this . webcamConfig . resizeWidth * 1.0 / this . webcamVideoElement . width ; const heightCroppingRatio = this . webcamConfig . resizeHeight * 1.0 / this . webcamVideoElement . height ; const widthCropStart = ( 1 - widthCroppingRatio ) / 2 ; const heightCropStart = ( 1 - heightCroppingRatio ) / 2 ; const widthCropEnd = widthCropStart + widthCroppingRatio ; const heightCropEnd = heightCroppingRatio + heightCropStart ; this . cropBox = tensor2d ( [ heightCropStart , widthCropStart , heightCropEnd , widthCropEnd ] , [ 1 , 4 ] ) ; } else { this . cropBox = tensor2d ( [ 0 , 0 , 1 , 1 ] , [ 1 , 4 ] ) ; } } } summary ( ) { return `webcam` ; } static async create ( webcamVideoElement ? , webcamConfig = { } ) { if ( ENV . get ( 'IS_NODE' ) ) { throw new Error ( 'tf.data.webcam is only supported in browser environment.' ) ; } if ( ! webcamVideoElement ) { webcamVideoElement = document . createElement ( 'video' ) ; if ( ! webcamConfig . resizeWidth || ! webcamConfig . resizeHeight ) { throw new Error ( 'Please provide webcam video element, or resizeWidth and ' + 'resizeHeight to create a hidden video element.' ) ; } webcamVideoElement . width = webcamConfig . resizeWidth ; webcamVideoElement . height = webcamConfig . resizeHeight ; } const webcamIterator = new WebcamIterator ( webcamVideoElement , webcamConfig ) ; await webcamIterator . start ( ) ; return webcamIterator ; } async start ( ) { if ( this . webcamConfig . facingMode ) { util . assert ( ( this . webcamConfig . facingMode === 'user' ) || ( this . webcamConfig . facingMode === 'environment' ) , ( ) => `Invalid webcam facing mode: ${ this . webcamConfig . facingMode }. ` + `Please provide 'user' or 'environment'` ) ; } try { this . stream = await navigator . mediaDevices . getUserMedia ( { video : { deviceId : this . webcamConfig . deviceId , facingMode : this . webcamConfig . facingMode ? this . webcamConfig . facingMode : 'user' , width : this . webcamVideoElement . width , height : this . webcamVideoElement . height } } ) ; } catch ( e ) { e . message = `Error thrown while initializing video stream: ${ e . message }` ; throw e ; } if ( ! this . stream ) { throw new Error ( 'Could not obtain video from webcam.' ) ; } try { this . webcamVideoElement . srcObject = this . stream ; } catch ( error ) { console . log ( error ) ; this . webcamVideoElement . src = window . URL . createObjectURL ( this . stream ) ; } this . webcamVideoElement . play ( ) ; this . isClosed = false ; return new Promise < void > ( resolve => { this . webcamVideoElement . onloadedmetadata = ( ) => { resolve ( ) ; } ; } ) ; } async next ( ) { if ( this . isClosed ) { return { value : null , done : true } ; } let img ; try { img = browser . fromPixels ( this . webcamVideoElement ) ; } catch ( e ) { throw new Error ( `Error thrown converting video to pixels: ${ JSON . stringify ( e ) }` ) ; } if ( this . resize ) { try { return { value : this . cropAndResizeFrame ( img ) , done : false } ; } catch ( e ) { throw new Error ( `Error thrown cropping the video: ${ e . message }` ) ; } finally { img . dispose ( ) ; } } else { return { value : img , done : false } ; } } private needToResize ( ) { if ( this . webcamConfig . resizeWidth && this . webcamConfig . resizeHeight && ( this . webcamVideoElement . width !== this . webcamConfig . resizeWidth || this . webcamVideoElement . height !== this . webcamConfig . resizeHeight ) ) { return true ; } return false ; } cropAndResizeFrame ( img ) { return tidy ( ( ) => { const expandedImage = img . toFloat ( ) . expandDims ( 0 ) ; let resizedImage ; resizedImage = image . cropAndResize ( expandedImage , this . cropBox , this . cropBoxInd , this . cropSize , 'bilinear' ) ; const shape = resizedImage . shape ; return resizedImage . reshape ( shape . slice ( 1 ) as [ number , number , number ] ) ; } ) ; } async capture ( ) { return ( await this . next ( ) ) . value ; } stop ( ) { const tracks = this . stream . getTracks ( ) ; tracks . forEach ( track => track . stop ( ) ) ; try { this . webcamVideoElement . srcObject = null ; } catch ( error ) { console . log ( error ) ; this . webcamVideoElement . src = null ; } this . isClosed = true ; } toArray ( ) { throw new Error ( 'Can not convert infinite video stream to array.' ) ; } }