license import * as tf from '@tensorflow/tfjs' ; const STYLENET_URL = 'https://cdn.jsdelivr.net/gh/reiinakano/arbitrary-image-stylization-tfjs@master/saved_model_style_js/model.json' ; const TRANSFORMNET_URL = 'https://cdn.jsdelivr.net/gh/reiinakano/arbitrary-image-stylization-tfjs@master/saved_model_transformer_separable_js/model.json' ; export class StyleTranfer { private styleNet ? : tf . GraphModel ; private transformNet ? : tf . GraphModel ; constructor ( ) { } async init ( ) { await Promise . all ( [ this . loadStyleModel ( ) , this . loadTransformerModel ( ) ] ) ; await this . warmup ( ) ; } async loadStyleModel ( ) { if ( this . styleNet == null ) { this . styleNet = await tf . loadGraphModel ( STYLENET_URL ) ; console . log ( 'stylenet loaded' ) ; } } async loadTransformerModel ( ) { if ( this . transformNet == null ) { this . transformNet = await tf . loadGraphModel ( TRANSFORMNET_URL ) ; console . log ( 'transformnet loaded' ) ; } } async warmup ( ) { const input = tf . randomNormal ( [ 320 , 240 , 3 ] ) ; const res = this . stylize ( input , input ) ; await res . data ( ) ; tf . dispose ( [ input , res ] ) ; } param style private predictStyleParameters ( styleImage ) { return tf . tidy ( ( ) => { if ( this . styleNet == null ) { throw new Error ( 'Stylenet not loaded' ) ; } return this . styleNet . predict ( styleImage . toFloat ( ) . div ( tf . scalar ( 255 ) ) . expandDims ( ) ) ; } ) as tf . Tensor4D ; } param content param bottleneck private produceStylized ( contentImage , bottleneck ) { return tf . tidy ( ( ) => { if ( this . transformNet == null ) { throw new Error ( 'Transformnet not loaded' ) ; } const input = contentImage . toFloat ( ) . div ( tf . scalar ( 255 ) ) . expandDims ( ) ; const image = this . transformNet . predict ( [ input , bottleneck ] ) as tf . Tensor4D ; return image . mul ( 255 ) . squeeze ( ) ; } ) ; } public stylize ( styleImage , contentImage ) { const start = Date . now ( ) ; console . log ( styleImage . shape , contentImage . shape ) ; const styleRepresentation = this . predictStyleParameters ( styleImage ) ; const stylized = this . produceStylized ( contentImage , styleRepresentation ) ; tf . dispose ( [ styleRepresentation ] ) ; const end = Date . now ( ) ; console . log ( 'stylization scheduled' , end - start ) ; return stylized ; } }