[memory leak] fixes memory leaks (#684)

* fix memory leak in bodypix

* fixes memory leak in unet

* fixes memory leak in style transfer

* fix memory leak in pix2pix

* fixes memory leaks in kmeans

* fixes memory leaks in DCGAN & adds options to constructor

* fix memory leak in cvae

* fixes mem leak in sentiment

* fixes unet - array3DToImage destroys tensor, order matters

Author: joeyklee

src/BodyPix/index.js

@@ -21,7 +21,8 @@ const DEFAULTS = {
     "multiplier": 0.75,
     "outputStride": 16,
     "segmentationThreshold": 0.5,
-    "palette": BODYPIX_PALETTE
+    "palette": BODYPIX_PALETTE,
+    "returnTensors": false,
 }
 
 class BodyPix {
@@ -40,7 +41,8 @@ class BodyPix {
             multiplier: options.multiplier || DEFAULTS.multiplier,
             outputStride: options.outputStride || DEFAULTS.outputStride,
             segmentationThreshold: options.segmentationThreshold || DEFAULTS.segmentationThreshold,
-            palette: options.palette || DEFAULTS.palette
+            palette: options.palette || DEFAULTS.palette,
+            returnTensors: options.returnTensors || DEFAULTS.returnTensors
         }
 
         this.ready = callCallback(this.loadModel(), callback);
@@ -126,7 +128,6 @@ class BodyPix {
         this.config.outputStride = segmentationOptions.outputStride || this.config.outputStride;
         this.config.segmentationThreshold = segmentationOptions.segmentationThreshold || this.config.segmentationThreshold;
 
-        // const segmentation = await this.model.estimatePersonSegmentation(imgToSegment, this.config.outputStride, this.config.segmentationThreshold)
         const segmentation = await this.model.estimatePartSegmentation(imgToSegment, this.config.outputStride, this.config.segmentationThreshold);
 
         const bodyPartsMeta = this.bodyPartsSpec(this.config.palette);
@@ -153,13 +154,12 @@ class BodyPix {
         result.raw.personMask = bp.toMaskImageData(segmentation, false);
         result.raw.partMask = bp.toColoredPartImageData(segmentation, colorsArray);
 
-        let normTensor = await tf.browser.fromPixels(imgToSegment);
-
         const {
             personMask,
             backgroundMask,
             partMask,
         } = tf.tidy(() => {
+            let normTensor = tf.browser.fromPixels(imgToSegment);
             // create a tensor from the input image
             const alpha = tf.ones([segmentation.height, segmentation.width, 1]).tile([1, 1, 1]).mul(255)
             normTensor = normTensor.concat(alpha, 2)
@@ -198,6 +198,17 @@ class BodyPix {
             result.partMask = await this.convertToP5Image(partMaskPixels, segmentation.width, segmentation.height)
         }
 
+        if (!this.config.returnTensors) {
+            personMask.dispose();
+            backgroundMask.dispose();
+            partMask.dispose();
+        } else {
+            // return tensors
+            result.tensor.personMask = personMask;
+            result.tensor.backgroundMask = backgroundMask;
+            result.tensor.partMask = partMask;
+        }
+
         return result;
 
     }
@@ -309,12 +320,11 @@ class BodyPix {
         // result.backgroundMask = bgMaskCanvas;
         // result.featureMask = featureMaskCanvas;
 
-        let normTensor = await tf.browser.fromPixels(imgToSegment);
-
         const {
             personMask,
             backgroundMask
         } = tf.tidy(() => {
+            let normTensor = tf.browser.fromPixels(imgToSegment);
             // create a tensor from the input image
             const alpha = tf.ones([segmentation.height, segmentation.width, 1]).tile([1, 1, 1]).mul(255)
             normTensor = normTensor.concat(alpha, 2)
@@ -337,16 +347,25 @@ class BodyPix {
         const personMaskPixels = await tf.browser.toPixels(personMask);
         const bgMaskPixels = await tf.browser.toPixels(backgroundMask);
 
-        // otherwise, return the pixels 
-        result.personMask = personMaskPixels;
-        result.backgroundMask = bgMaskPixels;
-
         // if p5 exists, convert to p5 image
         if (p5Utils.checkP5()) {
             result.personMask = await this.convertToP5Image(personMaskPixels, segmentation.width, segmentation.height)
             result.backgroundMask = await this.convertToP5Image(bgMaskPixels, segmentation.width, segmentation.height)
+        } else {
+            // otherwise, return the pixels 
+            result.personMask = personMaskPixels;
+            result.backgroundMask = bgMaskPixels;
         }
 
+        if (!this.config.returnTensors) {
+            personMask.dispose();
+            backgroundMask.dispose();
+        } else {
+            result.tensor.personMask = personMask;
+            result.tensor.backgroundMask = backgroundMask;
+        }
+
+
         return result;
 
     }

src/CVAE/index.js

@@ -97,8 +97,10 @@ class Cvae {
       const temp = this.model.predict([this.latentDim, input]);
       return temp.reshape([temp.shape[1], temp.shape[2], temp.shape[3]]);
     });
+  
 
     const raws = await tf.browser.toPixels(res);
+    res.dispose();
 
     const canvas = document.createElement('canvas'); // consider using offScreneCanvas
     const ctx = canvas.getContext('2d');

src/DCGAN/index.js

@@ -27,12 +27,15 @@ class DCGANBase {
      * @param {modelName} modelName - The name of the model to use.
      * @param {function} readyCb - A callback to be called when the model is ready.
      */
-    constructor(modelPath, callback) {
+    constructor(modelPath, options, callback) {
         this.model = {};
         this.modelPath = modelPath;
         this.modelInfo = {};
         this.modelPathPrefix = '';
         this.modelReady = false;
+        this.config = {
+            returnTensors: options.returnTensors || false,
+        }
         this.ready = callCallback(this.loadModel(), callback);
     }
 
@@ -98,6 +101,7 @@ class DCGANBase {
      * @return {object} includes blob, raw, and tensor. if P5 exists, then a p5Image
      */
     async generateInternal(latentVector) {
+        
         const {
             modelLatentDim
         } = this.modelInfo;
@@ -119,12 +123,19 @@ class DCGANBase {
         const result = {};
         result.blob = blob;
         result.raw = raw;
-        result.tensor = imageTensor;
+        
 
         if (p5Utils.checkP5()) {
             result.image = p5Image;
         }
 
+        if(!this.config.returnTensors){
+            result.tensor = null;
+            imageTensor.dispose();
+        } else {
+            result.tensor = imageTensor;
+        }
+
         return result;
 
     }
@@ -138,7 +149,9 @@ class DCGANBase {
 
 }
 
-const DCGAN = (modelPath, cb) => {
+const DCGAN = (modelPath, optionsOrCb, cb) => {
+    let callback;
+    let options = {};
     if (typeof modelPath !== 'string') {
         throw new Error(`Please specify a path to a "manifest.json" file: \n
          "models/face/manifest.json" \n\n
@@ -151,10 +164,17 @@ const DCGAN = (modelPath, cb) => {
          }
          `);
     }
+
+    if(typeof optionsOrCb === 'function'){
+        callback = optionsOrCb;
+    } else if (typeof optionsOrCb === 'object'){
+        options = optionsOrCb;
+        callback = cb;
+    }
 
 
-    const instance = new DCGANBase(modelPath, cb);
-    return cb ? instance : instance.ready;
+    const instance = new DCGANBase(modelPath, options, callback);
+    return callback ? instance : instance.ready;
 
 }