add onProgress for loaders (#180)

tensorspace-team · Feb 23, 2019 · 068097e · 068097e
1 parent f01f8f4
commit 068097e
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 5 deletions.
diff --git a/src/loader/KerasLoader.js b/src/loader/KerasLoader.js
@@ -62,8 +62,8 @@ KerasLoader.prototype = Object.assign( Object.create( Loader.prototype ), {
 	 */
 
 	load: async function() {
-
-		const loadedModel = await tf.loadLayersModel( this.url );
+		
+		const loadedModel = await tf.loadLayersModel( this.url, this.tfjsLoadOption );
 
 		this.model.resource = loadedModel;
 

diff --git a/src/loader/Loader.js b/src/loader/Loader.js
@@ -38,6 +38,22 @@ function Loader( model, config ) {
 	 */
 
 	this.onCompleteCallback = undefined;
+
+	/**
+	 * Store callback function fired periodically before load process is completed.
+	 *
+	 * @type { function }
+	 */
+
+	this.onProgressCallBack = undefined;
+
+	/**
+	 * Tfjs LoadOption object.
+	 *
+	 * @type { Object }
+	 */
+
+	this.tfjsLoadOption = {};
 
 	// Load loader's basic configuration.
 
@@ -56,7 +72,16 @@ Loader.prototype = {
 	loadLoaderConfig: function( config ) {
 
 		if ( this.config !== undefined )  {
-
+
+			// If onProgress is defined by user, store it.
+
+			if ( config.onProgress !== undefined ) {
+
+				this.onProgressCallBack = config.onProgress;
+				this.tfjsLoadOption.onProgress = config.onProgress;
+
+			}
+
 			// If onComplete callback is defined by user, store it.
 
 			if ( config.onComplete !== undefined ) {

diff --git a/src/loader/TfLoader.js b/src/loader/TfLoader.js
@@ -72,7 +72,7 @@ TfLoader.prototype = Object.assign( Object.create( Loader.prototype ), {
 
 	load: async function() {
 
-		const loadedModel = await tf.loadGraphModel( this.modelUrl );
+		const loadedModel = await tf.loadGraphModel( this.modelUrl, this.tfjsLoadOption );
 
 		this.model.resource = loadedModel;
 

diff --git a/src/loader/TfjsLoader.js b/src/loader/TfjsLoader.js
@@ -63,7 +63,7 @@ TfjsLoader.prototype = Object.assign( Object.create( Loader.prototype ), {
 
 	load: async function() {
 
-		const loadedModel = await tf.loadLayersModel( this.url );
+		const loadedModel = await tf.loadLayersModel( this.url, this.tfjsLoadOption );
 
 		this.model.resource = loadedModel;