Reduce memory consumption when writing recorded stimulus

QDSpy_core_presenter.py

@@ -20,7 +20,6 @@
 __author__ = "code@eulerlab.de"
 
 import os
-import pickle
 import platform
 import numpy as np
 import PIL
@@ -698,6 +697,7 @@ def onDraw(self):
                 if self.nFrTotal % self.Conf.rec_f_downsample_t == 0:
                     stimframe = self.View.grabStimFrame()
                     self.recordedStim.append(stimframe)
+                    self.recordedMark.append(self.Stim.cScMarkList[self.iSc])
 
         # Keep track of refresh duration
         if self.Conf.isTrackTime:
@@ -767,6 +767,7 @@ def run(self):
         if self.recordStim:
             self.View.prepareGrabStim()
             self.recordedStim = []
+            self.recordedMark = []
             self.recordedStimName = self.Stim.nameStr
 
         self.Stage.logData()
@@ -857,43 +858,46 @@ def finish(self):
     @staticmethod
     def stim_to_pil_image(
         image, f_downsample: int = 1
-    ) -> PIL.Image.Image:
-        """Convert a stimulus frame into a PIL image
+    ) -> np.ndarray:
+        """Convert a stimulus frame into a numpy.array
         """
         img_data = image.get_data()
-
         pil_image = PIL.Image.new(mode="RGBA", size=(image.width, image.height))
         pil_image.frombytes(img_data)
 
         if f_downsample > 1:
             pil_image = pil_image.resize(
-                tuple(s // f_downsample for s in pil_image.size)
+                tuple(s // f_downsample for s in pil_image.size),
+                resample=PIL.Image.Resampling.HAMMING,
             )
 
         pil_image = pil_image.convert("RGB")
         pil_image = pil_image.transpose(PIL.Image.Transpose.FLIP_TOP_BOTTOM)
-        return pil_image
+
+        frame_array = np.array(pil_image, dtype=np.uint8)
+        return frame_array
 
     @staticmethod
     def adapt_stimulus_recording_to_setup(
-        stimulus_stack: np.array, setup_id: int
-    ) -> np.array:
-        """Tweak stimulus according to 
+        stimulus_stack: np.ndarray, setup_id: int
+    ) -> None:
+        """Tweak stimulus inplace according to 
         https://cin-10.medizin.uni-tuebingen.de/eulerwiki/index.php/Orientation
         stimulus_stack.shape: frame, y, x, color
         """
-        # For both setups the x and y plane is swapped
+        # In-place manipulation by looping over frames (first axis)
         if setup_id == 1:
             # Swap x and y
-            stimulus_stack = stimulus_stack.transpose(0, 2, 1, 3)
+            for i in range(stimulus_stack.shape[0]):
+                stimulus_stack[i] = np.transpose(stimulus_stack[i], (1, 0, 2))
         elif setup_id == 3:
-            # Swap x and y
-            stimulus_stack = stimulus_stack.transpose(0, 2, 1, 3)
-            # flip direction in y-axis
-            stimulus_stack = np.flip(stimulus_stack, axis=1)
+            # Swap x and y, then flip direction in y-axis
+            for i in range(stimulus_stack.shape[0]):
+                frame = np.transpose(stimulus_stack[i], (1, 0, 2))
+                frame = np.flip(frame, axis=0)
+                stimulus_stack[i] = frame
         else:
             raise ValueError(f"Unknown setup: {setup_id=}")
-        return stimulus_stack
 
     def save_stim_to_file(self) -> None:
         """Save (downsampled) stimulus to file 
@@ -905,21 +909,37 @@ def save_stim_to_file(self) -> None:
         if not os.path.isdir(stim_folder):
             os.mkdir(stim_folder)
 
-        pil_image_array = [
-            self.stim_to_pil_image(s, f_downsample=self.Conf.rec_f_downsample_x)
-            for s in self.recordedStim
-        ]
-        recorded_stimulus_stack = np.stack(pil_image_array)
+        n_frames = len(self.recordedStim)
+        height = self.recordedStim[0].height // self.Conf.rec_f_downsample_x
+        width = self.recordedStim[0].width // self.Conf.rec_f_downsample_x
+        channels = 3
+        # Pre-allocate numpy array for all frames to avoid memory spikes
+        recorded_stimulus_np = np.empty(
+            (n_frames, height, width, channels), dtype=np.uint8
+        )
+
+        # Process frames one-by-one and fill directly into pre-allocated array
+        # We also set the processed recordedStim frames to None to reduce memory consumption
+        for i in range(n_frames):
+            frame_array = self.stim_to_pil_image(
+                self.recordedStim[i], f_downsample=self.Conf.rec_f_downsample_x
+            )
+            recorded_stimulus_np[i] = frame_array
+            self.recordedStim[i] = None
+
+        self.recordedStim = []
         if self.Conf.rec_setup_id is not None:
-            recorded_stimulus_stack = self.adapt_stimulus_recording_to_setup(
-                recorded_stimulus_stack, self.Conf.rec_setup_id
+            self.adapt_stimulus_recording_to_setup(
+                recorded_stimulus_np, self.Conf.rec_setup_id
             )
 
-        file_name = f"{stim_folder}/{self.recordedStimName}.pickle"
-        with open(file_name, "wb") as file:
-            pickle.dump(recorded_stimulus_stack, file, protocol=pickle.HIGHEST_PROTOCOL)
+        file_path = fsu.getJoinedPath(stim_folder, f"{self.recordedStimName}.npy")
+        np.save(file_path, recorded_stimulus_np, allow_pickle=False)
+        assert n_frames == len(self.recordedMark)
+        file_path = fsu.getJoinedPath(stim_folder, f"{self.recordedStimName}_mark.npy")
+        np.save(file_path, self.recordedMark, allow_pickle=False)
 
-        Log.write("DEBUG", f"Successfully saved stimulus recording to {file_name}")
+        Log.write("DEBUG", f"Successfully saved stimulus recording to {file_path}")
 
     # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     def prepare(self, _Stim, _Sync=None, _vol=0):