Fixes Memory Leak

apps/model/model.js

@@ -372,6 +372,8 @@ function camicStopDraw(e) {
       if (args) {
         runPredict(args.status);
       }
+      var memory = tf.memory();
+      console.log(memory);
       $UI.modelPanel.setPosition(box.rect.x,box.rect.y,box.rect.width,box.rect.height);
       $UI.modelPanel.open(args);
 
@@ -471,10 +473,10 @@ function runPredict(key) {
     self.showProgress("Model loaded...");
 
     // Warmup the model before using real data.
+    tf.tidy(()=>{
     const warmupResult = model.predict(tf.zeros([1, image_size, image_size, input_channels]));
-    warmupResult.dataSync();
-    warmupResult.dispose();
     console.log("Model ready");
+    });
 
     let temp = document.querySelector('#dummy');
     temp.height = step;
@@ -509,7 +511,7 @@ function runPredict(key) {
         fullResCvs.getContext('2d').drawImage(l_img, 0, 0);
 
         let imgData = fullResCvs.getContext('2d').getImageData(0,0,fullResCvs.width,fullResCvs.height);
-
+        tf.tidy(()=>{
         const img = tf.browser.fromPixels(imgData).toFloat();
         let img2;
         if (input_channels == 1) {
@@ -544,7 +546,7 @@ function runPredict(key) {
           normalized = img2.sub(mean).div(std);
         }    
         let batched = normalized.reshape([1, image_size, image_size, input_channels]);
-        let values = await model.predict(batched).data();
+        let values =model.predict(batched).dataSync();
 
         values.forEach((e) => {
           csvContent += e.toString() + ",";
@@ -555,6 +557,7 @@ function runPredict(key) {
         // Retrieving the top class
 
         dx += step;
+      });
       }
       dy += step;
     }

apps/segment/segment.js

@@ -456,13 +456,16 @@ function camicStopDraw(e) {
       } else {
         segmentModel(args.status);
       }
+      var memory = tf.memory();
+      console.log(memory);
       $UI.segmentPanel.setPosition(box.rect.x,box.rect.y,box.rect.width,box.rect.height);
       $UI.segmentPanel.open(args);
 
       // close
       canvasDraw.clear();
     }
 
+
   } else {
     console.error('Could not get feature collection.')
   }
@@ -682,12 +685,12 @@ async function segmentModel(key) {
 
     model = await tf.loadLayersModel(IDB_URL + key);
     console.log('Model Loaded');
-
+
+    tf.tidy(()=>{
     // Warmup the model before using real data.
     const warmupResult = model.predict(tf.zeros([1, image_size, image_size, input_channels]));
-    warmupResult.dataSync();
-    warmupResult.dispose();
     self.showProgress("Model loaded...");
+    });
 
     let fullResCvs = self.__fullsrc;
     fullResCvs.height = step;
@@ -722,7 +725,8 @@ async function segmentModel(key) {
 
         // dummy.getContext('2d').drawImage(img, dx, dy);
         let imgData = fullResCvs.getContext('2d').getImageData(0,0,fullResCvs.width,fullResCvs.height);
-
+        let val;
+        tf.tidy(()=>{
         const img = tf.browser.fromPixels(imgData).toFloat();
         let img2;
         if (input_channels == 1) {
@@ -759,12 +763,13 @@ async function segmentModel(key) {
 
 
         let batched = normalized.reshape([1, image_size, image_size, input_channels]);
-        let values = await model.predict(batched).data();
+        let values = model.predict(batched).dataSync();
         values = Array.from(values);
         //scale values
         values = values.map(x => x * 255)
-        let val = new Array();
+        val = new Array();
         while (values.length > 0) val.push(values.splice(0, image_size));
+        });
         await tf.browser.toPixels(val, temp);
         finalRes.getContext('2d').drawImage(temp, dx, dy);