diff --git a/QDSpy_core_presenter.py b/QDSpy_core_presenter.py index d332883..05a0849 100644 --- a/QDSpy_core_presenter.py +++ b/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):