Merge pull request #22 from piximi/issue/18-dispose-tensors

[fix] disposes tensors when reseting data

src/components/dialogs/NewProjectDialog/NewProjectDialog.tsx

@@ -35,7 +35,6 @@ export const NewProjectDialog = ({ onClose, open }: NewProjectDialogProps) => {
       })
     );
     dispatch(dataSlice.actions.resetData());
-
     dispatch(classifierSlice.actions.resetClassifier());
 
     closeDialog();

src/store/data/dataSlice.ts

@@ -72,6 +72,12 @@ export const dataSlice = createSlice({
   initialState: initialState,
   reducers: {
     resetData: (state) => {
+      Object.values(state.things.entities).forEach((entity) => {
+        dispose(entity.saved.data as unknown as TensorContainer);
+        if ("colors" in entity.saved) {
+          dispose(entity.saved.colors as unknown as TensorContainer);
+        }
+      });
       return initialState();
     },
     initializeState(
@@ -849,7 +855,7 @@ export const dataSlice = createSlice({
     ) {
       if (!action.payload.preparedByListener) return;
       if (!("thingIds" in action.payload)) return;
-      const { thingIds, disposeColorTensors, isPermanent } = action.payload;
+      const { thingIds, isPermanent } = action.payload;
       const imageChanges: Record<
         string,
         {
@@ -885,10 +891,8 @@ export const dataSlice = createSlice({
               const kind = state.kinds.entities[thingKind];
               const category = state.categories.entities[thingCategoryId];
               if (isPermanent) {
-                if (disposeColorTensors) {
-                  dispose(containedThing.saved.data as TensorContainer);
-                  dispose(containedThing.changes as TensorContainer);
-                }
+                dispose(containedThing.saved.data as TensorContainer);
+                dispose(containedThing.changes as TensorContainer);
 
                 /* UPDATE KIND'S CONTAINING LIST */
                 mutatingFilter(
@@ -952,10 +956,8 @@ export const dataSlice = createSlice({
         const kind = state.kinds.entities[thingKind];
         const category = state.categories.entities[thingCategoryId];
         if (isPermanent) {
-          if (disposeColorTensors) {
-            dispose(thingEntity.saved.data as TensorContainer);
-            dispose(thingEntity.changes as TensorContainer);
-          }
+          dispose(thingEntity.saved.data as TensorContainer);
+          dispose(thingEntity.changes as TensorContainer);
 
           /* UPDATE KIND'S CONTAINING LIST */
 

src/utils/common/helpers.ts

@@ -441,6 +441,7 @@ export const getPropertiesFromImageSync = (
   });
   const objSrc = objectImage.getCanvas().toDataURL();
   const data = tfImage.cropAndResize(image.data, box, [0], [height, width]);
+  box.dispose();
 
   return {
     data: data,

src/utils/common/tensorHelpers.ts

@@ -378,22 +378,22 @@ export const findMinMaxs = async <T extends Tensor3D | Tensor4D>(
 ): Promise<[number[], number[]]> => {
   let mins: number[];
   let maxs: number[];
+  let minTensor: Tensor1D;
+  let maxTensor: Tensor1D;
 
   if (imageTensor.rank === 3) {
-    mins = await tidy(
-      () => (imageTensor as Tensor3D).min([0, 1]) as Tensor1D
-    ).array();
-    maxs = await tidy(
-      () => (imageTensor as Tensor3D).max([0, 1]) as Tensor1D
-    ).array();
+    minTensor = (imageTensor as Tensor3D).min([0, 1]);
+    mins = minTensor.arraySync();
+    maxTensor = (imageTensor as Tensor3D).max([0, 1]);
+    maxs = maxTensor.arraySync();
   } else {
-    mins = await tidy(
-      () => (imageTensor as Tensor4D).min([0, 1, 2]) as Tensor1D
-    ).array();
-    maxs = await tidy(
-      () => (imageTensor as Tensor4D).max([0, 1, 2]) as Tensor1D
-    ).array();
+    minTensor = (imageTensor as Tensor4D).min([0, 1, 2]);
+    mins = minTensor.arraySync();
+    maxTensor = (imageTensor as Tensor4D).max([0, 1, 2]);
+    maxs = maxTensor.arraySync();
   }
+  minTensor.dispose();
+  maxTensor.dispose();
 
   opts.disposeImageTensor && imageTensor.dispose();
   return [mins, maxs];

src/views/Application/Application.tsx

@@ -47,6 +47,7 @@ export const Application = () => {
     });
 
     dispatch(dataSlice.actions.initializeState({ data: dataState }));
+
     dispatch(projectSlice.actions.setProjectImageChannels({ channels: 3 }));
     setHasLoaded(true);
   }, [dispatch]);